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

src/lerobot/rl/gym_manipulator.py

@@ -659,6 +659,7 @@ def control_loop(
     episode_step = 0
     episode_start_time = time.perf_counter()
 
+    try:
         while episode_idx < cfg.dataset.num_episodes_to_record:
             step_start_time = time.perf_counter()
 
@@ -667,7 +668,6 @@ def control_loop(
             if use_gripper:
                 neutral_action = torch.cat([neutral_action, torch.tensor([0.0])])  # Gripper stay
 
-        # Use the new step function
             transition = step_env_and_process_transition(
                 env=env,
                 transition=transition,
@@ -684,7 +684,6 @@ def control_loop(
                     for k, v in transition[TransitionKey.OBSERVATION].items()
                     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(
                     "teleop_action", transition[TransitionKey.ACTION]
                 )
@@ -695,8 +694,12 @@ def control_loop(
                     DONE: np.array([terminated or truncated], dtype=bool),
                 }
                 if use_gripper:
-                discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get("discrete_penalty", 0.0)
-                frame["complementary_info.discrete_penalty"] = np.array([discrete_penalty], dtype=np.float32)
+                    discrete_penalty = transition[TransitionKey.COMPLEMENTARY_DATA].get(
+                        "discrete_penalty", 0.0
+                    )
+                    frame["complementary_info.discrete_penalty"] = np.array(
+                        [discrete_penalty], dtype=np.float32
+                    )
 
                 if dataset is not None:
                     frame["task"] = cfg.dataset.task
@@ -732,6 +735,10 @@ def control_loop(
 
             # Maintain fps timing
             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:
         logging.info("Finalizing dataset before pushing to hub")