Skip to content

Update lerobot recorder #145

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yuecideng merged 10 commits into from
Feb 21, 2026
Merged

Update lerobot recorder #145

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
5a18d94
wip
yuecideng Feb 19, 2026
78c6495
wip
yuecideng Feb 20, 2026
3f2b303
Merge remote-tracking branch 'origin/main' into yueci/physics-step
yuecideng Feb 20, 2026
536c17a
wip
yuecideng Feb 20, 2026
2cac45f
wip
yuecideng Feb 20, 2026
3d58d98
Merge branch 'main' into yueci/update-lerobot-recorder
yuecideng Feb 20, 2026
b659ad1
wip
yuecideng Feb 21, 2026
a1c2e4b
wip
yuecideng Feb 21, 2026
f2f9fd0
wip
yuecideng Feb 21, 2026
4926e1f
wip
yuecideng Feb 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 1 addition & 7 deletions configs/gym/agent/pour_water_agent/fast_gym_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,14 +251,8 @@
"mode": "save",
"params": {
"robot_meta": {
"arm_dofs": 12,
"control_freq": 25,
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
"observation": {
"vision": {},
"states": ["qpos"]
},
"min_len_steps": 5
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
},
"instruction": {
"lang": "Pour water from the bottle into the mug."
Expand Down
8 changes: 1 addition & 7 deletions configs/gym/agent/rearrangement_agent/fast_gym_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,14 +236,8 @@
"mode": "save",
"params": {
"robot_meta": {
"arm_dofs": 12,
"control_freq": 25,
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
"observation": {
"vision": {},
"states": ["qpos"]
},
"min_len_steps": 125
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
},
"instruction": {
"lang": "Place the spoon and fork neatly into the plate on the table."
Expand Down
14 changes: 1 addition & 13 deletions configs/gym/pour_water/gym_config.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -264,20 +264,8 @@
"params": {
"robot_meta": {
"robot_type": "CobotMagic",
"arm_dofs": 12,
"control_freq": 25,
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
"observation": {
"vision": {
"cam_high": ["mask"],
"cam_right_wrist": ["mask"],
"cam_left_wrist": ["mask"]
},
"states": ["qpos"],
"exteroception": ["cam_high", "cam_right_wrist", "cam_left_wrist"]
},
"action": "qpos_with_eef_pose",
"min_len_steps": 5
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
},
"instruction": {
"lang": "Pour water from bottle to cup"
Expand Down
12 changes: 1 addition & 11 deletions configs/gym/pour_water/gym_config_simple.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -213,18 +213,8 @@
"params": {
"robot_meta": {
"robot_type": "CobotMagic",
"arm_dofs": 12,
"control_freq": 25,
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"],
"observation": {
"vision": {
"cam_high": ["mask"]
},
"states": ["qpos"],
"exteroception": ["cam_high", "cam_right_wrist", "cam_left_wrist"]
},
"action": "qpos_with_eef_pose",
"min_len_steps": 5
"control_parts": ["left_arm", "left_eef", "right_arm", "right_eef"]
},
"instruction": {
"lang": "Pour water from bottle to cup"
Expand Down
11 changes: 1 addition & 10 deletions configs/gym/special/simple_task_ur10.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,16 +33,7 @@
"params": {
"robot_meta": {
"robot_type": "UR10",
"arm_dofs": 6,
"control_freq": 3,
"observation": {
"vision": {
"cam_high": []
},
"states": ["qpos"]
},
"action": "qpos",
"min_len_steps": 5
"control_freq": 3
},
"instruction": {
"lang": "Acting with Oscillatory motion"
Expand Down
25 changes: 12 additions & 13 deletions embodichain/lab/gym/envs/base_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,21 +164,15 @@ def device(self) -> torch.device:

@cached_property
def single_observation_space(self) -> gym.spaces.Space:
if self.num_envs == 1:
return gym_utils.convert_observation_to_space(self._init_raw_obs)
else:
return gym_utils.convert_observation_to_space(
self._init_raw_obs, unbatched=True
)
return gym_utils.convert_observation_to_space(
self._init_raw_obs, unbatched=True
)

@cached_property
def observation_space(self) -> gym.spaces.Space:
if self.num_envs == 1:
return self.single_observation_space
else:
return gym.vector.utils.batch_space(
self.single_observation_space, n=self.num_envs
)
return gym_utils.convert_observation_to_space(
self._init_raw_obs, unbatched=False
)

@cached_property
def flattened_observation_space(self) -> gym.spaces.Box:
Expand Down Expand Up @@ -208,6 +202,11 @@ def action_space(self) -> gym.spaces.Space:
def elapsed_steps(self) -> Union[int, torch.Tensor]:
return self._elapsed_steps

@property
def has_sensors(self) -> bool:
"""Return whether the environment has sensors."""
return len(self.sensors) > 0

def get_sensor(self, name: str, **kwargs) -> BaseSensor:
"""Get the sensor instance by name.

Expand Down Expand Up @@ -576,10 +575,10 @@ def reset(
self.sim.reset_objects_state(
env_ids=reset_ids, excluded_uids=self._detached_uids_for_reset
)
self._elapsed_steps[reset_ids] = 0

# Reset hook for user to perform any custom reset logic.
self._initialize_episode(reset_ids, **options)
self._elapsed_steps[reset_ids] = 0
Comment on lines 579 to +581
Copy link

Copilot AI Feb 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The reset of _elapsed_steps was moved to after _initialize_episode. This changes the semantics: previously, elapsed_steps would be 0 during _initialize_episode for the reset environments. Now it retains the old value during _initialize_episode and only gets reset afterwards. This could affect any logic in _initialize_episode that depends on elapsed_steps being 0 for newly reset environments.

Suggested change
# Reset hook for user to perform any custom reset logic.
self._initialize_episode(reset_ids, **options)
self._elapsed_steps[reset_ids] = 0
# Reset elapsed steps for the environments being reset before initialization.
self._elapsed_steps[reset_ids] = 0
# Reset hook for user to perform any custom reset logic.
self._initialize_episode(reset_ids, **options)

Copilot uses AI. Check for mistakes.

return self.get_obs(**options), self.get_info(**options)

Expand Down
9 changes: 4 additions & 5 deletions embodichain/lab/gym/envs/embodied_env.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,6 +178,9 @@ def __init__(self, cfg: EmbodiedEnvCfg, **kwargs):

super().__init__(cfg, **kwargs)

if self.cfg.dataset and not self.cfg.filter_dataset_saving:
self.dataset_manager = DatasetManager(self.cfg.dataset, self)

self.episode_obs_buffer: Dict[int, List[EnvObs]] = {
i: [] for i in range(self.num_envs)
}
Expand Down Expand Up @@ -208,9 +211,6 @@ def _init_sim_state(self, **kwargs):
if self.cfg.rewards:
self.reward_manager = RewardManager(self.cfg.rewards, self)

if self.cfg.dataset and not self.cfg.filter_dataset_saving:
self.dataset_manager = DatasetManager(self.cfg.dataset, self)

def _apply_functor_filter(self) -> None:
"""Apply functor filters to the environment components based on configuration.

Expand Down Expand Up @@ -397,6 +397,7 @@ def _initialize_episode(
]

if successful_env_ids:

Copy link

Copilot AI Feb 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The extra blank line at line 400 appears to be an unintentional formatting change. Consider removing it to maintain consistent code formatting.

Suggested change

Copilot uses AI. Check for mistakes.
# Convert back to tensor if needed
successful_env_ids_tensor = torch.tensor(
successful_env_ids, device=self.device
Expand All @@ -405,8 +406,6 @@ def _initialize_episode(
mode="save",
env_ids=successful_env_ids_tensor,
)
Copy link

Copilot AI Feb 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The removed log message "No successful episodes to save" provided useful debugging information when save operations are triggered but no data needs to be saved. Consider adding equivalent logging in the parent if-statement to maintain visibility into save operations, or document why this logging was removed.

Suggested change
)
)
else:
# Log when a save is requested but there are no successful episodes
logger.debug(
"Save requested but no successful episodes to save for env_ids=%s",
env_ids_to_process,
)

Copilot uses AI. Check for mistakes.
else:
logger.log_warning("No successful episodes to save.")

# Clear episode buffers and reset success status for environments being reset
for env_id in env_ids_to_process:
Expand Down
Loading
Loading