From ba6a6a7eb5191bac4c5f70f12313185b15e3c350 Mon Sep 17 00:00:00 2001
From: ZiyueWang25 <wfuymu@gmail.com>
Date: Tue, 3 Oct 2023 13:08:01 -0700
Subject: [PATCH] fix dict env observation space

---
 tests/data/test_wrappers.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tests/data/test_wrappers.py b/tests/data/test_wrappers.py
index cfde9dbcc..279c72079 100644
--- 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):