Add partial support for dictionary observation spaces (bc, density)

setup.py

@@ -17,7 +17,7 @@
 ATARI_REQUIRE = [
     "seals[atari]~=0.2.1",
 ]
-PYTYPE = ["pytype==2023.9.27"] if IS_NOT_WINDOWS else []
+PYTYPE = ["pytype==2022.7.26"] if IS_NOT_WINDOWS else []
 
 # Note: the versions of the test and doc requirements should be tightly pinned to known
 #   working versions to make our CI/CD pipeline as stable as possible.

src/imitation/algorithms/dagger.py

@@ -508,7 +508,7 @@ def create_trajectory_collector(self) -> InteractiveTrajectoryCollector:
         beta = self.beta_schedule(self.round_num)
         collector = InteractiveTrajectoryCollector(
             venv=self.venv,
-            get_robot_acts=lambda obs: self.bc_trainer.policy.predict(obs)[0],
+            get_robot_acts=lambda acts: self.bc_trainer.policy.predict(acts)[0],
             beta=beta,
             save_dir=save_dir,
             rng=self.rng,

src/imitation/algorithms/density.py

@@ -145,14 +145,15 @@ def _get_demo_from_batch(
             )
 
         assert act_b.shape[1:] == self.venv.action_space.shape
-
+        ob_space = self.venv.observation_space
         if isinstance(obs_b, types.DictObs):
-            exp_shape = {k: v.shape for k, v in self.venv.observation_space.items()}  # type: ignore[attr-defined] # noqa: E501
-
+            exp_shape = {
+                k: v.shape for k, v in ob_space.items()  # type: ignore[attr-defined]
+            }
             obs_shape = {k: v.shape[1:] for k, v in obs_b.items()}
             assert exp_shape == obs_shape, f"Expected {exp_shape}, got {obs_shape}"
         else:
-            assert obs_b.shape[1:] == self.venv.observation_space.shape
+            assert obs_b.shape[1:] == ob_space.shape
         assert len(act_b) == len(obs_b)
         if next_obs_b is not None:
             assert next_obs_b.shape == obs_b.shape

src/imitation/data/rollout.py

@@ -490,10 +490,9 @@ def generate_trajectories(
                 assert v.shape is not None
                 exp_obs[k] = (n_steps + 1,) + v.shape
         else:
-            assert venv.observation_space.shape is not None
-            exp_obs = (
-                n_steps + 1,
-            ) + venv.observation_space.shape  # type: ignore[assignment]
+            obs_space_shape = venv.observation_space.shape
+            assert obs_space_shape is not None
+            exp_obs = (n_steps + 1,) + obs_space_shape  # type: ignore[assignment]
         real_obs = trajectory.obs.shape
         assert real_obs == exp_obs, f"expected shape {exp_obs}, got {real_obs}"
         assert venv.action_space.shape is not None

src/imitation/data/wrappers.py

@@ -1,18 +1,14 @@
 """Environment wrappers for collecting rollouts."""
 
-from typing import List, Optional, Sequence, Tuple, Dict, Union
+from typing import List, Optional, Sequence, Tuple
 
 import gymnasium as gym
-from gymnasium.core import Env
 import numpy as np
 import numpy.typing as npt
 from stable_baselines3.common.vec_env import VecEnv, VecEnvWrapper
 
 from imitation.data import rollout, types
 
-# The key for human readable data in the observation.
-HR_OBS_KEY = "HR_OBS"
-
 
 class BufferingWrapper(VecEnvWrapper):
     """Saves transitions of underlying VecEnv.
@@ -174,7 +170,7 @@ def pop_transitions(self) -> types.TransitionsWithRew:
 
 
 class RolloutInfoWrapper(gym.Wrapper):
-    """Adds the entire episode's rewards and observations to `info` at episode end.
+    """Add the entire episode's rewards and observations to `info` at episode end.
 
     Whenever done=True, `info["rollouts"]` is a dict with keys "obs" and "rews", whose
     corresponding values hold the NumPy arrays containing the raw observations and
@@ -210,56 +206,3 @@ def step(self, action):
                 "rews": np.stack(self._rews),
             }
         return obs, rew, terminated, truncated, info
-
-
-class HumanReadableWrapper(gym.Wrapper):
-    """Adds human-readable observation to `obs` at every step."""
-
-    def __init__(self, env: Env, original_obs_key: str = "ORI_OBS"):
-        """Builds HumanReadableWrapper
-
-        Args:
-            env: Environment to wrap.
-            original_obs_key: The key for original observation if the original
-                observation is not in dict format.
-        """
-        env.render_mode = "rgb_array"
-        self._original_obs_key = original_obs_key
-        super().__init__(env)
-
-    def _add_hr_obs(
-        self, obs: Union[np.ndarray, Dict[str, np.ndarray]]
-    ) -> Dict[str, np.ndarray]:
-        """Adds human-readable observation to obs.
-
-        Transforms obs into dictionary if it is not already, and adds the human-readable
-        observation from `env.render()` under the key HR_OBS_KEY.
-
-        Args:
-            obs: Observation from environment.
-
-        Returns:
-            Observation dictionary with the human-readable data
-
-        Raises:
-            KeyError: When the key HR_OBS_KEY already exists in the observation
-                dictionary.
-        """
-        assert self.env.render_mode is not None
-        assert self.env.render_mode == "rgb_array"
-        hr_obs = self.env.render()
-        if not isinstance(obs, Dict):
-            obs = {self._original_obs_key: obs}
-
-        if HR_OBS_KEY in obs:
-            raise KeyError(f"{HR_OBS_KEY!r} already exists in observation dict")
-        obs[HR_OBS_KEY] = hr_obs
-        return obs
-
-    def reset(self, **kwargs):
-        obs, info = super().reset(**kwargs)
-        return self._add_hr_obs(obs), info
-
-    def step(self, action):
-        obs, rew, terminated, truncated, info = self.env.step(action)
-        return self._add_hr_obs(obs), rew, terminated, truncated, info

tests/algorithms/conftest.py

@@ -1,9 +1,11 @@
 """Fixtures common across algorithm tests."""
 from typing import Sequence
 
+import gymnasium as gym
 import pytest
+from stable_baselines3.common import envs
 from stable_baselines3.common.policies import BasePolicy
-from stable_baselines3.common.vec_env import VecEnv
+from stable_baselines3.common.vec_env import DummyVecEnv, VecEnv
 
 from imitation.algorithms import bc
 from imitation.data.types import TrajectoryWithRew
@@ -109,3 +111,22 @@ def pendulum_single_venv(rng) -> VecEnv:
         post_wrappers=[lambda env, _: RolloutInfoWrapper(env)],
         rng=rng,
     )
+
+
+# TODO(GH#794): Remove after https://github.com/DLR-RM/stable-baselines3/pull/1676
+# merged and released.
+class FloatReward(gym.RewardWrapper):
+    """Typecasts reward to a float."""
+
+    def reward(self, reward):
+        return float(reward)
+
+
+@pytest.fixture
+def multi_obs_venv() -> VecEnv:
+    def make_env():
+        env = envs.SimpleMultiObsEnv(channel_last=False)
+        env = FloatReward(env)
+        return RolloutInfoWrapper(env)
+
+    return DummyVecEnv([make_env, make_env])

tests/algorithms/test_base.py

@@ -41,7 +41,10 @@ def test_check_fixed_horizon_flag(custom_logger):
 
 
 def _make_and_iterate_loader(*args, **kwargs):
-    loader = base.make_data_loader(*args, **kwargs)
+    # our pytype version doesn't understand optional arguments in TypedDict
+    # this is fixed in 2023.04.11, but we require 2022.7.26
+    # See https://github.com/google/pytype/issues/1195
+    loader = base.make_data_loader(*args, **kwargs)  # pytype: disable=wrong-arg-types
     for batch in loader:
         pass
 

tests/algorithms/test_bc.py

@@ -10,7 +10,6 @@
 import numpy as np
 import pytest
 import torch as th
-from stable_baselines3.common import envs as sb_envs
 from stable_baselines3.common import evaluation
 from stable_baselines3.common import policies as sb_policies
 from stable_baselines3.common import vec_env
@@ -291,44 +290,29 @@ def test_that_policy_reconstruction_preserves_parameters(
         th.testing.assert_close(original, reconstructed)
 
 
-# TODO(GH#794): Remove after https://github.com/DLR-RM/stable-baselines3/pull/1676
-# merged and released.
-class FloatReward(gym.RewardWrapper):
-    """Typecasts reward to a float."""
-
-    def reward(self, reward):
-        return float(reward)
-
-
-def test_dict_space():
-    def make_env():
-        env = sb_envs.SimpleMultiObsEnv(channel_last=False)
-        env = FloatReward(env)
-        return RolloutInfoWrapper(env)
-
-    env = vec_env.DummyVecEnv([make_env, make_env])
-
+def test_dict_space(multi_obs_venv: vec_env.VecEnv):
     # multi-input policy to accept dict observations
+    assert isinstance(multi_obs_venv.observation_space, gym.spaces.Dict)
     policy = sb_policies.MultiInputActorCriticPolicy(
-        env.observation_space,
-        env.action_space,
+        multi_obs_venv.observation_space,
+        multi_obs_venv.action_space,
         lambda _: 0.001,
     )
     rng = np.random.default_rng()
 
     # sample random transitions
     rollouts = rollout.rollout(
         policy=None,
-        venv=env,
+        venv=multi_obs_venv,
         sample_until=rollout.make_sample_until(min_timesteps=None, min_episodes=50),
         rng=rng,
         unwrap=True,
     )
     transitions = rollout.flatten_trajectories(rollouts)
     bc_trainer = bc.BC(
-        observation_space=env.observation_space,
+        observation_space=multi_obs_venv.observation_space,
         policy=policy,
-        action_space=env.action_space,
+        action_space=multi_obs_venv.action_space,
         rng=rng,
         demonstrations=transitions,
     )

tests/algorithms/test_density_baselines.py

@@ -7,13 +7,11 @@
 import numpy as np
 import pytest
 import stable_baselines3
-from stable_baselines3.common import envs as sb_envs
 from stable_baselines3.common import policies, vec_env
 
 from imitation.algorithms.density import DensityAlgorithm, DensityType
 from imitation.data import rollout, types
 from imitation.data.types import TrajectoryWithRew
-from imitation.data.wrappers import RolloutInfoWrapper
 from imitation.policies.base import RandomPolicy
 from imitation.testing import reward_improvement
 
@@ -172,27 +170,12 @@ def test_density_trainer_raises(
         density_trainer.set_demonstrations("foo")  # type: ignore[arg-type]
 
 
-# TODO(GH#794): Remove after https://github.com/DLR-RM/stable-baselines3/pull/1676
-# merged and released.
-class FloatReward(gym.RewardWrapper):
-    """Typecasts reward to a float."""
-
-    def reward(self, reward):
-        return float(reward)
-
-
-def test_dict_space():
-    def make_env():
-        env = sb_envs.SimpleMultiObsEnv(channel_last=False)
-        env = FloatReward(env)
-        return RolloutInfoWrapper(env)
-
-    venv = vec_env.DummyVecEnv([make_env, make_env])
-
+def test_dict_space(multi_obs_venv: vec_env.VecEnv):
     # multi-input policy to accept dict observations
+    assert isinstance(multi_obs_venv.observation_space, gym.spaces.Dict)
     rl_algo = stable_baselines3.PPO(
         policies.MultiInputActorCriticPolicy,
-        venv,
+        multi_obs_venv,
         n_steps=10,  # small value to make test faster
         n_epochs=2,  # small value to make test faster
     )
@@ -202,14 +185,14 @@ def make_env():
     sample_until = rollout.make_min_episodes(15)
     rollouts = rollout.rollout(
         policy=None,
-        venv=venv,
+        venv=multi_obs_venv,
         sample_until=sample_until,
         rng=rng,
     )
     density_trainer = DensityAlgorithm(
         demonstrations=rollouts,
         kernel="gaussian",
-        venv=venv,
+        venv=multi_obs_venv,
         rl_algo=rl_algo,
         kernel_bandwidth=0.2,
         standardise_inputs=True,

tests/data/test_rollout.py

@@ -423,5 +423,6 @@ def test_dictionary_observations(rng):
     )
     for traj in trajs:
         assert isinstance(traj.obs, types.DictObs)
-        assert venv.observation_space.contains(obs)
+        for obs in traj.obs:
+            assert venv.observation_space.contains(dict(obs.items()))
         np.testing.assert_allclose(traj.obs.get("a") / 2, traj.obs.get("b"))