mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-27 19:56:09 +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
|
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):
|
||||||
|
|||||||
@@ -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``.
|
||||||
|
|||||||
Reference in New Issue
Block a user