mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-27 11:46:04 +00:00
fix(robomme): expose episode language instruction
This commit is contained in:
@@ -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):
|
||||
|
||||
@@ -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``.
|
||||
|
||||
Reference in New Issue
Block a user