fix aloha pixels env test

2024-04-07 17:20:54 +00:00 · 2024-04-07 17:20:54 +00:00 · e1ac5dc62f
parent 4371a5570d
commit e1ac5dc62f
1 changed files with 9 additions and 2 deletions
--- a/lerobot/common/envs/aloha/env.py
+++ b/lerobot/common/envs/aloha/env.py
@ -48,8 +48,15 @@ class AlohaEnv(gym.Env):
                dtype=np.float64,
            )
        elif self.obs_type == "pixels":
-            self.observation_space = spaces.Box(
-                low=0, high=255, shape=(self.observation_height, self.observation_width, 3), dtype=np.uint8
+            self.observation_space = spaces.Dict(
+                {
+                    "top": spaces.Box(
+                        low=0,
+                        high=255,
+                        shape=(self.observation_height, self.observation_width, 3),
+                        dtype=np.uint8,
+                    )
+                }
            )
        elif self.obs_type == "pixels_agent_pos":
            self.observation_space = spaces.Dict(