Dictionary Observations

docs/misc/changelog.rst

@@ -8,6 +8,7 @@ Pre-Release 0.11.0a1 (WIP)
 
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
+- Breaks HER as HER needs to be updated to use the new dictionary observations
 - ``evaluate_policy`` now returns rewards/episode lengths from a ``Monitor`` wrapper if one is present,
   this allows to return the unnormalized reward in the case of Atari games for instance.
 - Renamed ``common.vec_env.is_wrapped`` to ``common.vec_env.is_vecenv_wrapped`` to avoid confusion
@@ -19,13 +20,16 @@ New Features:
   automatic check for image spaces.
 - ``VecFrameStack`` now has a ``channels_order`` argument to tell if observations should be stacked
   on the first or last observation dimension (originally always stacked on last).
+- Add support for dictionary observations in both RolloutBuffer (need to be tested in ReplayBuffer)
+- Added simple 4x4 and 9room test environments
 - Added ``common.env_util.is_wrapped`` and ``common.env_util.unwrap_wrapper`` functions for checking/unwrapping
   an environment for specific wrapper.
 - Added ``env_is_wrapped()`` method for ``VecEnv`` to check if its environments are wrapped
   with given Gym wrappers.
 - Added ``monitor_kwargs`` parameter to ``make_vec_env`` and ``make_atari_env``
 - Wrap the environments automatically with a ``Monitor`` wrapper when possible.
 
+
 Bug Fixes:
 ^^^^^^^^^^
 - Fixed bug where code added VecTranspose on channel-first image environments (thanks @qxcv)

multi_input_tests.py

@@ -0,0 +1,87 @@
+import argparse
+import gym
+import numpy as np
+
+from stable_baselines3 import PPO, SAC
+from stable_baselines3.common.policies import MultiInputActorCriticPolicy
+from stable_baselines3.common.vec_env import (
+    DummyVecEnv,
+    VecFrameStack,
+    VecTransposeImage,
+)
+
+from stable_baselines3.common.multi_input_envs import (
+    SimpleMultiObsEnv,
+    NineRoomMultiObsEnv,
+)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Runs the multi_input_tests script")
+    parser.add_argument(
+        "--timesteps",
+        type=int,
+        default=30000,
+        help="Number of timesteps to train for (default: 20000)",
+    )
+    parser.add_argument(
+        "--num_envs",
+        type=int,
+        default=10,
+        help="Number of environments to use (default: 10)",
+    )
+    parser.add_argument(
+        "--frame_stacks",
+        type=int,
+        default=1,
+        help="Number of stacked frames to use (default: 4)",
+    )
+    parser.add_argument(
+        "--room9",
+        action="store_true",
+        help="If true, uses more complex 9 room environment",
+    )
+    args = parser.parse_args()
+
+    ENV_CLS = NineRoomMultiObsEnv if args.room9 else SimpleMultiObsEnv
+
+    make_env = lambda: ENV_CLS(random_start=True)
+
+    env = DummyVecEnv([make_env for i in range(args.num_envs)])
+    if args.frame_stacks > 1:
+        env = VecFrameStack(env, n_stack=args.frame_stacks)
+
+    model = PPO(MultiInputActorCriticPolicy, env)
+
+    model.learn(args.timesteps)
+    env.close()
+    print("Done training, starting testing")
+
+    make_env = lambda: ENV_CLS(random_start=False)
+    test_env = DummyVecEnv([make_env])
+    if args.frame_stacks > 1:
+        test_env = VecFrameStack(test_env, n_stack=args.frame_stacks)
+
+    obs = test_env.reset()
+    num_episodes = 1
+    trajectories = [[]]
+    i_step, i_episode = 0, 0
+    while i_episode < num_episodes:
+        action, _states = model.predict(obs, deterministic=False)
+        obs, reward, done, info = test_env.step(action)
+        test_env.render()
+        trajectories[-1].append((test_env.get_attr("state")[0], action[0]))
+
+        i_step += 1
+
+        if done[0]:
+            if info[0]["got_to_end"]:
+                print(f"Episode {i_episode} : Got to end in {i_step} steps")
+            else:
+                print(f"Episode {i_episode} : Did not get to end")
+            obs = test_env.reset()
+            i_step = 0
+            trajectories.append([])
+            i_episode += 1
+
+    test_env.close()

stable_baselines3/common/base_class.py

@@ -197,7 +197,9 @@ def _wrap_env(env: GymEnv, verbose: int = 0, monitor_wrapper: bool = True) -> Ve
             env = VecTransposeImage(env)
 
         # check if wrapper for dict support is needed when using HER
-        if isinstance(env.observation_space, gym.spaces.dict.Dict):
+        if isinstance(env.observation_space, gym.spaces.dict.Dict) and set(env.observation_space.spaces.keys()) == set(
+            ["observation", "desired_goal", "achieved_goal"]
+        ):
             env = ObsDictWrapper(env)
 
         return env

stable_baselines3/common/buffers.py

@@ -13,7 +13,12 @@
     psutil = None
 
 from stable_baselines3.common.preprocessing import get_action_dim, get_obs_shape
-from stable_baselines3.common.type_aliases import ReplayBufferSamples, RolloutBufferSamples
+from stable_baselines3.common.type_aliases import (
+    DictReplayBufferSamples,
+    DictRolloutBufferSamples,
+    ReplayBufferSamples,
+    RolloutBufferSamples,
+)
 from stable_baselines3.common.vec_env import VecNormalize
 
 
@@ -42,6 +47,7 @@ def __init__(
         self.observation_space = observation_space
         self.action_space = action_space
         self.obs_shape = get_obs_shape(observation_space)
+        self.is_dict_data = isinstance(self.observation_space, spaces.Dict)
         self.action_dim = get_action_dim(action_space)
         self.pos = 0
         self.full = False
@@ -130,7 +136,8 @@ def to_torch(self, array: np.ndarray, copy: bool = True) -> th.Tensor:
 
     @staticmethod
     def _normalize_obs(
-        obs: Union[np.ndarray, Dict[str, np.ndarray]], env: Optional[VecNormalize] = None
+        obs: Union[np.ndarray, Dict[str, np.ndarray]],
+        env: Optional[VecNormalize] = None,
     ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
         if env is not None:
             return env.normalize_obs(obs)
@@ -177,20 +184,50 @@ def __init__(
             mem_available = psutil.virtual_memory().available
 
         self.optimize_memory_usage = optimize_memory_usage
-        self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=observation_space.dtype)
+
+        if self.is_dict_data:
+            self.observations = {
+                key: np.zeros((self.buffer_size, self.n_envs) + _obs_shape) for key, _obs_shape in self.obs_shape.items()
+            }
+        else:
+            self.observations = np.zeros(
+                (self.buffer_size, self.n_envs) + self.obs_shape,
+                dtype=observation_space.dtype,
+            )
         if optimize_memory_usage:
             # `observations` contains also the next observation
             self.next_observations = None
         else:
-            self.next_observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=observation_space.dtype)
+            if self.is_dict_data:
+                self.next_observations = {
+                    key: np.zeros((self.buffer_size, self.n_envs) + _obs_shape) for key, _obs_shape in self.obs_shape.items()
+                }
+            else:
+                self.next_observations = np.zeros(
+                    (self.buffer_size, self.n_envs) + self.obs_shape,
+                    dtype=observation_space.dtype,
+                )
         self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=action_space.dtype)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.dones = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
 
         if psutil is not None:
-            total_memory_usage = self.observations.nbytes + self.actions.nbytes + self.rewards.nbytes + self.dones.nbytes
+            obs_nbytes = 0
+            if self.is_dict_data:
+                for key, obs in self.observations.items():
+                    obs_nbytes += obs.nbytes
+            else:
+                obs_nbytes = self.observations.nbytes
+
+            total_memory_usage = obs_nbytes + self.actions.nbytes + self.rewards.nbytes + self.dones.nbytes
             if self.next_observations is not None:
-                total_memory_usage += self.next_observations.nbytes
+                next_obs_nbytes = 0
+                if self.is_dict_data:
+                    for key, obs in self.observations.items():
+                        next_obs_nbytes += obs.nbytes
+                else:
+                    next_obs_nbytes = self.next_observations.nbytes
+                total_memory_usage += next_obs_nbytes
 
             if total_memory_usage > mem_available:
                 # Convert to GB
@@ -201,13 +238,34 @@ def __init__(
                     f"replay buffer {total_memory_usage:.2f}GB > {mem_available:.2f}GB"
                 )
 
-    def add(self, obs: np.ndarray, next_obs: np.ndarray, action: np.ndarray, reward: np.ndarray, done: np.ndarray) -> None:
+    def add(
+        self,
+        obs: Union[np.ndarray, dict],
+        next_obs: np.ndarray,
+        action: np.ndarray,
+        reward: np.ndarray,
+        done: np.ndarray,
+    ) -> None:
         # Copy to avoid modification by reference
-        self.observations[self.pos] = np.array(obs).copy()
+
+        if self.is_dict_data:
+            for key in self.observations.keys():
+                self.observations[key][self.pos] = np.array(obs[key]).copy()
+        else:
+            self.observations[self.pos] = np.array(obs).copy()
+
         if self.optimize_memory_usage:
-            self.observations[(self.pos + 1) % self.buffer_size] = np.array(next_obs).copy()
+            if self.is_dict_data:
+                for key in self.observations.keys():
+                    self.observations[key][(self.pos + 1) % self.buffer_size] = np.array(next_obs[key]).copy()
+            else:
+                self.observations[(self.pos + 1) % self.buffer_size] = np.array(next_obs).copy()
         else:
-            self.next_observations[self.pos] = np.array(next_obs).copy()
+            if self.is_dict_data:
+                for key in self.next_observations.keys():
+                    self.next_observations[key][self.pos] = np.array(next_obs[key]).copy()
+            else:
+                self.next_observations[self.pos] = np.array(next_obs).copy()
 
         self.actions[self.pos] = np.array(action).copy()
         self.rewards[self.pos] = np.array(reward).copy()
@@ -241,6 +299,35 @@ def sample(self, batch_size: int, env: Optional[VecNormalize] = None) -> ReplayB
         return self._get_samples(batch_inds, env=env)
 
     def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> ReplayBufferSamples:
+
+        if self.is_dict_data:
+            if self.optimize_memory_usage:
+                next_obs = {
+                    key: self.to_torch(
+                        self._normalize_obs(
+                            obs[(batch_inds + 1) % self.buffer_size, 0, :],
+                            env,
+                        )
+                    )
+                    for key, obs in self.observations.items()
+                }
+            else:
+                next_obs = {
+                    key: self.to_torch(self._normalize_obs(obs[batch_inds, 0, :], env))
+                    for key, obs in self.next_observations.items()
+                }
+
+            normalized_obs = {
+                key: self.to_torch(self._normalize_obs(obs[batch_inds, 0, :], env)) for key, obs in self.observations.items()
+            }
+            return DictReplayBufferSamples(
+                observations=normalized_obs,
+                actions=self.to_torch(self.actions[batch_inds]),
+                next_observations=next_obs,
+                dones=self.to_torch(self.dones[batch_inds]),
+                returns=self.to_torch(self._normalize_reward(self.rewards[batch_inds], env)),
+            )
+
         if self.optimize_memory_usage:
             next_obs = self._normalize_obs(self.observations[(batch_inds + 1) % self.buffer_size, 0, :], env)
         else:
@@ -293,13 +380,24 @@ def __init__(
         super(RolloutBuffer, self).__init__(buffer_size, observation_space, action_space, device, n_envs=n_envs)
         self.gae_lambda = gae_lambda
         self.gamma = gamma
-        self.observations, self.actions, self.rewards, self.advantages = None, None, None, None
+        self.observations, self.actions, self.rewards, self.advantages = (
+            None,
+            None,
+            None,
+            None,
+        )
         self.returns, self.dones, self.values, self.log_probs = None, None, None, None
         self.generator_ready = False
         self.reset()
 
     def reset(self) -> None:
-        self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=np.float32)
+
+        if self.is_dict_data:
+            self.observations = {}
+            for (key, obs_input_shape) in self.obs_shape.items():
+                self.observations[key] = np.zeros((self.buffer_size, self.n_envs) + obs_input_shape, dtype=np.float32)
+        else:
+            self.observations = np.zeros((self.buffer_size, self.n_envs) + self.obs_shape, dtype=np.float32)
         self.actions = np.zeros((self.buffer_size, self.n_envs, self.action_dim), dtype=np.float32)
         self.rewards = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
         self.returns = np.zeros((self.buffer_size, self.n_envs), dtype=np.float32)
@@ -342,7 +440,13 @@ def compute_returns_and_advantage(self, last_values: th.Tensor, dones: np.ndarra
         self.returns = self.advantages + self.values
 
     def add(
-        self, obs: np.ndarray, action: np.ndarray, reward: np.ndarray, done: np.ndarray, value: th.Tensor, log_prob: th.Tensor
+        self,
+        obs: Union[np.ndarray, dict],
+        action: np.ndarray,
+        reward: np.ndarray,
+        done: np.ndarray,
+        value: th.Tensor,
+        log_prob: th.Tensor,
     ) -> None:
         """
         :param obs: Observation
@@ -358,7 +462,11 @@ def add(
             # Reshape 0-d tensor to avoid error
             log_prob = log_prob.reshape(-1, 1)
 
-        self.observations[self.pos] = np.array(obs).copy()
+        if self.is_dict_data:
+            for key in self.observations.keys():
+                self.observations[key][self.pos] = np.array(obs[key]).copy()
+        else:
+            self.observations[self.pos] = np.array(obs).copy()
         self.actions[self.pos] = np.array(action).copy()
         self.rewards[self.pos] = np.array(reward).copy()
         self.dones[self.pos] = np.array(done).copy()
@@ -373,7 +481,15 @@ def get(self, batch_size: Optional[int] = None) -> Generator[RolloutBufferSample
         indices = np.random.permutation(self.buffer_size * self.n_envs)
         # Prepare the data
         if not self.generator_ready:
-            for tensor in ["observations", "actions", "values", "log_probs", "advantages", "returns"]:
+
+            _tensor_names = ["actions", "values", "log_probs", "advantages", "returns"]
+            if self.is_dict_data:
+                for key, obs in self.observations.items():
+                    self.observations[key] = self.swap_and_flatten(obs)
+            else:
+                _tensor_names.append("observations")
+
+            for tensor in _tensor_names:
                 self.__dict__[tensor] = self.swap_and_flatten(self.__dict__[tensor])
             self.generator_ready = True
 
@@ -387,6 +503,16 @@ def get(self, batch_size: Optional[int] = None) -> Generator[RolloutBufferSample
             start_idx += batch_size
 
     def _get_samples(self, batch_inds: np.ndarray, env: Optional[VecNormalize] = None) -> RolloutBufferSamples:
+        if self.is_dict_data:
+            return DictRolloutBufferSamples(
+                observations={key: self.to_torch(obs[batch_inds]) for (key, obs) in self.observations.items()},
+                actions=self.to_torch(self.actions[batch_inds]),
+                old_values=self.to_torch(self.values[batch_inds].flatten()),
+                old_log_prob=self.to_torch(self.log_probs[batch_inds].flatten()),
+                advantages=self.to_torch(self.advantages[batch_inds].flatten()),
+                returns=self.to_torch(self.returns[batch_inds].flatten()),
+            )
+
         data = (
             self.observations[batch_inds],
             self.actions[batch_inds],