is_headless

lerobot/scripts/control_robot.py

@@ -84,6 +84,7 @@ import concurrent.futures
 import json
 import logging
 import os
+import platform
 import shutil
 import time
 from contextlib import nullcontext
@@ -94,7 +95,6 @@ import tqdm
 from huggingface_hub import create_branch
 from omegaconf import DictConfig
 from PIL import Image
-from pynput import keyboard
 
 # from safetensors.torch import load_file, save_file
 from lerobot.common.datasets.compute_stats import compute_stats
@@ -270,6 +270,19 @@ def record_dataset(
     rerecord_episode = False
     stop_recording = False
 
+    def is_headless():
+        if platform.system() == "Linux":
+            display = os.environ.get("DISPLAY")
+            if display is None or display == "":
+                return True
+        return False
+
+    # Only import pynput if not in a headless environment
+    if is_headless():
+        logging.info("Headless environment detected. Keyboard input will not be available.")
+    else:
+        from pynput import keyboard
+
         def on_press(key):
             nonlocal exit_early, rerecord_episode, stop_recording
             try:
@@ -408,8 +421,10 @@ def record_dataset(
             if stop_recording or episode_index == num_episodes:
                 logging.info("Done recording")
                 os.system('say "Done recording"')
-                logging.info("Waiting for threads writing the images on disk to terminate...")
+                if not is_headless():
                     listener.stop()
+
+                logging.info("Waiting for threads writing the images on disk to terminate...")
                 for _ in tqdm.tqdm(
                     concurrent.futures.as_completed(futures), total=len(futures), desc="Writting images"
                 ):