leor-c
diff --git a/‎stable_baselines3/common/vec_env/obs_dict_wrapper.py‎
Lines changed: 68 additions & 0 deletions b/‎stable_baselines3/common/vec_env/obs_dict_wrapper.py‎
Lines changed: 68 additions & 0 deletions
@@ -0,0 +1,68 @@
+from typing import Dict
+
+import numpy as np
+from gym import spaces
+
+from stable_baselines3.common.vec_env import VecEnv, VecEnvWrapper
+
+
+class ObsDictWrapper(VecEnvWrapper):
+    """
+    Wrapper for a VecEnv which overrides the observation space for Hindsight Experience Replay to support dict observations.
+
+    :param env: The vectorized environment to wrap.
+    """
+
+    def __init__(self, venv: VecEnv):
+        super(ObsDictWrapper, self).__init__(venv, venv.observation_space, venv.action_space)
+
+        self.venv = venv
+
+        self.spaces = list(venv.observation_space.spaces.values())
+
+        # get dimensions of observation and goal
+        if isinstance(self.spaces[0], spaces.Discrete):
+            self.obs_dim = 1
+            self.goal_dim = 1
+        else:
+            self.obs_dim = venv.observation_space.spaces["observation"].shape[0]
+            self.goal_dim = venv.observation_space.spaces["achieved_goal"].shape[0]
+
+        # new observation space with concatenated observation and (desired) goal
+        # for the different types of spaces
+        if isinstance(self.spaces[0], spaces.Box):
+            low_values = np.concatenate(
+                [venv.observation_space.spaces["observation"].low, venv.observation_space.spaces["desired_goal"].low]
+            )
+            high_values = np.concatenate(
+                [venv.observation_space.spaces["observation"].high, venv.observation_space.spaces["desired_goal"].high]
+            )
+            self.observation_space = spaces.Box(low_values, high_values, dtype=np.float32)
+        elif isinstance(self.spaces[0], spaces.MultiBinary):
+            total_dim = self.obs_dim + self.goal_dim
+            self.observation_space = spaces.MultiBinary(total_dim)
+        elif isinstance(self.spaces[0], spaces.Discrete):
+            dimensions = [venv.observation_space.spaces["observation"].n, venv.observation_space.spaces["desired_goal"].n]
+            self.observation_space = spaces.MultiDiscrete(dimensions)
+        else:
+            raise NotImplementedError(f"{type(self.spaces[0])} space is not supported")
+
+    def reset(self):
+        return self.venv.reset()
+
+    def step_wait(self):
+        return self.venv.step_wait()
+
+    @staticmethod
+    def convert_dict(
+        observation_dict: Dict[str, np.ndarray], observation_key: str = "observation", goal_key: str = "desired_goal"
+    ) -> np.ndarray:
+        """
+        Concatenate observation and (desired) goal of observation dict.
+
+        :param observation_dict: Dictionary with observation.
+        :param observation_key: Key of observation in dicitonary.
+        :param goal_key: Key of (desired) goal in dicitonary.
+        :return: Concatenated observation.
+        """
+        return np.concatenate([observation_dict[observation_key], observation_dict[goal_key]], axis=-1)