fix(robomme): expose episode language instruction

This commit is contained in:
Pepijn
2026-07-15 19:58:06 +02:00
parent 5c453c3786
commit 4627900be2
2 changed files with 27 additions and 1 deletions
+8
View File
@@ -63,6 +63,8 @@ class RoboMMEGymEnv(gym.Env):
from robomme.env_record_wrapper import BenchmarkEnvBuilder from robomme.env_record_wrapper import BenchmarkEnvBuilder
self._task = task self._task = task
self.task = task
self.task_description = task
self._action_space_type = action_space_type self._action_space_type = action_space_type
self._dataset = dataset self._dataset = dataset
self._episode_idx = episode_idx self._episode_idx = episode_idx
@@ -105,6 +107,12 @@ class RoboMMEGymEnv(gym.Env):
) )
obs, info = self._env.reset() obs, info = self._env.reset()
self._last_raw_obs = obs self._last_raw_obs = obs
task_goal = info.get("task_goal")
# RoboMME returns [simple_subgoal, grounded_subgoal]. The published
# LeRobot training dataset uses the simple subgoal as its `task` text.
if isinstance(task_goal, (list, tuple)):
task_goal = task_goal[0] if task_goal else ""
self.task_description = str(task_goal or self._task)
return self._convert_obs(obs), self._convert_info(info) return self._convert_obs(obs), self._convert_info(info)
def step(self, action): def step(self, action):
+19 -1
View File
@@ -44,7 +44,10 @@ def _install_robomme_stub():
"joint_state_list": [np.zeros(7, dtype=np.float32)], "joint_state_list": [np.zeros(7, dtype=np.float32)],
"gripper_state_list": [np.zeros(2, dtype=np.float32)], "gripper_state_list": [np.zeros(2, dtype=np.float32)],
} }
env.reset.return_value = (obs, {"status": "ongoing", "task_goal": "pick the cube"}) env.reset.return_value = (
obs,
{"status": "ongoing", "task_goal": ["pick the cube", "pick the blue cube"]},
)
env.step.return_value = (obs, 0.0, False, False, {"status": "ongoing", "task_goal": ""}) env.step.return_value = (obs, 0.0, False, False, {"status": "ongoing", "task_goal": ""})
return env return env
@@ -116,6 +119,21 @@ def test_robomme_features_action_dim_ee_pose():
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def test_reset_exposes_episode_task_description():
"""VLA evaluation receives the episode-specific language instruction."""
_install_robomme_stub()
try:
from lerobot.envs.robomme import RoboMMEGymEnv
env = RoboMMEGymEnv(task="PickXtimes")
env.reset()
assert env.task == "PickXtimes"
assert env.task_description == "pick the cube"
finally:
_uninstall_robomme_stub()
def test_convert_obs_list_format(): def test_convert_obs_list_format():
"""_convert_obs takes the last element from list-format obs fields and """_convert_obs takes the last element from list-format obs fields and
emits a nested ``pixels`` dict (image, wrist_image) plus ``agent_pos``. emits a nested ``pixels`` dict (image, wrist_image) plus ``agent_pos``.