Raise error when observation keys and observation space keys don't match (#1047)

* Raise error when observation keys and observation space keys don't match

* Print the difference in keys

* Update changelog

docs/misc/changelog.rst

@@ -15,6 +15,7 @@ New Features:
 - Support logging hyperparameters to tensorboard (@timothe-chaumont)
 - Added checkpoints for replay buffer and ``VecNormalize`` statistics (@anand-bala)
 - Added option for ``Monitor`` to append to existing file instead of overriding (@sidney-tio)
+- The env checker now raises an error when using dict observation spaces and observation keys don't match observation space keys
 
 SB3-Contrib
 ^^^^^^^^^^^

stable_baselines3/common/env_checker.py

@@ -143,6 +143,13 @@ def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action
 
     if isinstance(observation_space, spaces.Dict):
         assert isinstance(obs, dict), "The observation returned by `reset()` must be a dictionary"
+
+        if not obs.keys() == observation_space.spaces.keys():
+            raise AssertionError(
+                "The observation keys returned by `reset()` must match the observation "
+                f"space keys: {obs.keys()} != {observation_space.spaces.keys()}"
+            )
+
         for key in observation_space.spaces.keys():
             try:
                 _check_obs(obs[key], observation_space.spaces[key], "reset")
@@ -162,6 +169,13 @@ def _check_returned_values(env: gym.Env, observation_space: spaces.Space, action
 
     if isinstance(observation_space, spaces.Dict):
         assert isinstance(obs, dict), "The observation returned by `step()` must be a dictionary"
+
+        if not obs.keys() == observation_space.spaces.keys():
+            raise AssertionError(
+                "The observation keys returned by `step()` must match the observation "
+                f"space keys: {obs.keys()} != {observation_space.spaces.keys()}"
+            )
+
         for key in observation_space.spaces.keys():
             try:
                 _check_obs(obs[key], observation_space.spaces[key], "step")

tests/test_envs.py

@@ -198,6 +198,14 @@ def test_common_failures_reset():
     check_reset_assert_error(env, (env.observation_space.sample(), False))
 
     env = SimpleMultiObsEnv()
+
+    # Observation keys and observation space keys must match
+    wrong_obs = env.observation_space.sample()
+    wrong_obs.pop("img")
+    check_reset_assert_error(env, wrong_obs)
+    wrong_obs = {**env.observation_space.sample(), "extra_key": None}
+    check_reset_assert_error(env, wrong_obs)
+
     obs = env.reset()
 
     def wrong_reset(self):
@@ -249,6 +257,14 @@ def test_common_failures_step():
     check_step_assert_error(env, (env.observation_space.sample(), 0.0, 1, {}))
 
     env = SimpleMultiObsEnv()
+
+    # Observation keys and observation space keys must match
+    wrong_obs = env.observation_space.sample()
+    wrong_obs.pop("img")
+    check_step_assert_error(env, (wrong_obs, 0.0, False, {}))
+    wrong_obs = {**env.observation_space.sample(), "extra_key": None}
+    check_step_assert_error(env, (wrong_obs, 0.0, False, {}))
+
     obs = env.reset()
 
     def wrong_step(self, action):