Add support for MultiDiscrete/MultiBinary action spaces

[araffin]

• distributions.py need to be updated (and maybe ppo/a2c) with MultiCategorical and Bernoulli distributions

• the envs from identity_env.py should help to create tests

@rolandgvc is working on it

[araffin]

Flattened seems the easiest and cleanest way, no?

[Miffyli]

Hmm where did the question go disappear? o: .I would like to comment on this after taking a good look at how it could be done. A multi-discrete distribution returning flattened output seems bit convoluted at first sight, but I will comment better when there is an example code suggestion.

[rolandgvc]

@Miffyli deleted it before @araffin answered, as I though it was obvious =) here is the suggestion:

class MultiCategoricalDistribution(Distribution):
    """
    MultiCategorical distribution for multi discrete actions.
    
    :param action_dims: ([int]) List of sizes of discrete action spaces.
    """
    def __init__(self, action_dims: [int]):
        super(MultiCategoricalDistribution, self).__init__()
        self.action_dims = action_dims
        self.distributions = None

    def proba_distribution_net(self, latent_dim: int) -> nn.Module:
        """
        Create the layer that represents the distribution:
        it will be the logits of the Categorical distribution.
        You can then get probabilities using a softmax.

        :param latent_dim: (int) Dimension of the last layer
            of the policy network (before the action layer)
        :return: (nn.Linear)
        """
        action_logits = nn.Linear(latent_dim, np.sum(self.action_dims))
        return action_logits 

    def proba_distribution(self, action_logits: th.Tensor) -> 'MultiCategoricalDistribution':
        reshaped_logits = action_logits.split(self.action_dims)
        self.distributions = [Categorical(logits=l) for l in reshaped_logits]
        return self

[rolandgvc]

Progress so far (tested):

class MultiCategoricalDistribution(Distribution):
    """
    MultiCategorical distribution for multi discrete actions.

    :param action_dims: ([int]) List of sizes of discrete action spaces.
    """
    def __init__(self, action_dims: [int]):
        super(MultiCategoricalDistribution, self).__init__()
        self.action_dims = action_dims
        self.distributions = None

    def proba_distribution_net(self, latent_dim: int) -> nn.Module:
        """
        Create the layer that represents the distribution:
        it will be the logits of the Categorical distribution.
        You can then get probabilities using a softmax.

        :param latent_dim: (int) Dimension of the last layer
            of the policy network (before the action layer)
        :return: (nn.Linear)
        """
        action_logits = nn.Linear(latent_dim, np.sum(self.action_dims))
        return action_logits 

    def proba_distribution(self, action_logits: th.Tensor) -> 'MultiCategoricalDistribution':
        reshaped_logits = th.split(action_logits, self.action_dims, dim=1)
        self.distributions = [Categorical(logits=l) for l in reshaped_logits]
        return self

    def mode(self) -> th.Tensor:
        return th.stack([th.argmax(d.probs, dim=1) for d in self.distributions])

    def sample(self) -> th.Tensor:
        return th.stack([d.sample() for d in self.distributions])

    def entropy(self) -> th.Tensor:
        return sum([d.entropy() for d in self.distributions])

    def actions_from_params(self, action_logits: th.Tensor,
                            deterministic: bool = False) -> th.Tensor:
        # Update the proba distribution
        self.proba_distribution(action_logits)
        return self.get_actions(deterministic=deterministic)

    def log_prob_from_params(self, action_logits: th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
        actions = self.actions_from_params(action_logits)
        log_prob = self.log_prob(actions)
        return actions, log_prob

    def log_prob(self, actions: th.Tensor) -> th.Tensor:
        return sum(d.log_prob(x) for d, x in zip(self.distributions, th.unbind(actions)))

Let me know if you have any design suggestions @Miffyli @araffin

[araffin]

I think it would be better if you open a draft pull request ;)

[rolandgvc]

@araffin can we just use the .shape attribute for multi spaces here?

def get_obs_shape(observation_space: spaces.Space) -> Tuple[int, ...]:
    """
    Get the shape of the observation (useful for the buffers).

    :param observation_space: (spaces.Space)
    :return: (Tuple[int, ...])
    """
    if isinstance(observation_space, spaces.Box):
        return observation_space.shape
    elif isinstance(observation_space, spaces.Discrete):
        # Observation is an int
        return 1
    elif isinstance(observation_space, spaces.MultiDiscrete):
        return observation_space.shape
    elif isinstance(observation_space, spaces.MultiBinary):
        return observation_space.shape
    else:
        raise NotImplementedError()

[araffin]

@rolandgvc this is for another issue (this one: #4 ) no?

Looking at the source, this won't work as obs.shape is not defined for MultiBinary, and same for multi discrete

As Gym is not documented, I really recommend to read the source code.