[BUG] action masking does not work with VecEnv and MultiDiscrete action space about stable-baselines3-contrib

[araffin]

Hello,
best is to start with a working example:

stable-baselines3-contrib/sb3_contrib/common/envs/invalid_actions_env.py

Line 39 in 75b2de1

class InvalidActionEnvMultiDiscrete(IdentityEnv):

that being said, there might be a bug too.
Tagging @kronion and @vwxyzjn as they actually worked with it.

[kronion]

It might be a bug, but it's hard to say from the description. Could you share the code to reproduce? And could you show an example of how the mask is being split weirdly? My initial impression is that the (128, 360) shape is intended because each row corresponds to an env in the vecenv.

[araffin]

BUT: the shape of the mask is not (360,) or (1,360) but instead it is (128, 360)

this actually looks good to me, we need to retrieve one mask per env.
Does it produce an error?
if so, please provide a minimal example to reproduce the issue and provide the traceback.

(fyi I think that we expect 1D mask from the env even for multi discrete (see #80 (comment)), it will be reshaped by the algorithm afterward)