def get_action_dist(action_space, config, dist_type=None): """Returns action distribution class and size for the given action space. Args: action_space (Space): Action space of the target gym env. config (dict): Optional model config. dist_type (str): Optional identifier of the action distribution. Returns: dist_class (ActionDistribution): Python class of the distribution. dist_dim (int): The size of the input vector to the distribution. """ config = config or MODEL_DEFAULTS if isinstance(action_space, gym.spaces.Box): if len(action_space.shape) > 1: raise ValueError( "Action space has multiple dimensions " "{}. ".format(action_space.shape) + "Consider reshaping this into a single dimension, " "using a Tuple action space, or the multi-agent API.") if dist_type is None: dist = DiagGaussian if config.get("squash_to_range"): dist = squash_to_range(dist, action_space.low, action_space.high) return dist, action_space.shape[0] * 2 elif dist_type == "deterministic": return Deterministic, action_space.shape[0] elif isinstance(action_space, gym.spaces.Discrete): return Categorical, action_space.n elif isinstance(action_space, gym.spaces.Tuple): child_dist = [] input_lens = [] for action in action_space.spaces: dist, action_size = ModelCatalog.get_action_dist( action, config) child_dist.append(dist) input_lens.append(action_size) return partial( MultiActionDistribution, child_distributions=child_dist, action_space=action_space, input_lens=input_lens), sum(input_lens) raise NotImplementedError("Unsupported args: {} {}".format( action_space, dist_type))
def get_action_dist(action_space, config, dist_type=None): """Returns action distribution class and size for the given action space. Args: action_space (Space): Action space of the target gym env. config (dict): Optional model config. dist_type (str): Optional identifier of the action distribution. Returns: dist_class (ActionDistribution): Python class of the distribution. dist_dim (int): The size of the input vector to the distribution. """ # TODO(ekl) are list spaces valid? if isinstance(action_space, list): action_space = gym.spaces.Tuple(action_space) config = config or {} if isinstance(action_space, gym.spaces.Box): if dist_type is None: dist = DiagGaussian if config.get("squash_to_range"): dist = squash_to_range(dist, action_space.low, action_space.high) return dist, action_space.shape[0] * 2 elif dist_type == 'deterministic': return Deterministic, action_space.shape[0] elif isinstance(action_space, gym.spaces.Discrete): return Categorical, action_space.n elif isinstance(action_space, gym.spaces.Tuple): child_dist = [] input_lens = [] for action in action_space.spaces: dist, action_size = ModelCatalog.get_action_dist( action, config) child_dist.append(dist) input_lens.append(action_size) return partial(MultiActionDistribution, child_distributions=child_dist, action_space=action_space, input_lens=input_lens), sum(input_lens) raise NotImplementedError("Unsupported args: {} {}".format( action_space, dist_type))