Added a compatibility to an obs with multiple informations in aec env

supersuit/lambda_wrappers/observation_lambda.py

@@ -45,6 +45,8 @@ def _check_wrapper_params(self):
     def observation_space(self, agent):
         if self.change_obs_space_fn is None:
             space = self.env.observation_space(agent)
+            if isinstance(space, gymnasium.spaces.Dict):
+                space = space["observation"]
             try:
                 trans_low = self.change_observation_fn(space.low, space, agent)
                 trans_high = self.change_observation_fn(space.high, space, agent)
@@ -57,13 +59,17 @@ def observation_space(self, agent):
             return Box(low=new_low, high=new_high, dtype=new_low.dtype)
         else:
             old_obs_space = self.env.observation_space(agent)
+            if isinstance(old_obs_space, gymnasium.spaces.Dict):
+                old_obs_space = old_obs_space["observation"]
             try:
                 return self.change_obs_space_fn(old_obs_space, agent)
             except TypeError:
                 return self.change_obs_space_fn(old_obs_space)
 
     def _modify_observation(self, agent, observation):
         old_obs_space = self.env.observation_space(agent)
+        if isinstance(old_obs_space, gymnasium.spaces.Dict):
+            old_obs_space = old_obs_space["observation"]
         try:
             return self.change_observation_fn(observation, old_obs_space, agent)
         except TypeError:

supersuit/utils/base_aec_wrapper.py

@@ -35,6 +35,9 @@ def observe(self, agent):
         obs = super().observe(
             agent
         )  # problem is in this line, the obs is sometimes a different size from the obs space
+        # In the case of more information in the obs
+        if isinstance(obs, dict):
+            obs = obs["observation"]
         observation = self._modify_observation(agent, obs)
         return observation