From 93e84ffb49e33a066b44874725a1ce25f78a35ed Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Fri, 17 Apr 2026 14:24:41 +0100
Subject: [PATCH] fix(robomme): nest `pixels` as a dict so
 preprocess_observation picks it up

`_convert_obs` was returning flat keys (`pixels/image`,
`pixels/wrist_image`). `preprocess_observation()` in envs/utils.py
keys off the top-level `"pixels"` entry and, not finding it,
silently dropped every image from the batch. The policy then saw
zero image features and raised

    ValueError: All image features are missing from the batch.

Match the LIBERO layout: return
`{"pixels": {"image": ..., "wrist_image": ...}, "agent_pos": ...}`
and declare the same shape in `observation_space`.

Made-with: Cursor
---
 src/lerobot/envs/robomme.py | 15 +++++++++++----
 1 file changed, 11 insertions(+), 4 deletions(-)

diff --git a/src/lerobot/envs/robomme.py b/src/lerobot/envs/robomme.py
index 7a0a76ed1..fdfe811fb 100644
--- a/src/lerobot/envs/robomme.py
+++ b/src/lerobot/envs/robomme.py
@@ -79,10 +79,18 @@ class RoboMMEGymEnv(gym.Env):
 
         action_dim = 8 if action_space_type == "joint_angle" else 7
         self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(action_dim,), dtype=np.float32)
+        # `pixels` must be a nested Dict so `preprocess_observation()` in
+        # envs/utils.py picks it up and maps each camera to
+        # `observation.images.<cam>`. A flat layout (`pixels/image`,
+        # `pixels/wrist_image`) silently drops every image from the batch.
         self.observation_space = spaces.Dict(
             {
-                "pixels/image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
-                "pixels/wrist_image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
+                "pixels": spaces.Dict(
+                    {
+                        "image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
+                        "wrist_image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8),
+                    }
+                ),
                 "agent_pos": spaces.Box(-np.inf, np.inf, shape=(8,), dtype=np.float32),
             }
         )
@@ -146,8 +154,7 @@ class RoboMMEGymEnv(gym.Env):
         state = np.concatenate([joint, gripper])
 
         return {
-            "pixels/image": front_rgb,
-            "pixels/wrist_image": wrist_rgb,
+            "pixels": {"image": front_rgb, "wrist_image": wrist_rgb},
             "agent_pos": state,
         }