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
This commit is contained in:
Pepijn
2026-04-17 14:24:41 +01:00
parent 0f03129536
commit 93e84ffb49
+11 -4
View File
@@ -79,10 +79,18 @@ class RoboMMEGymEnv(gym.Env):
action_dim = 8 if action_space_type == "joint_angle" else 7 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) 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( self.observation_space = spaces.Dict(
{ {
"pixels/image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8), "pixels": spaces.Dict(
"pixels/wrist_image": spaces.Box(0, 255, shape=(256, 256, 3), dtype=np.uint8), {
"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), "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]) state = np.concatenate([joint, gripper])
return { return {
"pixels/image": front_rgb, "pixels": {"image": front_rgb, "wrist_image": wrist_rgb},
"pixels/wrist_image": wrist_rgb,
"agent_pos": state, "agent_pos": state,
} }