Add rgb observation to dagger

examples/train_dagger_atari_interactive_policy.py

@@ -7,28 +7,48 @@
 
 import gymnasium as gym
 import numpy as np
-from stable_baselines3.common import vec_env
+import torch as th
+from stable_baselines3.common import torch_layers, vec_env
 
 from imitation.algorithms import bc, dagger
-from imitation.policies import interactive
+from imitation.data import wrappers as data_wrappers
+from imitation.policies import base as policy_base
+from imitation.policies import interactive, obs_update_wrapper
+
+
+def lr_schedule(_: float):
+    # Set lr_schedule to max value to force error if policy.optimizer
+    # is used by mistake (should use self.optimizer instead).
+    return th.finfo(th.float32).max
+
 
 if __name__ == "__main__":
     rng = np.random.default_rng(0)
 
-    env = vec_env.DummyVecEnv([lambda: gym.wrappers.TimeLimit(gym.make("Pong-v4"), 10)])
-    env.seed(0)
+    env = gym.make("PongNoFrameskip-v4", render_mode="rgb_array")
+    hr_env = data_wrappers.HumanReadableWrapper(env)
+    venv = vec_env.DummyVecEnv([lambda: hr_env])
+    venv.seed(0)
 
-    expert = interactive.AtariInteractivePolicy(env)
+    expert = interactive.AtariInteractivePolicy(venv)
+    policy = policy_base.FeedForward32Policy(
+        observation_space=env.observation_space,
+        action_space=env.action_space,
+        lr_schedule=lr_schedule,
+        features_extractor_class=torch_layers.FlattenExtractor,
+    )
+    wrapped_policy = obs_update_wrapper.RemoveHR(policy, lr_schedule=lr_schedule)
 
     bc_trainer = bc.BC(
         observation_space=env.observation_space,
         action_space=env.action_space,
+        policy=wrapped_policy,
         rng=rng,
     )
 
     with tempfile.TemporaryDirectory(prefix="dagger_example_") as tmpdir:
         dagger_trainer = dagger.SimpleDAggerTrainer(
-            venv=env,
+            venv=venv,
             scratch_dir=tmpdir,
             expert_policy=expert,
             bc_trainer=bc_trainer,

src/imitation/algorithms/dagger.py

@@ -11,7 +11,7 @@
 import os
 import pathlib
 import uuid
-from typing import Any, Callable, List, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
 import numpy as np
 import torch as th
@@ -161,13 +161,16 @@ class InteractiveTrajectoryCollector(vec_env.VecEnvWrapper):
     """
 
     traj_accum: Optional[rollout.TrajectoryAccumulator]
-    _last_obs: Optional[np.ndarray]
+    _last_obs: Optional[Union[Dict[str, np.ndarray], np.ndarray]]
     _last_user_actions: Optional[np.ndarray]
 
     def __init__(
         self,
         venv: vec_env.VecEnv,
-        get_robot_acts: Callable[[np.ndarray], np.ndarray],
+        get_robot_acts: Callable[
+            [Union[Dict[str, np.ndarray], np.ndarray]],
+            np.ndarray,
+        ],
         beta: float,
         save_dir: types.AnyPath,
         rng: np.random.Generator,
@@ -213,16 +216,20 @@ def seed(self, seed: Optional[int] = None) -> List[Optional[int]]:
         self.rng = np.random.default_rng(seed=seed)
         return list(self.venv.seed(seed))
 
-    def reset(self) -> np.ndarray:
+    def reset(self) -> Union[np.ndarray, Dict[str, np.ndarray]]:
         """Resets the environment.
 
         Returns:
             obs: first observation of a new trajectory.
         """
         self.traj_accum = rollout.TrajectoryAccumulator()
         obs = self.venv.reset()
-        assert isinstance(obs, np.ndarray)
-        for i, ob in enumerate(obs):
+        assert isinstance(
+            obs,
+            (np.ndarray, dict),
+        ), "Tuple observations are not supported."
+        dictobs = types.maybe_wrap_in_dictobs(obs)
+        for i, ob in enumerate(dictobs):
             self.traj_accum.add_step({"obs": ob}, key=i)
         self._last_obs = obs
         self._is_reset = True
@@ -256,7 +263,9 @@ def step_async(self, actions: np.ndarray) -> None:
 
         mask = self.rng.uniform(0, 1, size=(self.num_envs,)) > self.beta
         if np.sum(mask) != 0:
-            actual_acts[mask] = self.get_robot_acts(self._last_obs[mask])
+            last_obs = types.maybe_wrap_in_dictobs(self._last_obs)
+            obs_for_robot = types.maybe_unwrap_dictobs(last_obs[mask])
+            actual_acts[mask] = self.get_robot_acts(obs_for_robot)
 
         self._last_user_actions = actions
         self.venv.step_async(actual_acts)
@@ -270,9 +279,13 @@ def step_wait(self) -> VecEnvStepReturn:
             Observation, reward, dones (is terminal?) and info dict.
         """
         next_obs, rews, dones, infos = self.venv.step_wait()
-        assert isinstance(next_obs, np.ndarray)
         assert self.traj_accum is not None
         assert self._last_user_actions is not None
+        assert isinstance(
+            next_obs,
+            (np.ndarray, dict),
+        ), "Tuple observations are not supported."
+
         self._last_obs = next_obs
         fresh_demos = self.traj_accum.add_steps_and_auto_finish(
             obs=next_obs,
@@ -508,7 +521,7 @@ def create_trajectory_collector(self) -> InteractiveTrajectoryCollector:
         beta = self.beta_schedule(self.round_num)
         collector = InteractiveTrajectoryCollector(
             venv=self.venv,
-            get_robot_acts=lambda acts: self.bc_trainer.policy.predict(acts)[0],
+            get_robot_acts=lambda obs: self.bc_trainer.policy.predict(obs)[0],
             beta=beta,
             save_dir=save_dir,
             rng=self.rng,
@@ -550,7 +563,7 @@ def save_trainer(self) -> Tuple[pathlib.Path, pathlib.Path]:
 
 
 class SimpleDAggerTrainer(DAggerTrainer):
-    """Simpler subclass of DAggerTrainer for training with synthetic feedback."""
+    """Simpler subclass of DAggerTrainer for training with feedback."""
 
     def __init__(
         self,
@@ -571,7 +584,7 @@ def __init__(
                 simultaneously for that timestep.
             scratch_dir: Directory to use to store intermediate training
                 information (e.g. for resuming training).
-            expert_policy: The expert policy used to generate synthetic demonstrations.
+            expert_policy: The expert policy used to generate demonstrations.
             rng: Random state to use for the random number generator.
             expert_trajs: Optional starting dataset that is inserted into the round 0
                 dataset.

src/imitation/data/types.py

@@ -94,7 +94,7 @@ def dict_len(self):
 
     def __getitem__(
         self,
-        key: Union[int, slice, Tuple[Union[int, slice], ...]],
+        key: Union[int, slice, Tuple[Union[int, slice], ...], np.ndarray],
     ) -> "DictObs":
         """Indexes or slices each array.
 

src/imitation/data/wrappers.py

@@ -1,14 +1,18 @@
 """Environment wrappers for collecting rollouts."""
 
-from typing import List, Optional, Sequence, Tuple
+from typing import Dict, List, Optional, Sequence, Tuple, Union
 
 import gymnasium as gym
 import numpy as np
 import numpy.typing as npt
+from gymnasium.core import Env
 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.
@@ -170,7 +174,7 @@ def pop_transitions(self) -> types.TransitionsWithRew:
 
 
 class RolloutInfoWrapper(gym.Wrapper):
-    """Add the entire episode's rewards and observations to `info` at episode end.
+    """Adds 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
@@ -206,3 +210,54 @@ def step(self, action):
                 "rews": np.stack(self._rews),
             }
         return obs, rew, terminated, truncated, info
+
+
+class HumanReadableWrapper(gym.ObservationWrapper):
+    """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.
+
+        Raises:
+            ValueError: If `env.render_mode` is not "rgb_array".
+
+        """
+        if env.render_mode != "rgb_array":
+            raise ValueError(
+                "HumanReadableWrapper requires render_mode='rgb_array', "
+                f"got {env.render_mode!r}",
+            )
+        self._original_obs_key = original_obs_key
+        super().__init__(env)
+
+    def observation(
+        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.
+        """
+        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] = self.env.render()  # type: ignore[assignment]
+        return obs

src/imitation/policies/interactive.py

@@ -11,6 +11,7 @@
 from stable_baselines3.common import vec_env
 
 import imitation.policies.base as base_policies
+from imitation.data import wrappers
 from imitation.util import util
 
 
@@ -64,9 +65,6 @@ def _choose_action(
         if self.clear_screen_on_query:
             util.clear_screen()
 
-        if isinstance(obs, dict):
-            raise ValueError("Dictionary observations are not supported here")
-
         context = self._render(obs)
         key = self._get_input_key()
         self._clean_up(context)
@@ -87,7 +85,10 @@ def _get_input_key(self) -> str:
         return key
 
     @abc.abstractmethod
-    def _render(self, obs: np.ndarray) -> Optional[object]:
+    def _render(
+        self,
+        obs: Union[np.ndarray, Dict[str, np.ndarray]],
+    ) -> Optional[object]:
         """Renders an observation, optionally returns a context for later cleanup."""
 
     def _clean_up(self, context: object) -> None:
@@ -97,7 +98,7 @@ def _clean_up(self, context: object) -> None:
 class ImageObsDiscreteInteractivePolicy(DiscreteInteractivePolicy):
     """DiscreteInteractivePolicy that renders image observations."""
 
-    def _render(self, obs: np.ndarray) -> plt.Figure:
+    def _render(self, obs: Union[np.ndarray, Dict[str, np.ndarray]]) -> plt.Figure:
         img = self._prepare_obs_image(obs)
 
         fig, ax = plt.subplots()
@@ -110,9 +111,16 @@ def _render(self, obs: np.ndarray) -> plt.Figure:
     def _clean_up(self, context: plt.Figure) -> None:
         plt.close(context)
 
-    def _prepare_obs_image(self, obs: np.ndarray) -> np.ndarray:
+    def _prepare_obs_image(
+        self,
+        obs: Union[np.ndarray, Dict[str, np.ndarray]],
+    ) -> np.ndarray:
         """Applies any required observation processing to get an image to show."""
-        return obs
+        if not isinstance(obs, Dict):
+            return obs
+        if wrappers.HR_OBS_KEY not in obs:
+            raise KeyError(f"Observation does not contain {wrappers.HR_OBS_KEY!r}")
+        return obs[wrappers.HR_OBS_KEY]
 
 
 ATARI_ACTION_NAMES_TO_KEYS = {