[Bug] get_obs_shape returns wrong shape for MultiBinary spaces

Important Note: We do not do technical support, nor consulting and don’t answer personal questions per email. Please post your question on the RL Discord, Reddit or Stack Overflow in that case.

If your issue is related to a custom gym environment, please use the custom gym env template.

🐛 Bug

stable_baselines3.common.preprocessing.get_obs_shape returns the wrong shape when a MultiBinary spaces is multi-dimensions.

To Reproduce

Steps to reproduce the behavior.

Please try to provide a minimal example to reproduce the bug. Error messages and stack traces are also helpful.

from gym import spaces
from stable_baselines3.common.preprocessing import get_obs_shape


test_multi = spaces.MultiBinary([5, 4, 5])

get_obs_shape(test_multi)

    151     elif isinstance(observation_space, spaces.MultiBinary):
    152         # Number of binary features
--> 153         return (int(observation_space.n),)
    154     elif isinstance(observation_space, spaces.Dict):
    155         return {key: get_obs_shape(subspace) for (key, subspace) in observation_space.spaces.items()}

TypeError: int() argument must be a string, a bytes-like object or a number, not 'list'

Expected behavior

get_obs_shape should returnobservation_space.shape not int(observation_space.n)

### System Info

Describe the characteristic of your environment:

• Describe how the library was installed (pip, docker, source, …): pip
• GPU models and configuration: none
• Python version: 3.7.
• PyTorch version: 1.9.0
• Gym version" 0.18.3

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)

[ylchan87]

Despite the doc, turns out gym itself also dislike N-Dim MultiBinary,

def flatdim(space):
    ....
    elif isinstance(space, MultiBinary):
>        return int(space.n)

E           TypeError: int() argument must be a string, a bytes-like object or a number, not 'list'

https://github.com/openai/gym/blob/4ede9280f9c477f1ca09929d10cdc1e1ba1129f1/gym/spaces/utils.py#L28

So I guess a wrapper would be the solution, not a temporary one

[Miffyli]

I think the reshaping will be done regardless, so making a quick wrapper (for now, before official support) would be sufficient. Something along lines of (not tested, just a sketch):

class MultiBinaryFlattenWrapper(gym.Wrapper):
    def __init__(self, env):
        super().__init__(env)
        self.original_shape = self.observation_space.n
        assert isinstance(self.original_shape, tuple) and len(self.original_shape) > 1
        self.observation_space = gym.spaces.MultiBinary(n=np.product(self.observation_space.n))

    def step(self, action):
        action = action.reshape(self.original_shape)
        return self.env.step(action)