Skip to content

Commit ba8895f

Browse files
committed
Add support to start stop episodes with keyboard controls
1 parent 8e88118 commit ba8895f

File tree

2 files changed

+36
-17
lines changed

2 files changed

+36
-17
lines changed

piper_sdk

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Subproject commit 31d1a509ee50467b2051f1c0d6e3267b3e7dddd4

src/lerobot/utils/control_utils.py

Lines changed: 35 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818

1919

2020
import logging
21+
import threading
2122
import traceback
2223
from contextlib import nullcontext
2324
from copy import copy
@@ -131,6 +132,8 @@ def init_keyboard_listener():
131132
# Allow to exit early while recording an episode or resetting the environment,
132133
# by tapping the right arrow key '->'. This might require a sudo permission
133134
# to allow your terminal to monitor keyboard events.
135+
# Initialize simple text-based keyboard listener to avoid sudo permission issues.
136+
# Uses input() in a background thread - no special permissions required.
134137
events = {}
135138
events["exit_early"] = False
136139
events["rerecord_episode"] = False
@@ -143,27 +146,42 @@ def init_keyboard_listener():
143146
listener = None
144147
return listener, events
145148

146-
# Only import pynput if not in a headless environment
147-
from pynput import keyboard
149+
print("Text-based keyboard controls enabled:")
150+
print(" 'n' + Enter: Next/forward command (simulate right arrow) - Exit current loop")
151+
print(" 'b' + Enter: Back command (simulate left arrow) - Re-record episode")
152+
print(" 's' + Enter: Stop recording completely")
148153

149-
def on_press(key):
154+
def input_listener():
150155
try:
151-
if key == keyboard.Key.right:
152-
print("Right arrow key pressed. Exiting loop...")
153-
events["exit_early"] = True
154-
elif key == keyboard.Key.left:
155-
print("Left arrow key pressed. Exiting loop and rerecord the last episode...")
156-
events["rerecord_episode"] = True
157-
events["exit_early"] = True
158-
elif key == keyboard.Key.esc:
159-
print("Escape key pressed. Stopping data recording...")
160-
events["stop_recording"] = True
161-
events["exit_early"] = True
156+
while not events["stop_recording"]:
157+
try:
158+
user_input = input().strip().lower()
159+
if user_input == 'n':
160+
print("Next/forward command received. Exiting loop...")
161+
events["exit_early"] = True
162+
elif user_input == 'b':
163+
print("Back command received. Exiting loop and re-record episode...")
164+
events["rerecord_episode"] = True
165+
events["exit_early"] = True
166+
elif user_input == 's':
167+
print("Stop command received. Stopping data recording...")
168+
events["stop_recording"] = True
169+
events["exit_early"] = True
170+
except EOFError:
171+
# Handle case where input is not available (e.g., piped input)
172+
break
162173
except Exception as e:
163-
print(f"Error handling key press: {e}")
174+
logging.debug(f"Input listener error: {e}")
164175

165-
listener = keyboard.Listener(on_press=on_press)
166-
listener.start()
176+
input_thread = threading.Thread(target=input_listener, daemon=True)
177+
input_thread.start()
178+
179+
# Return a simple object to maintain compatibility
180+
class SimpleListener:
181+
def stop(self):
182+
events["stop_recording"] = True
183+
184+
listener = SimpleListener()
167185

168186
return listener, events
169187

0 commit comments

Comments
 (0)