Supersuit Wrapper with Action Mask

[YiqingDing]

Hi, I am pretty new to this and I am trying to use Supersuit multiagent wrappers (both pad_observations_v0 and pad_action_space_v0) with action masking. Once I applied both wrappers to the environment (sequentially) and ran env.last(), I received the following error:

Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/supersuit/lambda_wrappers/observation_lambda.py", line 68, in _modify_observation
    return self.change_observation_fn(observation, old_obs_space, agent)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: pad_observations_v0.<locals>.<lambda>() takes 2 positional arguments but 3 were given

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/Base Code/test.py", line 40, in <module>
    env.last()
  File "/usr/local/lib/python3.11/site-packages/pettingzoo/utils/env.py", line 186, in last
    observation = self.observe(agent) if observe else None
                  ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/supersuit/utils/base_aec_wrapper.py", line 35, in observe
    obs = super().observe(
          ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/pettingzoo/utils/wrappers/order_enforcing.py", line 101, in observe
    return super().observe(agent)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/pettingzoo/utils/wrappers/base.py", line 41, in observe
    return self.env.observe(agent)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/supersuit/utils/base_aec_wrapper.py", line 38, in observe
    observation = self._modify_observation(agent, obs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/supersuit/lambda_wrappers/observation_lambda.py", line 70, in _modify_observation
    return self.change_observation_fn(observation, old_obs_space)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/supersuit/multiagent_wrappers/padding_wrappers.py", line 38, in <lambda>
    lambda obs, obs_space: homogenize_ops.homogenize_observations(
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/supersuit/utils/action_transforms/homogenize_ops.py", line 99, in homogenize_observations
    return pad_to(obs, obs_space.shape, 0)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/supersuit/utils/action_transforms/homogenize_ops.py", line 34, in pad_to
    old_shape = arr.shape
                ^^^^^^^^^
AttributeError: 'dict' object has no attribute 'shape'

It appears that when obs = self.observe() returns a dictionary (with action mask), supersuit's wrappers passed the entire dictionary instead of the observation keyed value (obs['observation']), correct? Is there anyway to use supersuit with action masking in this case? Thanks!