From 8699a28be0c0db33120c35ce85c9c72f5783f2a1 Mon Sep 17 00:00:00 2001
From: Jannik Grothusen <56967823+J4nn1K@users.noreply.github.com>
Date: Tue, 25 Feb 2025 19:28:26 +0100
Subject: [PATCH] [QOL] Enable teleoperation during environment reset (#725)

---
 lerobot/common/robot_devices/control_utils.py | 23 +++++++------------
 lerobot/scripts/control_robot.py              |  2 +-
 2 files changed, 9 insertions(+), 16 deletions(-)

diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py
index 6c97d0cb..d2361a64 100644
--- a/lerobot/common/robot_devices/control_utils.py
+++ b/lerobot/common/robot_devices/control_utils.py
@@ -12,7 +12,6 @@ from functools import cache
 
 import cv2
 import torch
-import tqdm
 from deepdiff import DeepDiff
 from termcolor import colored
 
@@ -276,24 +275,18 @@ def control_loop(
             break
 
 
-def reset_environment(robot, events, reset_time_s):
+def reset_environment(robot, events, reset_time_s, fps):
     # TODO(rcadene): refactor warmup_record and reset_environment
-    # TODO(alibets): allow for teleop during reset
     if has_method(robot, "teleop_safety_stop"):
         robot.teleop_safety_stop()
 
-    timestamp = 0
-    start_vencod_t = time.perf_counter()
-
-    # Wait if necessary
-    with tqdm.tqdm(total=reset_time_s, desc="Waiting") as pbar:
-        while timestamp < reset_time_s:
-            time.sleep(1)
-            timestamp = time.perf_counter() - start_vencod_t
-            pbar.update(1)
-            if events["exit_early"]:
-                events["exit_early"] = False
-                break
+    control_loop(
+        robot=robot,
+        control_time_s=reset_time_s,
+        events=events,
+        fps=fps,
+        teleoperate=True,
+    )
 
 
 def stop_recording(robot, listener, display_cameras):
diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py
index 32f3b181..ab5d0e8a 100644
--- a/lerobot/scripts/control_robot.py
+++ b/lerobot/scripts/control_robot.py
@@ -299,7 +299,7 @@ def record(
             (recorded_episodes < cfg.num_episodes - 1) or events["rerecord_episode"]
         ):
             log_say("Reset the environment", cfg.play_sounds)
-            reset_environment(robot, events, cfg.reset_time_s)
+            reset_environment(robot, events, cfg.reset_time_s, cfg.fps)
 
         if events["rerecord_episode"]:
             log_say("Re-record episode", cfg.play_sounds)