Merge ea9fa69435 into 1c873df5c0

lerobot/common/robot_devices/control_utils.py

@@ -26,6 +26,7 @@ from functools import cache
 
 import cv2
 import torch
+import tqdm
 from deepdiff import DeepDiff
 from termcolor import colored
 
@@ -287,14 +288,31 @@ def reset_environment(robot, events, reset_time_s, fps):
     # TODO(rcadene): refactor warmup_record and reset_environment
     if has_method(robot, "teleop_safety_stop"):
         robot.teleop_safety_stop()
+    timestamp = 0
+    start_vencod_t = time.perf_counter()
+    if "next.reward" in events:
+        events["next.reward"] = 0
 
-    control_loop(
-        robot=robot,
-        control_time_s=reset_time_s,
-        events=events,
-        fps=fps,
-        teleoperate=True,
-    )
+    # Wait if necessary
+    with tqdm.tqdm(total=reset_time_s, desc="Waiting") as pbar:
+        last_update = 0  # Track the last update time
+        while timestamp < reset_time_s:
+            start_loop_t = time.perf_counter()
+            robot.teleop_step(record_data=False)
+            timestamp = time.perf_counter() - start_vencod_t
+
+            # Update progress bar every second
+            current_second = int(timestamp)
+            if current_second > last_update:
+                pbar.update(current_second - last_update)
+                last_update = current_second
+
+            if events["exit_early"]:
+                events["exit_early"] = False
+                break
+
+            dt_s = time.perf_counter() - start_loop_t
+            busy_wait(1 / fps - dt_s)
 
 
 def stop_recording(robot, listener, display_cameras):