From ac85a0d17eba4ca670442127873fd1669c5fea24 Mon Sep 17 00:00:00 2001 From: tangger Date: Wed, 7 May 2025 21:05:28 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=8E=A7=E5=88=B6=E6=8C=89?= =?UTF-8?q?=E9=92=AE=EF=BC=9B=E5=A2=9E=E5=8A=A0enter=E4=BB=A5=E7=A1=AE?= =?UTF-8?q?=E8=AE=A4=E5=BD=93=E5=89=8D=E8=BD=A8=E8=BF=B9=EF=BC=9B=E5=A2=9E?= =?UTF-8?q?=E5=8A=A0backspace=E4=BB=A5=E4=B8=A2=E5=BC=83=E5=BD=93=E5=89=8D?= =?UTF-8?q?=E8=BD=A8=E8=BF=B9=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- lerobot/common/robot_devices/control_utils.py | 19 +++++++++------- lerobot/scripts/control_robot.py | 22 +++++++++++++++++++ 2 files changed, 33 insertions(+), 8 deletions(-) diff --git a/lerobot/common/robot_devices/control_utils.py b/lerobot/common/robot_devices/control_utils.py index e1259ec..b7d1c79 100644 --- a/lerobot/common/robot_devices/control_utils.py +++ b/lerobot/common/robot_devices/control_utils.py @@ -121,36 +121,39 @@ def init_keyboard_listener(): events["exit_early"] = False events["rerecord_episode"] = False events["stop_recording"] = False - + events["confirm_save"] = False + events["discard_episode"] = False if is_headless(): logging.warning( "Headless environment detected. On-screen cameras display and keyboard inputs will not be available." ) listener = None return listener, events - # Only import pynput if not in a headless environment from pynput import keyboard - def on_press(key): try: if key == keyboard.Key.right: - print("Right arrow key pressed. Exiting loop...") + print("Right arrow key pressed. Exiting current recording...") events["exit_early"] = True elif key == keyboard.Key.left: - print("Left arrow key pressed. Exiting loop and rerecord the last episode...") + print("Left arrow key pressed. Interrupting and preparing to rerecord...") events["rerecord_episode"] = True events["exit_early"] = True elif key == keyboard.Key.esc: - print("Escape key pressed. Stopping data recording...") + print("Escape key pressed. Stopping data recording session...") events["stop_recording"] = True events["exit_early"] = True + elif key == keyboard.Key.enter: + print("Enter key pressed. Confirming and saving current episode...") + events["confirm_save"] = True + elif key == keyboard.Key.backspace: + print("Back key pressed. Discarding completed episode...") + events["discard_episode"] = True except Exception as e: print(f"Error handling key press: {e}") - listener = keyboard.Listener(on_press=on_press) listener.start() - return listener, events diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py index 6be27b7..111864c 100644 --- a/lerobot/scripts/control_robot.py +++ b/lerobot/scripts/control_robot.py @@ -291,6 +291,28 @@ def record( single_task=cfg.single_task, ) + # 录制完成后等待用户确认是否保存 + if not events["stop_recording"] and not events["rerecord_episode"]: + log_say("Episode recorded. Press Enter to save, Backspace to discard, Esc to end session", cfg.play_sounds) + print("Confirmation controls: Enter to save, Backspace to discard, Esc to end recording session") + + # 等待用户确认是否保存该段录制 + waiting_for_confirmation = True + events["confirm_save"] = False + events["discard_episode"] = False + + while waiting_for_confirmation and not events["stop_recording"]: + time.sleep(0.1) + if events["confirm_save"]: + log_say("Saving episode", cfg.play_sounds) + waiting_for_confirmation = False + elif events["discard_episode"]: + log_say("Discarding episode", cfg.play_sounds) + events["rerecord_episode"] = True + waiting_for_confirmation = False + elif events["stop_recording"]: + waiting_for_confirmation = False + # Execute a few seconds without recording to give time to manually reset the environment # Current code logic doesn't allow to teleoperate during this time. # TODO(rcadene): add an option to enable teleoperation during reset