From 4627900be2a264552567eed18630ad22487169e1 Mon Sep 17 00:00:00 2001
From: Pepijn <pepijn@huggingface.co>
Date: Wed, 15 Jul 2026 19:58:06 +0200
Subject: [PATCH] fix(robomme): expose episode language instruction

---
 src/lerobot/envs/robomme.py |  8 ++++++++
 tests/test_robomme_env.py   | 20 +++++++++++++++++++-
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/lerobot/envs/robomme.py b/src/lerobot/envs/robomme.py
index 69d665bd4..a59389f5a 100644
--- a/src/lerobot/envs/robomme.py
+++ b/src/lerobot/envs/robomme.py
@@ -63,6 +63,8 @@ class RoboMMEGymEnv(gym.Env):
         from robomme.env_record_wrapper import BenchmarkEnvBuilder
 
         self._task = task
+        self.task = task
+        self.task_description = task
         self._action_space_type = action_space_type
         self._dataset = dataset
         self._episode_idx = episode_idx
@@ -105,6 +107,12 @@ class RoboMMEGymEnv(gym.Env):
         )
         obs, info = self._env.reset()
         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)
 
     def step(self, action):
diff --git a/tests/test_robomme_env.py b/tests/test_robomme_env.py
index 20646430a..4c344574f 100644
--- a/tests/test_robomme_env.py
+++ b/tests/test_robomme_env.py
@@ -44,7 +44,10 @@ def _install_robomme_stub():
                 "joint_state_list": [np.zeros(7, 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": ""})
             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():
     """_convert_obs takes the last element from list-format obs fields and
     emits a nested ``pixels`` dict (image, wrist_image) plus ``agent_pos``.