admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-26 11:16:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(robomme): nest pixels as a dict so preprocess_observation picks it up

Browse Source 
`_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
1 changed files with 11 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/envs/robomme.py
+11 -4
View File
@@ -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,
}