admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 10:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: add try/finally to control_loop to ensure image writer cleanup on exit

Browse Source
This commit is contained in:
Khalil Meftah
2026-04-14 17:54:35 +02:00
parent da6e36fd03
commit 79a9ebdaa6
1 changed files with 72 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/rl/gym_manipulator.py
+11 -4
View File
@@ -659,6 +659,7 @@ def control_loop(
episode_step = 0 episode_step = 0
episode_start_time = time.perf_counter() episode_start_time = time.perf_counter()
try:
while episode_idx < cfg.dataset.num_episodes_to_record: while episode_idx < cfg.dataset.num_episodes_to_record:
step_start_time = time.perf_counter() step_start_time = time.perf_counter()
@@ -667,7 +668,6 @@ def control_loop(
if use_gripper: if use_gripper:
neutral_action = torch.cat([neutral_action, torch.tensor([0.0])]) # Gripper stay neutral_action = torch.cat([neutral_action, torch.tensor([0.0])]) # Gripper stay
# Use the new step function
transition = step_env_and_process_transition( transition = step_env_and_process_transition(
env=env, env=env,
transition=transition, transition=transition,
@@ -684,7 +684,6 @@ def control_loop(
for k, v in transition[TransitionKey.OBSERVATION].items() for k, v in transition[TransitionKey.OBSERVATION].items()
if isinstance(v, torch.Tensor) if isinstance(v, torch.Tensor)
} }
# Use teleop_action if available, otherwise use the action from the transition
action_to_record = transition[TransitionKey.COMPLEMENTARY_DATA].get( action_to_record = transition[TransitionKey.COMPLEMENTARY_DATA].get(
"teleop_action", transition[TransitionKey.ACTION] "teleop_action", transition[TransitionKey.ACTION]
) )
@@ -695,8 +694,12 @@ def control_loop(
DONE: np.array([terminated or truncated], dtype=bool), DONE: np.array([terminated or truncated], dtype=bool),
} }
if use_gripper: if use_gripper:
discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0) discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get(
frame["complementary_info.discrete_penalty"] = np.array([discrete_penalty], dtype=np.float32) "discrete_penalty", 0.0
)
frame["complementary_info.discrete_penalty"] = np.array(
[discrete_penalty], dtype=np.float32
)
if dataset is not None: if dataset is not None:
frame["task"] = cfg.dataset.task frame["task"] = cfg.dataset.task
@@ -732,6 +735,10 @@ def control_loop(
# Maintain fps timing # Maintain fps timing
precise_sleep(max(dt - (time.perf_counter() - step_start_time), 0.0)) precise_sleep(max(dt - (time.perf_counter() - step_start_time), 0.0))
finally:
if dataset is not None and dataset.writer is not None and dataset.writer.image_writer is not None:
logging.info("Waiting for image writer to finish...")
dataset.writer.image_writer.stop()
if dataset is not None and cfg.dataset.push_to_hub: if dataset is not None and cfg.dataset.push_to_hub:
logging.info("Finalizing dataset before pushing to hub") logging.info("Finalizing dataset before pushing to hub")