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