fix dict env observation space

HumanCompatibleAI · Oct 3, 2023 · ba6a6a7 · ba6a6a7
1 parent 6e5c3e8
commit ba6a6a7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tests/data/test_wrappers.py b/tests/data/test_wrappers.py
@@ -61,7 +61,10 @@ class _CountingDictEnv(_CountingEnv):  # pragma: no cover
     def __init__(self, episode_length=5):
         super().__init__(episode_length)
         self.observation_space = gym.spaces.Dict(
-            spaces={"t": gym.spaces.Box(low=0, high=np.inf, shape=())},
+            spaces={
+                "t": gym.spaces.Box(low=0, high=np.inf, shape=()),
+                "2t": gym.spaces.Box(low=0, high=np.inf, shape=()),
+            },
         )
 
     def reset(self, *, seed: Optional[int] = None, options: Optional[dict] = None):