diff --git a/docs/source/robomme.mdx b/docs/source/robomme.mdx index 522739b8f..d2d4e5478 100644 --- a/docs/source/robomme.mdx +++ b/docs/source/robomme.mdx @@ -109,15 +109,9 @@ dataset = LeRobotDataset("lerobot/robomme") ### Feature key alignment (training) -The env wrapper exposes `front_rgb` and `wrist_rgb` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `obs_images.front` and `obs_images.wrist` for the policy. +The env wrapper exposes `pixels/image` and `pixels/wrist_image` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `observation.images.image` and `observation.images.wrist_image` for the policy. State is exposed as `agent_pos` and maps to `observation.state`. -The dataset uses `image` and `wrist_image` as column names. When fine-tuning from the dataset, you may need to rename columns to match your policy's expected input keys: - -```python -dataset = LeRobotDataset("lerobot/robomme") -# rename dataset columns if needed: -# dataset.hf_dataset = dataset.hf_dataset.rename_column("image", "front_rgb") -``` +The dataset's `image` and `wrist_image` columns already align with the policy input keys, so no renaming is needed when fine-tuning. ## Action Spaces diff --git a/tests/test_robomme_env.py b/tests/test_robomme_env.py index 840a8bdcd..20646430a 100644 --- a/tests/test_robomme_env.py +++ b/tests/test_robomme_env.py @@ -118,7 +118,11 @@ def test_robomme_features_action_dim_ee_pose(): def test_convert_obs_list_format(): """_convert_obs takes the last element from list-format obs fields and - emits the lerobot-canonical keys (pixels/image, pixels/wrist_image, agent_pos).""" + emits a nested ``pixels`` dict (image, wrist_image) plus ``agent_pos``. + + The nested layout is required so ``preprocess_observation()`` in + ``envs/utils.py`` maps each camera to ``observation.images.``. + """ _install_robomme_stub() try: from lerobot.envs.robomme import RoboMMEGymEnv @@ -138,8 +142,8 @@ def test_convert_obs_list_format(): } result = env._convert_obs(obs_raw) - np.testing.assert_array_equal(result["pixels/image"], front) - np.testing.assert_array_equal(result["pixels/wrist_image"], wrist) + np.testing.assert_array_equal(result["pixels"]["image"], front) + np.testing.assert_array_equal(result["pixels"]["wrist_image"], wrist) assert result["agent_pos"].shape == (8,) np.testing.assert_array_almost_equal(result["agent_pos"][:7], joints) assert result["agent_pos"][7] == gripper[0] @@ -163,8 +167,8 @@ def test_convert_obs_array_format(): "gripper_state_list": np.zeros(2, dtype=np.float32), } result = env._convert_obs(obs_raw) - assert result["pixels/image"].shape == (256, 256, 3) - assert result["pixels/wrist_image"].shape == (256, 256, 3) + assert result["pixels"]["image"].shape == (256, 256, 3) + assert result["pixels"]["wrist_image"].shape == (256, 256, 3) assert result["agent_pos"].shape == (8,) finally: _uninstall_robomme_stub()