WIP - working role-relative obs wrapper

diambra/arena/wrappers/observation.py

@@ -5,7 +5,7 @@
 from collections.abc import Mapping
 import cv2  # pytype:disable=import-error
 cv2.ocl.setUseOpenCL(False)
-from diambra.engine import model
+from diambra.engine import Roles
 
 # Env Wrappers classes
 class GrayscaleFrame(gym.ObservationWrapper):
@@ -274,17 +274,59 @@ def _obs_normalization_func(self, observation, observation_space):
 
         return observation
 
-class RoleRelativeObservation(gym.ObservationWrapper):
+class RoleRelativeObservation(gym.Wrapper):
     def __init__(self, env):
-        gym.ObservationWrapper.__init__(self, env)
+        gym.Wrapper.__init__(self, env)
+
+        new_observation_space = {}
+        if self.unwrapped.env_settings.n_players == 1:
+            for k, v in self.observation_space.items():
+                if not isinstance(v, gym.spaces.Dict):
+                    new_observation_space[k] = v
+            new_observation_space["own"] = self.observation_space["P1"]
+            new_observation_space["opp"] = self.observation_space["P1"]
+        else:
+            for k, v in self.observation_space.items():
+                if not isinstance(v, gym.spaces.Dict) or k.startswith("agent_"):
+                    new_observation_space[k] = v
+            for idx in range(self.unwrapped.env_settings.n_players):
+                new_observation_space["agent_{}".format(idx)]["own"] = self.observation_space["P1"]
+                new_observation_space["agent_{}".format(idx)]["opp"] = self.observation_space["P1"]
 
-        self.observation_space["own"] = self.observation_space["P1"]
-        self.observation_space["opp"] = self.observation_space["P1"]
-        del self.observation_space["P1"]
-        del self.observation_space["P2"]
+        self.observation_space = gym.spaces.Dict(new_observation_space)
 
-    def observation(self, observation):
-        return None
+    def reset(self, **kwargs):
+        obs, info = self.env.reset(**kwargs)
+        if self.unwrapped.env_settings.n_players == 1:
+            def _process_obs_1p(observation):
+                new_observation = {}
+                role_name = Roles.Name(info["settings"].episode_settings.player_settings[0].role)
+                opponent_role_name = "P2" if role_name == "P1" else "P1"
+                for k, v in observation.items():
+                    if not isinstance(v, dict):
+                        new_observation[k] = v
+                new_observation["own"] = observation[role_name]
+                new_observation["opp"] = observation[opponent_role_name]
+                return new_observation
+            self._process_obs = _process_obs_1p
+        else:
+            def _process_obs_2p(observation):
+                new_observation = {}
+                for k, v in observation.items():
+                    if not isinstance(v, dict) or k.startswith("agent_"):
+                        new_observation[k] = v
+                for idx in range(self.unwrapped.env_settings.n_players):
+                    role_name = Roles.Name(info["settings"].episode_settings.player_settings[idx].role)
+                    opponent_role_name = "P2" if role_name == "P1" else "P1"
+                    new_observation["agent_{}".format(idx)]["own"] = observation[role_name]
+                    new_observation["agent_{}".format(idx)]["opp"] = observation[opponent_role_name]
+                return new_observation
+            self._process_obs = _process_obs_2p
+        return self._process_obs(obs), info
+
+    def step(self, action):
+        obs, reward, terminated, truncated, info = self.env.step(action)
+        return self._process_obs(obs), reward, terminated, truncated, info
 
     """
     def rename_key_recursive(dictionary, old_key, new_key):

examples/wrappers_options.py

@@ -64,7 +64,6 @@ def main():
     #wrappers_settings["exclude_image_scaling"] = True
     #wrappers_settings["process_discrete_binary"] = True
 
-    """
     # If to make the observation relative to the agent as a function to its role (P1 or P2) (deactivate by default)
     # i.e.:
     #  - In 1P environments, if the agent is P1 then the observation "P1" nesting level becomes "own" and "P2" becomes "opp"
@@ -74,7 +73,7 @@ def main():
     #    - Under "agent_1", "P1" nesting level becomes "opp" and "P2" becomes "own"
     wrappers_settings["role_relative_observation"] = True
 
-
+    """
     # Flattening observation dictionary and filtering
     # a sub-set of the RAM states
     wrappers_settings["flatten"] = True