diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index 6c97d0cb6e..d2361a6468 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -12,7 +12,6 @@ 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 32f3b18148..ab5d0e8ac2 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)