diff --git a/README.md b/README.md index fa3e9e1a3..f72952102 100644 --- a/README.md +++ b/README.md @@ -97,7 +97,7 @@ Training a policy is as simple as running a script configuration: ```bash lerobot-train \ - --policy=act \ + --policy.type=act \ --dataset.repo_id=lerobot/aloha_mobile_cabinet ``` @@ -136,6 +136,7 @@ Learn how to implement your own simulation environment or benchmark and distribu - **[X](https://x.com/LeRobotHF):** Follow us on X to stay up-to-date with the latest developments. - **[Robot Learning Tutorial](https://huggingface.co/spaces/lerobot/robot-learning-tutorial):** A free, hands-on course to learn robot learning using LeRobot. - **[T-Shirt Folding Experiment](https://huggingface.co/spaces/lerobot/robot-folding):** An end-to-end demonstration of folding t-shirts with LeRobot. +- **[LeLab](https://github.com/huggingface/leLab):** A web interface for LeRobot — teleoperate, calibrate, record datasets, replay, and train your SO arm from the browser, no CLI required. ## Citation diff --git a/docs/source/il_robots.mdx b/docs/source/il_robots.mdx index 53ae5af82..6a820e0db 100644 --- a/docs/source/il_robots.mdx +++ b/docs/source/il_robots.mdx @@ -390,9 +390,17 @@ Set the flow of data recording using command-line arguments: Control the data recording flow using keyboard shortcuts: -- Press **Right Arrow (`→`)**: Early stop the current episode or reset time and move to the next. -- Press **Left Arrow (`←`)**: Cancel the current episode and re-record it. -- Press **Escape (`ESC`)**: Immediately stop the session, encode videos, and upload the dataset. +- Press **Right Arrow (`→`)** or **`n`**: Early stop the current episode or reset time and move to the next. +- Press **Left Arrow (`←`)** or **`r`**: Cancel the current episode and re-record it. +- Press **Escape (`ESC`)** or **`q`**: Immediately stop the session, encode videos, and upload the dataset. + + + +These control-flow shortcuts work on **X11, Wayland, and headless/SSH** sessions. When a global keyboard backend isn't available (Wayland, a headless machine, or macOS without Accessibility permission), `lerobot-record` automatically reads the same keys from the terminal — launch it from an interactive terminal and keep it focused. You can also use the letter equivalents **`n`** (next, same as `→`), **`r`** (re-record, same as `←`) and **`q`** (quit, same as `ESC`). No `$DISPLAY` setup is required. + +This applies to the recording control flow only. Keyboard **teleoperation** (driving the robot with the keyboard) still needs a global key backend, so it works only on an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — not on Wayland or headless sessions. + + #### Tips for gathering data @@ -406,7 +414,7 @@ If you want to dive deeper into this important topic, you can check out the [blo #### Troubleshooting: -- On Linux, if the left and right arrow keys and escape key don't have any effect during data recording, make sure you've set the `$DISPLAY` environment variable. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). +- On Linux, the recording control-flow keys (arrow keys, Escape) work on X11, Wayland, and headless/SSH sessions as long as `lerobot-record` runs in an interactive terminal — no `$DISPLAY` setup is needed. If the keys have no effect, make sure you are in an interactive (TTY) terminal, not a piped/non-TTY session, and that it is focused; the letter equivalents `n` / `r` / `q` also work. Keyboard _teleoperation_ (as opposed to the recording control flow) still requires a global key backend — an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — and is unavailable on Wayland or headless machines. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). ## Visualize a dataset diff --git a/docs/source/lekiwi.mdx b/docs/source/lekiwi.mdx index 7e7c1a680..739073b65 100644 --- a/docs/source/lekiwi.mdx +++ b/docs/source/lekiwi.mdx @@ -319,7 +319,7 @@ If you want to dive deeper into this important topic, you can check out the [blo #### Troubleshooting: -- On Linux, if the left and right arrow keys and escape key don't have any effect during data recording, make sure you've set the `$DISPLAY` environment variable. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). +- On Linux, the recording control-flow keys (arrow keys, Escape) work on X11, Wayland, and headless/SSH sessions as long as you run the recording from an interactive terminal (keep it focused) — no `$DISPLAY` setup is needed; the letter equivalents `n` / `r` / `q` also work. Note that **keyboard teleoperation of the LeKiwi base** is different: it relies on a global key backend and therefore works only on an X11 session, a Windows desktop, or macOS with Accessibility/Input Monitoring granted — not on Wayland or headless machines. See [pynput limitations](https://pynput.readthedocs.io/en/latest/limitations.html#linux). ## Replay an episode diff --git a/docs/source/multi_gpu_training.mdx b/docs/source/multi_gpu_training.mdx index d7369e8f8..7c212364e 100644 --- a/docs/source/multi_gpu_training.mdx +++ b/docs/source/multi_gpu_training.mdx @@ -95,7 +95,7 @@ If you want to scale your hyperparameters when using multiple GPUs, you should d accelerate launch --num_processes=2 $(which lerobot-train) \ --optimizer.lr=2e-4 \ --dataset.repo_id=lerobot/pusht \ - --policy=act + --policy.type=act ``` **Training Steps Scaling:** @@ -110,9 +110,64 @@ accelerate launch --num_processes=2 $(which lerobot-train) \ --batch_size=8 \ --steps=50000 \ --dataset.repo_id=lerobot/pusht \ - --policy=act + --policy.type=act ``` +## Training Large Models with FSDP + +DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under +DDP either. For large models, use **FSDP** (Fully Sharded Data Parallel), which shards parameters, +gradients, and optimizer state across GPUs. See the [accelerate FSDP guide](https://huggingface.co/docs/accelerate/usage_guides/fsdp) for background. + +An example on how to launch LeRobot training with FSDP across 4 GPUs (1 machine): + +```bash +accelerate launch --config_file fsdp.yaml --num_processes=4 $(which lerobot-train) \ + --dataset.repo_id=${HF_USER}/my_dataset \ + --policy.type= \ + --output_dir=outputs/train/my_policy_fsdp +``` + +A minimal `fsdp.yaml` (FSDP1; shards params/grads/optimizer — ZeRO-3-equivalent): + +```yaml +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +mixed_precision: bf16 +num_machines: 1 +num_processes: 4 +fsdp_config: + fsdp_version: 1 + fsdp_sharding_strategy: FULL_SHARD # params + grads + optimizer (ZeRO-3) + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_transformer_layer_cls_to_wrap: # repeated block class to shard + fsdp_use_orig_params: true # required: optimizer is built pre-prepare + fsdp_state_dict_type: FULL_STATE_DICT +``` + +Set `fsdp_transformer_layer_cls_to_wrap` to your model's repeated transformer-block class so each +block is sharded as its own unit. `fsdp_use_orig_params: true` is required because LeRobot builds the +optimizer before `accelerator.prepare()`. + +### FSDP checkpoints + +LeRobot gathers the full state dict across all ranks and the main process writes it as a single +`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for: + +- **Checkpoints store fp32 weights.** Under mixed precision (`bf16`/`fp16`) FSDP keeps an fp32 master + copy, and the checkpoint saves it (~2× the bf16 size on disk) so training can resume consistently + with the fp32 optimizer state; `from_pretrained` casts back to the policy dtype on load. FSDP-specific + caveat: an fp32 checkpoint is materialized in full precision on the target device _before_ casting, + so loading it for inference on a tight GPU can OOM even when the bf16 model would fit — load on CPU + first, or cast `model.safetensors` to the deployment dtype offline. +- The sharded optimizer state is gathered into a full (world-size-independent) state dict and saved + alongside the model in the same `optimizer_state.safetensors` / `optimizer_param_groups.json` + format as single-GPU training, so **resume-from-checkpoint is supported** with `--resume=true`. + Resume reshards both the model and the optimizer state to the _current_ FSDP topology, so you can + resume an FSDP checkpoint on a different number of GPUs. Note that the data sampler is only + sample-exact when the world size and batch size match the original run (a warning is logged + otherwise); the optimizer/model state itself is unaffected. + ## Notes - The `--policy.use_amp` flag in `lerobot-train` is only used when **not** running with accelerate. When using accelerate, mixed precision is controlled by accelerate's configuration. diff --git a/docs/source/pi0fast.mdx b/docs/source/pi0fast.mdx index f7272acc5..15dff8071 100644 --- a/docs/source/pi0fast.mdx +++ b/docs/source/pi0fast.mdx @@ -96,7 +96,7 @@ lerobot-train \ --policy.type=pi0_fast \ --output_dir=./outputs/pi0fast_training \ --job_name=pi0fast_training \ - --policy.pretrained_path=lerobot/pi0_fast_base \ + --policy.pretrained_path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --policy.gradient_checkpointing=true \ --policy.chunk_size=10 \ @@ -187,7 +187,7 @@ lerobot-train \ --dataset.repo_id=lerobot/libero \ --output_dir=outputs/libero_pi0fast \ --job_name=libero_pi0fast \ - --policy.path=lerobot/pi0fast_base \ + --policy.path=lerobot/pi0fast-base \ --policy.dtype=bfloat16 \ --steps=100000 \ --save_freq=20000 \ diff --git a/examples/lekiwi/evaluate.py b/examples/lekiwi/evaluate.py index 3ddcd1f14..13bb6ac28 100644 --- a/examples/lekiwi/evaluate.py +++ b/examples/lekiwi/evaluate.py @@ -17,7 +17,7 @@ import logging import time -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.datasets import LeRobotDataset from lerobot.policies import make_pre_post_processors from lerobot.policies.act import ACTPolicy @@ -26,6 +26,7 @@ from lerobot.processor import make_default_processors from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, hw_to_dataset_features +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/lekiwi/record.py b/examples/lekiwi/record.py index 2c581f5ff..f62a9eb49 100644 --- a/examples/lekiwi/record.py +++ b/examples/lekiwi/record.py @@ -14,7 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset from lerobot.processor import make_default_processors from lerobot.robots.lekiwi import LeKiwiClient, LeKiwiClientConfig @@ -23,6 +22,7 @@ from lerobot.teleoperators.keyboard import KeyboardTeleop, KeyboardTeleopConfig from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import hw_to_dataset_features +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/examples/phone_to_so100/evaluate.py b/examples/phone_to_so100/evaluate.py index e859123d0..d1fb4de67 100644 --- a/examples/phone_to_so100/evaluate.py +++ b/examples/phone_to_so100/evaluate.py @@ -18,7 +18,7 @@ import logging import time from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics @@ -41,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/phone_to_so100/record.py b/examples/phone_to_so100/record.py index 87b8e49fd..612e94ab9 100644 --- a/examples/phone_to_so100/record.py +++ b/examples/phone_to_so100/record.py @@ -15,7 +15,6 @@ # limitations under the License. from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( @@ -39,6 +38,7 @@ from lerobot.teleoperators.phone.config_phone import PhoneOS from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction from lerobot.types import RobotAction, RobotObservation from lerobot.utils.feature_utils import combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/examples/so100_to_so100_EE/evaluate.py b/examples/so100_to_so100_EE/evaluate.py index 63def68d0..2a2022623 100644 --- a/examples/so100_to_so100_EE/evaluate.py +++ b/examples/so100_to_so100_EE/evaluate.py @@ -18,7 +18,7 @@ import logging import time from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener, predict_action +from lerobot.common.control_utils import predict_action from lerobot.configs import FeatureType, PolicyFeature from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics @@ -41,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import ( from lerobot.types import RobotAction, RobotObservation from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun, log_rerun_data diff --git a/examples/so100_to_so100_EE/record.py b/examples/so100_to_so100_EE/record.py index a0b92da3b..3706ee4f5 100644 --- a/examples/so100_to_so100_EE/record.py +++ b/examples/so100_to_so100_EE/record.py @@ -16,7 +16,6 @@ from lerobot.cameras.opencv import OpenCVCameraConfig -from lerobot.common.control_utils import init_keyboard_listener from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features from lerobot.model.kinematics import RobotKinematics from lerobot.processor import ( @@ -36,6 +35,7 @@ from lerobot.scripts.lerobot_record import record_loop from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig from lerobot.types import RobotAction, RobotObservation from lerobot.utils.feature_utils import combine_feature_dicts +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import init_rerun diff --git a/pyproject.toml b/pyproject.toml index 0dc86d7ff..3961ded19 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -140,7 +140,14 @@ av-dep = ["av>=15.0.0,<16.0.0"] pygame-dep = ["pygame>=2.5.1,<2.7.0"] # NOTE: 0.9.16 links against liburdfdom_sensor.so.4, which is unavailable on Ubuntu 24.04 # (noble ships urdfdom 3.x). Cap below 0.9.16 until system urdfdom 4.x is broadly available. -placo-dep = ["placo>=0.9.6,<0.9.16"] +# +# NOTE: placo pulls in pin (Pinocchio), whose binary wheels dlopen specific cmeel sonames +# (liburdfdom_sensor.so.4.0, libtinyxml2.so.10) but declare only `>=` floors on their cmeel +# packages. The 2026-05-21 major bumps (cmeel-urdfdom 6.0.0 -> .so.6, cmeel-tinyxml2 11.0.0 +# -> .so.11) ship newer sonames, so left unpinned the resolver grabs them and `import placo` +# fails at load with "liburdfdom_sensor.so.4.0: cannot open shared object file" (see #3755). +# There is no cmeel-urdfdom 5.x; <5 selects the 4.x ABI the placo/pin wheels are built against. +placo-dep = ["placo>=0.9.6,<0.9.16", "cmeel-urdfdom>=4,<5", "cmeel-tinyxml2<11"] transformers-dep = ["transformers>=5.4.0,<5.6.0"] grpcio-dep = ["grpcio>=1.73.1,<2.0.0", "protobuf>=6.31.1,<8.0.0"] accelerate-dep = ["accelerate>=1.14.0,<2.0.0"] diff --git a/src/lerobot/cameras/opencv/camera_opencv.py b/src/lerobot/cameras/opencv/camera_opencv.py index 3e92eaf06..b3c20e8dd 100644 --- a/src/lerobot/cameras/opencv/camera_opencv.py +++ b/src/lerobot/cameras/opencv/camera_opencv.py @@ -442,11 +442,12 @@ class OpenCVCamera(Camera): Stops on DeviceNotConnectedError, logs other errors and continues. """ - if self.stop_event is None: + stop_event = self.stop_event + if stop_event is None: raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.") failure_count = 0 - while not self.stop_event.is_set(): + while not stop_event.is_set(): try: raw_frame = self._read_from_hardware() processed_frame = self._postprocess_image(raw_frame) diff --git a/src/lerobot/cameras/realsense/camera_realsense.py b/src/lerobot/cameras/realsense/camera_realsense.py index e156e6d14..80008e9f9 100644 --- a/src/lerobot/cameras/realsense/camera_realsense.py +++ b/src/lerobot/cameras/realsense/camera_realsense.py @@ -471,11 +471,12 @@ class RealSenseCamera(Camera): Stops on DeviceNotConnectedError, logs other errors and continues. """ - if self.stop_event is None: + stop_event = self.stop_event + if stop_event is None: raise RuntimeError(f"{self}: stop_event is not initialized before starting read loop.") failure_count = 0 - while not self.stop_event.is_set(): + while not stop_event.is_set(): try: frame = self._read_from_hardware() color_frame_raw = frame.get_color_frame() diff --git a/src/lerobot/cameras/zmq/camera_zmq.py b/src/lerobot/cameras/zmq/camera_zmq.py index 1b0be5de6..f3df17814 100644 --- a/src/lerobot/cameras/zmq/camera_zmq.py +++ b/src/lerobot/cameras/zmq/camera_zmq.py @@ -246,11 +246,12 @@ class ZMQCamera(Camera): """ Internal loop run by the background thread for asynchronous reading. """ - if self.stop_event is None: + stop_event = self.stop_event + if stop_event is None: raise RuntimeError(f"{self}: stop_event is not initialized.") failure_count = 0 - while not self.stop_event.is_set(): + while not stop_event.is_set(): try: frame = self._read_from_hardware() capture_time = time.perf_counter() diff --git a/src/lerobot/common/control_utils.py b/src/lerobot/common/control_utils.py index ddaf77d26..e3130643d 100644 --- a/src/lerobot/common/control_utils.py +++ b/src/lerobot/common/control_utils.py @@ -17,12 +17,9 @@ from __future__ import annotations ######################################################################################## # Utilities ######################################################################################## -import logging import time -import traceback from contextlib import nullcontext from copy import copy -from functools import cache from typing import TYPE_CHECKING, Any import numpy as np @@ -43,34 +40,6 @@ from lerobot.robots import Robot from lerobot.types import PolicyAction -@cache -def is_headless(): - """ - Detects if the Python script is running in a headless environment (e.g., without a display). - - This function attempts to import `pynput`, a library that requires a graphical environment. - If the import fails, it assumes the environment is headless. The result is cached to avoid - re-running the check. - - Returns: - True if the environment is determined to be headless, False otherwise. - """ - try: - import pynput # noqa - - return False - except Exception: - print( - "Error trying to import pynput. Switching to headless mode. " - "As a result, the video stream from the cameras won't be shown, " - "and you won't be able to change the control flow with keyboards. " - "For more info, see traceback below.\n" - ) - traceback.print_exc() - print() - return True - - def predict_action( observation: dict[str, np.ndarray], policy: PreTrainedPolicy, @@ -122,59 +91,6 @@ def predict_action( return action -def init_keyboard_listener(): - """ - Initializes a non-blocking keyboard listener for real-time user interaction. - - This function sets up a listener for specific keys (right arrow, left arrow, escape) to control - the program flow during execution, such as stopping recording or exiting loops. It gracefully - handles headless environments where keyboard listening is not possible. - - Returns: - A tuple containing: - - The `pynput.keyboard.Listener` instance, or `None` if in a headless environment. - - A dictionary of event flags (e.g., `exit_early`) that are set by key presses. - """ - # Allow to exit early while recording an episode or resetting the environment, - # by tapping the right arrow key '->'. This might require a sudo permission - # to allow your terminal to monitor keyboard events. - events = {} - events["exit_early"] = False - events["rerecord_episode"] = False - events["stop_recording"] = 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...") - events["exit_early"] = True - elif key == keyboard.Key.left: - print("Left arrow key pressed. Exiting loop and rerecord the last episode...") - events["rerecord_episode"] = True - events["exit_early"] = True - elif key == keyboard.Key.esc: - print("Escape key pressed. Stopping data recording...") - events["stop_recording"] = True - events["exit_early"] = True - except Exception as e: - print(f"Error handling key press: {e}") - - listener = keyboard.Listener(on_press=on_press) - listener.start() - - return listener, events - - def sanity_check_dataset_name(repo_id, policy_cfg): """ Validates the dataset repository name against the presence of a policy configuration. diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index 2d23b4003..5ae593bb8 100644 --- a/src/lerobot/common/train_utils.py +++ b/src/lerobot/common/train_utils.py @@ -21,6 +21,7 @@ from torch.optim.lr_scheduler import LRScheduler from lerobot.configs.train import TrainPipelineConfig from lerobot.optim import ( load_optimizer_state, + load_optimizer_state_dict, load_scheduler_state, save_optimizer_state, save_scheduler_state, @@ -98,6 +99,8 @@ def save_checkpoint( postprocessor: PolicyProcessorPipeline | None = None, num_processes: int | None = None, batch_size: int | None = None, + model_state_dict: dict | None = None, + optim_state_dict: dict | None = None, ) -> None: """This function creates the following directory structure: @@ -127,9 +130,18 @@ def save_checkpoint( resume. Defaults to None (not recorded). batch_size (int | None, optional): Per-process batch size to record for sample-exact resume. Defaults to None (not recorded). + model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP, + where `policy.state_dict()` would return sharded tensors; the caller gathers it via a + cross-rank collective and passes it here so rank 0 can write it directly. It holds + FSDP's fp32 master weights and is saved as-is (the loader casts to the policy dtype on + read). When None (DDP / single-GPU), the model is saved the normal way. Defaults to None. + optim_state_dict: Pre-gathered full (unsharded) optimizer state dict. Required under FSDP + (gathered alongside `model_state_dict` via `gather_fsdp_state_dicts`); saved in the same + safetensors format as the single-GPU path. When None, `optimizer.state_dict()` is used. + Defaults to None. """ pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR - policy.save_pretrained(pretrained_dir) + policy.save_pretrained(pretrained_dir, state_dict=model_state_dict) cfg.save_pretrained(pretrained_dir) if cfg.peft is not None: # When using PEFT, policy.save_pretrained will only write the adapter weights + config, not the @@ -140,7 +152,13 @@ def save_checkpoint( if postprocessor is not None: postprocessor.save_pretrained(pretrained_dir) save_training_state( - checkpoint_dir, step, optimizer, scheduler, num_processes=num_processes, batch_size=batch_size + checkpoint_dir, + step, + optimizer, + scheduler, + num_processes=num_processes, + batch_size=batch_size, + optim_state_dict=optim_state_dict, ) @@ -151,6 +169,7 @@ def save_training_state( scheduler: LRScheduler | None = None, num_processes: int | None = None, batch_size: int | None = None, + optim_state_dict: dict | None = None, ) -> None: """ Saves the training step, optimizer state, scheduler state, and rng state. @@ -164,19 +183,21 @@ def save_training_state( Defaults to None. num_processes (int | None, optional): Distributed world size to record. Defaults to None. batch_size (int | None, optional): Per-process batch size to record. Defaults to None. + optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of + `optimizer.state_dict()` when provided. Defaults to None. """ save_dir = checkpoint_dir / TRAINING_STATE_DIR save_dir.mkdir(parents=True, exist_ok=True) save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size) save_rng_state(save_dir) if optimizer is not None: - save_optimizer_state(optimizer, save_dir) + save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) if scheduler is not None: save_scheduler_state(scheduler, save_dir) def load_training_state( - checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None + checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None, load_optimizer: bool = True ) -> tuple[int, Optimizer, LRScheduler | None]: """ Loads the training step, optimizer state, scheduler state, and rng state. @@ -186,6 +207,10 @@ def load_training_state( checkpoint_dir (Path): The checkpoint directory. Should contain a 'training_state' dir. optimizer (Optimizer): The optimizer to load the state_dict to. scheduler (LRScheduler | None): The scheduler to load the state_dict to (can be None). + load_optimizer (bool, optional): Whether to load the optimizer state from disk. Defaults to + True. Set to False under FSDP, where the sharded optimizer state must be loaded after + `accelerator.prepare()` via `load_fsdp_optimizer_state` (the optimizer is returned + untouched here). Raises: NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir @@ -200,8 +225,61 @@ def load_training_state( load_rng_state(training_state_dir) step = load_training_step(training_state_dir) - optimizer = load_optimizer_state(optimizer, training_state_dir) + if load_optimizer: + optimizer = load_optimizer_state(optimizer, training_state_dir) if scheduler is not None: scheduler = load_scheduler_state(scheduler, training_state_dir) return step, optimizer, scheduler + + +def gather_fsdp_state_dicts(model, optimizer) -> tuple[dict, dict]: + """Gather the full (unsharded) model and optimizer state dicts under FSDP. + + `model.state_dict()` and `FSDP.optim_state_dict(...)` are cross-rank collectives, so this must be + called on *every* rank with the prepared (FSDP-wrapped) `model` and `optimizer`. With + `rank0_only=True` and `offload_to_cpu=True`, every rank runs the all-gather but only rank 0 + materializes the full dicts (the others get empty dicts) and they are kept on CPU to bound GPU + memory. The returned optimizer state dict is keyed by parameter FQNs and is world-size + independent; `load_fsdp_optimizer_state` reshards it on resume. + + Returns: + (model_state_dict, optim_state_dict): full dicts on rank 0, empty dicts on other ranks. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + state_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + optim_cfg = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + model_state_dict = model.state_dict() + optim_state_dict = FSDP.optim_state_dict(model, optimizer) + return model_state_dict, optim_state_dict + + +def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None: + """Load the FSDP optimizer state (saved as safetensors) and reshard it into the optimizer. + + This is a cross-rank collective and must be called on every rank *after* `accelerator.prepare()` + with the prepared (FSDP-wrapped) `model` and `optimizer`. The saved state is the full, + world-size-independent optimizer state (keyed by parameter FQNs); `FSDP.optim_state_dict_to_load` + reshards it to the current FSDP topology, so resume on a different number of GPUs works. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + # Every rank reads the same full state from the (shared) checkpoint dir, so rank0_only=False. + full_osd = load_optimizer_state_dict(checkpoint_dir / TRAINING_STATE_DIR) + state_cfg = FullStateDictConfig(rank0_only=False) + optim_cfg = FullOptimStateDictConfig(rank0_only=False) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd) + optimizer.load_state_dict(sharded_osd) diff --git a/src/lerobot/common/wandb_utils.py b/src/lerobot/common/wandb_utils.py index b782cd751..c229b5eaa 100644 --- a/src/lerobot/common/wandb_utils.py +++ b/src/lerobot/common/wandb_utils.py @@ -180,24 +180,26 @@ class WandBLogger: self._wandb_custom_step_key.add(new_custom_key) self._wandb.define_metric(new_custom_key, hidden=True) + batch_data = {} for k, v in d.items(): + # Skip the custom step key here, it's added to the batch below. + if custom_step_key is not None and k == custom_step_key: + continue + if not isinstance(v, (int | float | str)): logging.warning( f'WandB logging of key "{k}" was ignored as its type "{type(v)}" is not handled by this wrapper.' ) continue - # Do not log the custom step key itself. - if self._wandb_custom_step_key is not None and k in self._wandb_custom_step_key: - continue + batch_data[f"{mode}/{k}"] = v + if batch_data: if custom_step_key is not None: - value_custom_step = d[custom_step_key] - data = {f"{mode}/{k}": v, f"{mode}/{custom_step_key}": value_custom_step} - self._wandb.log(data) - continue - - self._wandb.log(data={f"{mode}/{k}": v}, step=step) + batch_data[f"{mode}/{custom_step_key}"] = d[custom_step_key] + self._wandb.log(batch_data) + else: + self._wandb.log(data=batch_data, step=step) def log_video(self, video_path: str, step: int, mode: str = "train"): if mode not in {"train", "eval"}: diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index b809e71d9..648e03f33 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -73,8 +73,17 @@ class EvalConfig: # `use_async_envs` specifies whether to use asynchronous environments (multiprocessing). # Defaults to True; automatically downgraded to SyncVectorEnv when batch_size=1. use_async_envs: bool = True + # Whether to record eval rollouts as a LeRobot dataset on disk. + recording: bool = False + # If set, push recorded eval datasets to the Hub under this repo id (one repo per task, + # suffixed by task and env index). Requires recording=true. + recording_repo_id: str | None = None + # Whether the pushed recording repositories should be private. + recording_private: bool = False def __post_init__(self) -> None: + if self.recording_repo_id is not None and not self.recording: + raise ValueError("eval.recording_repo_id requires eval.recording=true.") if self.batch_size == 0: self.batch_size = self._auto_batch_size() if self.batch_size > self.n_episodes: diff --git a/src/lerobot/configs/policies.py b/src/lerobot/configs/policies.py index 91701af6d..b0f003519 100644 --- a/src/lerobot/configs/policies.py +++ b/src/lerobot/configs/policies.py @@ -79,6 +79,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno # Either the repo ID of a model hosted on the Hub or a path to a directory containing weights # saved using `Policy.save_pretrained`. If not provided, the policy is initialized from scratch. pretrained_path: Path | None = None + # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained model version. + pretrained_revision: str | None = None def __post_init__(self) -> None: if not self.device or not is_torch_device_available(self.device): diff --git a/src/lerobot/configs/rewards.py b/src/lerobot/configs/rewards.py index 7e99e7f71..92490bc9f 100644 --- a/src/lerobot/configs/rewards.py +++ b/src/lerobot/configs/rewards.py @@ -56,6 +56,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): device: str | None = None pretrained_path: str | None = None + # Optional Hub revision (commit hash, branch, or tag) to pin the pretrained reward model version. + pretrained_revision: str | None = None push_to_hub: bool = False repo_id: str | None = None diff --git a/src/lerobot/datasets/dataset_metadata.py b/src/lerobot/datasets/dataset_metadata.py index 39a1b6d2b..b496e4f65 100644 --- a/src/lerobot/datasets/dataset_metadata.py +++ b/src/lerobot/datasets/dataset_metadata.py @@ -15,6 +15,7 @@ # limitations under the License. import contextlib from collections.abc import Callable +from copy import deepcopy from pathlib import Path import numpy as np @@ -709,7 +710,7 @@ class LeRobotDatasetMetadata: obj.root.mkdir(parents=True, exist_ok=False) - features = {**features, **DEFAULT_FEATURES} + features = {**deepcopy(features), **DEFAULT_FEATURES} _validate_feature_names(features) obj.tasks = None diff --git a/src/lerobot/datasets/dataset_reader.py b/src/lerobot/datasets/dataset_reader.py index 59aaa40e5..d7289ac48 100644 --- a/src/lerobot/datasets/dataset_reader.py +++ b/src/lerobot/datasets/dataset_reader.py @@ -74,6 +74,8 @@ class DatasetReader: self.episodes = episodes self._tolerance_s = tolerance_s self._video_backend = video_backend + if image_transforms is not None and not callable(image_transforms): + raise TypeError("image_transforms must be callable or None.") self._image_transforms = image_transforms self._return_uint8 = return_uint8 @@ -86,6 +88,16 @@ class DatasetReader: check_delta_timestamps(delta_timestamps, meta.fps, tolerance_s) self.delta_indices = get_delta_indices(delta_timestamps, meta.fps) + def set_image_transforms(self, image_transforms: Callable | None) -> None: + """Replace the transform applied to visual observations.""" + if image_transforms is not None and not callable(image_transforms): + raise TypeError("image_transforms must be callable or None.") + self._image_transforms = image_transforms + + def clear_image_transforms(self) -> None: + """Remove the transform applied to visual observations.""" + self._image_transforms = None + def try_load(self) -> bool: """Attempt to load from local cache. Returns True if data is sufficient.""" try: diff --git a/src/lerobot/datasets/dataset_tools.py b/src/lerobot/datasets/dataset_tools.py index 91dc66af2..9aca859b4 100644 --- a/src/lerobot/datasets/dataset_tools.py +++ b/src/lerobot/datasets/dataset_tools.py @@ -27,6 +27,7 @@ import logging import shutil from collections.abc import Callable from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed +from copy import deepcopy from pathlib import Path import datasets @@ -1101,7 +1102,9 @@ def _copy_episodes_metadata_and_stats( if dst_meta.video_keys and src_dataset.meta.video_keys: for key in dst_meta.video_keys: if key in src_dataset.meta.features: - dst_meta.info.features[key]["info"] = src_dataset.meta.info.features[key].get("info", {}) + dst_meta.info.features[key]["info"] = deepcopy( + src_dataset.meta.info.features[key].get("info", {}) + ) write_info(dst_meta.info, dst_meta.root) diff --git a/src/lerobot/datasets/io_utils.py b/src/lerobot/datasets/io_utils.py index b6344942c..be94f3b3a 100644 --- a/src/lerobot/datasets/io_utils.py +++ b/src/lerobot/datasets/io_utils.py @@ -154,7 +154,7 @@ def cast_stats_to_numpy(stats: dict) -> dict[str, dict[str, np.ndarray]]: Returns: dict: The statistics dictionary with values cast to numpy arrays. """ - stats = {key: np.array(value) for key, value in flatten_dict(stats).items()} + stats = {key: np.atleast_1d(np.array(value)) for key, value in flatten_dict(stats).items()} return unflatten_dict(stats) diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index d0dcf087d..d1e65fef1 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -201,8 +201,6 @@ class LeRobotDataset(torch.utils.data.Dataset): super().__init__() self.repo_id = repo_id self._requested_root = Path(root) if root else None - self.reader = None - self.set_image_transforms(image_transforms) self.delta_timestamps = delta_timestamps self.tolerance_s = tolerance_s self.revision = revision if revision else CODEBASE_VERSION @@ -249,6 +247,7 @@ class LeRobotDataset(torch.utils.data.Dataset): image_transforms=image_transforms, return_uint8=self._return_uint8, ) + self.image_transforms = image_transforms # Load actual data if force_cache_sync or not self.reader.try_load(): @@ -505,15 +504,14 @@ class LeRobotDataset(torch.utils.data.Dataset): def set_image_transforms(self, image_transforms: Callable | None) -> None: """Replace the transform applied to visual observations.""" - if image_transforms is not None and not callable(image_transforms): - raise TypeError("image_transforms must be callable or None.") + self._ensure_reader().set_image_transforms(image_transforms) self.image_transforms = image_transforms - if self.reader is not None: - self.reader._image_transforms = image_transforms def clear_image_transforms(self) -> None: """Remove the transform applied to visual observations.""" - self.set_image_transforms(None) + if self.reader is not None: + self.reader.set_image_transforms(None) + self.image_transforms = None # ── Hub methods (stay on facade) ────────────────────────────────── diff --git a/src/lerobot/envs/utils.py b/src/lerobot/envs/utils.py index 6e6f352e9..8b9c4f94b 100644 --- a/src/lerobot/envs/utils.py +++ b/src/lerobot/envs/utils.py @@ -126,6 +126,26 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten if "camera_obs" in observations: return_observations[f"{OBS_STR}.camera_obs"] = observations["camera_obs"] + # Pass through any remaining ndarray/tensor keys not already handled above, + # so env plugins can expose extra observation keys via get_env_processors(). + _handled = {"pixels", "environment_state", "agent_pos", "robot_state", "policy", "camera_obs"} + for key, value in observations.items(): + if key in _handled: + continue + target = f"{OBS_STR}.{key}" + if target in return_observations: + continue + if isinstance(value, np.ndarray): + val = torch.from_numpy(value).float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + elif isinstance(value, Tensor): + val = value.float() + if val.dim() == 1: + val = val.unsqueeze(0) + return_observations[target] = val + return return_observations diff --git a/src/lerobot/optim/__init__.py b/src/lerobot/optim/__init__.py index 46676027b..2d564c25f 100644 --- a/src/lerobot/optim/__init__.py +++ b/src/lerobot/optim/__init__.py @@ -20,6 +20,7 @@ from .optimizers import ( SGDConfig as SGDConfig, XVLAAdamWConfig as XVLAAdamWConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from .schedulers import ( @@ -50,6 +51,7 @@ __all__ = [ "VQBeTSchedulerConfig", # State management "load_optimizer_state", + "load_optimizer_state_dict", "load_scheduler_state", "save_optimizer_state", "save_scheduler_state", diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 0bdd7a37e..0a462e1aa 100644 --- a/src/lerobot/optim/optimizers.py +++ b/src/lerobot/optim/optimizers.py @@ -27,7 +27,7 @@ from lerobot.utils.constants import ( OPTIMIZER_PARAM_GROUPS, OPTIMIZER_STATE, ) -from lerobot.utils.io_utils import deserialize_json_into_object, write_json +from lerobot.utils.io_utils import deserialize_json_into_object, load_json, write_json from lerobot.utils.utils import flatten_dict, unflatten_dict # Type alias for parameters accepted by optimizer build() methods. @@ -281,28 +281,37 @@ class MultiAdamConfig(OptimizerConfig): def save_optimizer_state( - optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], save_dir: Path + optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], + save_dir: Path, + optim_state_dict: dict | None = None, ) -> None: """Save optimizer state to disk. Args: optimizer: Either a single optimizer or a dictionary of optimizers. save_dir: Directory to save the optimizer state. + optim_state_dict: Pre-gathered optimizer state dict (for FSDP, where the sharded state must + be gathered across ranks first). If provided, it is saved directly instead of calling + ``optimizer.state_dict()``. Only supported for a single optimizer. Defaults to None. """ if isinstance(optimizer, dict): # Handle dictionary of optimizers + if optim_state_dict is not None: + raise ValueError("optim_state_dict is not supported for a dict of optimizers") for name, opt in optimizer.items(): optimizer_dir = save_dir / name optimizer_dir.mkdir(exist_ok=True, parents=True) _save_single_optimizer_state(opt, optimizer_dir) else: # Handle single optimizer - _save_single_optimizer_state(optimizer, save_dir) + _save_single_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) -def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None: +def _save_single_optimizer_state( + optimizer: torch.optim.Optimizer, save_dir: Path, optim_state_dict: dict | None = None +) -> None: """Save a single optimizer's state to disk.""" - state = optimizer.state_dict() + state = dict(optim_state_dict) if optim_state_dict is not None else optimizer.state_dict() param_groups = state.pop("param_groups") flat_state = flatten_dict(state) save_file(flat_state, save_dir / OPTIMIZER_STATE) @@ -356,3 +365,19 @@ def _load_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Pat optimizer.load_state_dict(loaded_state_dict) return optimizer + + +def load_optimizer_state_dict(save_dir: Path) -> dict: + """Read a saved optimizer state dict (safetensors + json) back into a plain dict. + + Unlike `load_optimizer_state`, this does not load into an optimizer and preserves the original + ``state`` keys verbatim (e.g. FSDP parameter FQNs, which are not integer-castable). It is used by + the FSDP resume path, where the full state must be resharded via `FSDP.optim_state_dict_to_load` + before being loaded into the (sharded) optimizer. + """ + flat_state = load_file(save_dir / OPTIMIZER_STATE) + state = unflatten_dict(flat_state) + return { + "state": state.get("state", {}), + "param_groups": load_json(save_dir / OPTIMIZER_PARAM_GROUPS), + } diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index 5651fbfb1..1432b68a5 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/src/lerobot/policies/act/modeling_act.py @@ -148,7 +148,7 @@ class ACTPolicy(PreTrainedPolicy): l1_loss = (abs_err * valid_mask).sum() / num_valid.clamp_min(1) loss_dict = {"l1_loss": l1_loss.item()} - if self.config.use_vae: + if self.config.use_vae and log_sigma_x2_hat is not None: # Calculate Dₖₗ(latent_pdf || standard_normal). Note: After computing the KL-divergence for # each dimension independently, we sum over the latent dimension to get the total # KL-divergence per batch element, then take the mean over the batch. diff --git a/src/lerobot/policies/diffusion/modeling_diffusion.py b/src/lerobot/policies/diffusion/modeling_diffusion.py index 9fbe1f703..8758a7e29 100644 --- a/src/lerobot/policies/diffusion/modeling_diffusion.py +++ b/src/lerobot/policies/diffusion/modeling_diffusion.py @@ -101,11 +101,23 @@ class DiffusionPolicy(PreTrainedPolicy): @torch.no_grad() def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: - """Predict a chunk of actions given environment observations.""" - # stack n latest observations from the queue - batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} - actions = self.diffusion.generate_actions(batch, noise=noise) + """Predict a chunk of actions given environment observations. + Supports two modes: + - Online (queues populated via select_action): stacks observations from internal queues. + - Offline (empty queues, e.g. dataloader batch): uses the batch directly. + """ + queues_populated = any(len(q) > 0 for q in self._queues.values()) + if queues_populated: + batch = {k: torch.stack(list(self._queues[k]), dim=1) for k in batch if k in self._queues} + else: + batch = dict(batch) + if self.config.image_features: + for key in self.config.image_features: + if batch[key].ndim == 4: + batch[key] = batch[key].unsqueeze(1) + batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4) + actions = self.diffusion.generate_actions(batch, noise=noise) return actions @torch.no_grad() diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index a42b38ba4..b82eaeb72 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -252,6 +252,7 @@ class ProcessorConfigKwargs(TypedDict, total=False): def make_pre_post_processors( policy_cfg: PreTrainedConfig, pretrained_path: str | None = None, + pretrained_revision: str | None = None, **kwargs: Unpack[ProcessorConfigKwargs], ) -> tuple[ PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], @@ -309,6 +310,7 @@ def make_pre_post_processors( overrides=kwargs.get("preprocessor_overrides", {}), to_transition=batch_to_transition, to_output=transition_to_batch, + revision=pretrained_revision, ) postprocessor = PolicyProcessorPipeline.from_pretrained( pretrained_model_name_or_path=pretrained_path, @@ -318,6 +320,7 @@ def make_pre_post_processors( overrides=kwargs.get("postprocessor_overrides", {}), to_transition=policy_action_to_transition, to_output=transition_to_policy_action, + revision=pretrained_revision, ) _reconnect_relative_absolute_steps(preprocessor, postprocessor) return preprocessor, postprocessor @@ -557,6 +560,7 @@ def make_policy( # Load a pretrained policy and override the config if needed (for example, if there are inference-time # hyperparameters that we want to vary). kwargs["pretrained_name_or_path"] = cfg.pretrained_path + kwargs["revision"] = cfg.pretrained_revision policy = policy_cls.from_pretrained(**kwargs) elif cfg.pretrained_path and cfg.use_peft: # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index a69487f3f..a7aabb3f3 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -23,7 +23,7 @@ from typing import TypedDict, TypeVar, Unpack import packaging import safetensors -from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download +from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from huggingface_hub.errors import HfHubHTTPError from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor @@ -129,10 +129,43 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): if not getattr(cls, "name", None): raise TypeError(f"Class {cls.__name__} must define 'name'") - def _save_pretrained(self, save_directory: Path) -> None: + def save_pretrained( + self, + save_directory: str | Path, + *, + state_dict: dict[str, Tensor] | None = None, + repo_id: str | None = None, + push_to_hub: bool = False, + card_kwargs: dict | None = None, + **push_to_hub_kwargs, + ) -> str | None: + """Save the policy to a directory (and optionally push to the Hub). + + Overrides `HubMixin.save_pretrained` to add a `state_dict` argument (mirroring + `transformers.PreTrainedModel.save_pretrained`). Under FSDP, `self.state_dict()` would + return sharded tensors, so the caller gathers the full state dict via a cross-rank + collective and passes it here for `_save_pretrained` to write directly. + """ + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + self._save_pretrained(save_directory, state_dict=state_dict) + if push_to_hub: + if repo_id is None: + repo_id = save_directory.name + return self.push_to_hub(repo_id=repo_id, card_kwargs=card_kwargs, **push_to_hub_kwargs) + return None + + def _save_pretrained(self, save_directory: Path, state_dict: dict[str, Tensor] | None = None) -> None: self.config._save_pretrained(save_directory) model_to_save = self.module if hasattr(self, "module") else self - save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + if state_dict is None: + save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + return + # A pre-gathered (e.g. FSDP full) state dict was supplied: write it directly. + # `save_torch_state_dict` discards shared-tensor duplicates just like `save_model` does; + # pin `max_shard_size` above the total size so the output stays a single `model.safetensors` + total_bytes = sum(t.numel() * t.element_size() for t in state_dict.values()) + save_torch_state_dict(state_dict, str(save_directory), max_shard_size=max(total_bytes, 1)) @classmethod def from_pretrained( @@ -270,6 +303,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self, cfg: TrainPipelineConfig, peft_model=None, + state_dict: dict[str, Tensor] | None = None, ): api = HfApi() repo_id = api.create_repo( @@ -287,7 +321,8 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): peft_model.save_pretrained(saved_path) self.config.save_pretrained(saved_path) else: - self.save_pretrained(saved_path) # Calls _save_pretrained and stores model tensors + # Calls _save_pretrained and stores model tensors + self.save_pretrained(saved_path, state_dict=state_dict) card = self.generate_model_card( cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg diff --git a/src/lerobot/rewards/factory.py b/src/lerobot/rewards/factory.py index 2d73ae575..fee90c211 100644 --- a/src/lerobot/rewards/factory.py +++ b/src/lerobot/rewards/factory.py @@ -124,6 +124,7 @@ def make_reward_model(cfg: RewardModelConfig, **kwargs) -> PreTrainedRewardModel if cfg.pretrained_path: kwargs["pretrained_name_or_path"] = cfg.pretrained_path + kwargs["revision"] = cfg.pretrained_revision reward_model = reward_cls.from_pretrained(**kwargs) else: reward_model = reward_cls(**kwargs) diff --git a/src/lerobot/rollout/strategies/dagger.py b/src/lerobot/rollout/strategies/dagger.py index 8791a5502..21d1e8e98 100644 --- a/src/lerobot/rollout/strategies/dagger.py +++ b/src/lerobot/rollout/strategies/dagger.py @@ -47,8 +47,6 @@ from __future__ import annotations import contextlib import enum import logging -import os -import sys import time from concurrent.futures import Future, ThreadPoolExecutor from threading import Event, Lock @@ -58,7 +56,6 @@ import numpy as np from lerobot.common.control_utils import ( follower_smooth_move_to, - is_headless, teleop_smooth_move_to, teleop_supports_feedback, ) @@ -66,7 +63,7 @@ from lerobot.datasets import VideoEncodingManager from lerobot.datasets.utils import DEFAULT_VIDEO_FILE_SIZE_IN_MB from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame -from lerobot.utils.import_utils import _pynput_available +from lerobot.utils.keyboard_input import create_key_listener from lerobot.utils.pedal import start_pedal_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say @@ -75,19 +72,6 @@ from ..configs import DAggerKeyboardConfig, DAggerPedalConfig, DAggerStrategyCon from ..context import RolloutContext from .core import RolloutStrategy, estimate_max_episode_seconds, safe_push_to_hub, send_next_action -PYNPUT_AVAILABLE = _pynput_available -keyboard = None -if PYNPUT_AVAILABLE: - try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard - except Exception as e: - PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") - logger = logging.getLogger(__name__) @@ -180,64 +164,36 @@ class DAggerEvents: def _init_dagger_keyboard(events: DAggerEvents, cfg: DAggerKeyboardConfig): - """Initialise keyboard listener with DAgger 3-key controls. + """Initialise a keyboard listener for DAgger's 3 controls. - Returns the pynput Listener (or ``None`` in headless mode or when - pynput is unavailable). + Backend selection (pynput on X11 / trusted-macOS / Windows, a terminal reader on + Wayland / headless TTY) is delegated to :func:`create_key_listener`. Returns the + listener (exposing ``stop()``) or ``None`` when no keyboard backend is usable. """ - if not PYNPUT_AVAILABLE or is_headless(): - logger.warning("Headless environment or pynput unavailable — keyboard controls disabled") - return None - - # Map config key names to pynput Key objects for special keys - special_keys = { - "space": keyboard.Key.space, - "tab": keyboard.Key.tab, - "enter": keyboard.Key.enter, - } - - def _resolve_key(key) -> str | None: - """Resolve a pynput key event to a config-comparable string.""" - if key == keyboard.Key.esc: - return "esc" - for name, pynput_key in special_keys.items(): - if key == pynput_key: - return name - if hasattr(key, "char") and key.char: - return key.char - return None - - # Build mapping: resolved key string -> DAgger event name + # Map config key names to DAgger event names. key_to_event = { cfg.pause_resume: "pause_resume", cfg.correction: "correction", } - def on_press(key): - try: - resolved = _resolve_key(key) - if resolved is None: - return - if resolved == "esc": - logger.info("Stop recording...") - events.stop_recording.set() - return - if resolved in key_to_event: - events.request_transition(key_to_event[resolved]) - if resolved == cfg.upload: - events.upload_requested.set() - except Exception as e: - logger.debug("Key error: %s", e) + def dispatch(name: str) -> None: + """Apply a resolved key name to the DAgger events.""" + if name == "esc": + logger.info("Stop recording...") + events.stop_recording.set() + return + if name in key_to_event: + events.request_transition(key_to_event[name]) + if name == cfg.upload: + events.upload_requested.set() - listener = keyboard.Listener(on_press=on_press) - listener.start() - logger.info( - "DAgger keyboard listener started (pause_resume='%s', correction='%s', upload='%s', ESC=stop)", - cfg.pause_resume, - cfg.correction, - cfg.upload, + return create_key_listener( + dispatch, + controls_help=( + f"pause_resume='{cfg.pause_resume}', correction='{cfg.correction}', " + f"upload='{cfg.upload}', ESC=stop" + ), ) - return listener def _init_dagger_pedal(events: DAggerEvents, cfg: DAggerPedalConfig): @@ -328,7 +284,7 @@ class DAggerStrategy(RolloutStrategy): logger.info("Stopping DAgger recording") log_say("Stopping DAgger recording", play_sounds) - if self._listener is not None and not is_headless(): + if self._listener is not None: logger.info("Stopping keyboard listener") self._listener.stop() diff --git a/src/lerobot/rollout/strategies/episodic.py b/src/lerobot/rollout/strategies/episodic.py index e925fb2ea..e70e66787 100644 --- a/src/lerobot/rollout/strategies/episodic.py +++ b/src/lerobot/rollout/strategies/episodic.py @@ -35,14 +35,13 @@ import time from lerobot.common.control_utils import ( follower_smooth_move_to, - init_keyboard_listener, - is_headless, teleop_smooth_move_to, teleop_supports_feedback, ) from lerobot.datasets import VideoEncodingManager from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say from lerobot.utils.visualization_utils import log_rerun_data @@ -307,7 +306,7 @@ class EpisodicStrategy(RolloutStrategy): log_say("Stop recording", play_sounds, blocking=True) - if not is_headless() and self._listener is not None: + if self._listener is not None: self._listener.stop() if ctx.data.dataset is not None: diff --git a/src/lerobot/rollout/strategies/highlight.py b/src/lerobot/rollout/strategies/highlight.py index baff70da7..385a9e2b6 100644 --- a/src/lerobot/rollout/strategies/highlight.py +++ b/src/lerobot/rollout/strategies/highlight.py @@ -18,17 +18,14 @@ from __future__ import annotations import contextlib import logging -import os -import sys import time from concurrent.futures import Future, ThreadPoolExecutor from threading import Event as ThreadingEvent, Lock -from lerobot.common.control_utils import is_headless from lerobot.datasets import VideoEncodingManager from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame -from lerobot.utils.import_utils import _pynput_available, require_package +from lerobot.utils.keyboard_input import create_key_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import log_say @@ -37,19 +34,6 @@ from ..context import RolloutContext from ..ring_buffer import RolloutRingBuffer from .core import RolloutStrategy, safe_push_to_hub, send_next_action -PYNPUT_AVAILABLE = _pynput_available -keyboard = None -if PYNPUT_AVAILABLE: - try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard - except Exception as e: - PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") - logger = logging.getLogger(__name__) @@ -72,7 +56,6 @@ class HighlightStrategy(RolloutStrategy): def __init__(self, config: HighlightStrategyConfig): super().__init__(config) - require_package("pynput", extra="pynput-dep") self._ring: RolloutRingBuffer | None = None self._listener = None self._save_requested = ThreadingEvent() @@ -234,30 +217,27 @@ class HighlightStrategy(RolloutStrategy): logger.info("Highlight strategy teardown complete") def _setup_keyboard(self, shutdown_event: ThreadingEvent) -> None: - """Set up keyboard listener for save and push keys.""" - if is_headless(): - logger.warning("Headless environment — highlight keys unavailable") - return + """Set up a keyboard listener for the save and push keys. - try: - save_key = self.config.save_key - push_key = self.config.push_key + Backend selection (pynput on X11 / trusted-macOS / Windows, a terminal reader on + Wayland / headless TTY) is delegated to :func:`create_key_listener`. + """ + save_key = self.config.save_key + push_key = self.config.push_key - def on_press(key): - with contextlib.suppress(Exception): - if hasattr(key, "char") and key.char == save_key: - self._save_requested.set() - elif hasattr(key, "char") and key.char == push_key: - self._push_requested.set() - elif key == keyboard.Key.esc: - self._save_requested.clear() - shutdown_event.set() + def dispatch(name: str) -> None: + """Apply a resolved key name to the highlight events.""" + if name == save_key: + self._save_requested.set() + elif name == push_key: + self._push_requested.set() + elif name == "esc": + self._save_requested.clear() + shutdown_event.set() - self._listener = keyboard.Listener(on_press=on_press) - self._listener.start() - logger.info("Keyboard listener started (save='%s', push='%s', ESC=stop)", save_key, push_key) - except ImportError: - logger.warning("pynput not available — keyboard listener disabled") + self._listener = create_key_listener( + dispatch, controls_help=f"save='{save_key}', push='{push_key}', ESC=stop" + ) def _background_push(self, dataset, cfg) -> None: """Queue a Hub push on the single-worker executor.""" diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index d45483d21..1ec4ea75f 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -72,8 +72,9 @@ from termcolor import colored from torch import Tensor, nn from tqdm import trange -from lerobot.configs import parser +from lerobot.configs import FeatureType, parser from lerobot.configs.eval import EvalPipelineConfig +from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.envs import ( check_env_attributes_and_types, close_envs, @@ -84,7 +85,7 @@ from lerobot.envs import ( from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.processor import PolicyProcessorPipeline from lerobot.types import PolicyAction -from lerobot.utils.constants import ACTION, DONE, OBS_STR, REWARD +from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.io_utils import write_video @@ -95,6 +96,65 @@ from lerobot.utils.utils import ( ) +def _env_features_to_dataset_features(env_features: dict) -> dict: + """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" + features = {} + for key, ft in env_features.items(): + shape = tuple(ft.shape) + if ft.type is FeatureType.VISUAL: + features[key] = {"dtype": "video", "shape": shape, "names": ["height", "width", "channel"]} + else: + features[key] = {"dtype": "float32", "shape": shape, "names": None} + features["next.reward"] = {"dtype": "float32", "shape": (1,), "names": None} + features["next.success"] = {"dtype": "bool", "shape": (1,), "names": None} + features["next.done"] = {"dtype": "bool", "shape": (1,), "names": None} + return features + + +def _build_raw_frame( + raw_obs: dict, + env_idx: int, + action: np.ndarray, + reward: float, + success: bool, + done: bool, + task: str, + env_features: dict, +) -> dict: + """Build a dataset frame from raw env observations for one env index. + + Keys in the frame match the keys in env_features so they align with the + dataset schema created by _env_features_to_dataset_features(). + """ + frame: dict[str, Any] = {} + for key in env_features: + if key == ACTION: + continue + if key.startswith("next."): + continue + if "pixels" in raw_obs and isinstance(raw_obs["pixels"], dict): + for cam_name, img in raw_obs["pixels"].items(): + candidate = f"{OBS_IMAGES}.{cam_name}" + if candidate == key: + frame[key] = img[env_idx] + if key in frame: + continue + if "pixels" in raw_obs and not isinstance(raw_obs["pixels"], dict) and key in ("pixels", OBS_IMAGE): + frame[key] = raw_obs["pixels"][env_idx] + continue + if key in raw_obs and isinstance(raw_obs[key], np.ndarray): + val = raw_obs[key][env_idx] + if val.dtype == np.float64: + val = val.astype(np.float32) + frame[key] = val + frame[ACTION] = action + frame["next.reward"] = np.atleast_1d(np.float32(reward)) + frame["next.success"] = np.atleast_1d(np.bool_(success)) + frame["next.done"] = np.atleast_1d(np.bool_(done)) + frame["task"] = task + return frame + + def rollout( env: gym.vector.VectorEnv, policy: PreTrainedPolicy, @@ -105,6 +165,10 @@ def rollout( seeds: list[int] | None = None, return_observations: bool = False, render_callback: Callable[[gym.vector.VectorEnv], None] | None = None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> dict: """Run a batched policy rollout once through a batch of environments. @@ -145,6 +209,33 @@ def rollout( if render_callback is not None: render_callback(env) + recording_datasets: list[LeRobotDataset] | None = None + raw_observation = None + task_desc = "" + if recording_dir is not None and env_features is not None: + features = _env_features_to_dataset_features(env_features) + fps = env.unwrapped.metadata.get("render_fps", 30) + recording_datasets = [] + multi_env = env.num_envs > 1 + base_repo_id = recording_repo_id or "eval_recording" + for i in range(env.num_envs): + root = str(recording_dir / f"env_{i}") if multi_env else str(recording_dir) + repo_id = f"{base_repo_id}_env_{i}" if multi_env else base_repo_id + recording_datasets.append( + LeRobotDataset.create( + repo_id=repo_id, + fps=fps, + features=features, + root=root, + use_videos=True, + ) + ) + raw_observation = deepcopy(observation) + try: + task_desc = list(env.call("task_description"))[0] + except (AttributeError, NotImplementedError): + task_desc = "" + all_observations = [] all_actions = [] all_rewards = [] @@ -162,80 +253,112 @@ def rollout( leave=False, ) check_env_attributes_and_types(env) - while not np.all(done) and step < max_steps: - # Numpy array to tensor and changing dictionary keys to LeRobot policy format. - observation = preprocess_observation(observation) - if return_observations: - all_observations.append(deepcopy(observation)) + try: + while not np.all(done) and step < max_steps: + # Numpy array to tensor and changing dictionary keys to LeRobot policy format. + observation = preprocess_observation(observation) + if return_observations: + all_observations.append(deepcopy(observation)) - # Infer "task" from sub-environments (prefer natural language description). - # env.call() works with both SyncVectorEnv and AsyncVectorEnv. - try: - observation["task"] = list(env.call("task_description")) - except (AttributeError, NotImplementedError): + # Infer "task" from sub-environments (prefer natural language description). + # env.call() works with both SyncVectorEnv and AsyncVectorEnv. try: - observation["task"] = list(env.call("task")) + observation["task"] = list(env.call("task_description")) except (AttributeError, NotImplementedError): - observation["task"] = [""] * env.num_envs + try: + observation["task"] = list(env.call("task")) + except (AttributeError, NotImplementedError): + observation["task"] = [""] * env.num_envs - # Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO) - observation = env_preprocessor(observation) + # Apply environment-specific preprocessing (e.g., LiberoProcessorStep for LIBERO) + observation = env_preprocessor(observation) - observation = preprocessor(observation) - with torch.inference_mode(): - action = policy.select_action(observation) - action = postprocessor(action) + observation = preprocessor(observation) + with torch.inference_mode(): + action = policy.select_action(observation) + action = postprocessor(action) - action_transition = {ACTION: action} - action_transition = env_postprocessor(action_transition) - action = action_transition[ACTION] + action_transition = {ACTION: action} + action_transition = env_postprocessor(action_transition) + action = action_transition[ACTION] - # Convert to CPU / numpy. - action_numpy: np.ndarray = action.to("cpu").numpy() - assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)" + # Convert to CPU / numpy. + action_numpy: np.ndarray = action.to("cpu").numpy() + assert action_numpy.ndim == 2, "Action dimensions should be (batch, action_dim)" - # Apply the next action. - observation, reward, terminated, truncated, info = env.step(action_numpy) - if render_callback is not None: - render_callback(env) + # Apply the next action. + observation, reward, terminated, truncated, info = env.step(action_numpy) + if render_callback is not None: + render_callback(env) - # VectorEnv stores is_success in `info["final_info"][env_index]["is_success"]`. "final_info" isn't - # available if none of the envs finished. - if "final_info" in info: - final_info = info["final_info"] - if not isinstance(final_info, dict): - raise RuntimeError( - "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " - "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + # VectorEnv stores is_success in `info["final_info"][env_index]["is_success"]`. "final_info" isn't + # available if none of the envs finished. + if "final_info" in info: + final_info = info["final_info"] + if not isinstance(final_info, dict): + raise RuntimeError( + "Unsupported `final_info` format: expected dict (Gymnasium >= 1.0). " + "You're likely using an older version of gymnasium (< 1.0). Please upgrade." + ) + successes = final_info["is_success"].tolist() + elif "is_success" in info: + is_success = info["is_success"] + successes = ( + is_success.tolist() + if hasattr(is_success, "tolist") + else [bool(is_success)] * env.num_envs ) - successes = final_info["is_success"].tolist() - elif "is_success" in info: - is_success = info["is_success"] - successes = ( - is_success.tolist() if hasattr(is_success, "tolist") else [bool(is_success)] * env.num_envs + else: + successes = [False] * env.num_envs + + if recording_datasets is not None and raw_observation is not None: + prev_done = done.copy() + for env_idx in range(env.num_envs): + if prev_done[env_idx]: + continue + frame = _build_raw_frame( + raw_observation, + env_idx, + action_numpy[env_idx], + reward[env_idx], + successes[env_idx], + bool(terminated[env_idx] | truncated[env_idx]), + task_desc, + recording_datasets[env_idx].features, + ) + recording_datasets[env_idx].add_frame(frame) + if terminated[env_idx] or truncated[env_idx]: + recording_datasets[env_idx].save_episode() + raw_observation = deepcopy(observation) + + # Keep track of which environments are done so far. + # Mark the episode as done if we reach the maximum step limit. + # This ensures that the rollout always terminates cleanly at `max_steps`, + # and allows logging/saving (e.g., videos) to be triggered consistently. + done = terminated | truncated | done + if step + 1 == max_steps: + done = np.ones_like(done, dtype=bool) + + all_actions.append(torch.from_numpy(action_numpy)) + all_rewards.append(torch.from_numpy(reward)) + all_dones.append(torch.from_numpy(done)) + all_successes.append(torch.tensor(successes)) + + step += 1 + running_success_rate = ( + einops.reduce(torch.stack(all_successes, dim=1), "b n -> b", "any").numpy().mean() ) - else: - successes = [False] * env.num_envs - - # Keep track of which environments are done so far. - # Mark the episode as done if we reach the maximum step limit. - # This ensures that the rollout always terminates cleanly at `max_steps`, - # and allows logging/saving (e.g., videos) to be triggered consistently. - done = terminated | truncated | done - if step + 1 == max_steps: - done = np.ones_like(done, dtype=bool) - - all_actions.append(torch.from_numpy(action_numpy)) - all_rewards.append(torch.from_numpy(reward)) - all_dones.append(torch.from_numpy(done)) - all_successes.append(torch.tensor(successes)) - - step += 1 - running_success_rate = ( - einops.reduce(torch.stack(all_successes, dim=1), "b n -> b", "any").numpy().mean() - ) - progbar.set_postfix({"running_success_rate": f"{running_success_rate.item() * 100:.1f}%"}) - progbar.update() + progbar.set_postfix({"running_success_rate": f"{running_success_rate.item() * 100:.1f}%"}) + progbar.update() + finally: + if recording_datasets is not None: + for ds in recording_datasets: + ds.finalize() + if recording_repo_id is not None: + if ds.num_episodes > 0: + ds.push_to_hub(private=recording_private) + else: + logging.warning("No episodes recorded for %s — skipping push to hub.", ds.repo_id) # Track the final observation. if return_observations: @@ -273,6 +396,10 @@ def eval_policy( videos_dir: Path | None = None, return_episode_data: bool = False, start_seed: int | None = None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> dict: """ Args: @@ -361,6 +488,10 @@ def eval_policy( seeds=list(seeds) if seeds else None, return_observations=return_episode_data, render_callback=render_frame if max_episodes_rendered > 0 else None, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) # Figure out where in each rollout sequence the first done condition was encountered (results after @@ -563,6 +694,10 @@ def eval_main(cfg: EvalPipelineConfig): # Create environment-specific preprocessor and postprocessor (e.g., for LIBERO environments) env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env, policy_cfg=cfg.policy) + recording_dir = Path(cfg.output_dir) / "recordings" if cfg.eval.recording else None + max_episodes_rendered = 0 if cfg.eval.recording else 10 + videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos" + with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(): info = eval_policy_all( envs=envs, @@ -572,10 +707,15 @@ def eval_main(cfg: EvalPipelineConfig): preprocessor=preprocessor, postprocessor=postprocessor, n_episodes=cfg.eval.n_episodes, - max_episodes_rendered=10, - videos_dir=Path(cfg.output_dir) / "videos", + max_episodes_rendered=max_episodes_rendered, + videos_dir=videos_dir, + return_episode_data=False, start_seed=cfg.seed, max_parallel_tasks=cfg.env.max_parallel_tasks, + recording_dir=recording_dir, + env_features=cfg.env.features if cfg.eval.recording else None, + recording_repo_id=cfg.eval.recording_repo_id, + recording_private=cfg.eval.recording_private, ) print("Overall Aggregated Metrics:") print(info["overall"]) @@ -618,6 +758,10 @@ def eval_one( videos_dir: Path | None, return_episode_data: bool, start_seed: int | None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ) -> TaskMetrics: """Evaluates one task_id of one suite using the provided vec env.""" @@ -635,6 +779,10 @@ def eval_one( videos_dir=task_videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) per_episode = task_result["per_episode"] @@ -661,6 +809,10 @@ def run_one( videos_dir: Path | None, return_episode_data: bool, start_seed: int | None, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, ): """ Run eval_one for a single (task_group, task_id, env). @@ -672,7 +824,13 @@ def run_one( task_videos_dir = videos_dir / f"{task_group}_{task_id}" task_videos_dir.mkdir(parents=True, exist_ok=True) - # Call the existing eval_one (assumed to return TaskMetrics-like dict) + task_recording_dir = None + task_repo_id = None + if recording_dir is not None and env_features is not None: + task_recording_dir = recording_dir / f"{task_group}_{task_id}" + if recording_repo_id is not None: + task_repo_id = f"{recording_repo_id}_{task_group}_{task_id}" + metrics = eval_one( env, policy=policy, @@ -685,8 +843,12 @@ def run_one( videos_dir=task_videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=task_recording_dir, + env_features=env_features, + recording_repo_id=task_repo_id, + recording_private=recording_private, ) - # ensure we always provide video_paths key to simplify accumulation + if max_episodes_rendered > 0: metrics.setdefault("video_paths", []) return task_group, task_id, metrics @@ -702,6 +864,10 @@ def eval_policy_all( n_episodes: int, *, max_episodes_rendered: int = 0, + recording_dir: Path | None = None, + env_features: dict | None = None, + recording_repo_id: str | None = None, + recording_private: bool = False, videos_dir: Path | None = None, return_episode_data: bool = False, start_seed: int | None = None, @@ -761,6 +927,10 @@ def eval_policy_all( videos_dir=videos_dir, return_episode_data=return_episode_data, start_seed=start_seed, + recording_dir=recording_dir, + env_features=env_features, + recording_repo_id=recording_repo_id, + recording_private=recording_private, ) if max_parallel_tasks <= 1: diff --git a/src/lerobot/scripts/lerobot_record.py b/src/lerobot/scripts/lerobot_record.py index 0deb54b90..4d5518c7c 100644 --- a/src/lerobot/scripts/lerobot_record.py +++ b/src/lerobot/scripts/lerobot_record.py @@ -96,11 +96,7 @@ from lerobot.cameras.opencv import OpenCVCameraConfig # noqa: F401 from lerobot.cameras.reachy2_camera import Reachy2CameraConfig # noqa: F401 from lerobot.cameras.realsense import RealSenseCameraConfig # noqa: F401 from lerobot.cameras.zmq import ZMQCameraConfig # noqa: F401 -from lerobot.common.control_utils import ( - init_keyboard_listener, - is_headless, - sanity_check_dataset_robot_compatibility, -) +from lerobot.common.control_utils import sanity_check_dataset_robot_compatibility from lerobot.configs import parser from lerobot.configs.dataset import DatasetRecordConfig from lerobot.datasets import ( @@ -155,6 +151,7 @@ from lerobot.teleoperators.keyboard import KeyboardTeleop from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts from lerobot.utils.import_utils import register_third_party_plugins +from lerobot.utils.keyboard_input import init_keyboard_listener from lerobot.utils.robot_utils import precise_sleep from lerobot.utils.utils import ( init_logging, @@ -508,7 +505,7 @@ def record( if teleop and teleop.is_connected: teleop.disconnect() - if not is_headless() and listener: + if listener is not None: listener.stop() if cfg.dataset.push_to_hub: diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 70a5e9e9d..45281dac9 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -34,8 +34,10 @@ from torch.optim import Optimizer from tqdm import tqdm from lerobot.common.train_utils import ( + gather_fsdp_state_dicts, get_step_checkpoint_dir, get_step_identifier, + load_fsdp_optimizer_state, load_training_batch_size, load_training_num_processes, load_training_state, @@ -189,6 +191,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): require_package("accelerate", extra="training") from accelerate import Accelerator + from accelerate.utils import DistributedDataParallelKwargs, DistributedType cfg.validate() @@ -197,8 +200,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes # We set find_unused_parameters=True to handle models with conditional computation if accelerator is None: - from accelerate.utils import DistributedDataParallelKwargs - ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training). @@ -345,6 +346,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): preprocessor, postprocessor = make_pre_post_processors( policy_cfg=cfg.policy, pretrained_path=processor_pretrained_path, + pretrained_revision=getattr(cfg.policy, "pretrained_revision", None), **processor_kwargs, ) @@ -370,7 +372,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): step = 0 # number of policy updates (forward + backward + optim) if cfg.resume: - step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler) + # Under FSDP the optimizer state is sharded and must be loaded after `accelerator.prepare()` + # (see load_fsdp_optimizer_state below), so skip the optimizer here and load it then. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + step, optimizer, lr_scheduler = load_training_state( + cfg.checkpoint_path, optimizer, lr_scheduler, load_optimizer=not is_fsdp + ) num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad) num_total_params = sum(p.numel() for p in policy.parameters()) @@ -460,6 +467,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( policy, optimizer, dataloader, lr_scheduler ) + + # FSDP optimizer state is sharded across ranks, so it can only be loaded once the optimizer and + # model are FSDP-wrapped (i.e. after `prepare`). Collective: every rank must participate. + if cfg.resume and accelerator.distributed_type == DistributedType.FSDP: + load_fsdp_optimizer_state(policy, optimizer, cfg.checkpoint_path) + dl_iter = cycle(dataloader) policy.train() @@ -558,6 +571,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): train_tracker.reset_averages() if cfg.save_checkpoint and is_saving_step: + # Under FSDP, gathering the full model + optimizer state dicts is a cross-rank collective, + # so all ranks must participate; rank 0 then writes the materialized dicts. For DDP / + # single-GPU the state dicts are saved the normal way inside save_checkpoint. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + if is_fsdp: + model_state_dict, optim_state_dict = gather_fsdp_state_dicts(policy, optimizer) + else: + model_state_dict, optim_state_dict = None, None if is_main_process: logging.info(f"Checkpoint policy after step {step}") checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step) @@ -572,6 +593,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): postprocessor=postprocessor, num_processes=accelerator.num_processes, batch_size=cfg.batch_size, + model_state_dict=model_state_dict, + optim_state_dict=optim_state_dict, ) update_last_checkpoint(checkpoint_dir) if wandb_logger: @@ -634,6 +657,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if eval_env: close_envs(eval_env) + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None if is_main_process: logging.info("End of training") @@ -643,7 +668,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if not cfg.is_reward_model_training and cfg.policy.use_peft: unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) else: - unwrapped_model.push_model_to_hub(cfg) + unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict) preprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id) diff --git a/src/lerobot/teleoperators/gamepad/gamepad_utils.py b/src/lerobot/teleoperators/gamepad/gamepad_utils.py index c1531ca84..22dbb7cca 100644 --- a/src/lerobot/teleoperators/gamepad/gamepad_utils.py +++ b/src/lerobot/teleoperators/gamepad/gamepad_utils.py @@ -18,6 +18,7 @@ import logging from typing import TYPE_CHECKING from lerobot.utils.import_utils import _hidapi_available, _pygame_available, require_package +from lerobot.utils.keyboard_input import pynput_can_capture from ..utils import TeleopEvents @@ -123,6 +124,15 @@ class KeyboardController(InputController): def start(self): """Start the keyboard listener.""" + if not pynput_can_capture(): + logging.warning( + "Keyboard control is unavailable in this environment. pynput cannot capture keys " + "on Wayland or headless machines, or on macOS without Accessibility / Input " + "Monitoring permission. Keyboard motion will be inactive." + ) + self.running = False + return + from pynput import keyboard def on_press(key): diff --git a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py index 801789bcb..872cc7a26 100644 --- a/src/lerobot/teleoperators/keyboard/teleop_keyboard.py +++ b/src/lerobot/teleoperators/keyboard/teleop_keyboard.py @@ -15,8 +15,6 @@ # limitations under the License. import logging -import os -import sys import time from queue import Queue from typing import Any @@ -24,6 +22,7 @@ from typing import Any from lerobot.types import RobotAction from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected from lerobot.utils.import_utils import _pynput_available, require_package +from lerobot.utils.keyboard_input import pynput_can_capture from ..teleoperator import Teleoperator from ..utils import TeleopEvents @@ -37,14 +36,10 @@ PYNPUT_AVAILABLE = _pynput_available keyboard = None if PYNPUT_AVAILABLE: try: - if ("DISPLAY" not in os.environ) and ("linux" in sys.platform): - logging.info("No DISPLAY set. Skipping pynput import.") - PYNPUT_AVAILABLE = False - else: - from pynput import keyboard + from pynput import keyboard except Exception as e: PYNPUT_AVAILABLE = False - logging.info(f"Could not import pynput: {e}") + logging.info("Could not import pynput keyboard backend: %s", e) class KeyboardTeleop(Teleoperator): @@ -88,7 +83,7 @@ class KeyboardTeleop(Teleoperator): @check_if_already_connected def connect(self) -> None: - if PYNPUT_AVAILABLE: + if PYNPUT_AVAILABLE and pynput_can_capture(): logging.info("pynput is available - enabling local keyboard listener.") self.listener = keyboard.Listener( on_press=self._on_press, @@ -96,7 +91,13 @@ class KeyboardTeleop(Teleoperator): ) self.listener.start() else: - logging.info("pynput not available - skipping local keyboard listener.") + logging.warning( + "Keyboard teleoperation is unavailable in this environment. pynput can only " + "capture key events on an X11 session (Linux), a Windows desktop, or macOS with " + "Accessibility / Input Monitoring granted - not on Wayland or headless machines. " + "This keyboard teleoperator will produce no actions; use an X11 session, a " + "gamepad, or a leader-arm teleoperator instead." + ) self.listener = None def calibrate(self) -> None: diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index 5dbce2c5b..b0d894c04 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -216,9 +216,15 @@ def register_third_party_plugins() -> None: This function uses `importlib.metadata` to find packages installed in the environment (including editable installs) starting with 'lerobot_robot_', 'lerobot_camera_', - 'lerobot_teleoperator_', or 'lerobot_policy_' and imports them. + 'lerobot_teleoperator_', 'lerobot_policy_', or 'lerobot_env_' and imports them. """ - prefixes = ("lerobot_robot_", "lerobot_camera_", "lerobot_teleoperator_", "lerobot_policy_") + prefixes = ( + "lerobot_robot_", + "lerobot_camera_", + "lerobot_teleoperator_", + "lerobot_policy_", + "lerobot_env_", + ) imported: list[str] = [] failed: list[str] = [] diff --git a/src/lerobot/utils/keyboard_input.py b/src/lerobot/utils/keyboard_input.py new file mode 100644 index 000000000..00c0f53ec --- /dev/null +++ b/src/lerobot/utils/keyboard_input.py @@ -0,0 +1,440 @@ +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Display-independent keyboard input for interactive controls. + +This module centralizes everything related to *discrete* keyboard controls +(end-episode-early, re-record, stop, and the rollout strategies' custom keys): + +* environment detection — :func:`is_headless`, :func:`is_wayland`, + :func:`pynput_can_capture` (the single predicate every call-site should use to + decide whether ``pynput`` can actually capture keys here); +* a shared key mapping — :func:`apply_recording_control`; and +* two interchangeable backends behind one ``(listener, events)`` contract: + the ``pynput`` global listener (X11 / trusted-macOS / Windows) and a + standard-library :class:`TerminalKeyListener` that reads the controlling TTY + (Wayland / headless-SSH-with-TTY / macOS without Accessibility permission). + +NOTE: *continuous* key-state teleoperation ("hold a key to keep moving") is +deliberately NOT served here. A terminal in cbreak mode delivers only key-down +bytes — there is no key-release event — so the held-key model cannot be +reproduced. Those teleoperators stay on ``pynput`` and use +:func:`pynput_can_capture` to warn instead of silently doing nothing. +""" + +from __future__ import annotations + +import atexit +import contextlib +import logging +import os +import platform +import select +import sys +import threading +import time +from collections.abc import Callable +from functools import cache +from typing import TYPE_CHECKING + +from .import_utils import _pynput_available + +logger = logging.getLogger(__name__) + +# POSIX-only terminal modules (absent on Windows, where the pynput backend is used). +if TYPE_CHECKING: + import termios + import tty + + _TERMIOS_AVAILABLE = True +else: + try: + import termios + import tty + + _TERMIOS_AVAILABLE = True + except ImportError: # POSIX-only modules; unavailable on Windows + termios = tty = None + _TERMIOS_AVAILABLE = False + +keyboard = None +if _pynput_available: + try: + from pynput import keyboard + except Exception as e: # e.g. no reachable X display on a headless Linux box + logger.info("Could not import pynput keyboard backend: %s", e) + + +@cache +def is_headless() -> bool: + """Return ``True`` when no display server is available. + + * Linux: headless when neither ``DISPLAY`` (X11) nor ``WAYLAND_DISPLAY`` is set. + * macOS / Windows: a display is always assumed to be present. A genuinely GUI-less + Mac/Windows CI host would be misclassified but it doesn't matter, because the + sys.stdin.isatty() gate returns None there regardless. + """ + if platform.system() == "Linux": + return not (os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")) + return False + + +@cache +def is_wayland() -> bool: + """Return ``True`` when running under a Wayland session. + + ``pynput`` relies on an X11 backend. Under Wayland it still imports (XWayland + is usually present and ``$DISPLAY`` is set) but cannot capture *global* + hotkeys, so the documented arrow/Esc shortcuts silently do nothing. This case + is invisible to :func:`is_headless`, hence the dedicated check. + """ + return os.environ.get("XDG_SESSION_TYPE", "").lower() == "wayland" or bool( + os.environ.get("WAYLAND_DISPLAY") + ) + + +@cache +def pynput_can_capture() -> bool: + """Return ``True`` when a ``pynput`` global listener can actually capture keys. + + This is the single predicate every keyboard call-site should use to choose + between the ``pynput`` backend and a fallback. It is intentionally + conservative: + + * Linux: only a real X11 session (a display is present *and* it is not Wayland). + * macOS: ``True`` here — Accessibility / Input-Monitoring permission + (``IS_TRUSTED``) can only be confirmed at runtime *after* starting a + listener, so :func:`init_keyboard_listener` refines this with + :func:`pynput_listener_is_trusted`. + * Windows: ``True`` (the low-level global hook needs no special permission). + + Always ``False`` when ``pynput`` is not installed. + """ + if not _pynput_available: + return False + if platform.system() == "Linux": + return not is_headless() and not is_wayland() + return True + + +def pynput_listener_is_trusted(listener, timeout_s: float = 1.0) -> bool: + """Best-effort check that a freshly started ``pynput`` listener can capture. + + On macOS, ``pynput`` sets ``listener.IS_TRUSTED`` on its *listener thread* + once the Quartz event tap is created; the class default is ``False``. We + therefore wait for the thread to either flip it ``True`` (trusted) or for a + short timeout to elapse (untrusted — it stays ``False`` forever). This biases + toward the common trusted case (returns as soon as the flag flips) and only + pays the full ``timeout_s`` on an already-broken untrusted machine. + + On non-macOS backends the attribute is absent and capture is assumed to work. + """ + if platform.system() != "Darwin": + return True + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + if getattr(listener, "IS_TRUSTED", False): + return True + time.sleep(0.02) + return bool(getattr(listener, "IS_TRUSTED", False)) + + +def apply_recording_control(control: str, events: dict) -> None: + """Apply a recording control-flow key press to the shared ``events`` dict. + + Centralizes the mapping so the ``pynput`` and terminal backends behave + identically. ``control`` is one of ``"right"`` (end the loop early), ``"left"`` + (re-record the last episode), or ``"esc"`` (stop recording). + """ + if control == "right": + print("Right arrow key pressed. Exiting loop...") + events["exit_early"] = True + elif control == "left": + print("Left arrow key pressed. Exiting loop and rerecord the last episode...") + events["rerecord_episode"] = True + events["exit_early"] = True + elif control == "esc": + print("Escape key pressed. Stopping data recording...") + events["stop_recording"] = True + events["exit_early"] = True + + +# Terminal arrow keys arrive as a 3-byte escape sequence whose *final* byte identifies +# the direction. Two encodings exist depending on the terminal's cursor-key mode — CSI +# ("ESC [ X") and SS3 ("ESC O X", common over SSH/tmux) — but both share the same final +# byte, so this single table decodes either. Looked up by TerminalKeyListener._parse; +# an unknown final byte yields None (sequence ignored). +_ARROW_FINAL_BYTES = {"A": "up", "B": "down", "C": "right", "D": "left"} + + +class TerminalKeyListener: + """Display-independent keyboard listener that reads keys from the controlling TTY. + + Used as the Wayland / headless / macOS-untrusted equivalent of the ``pynput`` + listener for *discrete* controls. It puts the terminal into cbreak mode with + echo disabled and reads bytes on a daemon thread, decoding them into logical + key names that are passed to ``on_key``: + + * arrow keys (``ESC [ C`` / ``ESC O C`` …) -> ``"right"`` / ``"left"`` / ``"up"`` / ``"down"`` + * a bare ``ESC`` -> ``"esc"`` + * Enter / Tab / Space / Backspace -> ``"enter"`` / ``"tab"`` / ``"space"`` / ``"backspace"`` + * any other printable byte -> that character (e.g. ``"n"``, ``"s"``) + + Only key-down events are produced (terminals have no key-release), so this is + suitable for discrete commands but NOT for continuous "hold-to-move" teleop. + + The terminal is restored on :meth:`stop` and also via an ``atexit`` hook, so a + crash or Ctrl-C never leaves the shell in a no-echo cbreak state. POSIX-only + (``termios`` / ``tty`` / ``select``); those modules are imported lazily so this + file stays importable on Windows (where ``pynput`` is used instead). + """ + + def __init__(self, on_key: Callable[[str], None]): + self._on_key = on_key + self._running = False + self._thread: threading.Thread | None = None + self._fd: int | None = None + self._old_attrs = None + + def _read_char(self, timeout: float) -> str | None: + """Return one character from stdin within ``timeout`` seconds, or ``None``.""" + if self._fd is None: + return None + ready, _, _ = select.select([self._fd], [], [], timeout) + if not ready: + return None + try: + data = os.read(self._fd, 1) + except OSError: + return None + if not data: + return None + return data.decode(errors="ignore") + + def _parse(self, ch: str) -> str | None: + """Decode one (possibly multi-byte) key starting at ``ch`` into a key name.""" + if ch == "\x1b": + # Possible CSI / SS3 escape sequence (arrow keys) or a bare ESC. Use + # short follow-up reads so a lone ESC is not mistaken for a sequence. + ch2 = self._read_char(timeout=0.02) + if ch2 is None: + return "esc" + if ch2 in ("[", "O"): + ch3 = self._read_char(timeout=0.02) + return _ARROW_FINAL_BYTES.get(ch3 or "") + # Some other escape sequence (e.g. Alt+key); ignore it. + return None + if ch in ("\r", "\n"): + return "enter" + if ch == "\t": + return "tab" + if ch == " ": + return "space" + if ch in ("\x7f", "\x08"): + return "backspace" + if ch.isprintable(): + return ch + return None + + def _run(self) -> None: + while self._running: + ch = self._read_char(timeout=0.05) + if ch is None: + continue + name = self._parse(ch) + if name is None: + continue + try: + self._on_key(name) + except Exception as e: # never let a handler error kill the reader thread + logger.debug("Terminal key handler error: %s", e) + + def start(self) -> None: + """Switch the terminal to cbreak mode (echo off) and read keys on a daemon thread. + + No-op when stdin is not a TTY (piped/redirected input) or on platforms + without ``termios`` (e.g. Windows), so non-interactive runs are unaffected. + """ + if not sys.stdin.isatty(): + return + if not _TERMIOS_AVAILABLE: # POSIX-only modules (e.g. unavailable on Windows) + logger.warning("Terminal keyboard input is not supported on this platform.") + return + + self._fd = sys.stdin.fileno() + self._old_attrs = termios.tcgetattr(self._fd) + tty.setcbreak(self._fd) + # Explicitly disable ECHO so arrow-key escape sequences (e.g. ^[[C) are not + # echoed as garbage into the recording terminal. (Independent of the + # version-specific behavior of tty.setcbreak.) + new_attrs = termios.tcgetattr(self._fd) + new_attrs[3] &= ~termios.ECHO # index 3 == lflags + termios.tcsetattr(self._fd, termios.TCSADRAIN, new_attrs) + # Safety net: restore the terminal even if stop() is never reached (crash). + atexit.register(self.stop) + + self._running = True + self._thread = threading.Thread(target=self._run, daemon=True) + self._thread.start() + + def stop(self) -> None: + """Stop the reader thread and restore the original terminal attributes. + + Idempotent: safe to call multiple times (e.g. explicitly and via atexit). + """ + self._running = False + thread = self._thread + if thread is not None: + thread.join(timeout=0.5) + self._thread = None + if self._fd is not None and self._old_attrs is not None and _TERMIOS_AVAILABLE: + try: + termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs) + finally: + self._old_attrs = None + with contextlib.suppress(Exception): + atexit.unregister(self.stop) + + +# Map pynput key objects to the same canonical names TerminalKeyListener emits, so a +# single dispatch works across both backends. Empty when pynput is unavailable. +if keyboard is not None: + _PYNPUT_KEY_NAMES = { + keyboard.Key.right: "right", + keyboard.Key.left: "left", + keyboard.Key.up: "up", + keyboard.Key.down: "down", + keyboard.Key.esc: "esc", + keyboard.Key.enter: "enter", + keyboard.Key.tab: "tab", + keyboard.Key.space: "space", + keyboard.Key.backspace: "backspace", + } +else: + _PYNPUT_KEY_NAMES = {} + + +def _resolve_pynput_key(key) -> str | None: + """Resolve a pynput key event to the canonical name TerminalKeyListener also emits. + + Special keys map through :data:`_PYNPUT_KEY_NAMES`; character keys fall back to their + ``.char`` (e.g. ``"n"``). Returns ``None`` for keys with no mapping and no character. + """ + name = _PYNPUT_KEY_NAMES.get(key) + if name is not None: + return name + # ``or None`` keeps the historical truthy-char semantics: an empty/None char is "no key". + return getattr(key, "char", None) or None + + +def create_key_listener(dispatch: Callable[[str], None], *, controls_help: str = ""): + """Start a keyboard listener that routes resolved key names to ``dispatch``. + + Shared backend selection used by recording and the rollout strategies: + + * the ``pynput`` global listener on X11 / trusted-macOS / Windows (on macOS the + listener's ``IS_TRUSTED`` flag is checked after start, and an untrusted listener is + stopped so the terminal backend is used instead); + * the stdlib :class:`TerminalKeyListener` on Wayland / headless sessions with a TTY; + * ``None`` when no backend is usable (non-interactive / piped runs). + + Both backends pass ``dispatch`` the same canonical key names ("right" / "left" / "up" / + "down" / "esc" / "enter" / "tab" / "space" / "backspace", or a character), so one + ``dispatch`` works regardless of backend. ``controls_help`` is an optional hint + appended to the log messages. + + Returns the listener (exposing ``.stop()``) or ``None``. + """ + suffix = f" ({controls_help})" if controls_help else "" + + if pynput_can_capture() and keyboard is not None: + + def on_press(key): + with contextlib.suppress(Exception): + name = _resolve_pynput_key(key) + if name is not None: + dispatch(name) + + listener = keyboard.Listener(on_press=on_press) + listener.start() + if pynput_listener_is_trusted(listener): + logger.info("Keyboard listener started%s.", suffix) + return listener + # macOS without Accessibility / Input-Monitoring permission: the listener never + # fires. Stop it and fall through to the terminal backend. + logger.warning( + "pynput keyboard listener is not trusted (missing macOS Accessibility / " + "Input Monitoring permission); falling back to terminal keyboard input." + ) + listener.stop() + + if sys.stdin.isatty(): + listener = TerminalKeyListener(dispatch) + listener.start() + logger.info("Using terminal keyboard input — keep this terminal focused%s.", suffix) + return listener + + logger.warning( + "Keyboard controls unavailable: no usable display (Wayland/headless) and stdin is " + "not an interactive terminal%s.", + suffix, + ) + return None + + +def init_keyboard_listener(): + """Initialize a non-blocking keyboard listener for interactive recording controls. + + Backend selection: + + * ``pynput`` global listener when :func:`pynput_can_capture` is true (real + X11, macOS, Windows). On macOS the listener's ``IS_TRUSTED`` flag is checked + after start; if the process lacks Accessibility / Input-Monitoring + permission, the listener is stopped and the terminal backend is used. + * a :class:`TerminalKeyListener` reading the controlling TTY when ``pynput`` + cannot capture (Wayland / headless-SSH / macOS-untrusted) *and* stdin is a TTY. + * otherwise no listener (non-interactive / piped runs) — recording relies on + the episode/reset timers (or Ctrl+C). + + Both backends accept the same controls: Right/Left/Esc, plus the single-byte letter + equivalents ``n`` (next), ``r`` (re-record) and ``q`` (quit). The letters are the most + reliable choice over high-latency SSH/VNC links, where arrow-key escape sequences can + be split, delayed, or intercepted by the terminal. + + Returns: + A tuple ``(listener, events)`` where ``listener`` exposes ``.stop()`` or is + ``None``, and ``events`` is the dict of flags (``exit_early``, + ``rerecord_episode``, ``stop_recording``) set by key presses. + """ + events = { + "exit_early": False, + "rerecord_episode": False, + "stop_recording": False, + } + + # Accept the single-byte letter equivalents n/r/q alongside the arrow/Esc keys: the + # letters are immune to the escape-sequence split/delay/interception that affects arrows + # over laggy SSH/VNC links. Case-insensitive so Shift+letter still works. + def on_key(name: str) -> None: + key = name.lower() + if key in ("right", "n"): + apply_recording_control("right", events) + elif key in ("left", "r"): + apply_recording_control("left", events) + elif key in ("esc", "q"): + apply_recording_control("esc", events) + # other keys (incl. up/down) are intentionally ignored + + listener = create_key_listener(on_key, controls_help="Right/Left/Esc, or n=next, r=re-record, q=quit") + return listener, events diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py index 19c314fd6..1d2fb1d55 100644 --- a/tests/datasets/test_datasets.py +++ b/tests/datasets/test_datasets.py @@ -51,7 +51,7 @@ from lerobot.robots import make_robot_from_config from lerobot.transforms import ImageTransforms, ImageTransformsConfig from lerobot.utils.constants import ACTION, DONE, OBS_IMAGES, OBS_STATE, OBS_STR, REWARD from lerobot.utils.feature_utils import hw_to_dataset_features -from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_REPO_ID +from tests.fixtures.constants import DUMMY_CHW, DUMMY_HWC, DUMMY_MOTOR_FEATURES, DUMMY_REPO_ID from tests.mocks.mock_robot import MockRobotConfig from tests.utils import require_x86_64_kernel @@ -133,6 +133,21 @@ def test_dataset_feature_with_forward_slash_raises_error(): ) +def test_create_does_not_mutate_input_features(tmp_path, empty_lerobot_dataset_factory): + # ``create`` must deep-copy features so a dataset built from another's features stays independent. + dataset = empty_lerobot_dataset_factory( + root=tmp_path / "ds1", features=DUMMY_MOTOR_FEATURES, use_videos=False + ) + dataset_copy = empty_lerobot_dataset_factory( + root=tmp_path / "ds2", features=dataset.meta.features, use_videos=False + ) + + original_shape = dataset.meta.info.features["state"]["shape"] + dataset_copy.meta.info.features["state"]["shape"] = (999,) + + assert dataset.meta.info.features["state"]["shape"] == original_shape + + def test_add_frame_missing_task(tmp_path, empty_lerobot_dataset_factory): features = {"state": {"dtype": "float32", "shape": (1,), "names": None}} dataset = empty_lerobot_dataset_factory(root=tmp_path / "test", features=features) diff --git a/tests/optim/test_optimizers.py b/tests/optim/test_optimizers.py index d18565562..5b480f70d 100644 --- a/tests/optim/test_optimizers.py +++ b/tests/optim/test_optimizers.py @@ -20,6 +20,7 @@ from lerobot.optim.optimizers import ( MultiAdamConfig, SGDConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from lerobot.utils.constants import ( @@ -65,6 +66,44 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path): torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict()) +def test_save_and_load_fsdp_optimizer_state_dict_roundtrip(tmp_path): + """The FSDP full optimizer state dict is keyed by parameter FQNs (dotted strings), not the + integer indices of the single-GPU path. Verify it survives the safetensors save -> read + round-trip used by the FSDP save/resume path (save_optimizer_state(optim_state_dict=...) then + load_optimizer_state_dict), which the flatten/unflatten "/" separator must not corrupt.""" + full_osd = { + "state": { + "model.layers.0.weight": { + "step": torch.tensor(3.0), + "exp_avg": torch.randn(4, 4), + "exp_avg_sq": torch.randn(4, 4), + }, + "model.layers.0.bias": { + "step": torch.tensor(3.0), + "exp_avg": torch.randn(4), + "exp_avg_sq": torch.randn(4), + }, + }, + "param_groups": [ + {"lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0, "params": [0, 1]} + ], + } + + save_optimizer_state( + torch.optim.Adam([torch.nn.Parameter(torch.randn(1))]), tmp_path, optim_state_dict=full_osd + ) + assert (tmp_path / OPTIMIZER_STATE).is_file() + assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file() + + loaded = load_optimizer_state_dict(tmp_path) + # FQN keys must be preserved verbatim (not int-cast, not split on their dots). + assert set(loaded["state"].keys()) == set(full_osd["state"].keys()) + for fqn, sub in full_osd["state"].items(): + for k, v in sub.items(): + torch.testing.assert_close(loaded["state"][fqn][k], v) + assert loaded["param_groups"] == full_osd["param_groups"] + + @pytest.fixture def base_params_dict(): return { diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py index e9388b3ed..285b87d4c 100644 --- a/tests/policies/test_policies.py +++ b/tests/policies/test_policies.py @@ -23,6 +23,7 @@ import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") +from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from packaging import version from safetensors.torch import load_file @@ -300,6 +301,29 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name: torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) +def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path): + """Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict.""" + policy_cls = get_policy_class("act") + policy_cfg = make_policy_config("act") + features = dataset_to_policy_features(dummy_dataset_metadata.features) + policy_cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} + policy_cfg.input_features = { + key: ft for key, ft in features.items() if key not in policy_cfg.output_features + } + policy = policy_cls(policy_cfg) + policy.to(policy_cfg.device) + + save_dir = tmp_path / "fsdp_state_dict" + policy.save_pretrained(save_dir, state_dict=policy.state_dict()) + + # A single, unsharded safetensors file (no sharded set + index). + assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file() + assert not (save_dir / f"{SAFETENSORS_SINGLE_FILE}.index.json").exists() + + loaded_policy = policy_cls.from_pretrained(save_dir, config=policy_cfg) + torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) + + @pytest.mark.parametrize("multikey", [True, False]) def test_multikey_construction(multikey: bool): """ diff --git a/tests/training/test_multi_gpu.py b/tests/training/test_multi_gpu.py index 638dc3131..d37f1e35d 100644 --- a/tests/training/test_multi_gpu.py +++ b/tests/training/test_multi_gpu.py @@ -58,7 +58,46 @@ def download_dataset(repo_id, episodes): print(f"Dataset {repo_id} downloaded successfully") -def run_accelerate_training(config_args, num_processes=4, temp_dir=None): +def _write_multi_gpu_config(f, num_processes): + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: MULTI_GPU\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("downcast_bf16: 'no'\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + + +def _write_fsdp_config(f, num_processes): + # FSDP1 with FULL_SHARD (ZeRO-3-equivalent) and FULL_STATE_DICT, matching + # docs/source/multi_gpu_training.mdx. ACT's repeated transformer blocks are the wrap units; + # fsdp_use_orig_params is required because LeRobot builds the optimizer before prepare(). + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: FSDP\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + f.write("fsdp_config:\n") + f.write(" fsdp_version: 1\n") + f.write(" fsdp_sharding_strategy: FULL_SHARD\n") + f.write(" fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n") + f.write(" fsdp_transformer_layer_cls_to_wrap: ACTEncoderLayer,ACTDecoderLayer\n") + f.write(" fsdp_use_orig_params: true\n") + f.write(" fsdp_state_dict_type: FULL_STATE_DICT\n") + + +def run_accelerate_training(config_args, num_processes=4, temp_dir=None, distributed_type="MULTI_GPU"): """ Helper function to run training with accelerate launch. @@ -66,6 +105,7 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): config_args: List of config arguments to pass to lerobot_train.py num_processes: Number of processes (GPUs) to use temp_dir: Temporary directory for outputs + distributed_type: "MULTI_GPU" (DDP) or "FSDP" — selects the generated accelerate config. Returns: subprocess.CompletedProcess result @@ -75,18 +115,10 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): # Write YAML config with open(config_path, "w") as f: - f.write("compute_environment: LOCAL_MACHINE\n") - f.write("distributed_type: MULTI_GPU\n") - f.write("mixed_precision: 'no'\n") - f.write(f"num_processes: {num_processes}\n") - f.write("use_cpu: false\n") - f.write("gpu_ids: all\n") - f.write("downcast_bf16: 'no'\n") - f.write("machine_rank: 0\n") - f.write("main_training_function: main\n") - f.write("num_machines: 1\n") - f.write("rdzv_backend: static\n") - f.write("same_network: true\n") + if distributed_type == "FSDP": + _write_fsdp_config(f, num_processes) + else: + _write_multi_gpu_config(f, num_processes) cmd = [ "accelerate", @@ -211,3 +243,66 @@ class TestMultiGPUTraining: # Verify optimizer state exists optimizer_state = training_state_dir / "optimizer_state.safetensors" assert optimizer_state.exists(), f"No optimizer state in checkpoint {checkpoint_dir}" + + def test_fsdp_optimizer_save_and_resume(self): + """ + Test that FSDP saves the (gathered) optimizer state and can resume from it. + + Trains a few steps under FSDP, verifies the gathered optimizer state is written next to the + rest of the training state, then resumes from the checkpoint for more steps and checks it + completes without shape/key errors in the FSDP optimizer load path. + """ + # Pre-download dataset to avoid race conditions + download_dataset("lerobot/pusht", episodes=[0]) + + with tempfile.TemporaryDirectory() as temp_dir: + output_dir = Path(temp_dir) / "outputs" + + config_args = [ + "--dataset.repo_id=lerobot/pusht", + "--dataset.episodes=[0]", + "--policy.type=act", + "--policy.device=cuda", + "--policy.push_to_hub=false", + f"--output_dir={output_dir}", + "--batch_size=4", + "--steps=10", + "--eval_freq=-1", + "--log_freq=5", + "--save_freq=10", + "--seed=42", + "--num_workers=0", + ] + + result = run_accelerate_training( + config_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert result.returncode == 0, ( + f"FSDP training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}" + ) + + # The gathered optimizer state must be written under FSDP (proves the save collective ran), + # in the same safetensors format as single-GPU training. + training_state_dir = output_dir / "checkpoints" / "last" / "training_state" + optimizer_state = training_state_dir / "optimizer_state.safetensors" + optimizer_param_groups = training_state_dir / "optimizer_param_groups.json" + assert optimizer_state.exists(), f"FSDP optimizer state not saved in {training_state_dir}" + assert optimizer_param_groups.exists(), ( + f"FSDP optimizer param groups not saved in {training_state_dir}" + ) + + # Resume from the checkpoint for more steps. A successful run proves load_fsdp_optimizer + # accepts the saved state and reshards it without shape/key errors. + resume_config = output_dir / "checkpoints" / "last" / "pretrained_model" / "train_config.json" + resume_args = [ + f"--config_path={resume_config}", + "--resume=true", + "--steps=20", + ] + resume_result = run_accelerate_training( + resume_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert resume_result.returncode == 0, ( + f"FSDP resume failed:\nSTDOUT:\n{resume_result.stdout}\n\nSTDERR:\n{resume_result.stderr}" + ) + assert "End of training" in resume_result.stdout or "End of training" in resume_result.stderr diff --git a/tests/utils/test_keyboard_input.py b/tests/utils/test_keyboard_input.py new file mode 100644 index 000000000..2f0dee889 --- /dev/null +++ b/tests/utils/test_keyboard_input.py @@ -0,0 +1,228 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the display-independent keyboard input helpers. + +These cover the parts most likely to regress: the environment-detection decision +table (the heart of the Wayland/headless fix), the macOS trust probe, the control +mapping, the terminal escape-sequence parsing, and backend selection. They require +neither ``pynput`` nor a real terminal. +""" + +import io +import platform +import sys + +import pytest + +import lerobot.utils.keyboard_input as ki +from lerobot.utils.keyboard_input import ( + TerminalKeyListener, + apply_recording_control, + create_key_listener, + init_keyboard_listener, + is_headless, + is_wayland, + pynput_can_capture, + pynput_listener_is_trusted, +) + + +@pytest.fixture(autouse=True) +def _clear_detection_caches(): + """The detection helpers are ``@cache``-decorated; clear around each test.""" + for fn in (is_headless, is_wayland, pynput_can_capture): + fn.cache_clear() + yield + for fn in (is_headless, is_wayland, pynput_can_capture): + fn.cache_clear() + + +def _set_platform(monkeypatch, name): + monkeypatch.setattr(platform, "system", lambda: name) + + +def _set_tty(monkeypatch, is_tty): + stdin = io.StringIO("") + stdin.isatty = lambda: is_tty + monkeypatch.setattr(sys, "stdin", stdin) + + +# --- Environment detection (the core of the fix) --------------------------- +@pytest.mark.parametrize( + ("system", "env", "expected"), + [ + ("Linux", {}, True), # no display server + ("Linux", {"DISPLAY": ":0"}, False), # X11 + ("Linux", {"WAYLAND_DISPLAY": "wayland-0"}, False), # Wayland + ("Darwin", {}, False), # display always assumed present + ], +) +def test_is_headless(monkeypatch, system, env, expected): + _set_platform(monkeypatch, system) + monkeypatch.delenv("DISPLAY", raising=False) + monkeypatch.delenv("WAYLAND_DISPLAY", raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert is_headless() is expected + + +@pytest.mark.parametrize( + ("env", "expected"), + [ + ({"XDG_SESSION_TYPE": "wayland"}, True), + ({"WAYLAND_DISPLAY": "wayland-0"}, True), + ({"XDG_SESSION_TYPE": "x11"}, False), + ({}, False), + ], +) +def test_is_wayland(monkeypatch, env, expected): + monkeypatch.delenv("XDG_SESSION_TYPE", raising=False) + monkeypatch.delenv("WAYLAND_DISPLAY", raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert is_wayland() is expected + + +@pytest.mark.parametrize( + ("system", "env", "pynput_available", "expected"), + [ + ("Linux", {"DISPLAY": ":0"}, True, True), # X11 + ("Linux", {"DISPLAY": ":0", "WAYLAND_DISPLAY": "wayland-0"}, True, False), # Wayland + ("Linux", {}, True, False), # headless + ("Darwin", {}, True, True), + ("Linux", {"DISPLAY": ":0"}, False, False), # pynput not installed + ], +) +def test_pynput_can_capture(monkeypatch, system, env, pynput_available, expected): + _set_platform(monkeypatch, system) + monkeypatch.setattr(ki, "_pynput_available", pynput_available) + for var in ("DISPLAY", "WAYLAND_DISPLAY", "XDG_SESSION_TYPE"): + monkeypatch.delenv(var, raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + assert pynput_can_capture() is expected + + +# --- macOS trust probe ------------------------------------------------------ +class _FakeListener: + def __init__(self, is_trusted): + self.IS_TRUSTED = is_trusted + + +def test_pynput_listener_is_trusted(monkeypatch): + _set_platform(monkeypatch, "Linux") + assert pynput_listener_is_trusted(_FakeListener(False)) is True # non-macOS: always assumed ok + _set_platform(monkeypatch, "Darwin") + assert pynput_listener_is_trusted(_FakeListener(False), timeout_s=0.05) is False + + +# --- Control mapping -------------------------------------------------------- +def test_apply_recording_control(): + events = {"exit_early": False, "rerecord_episode": False, "stop_recording": False} + apply_recording_control("left", events) + assert events == {"exit_early": True, "rerecord_episode": True, "stop_recording": False} + apply_recording_control("esc", events) + assert events["stop_recording"] is True + apply_recording_control("up", events) # unknown control -> no-op (no error) + + +# --- Terminal escape-sequence parsing (the tricky bit) ---------------------- +def _drive(listener, byte_seq): + """Run the listener's read loop over a scripted list of bytes (no real terminal).""" + script = list(byte_seq) + + def fake_read(timeout): + if script: + return script.pop(0) + listener._running = False + return None + + listener._read_char = fake_read + listener._running = True + listener._run() + + +@pytest.mark.parametrize( + ("byte_seq", "expected"), + [ + (["\x1b", "[", "C"], ["right"]), # CSI arrow + (["\x1b", "O", "D"], ["left"]), # SS3 arrow (e.g. over SSH/tmux) + (["\x1b"], ["esc"]), # bare ESC + (["\x1b", "[", "A"], ["up"]), # decoded even though the record handler ignores it + (["n"], ["n"]), # letter passthrough + ], +) +def test_terminal_parsing(byte_seq, expected): + collected = [] + _drive(TerminalKeyListener(collected.append), byte_seq) + assert collected == expected + + +# --- Backend selection ------------------------------------------------------ +def test_init_selects_terminal_when_pynput_cannot_capture(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) # avoid touching termios + listener, _ = init_keyboard_listener() + assert isinstance(listener, TerminalKeyListener) + + +def test_init_returns_none_without_tty(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=False) + listener, _ = init_keyboard_listener() + assert listener is None + + +@pytest.mark.parametrize( + ("key", "flag"), + [("right", "exit_early"), ("r", "rerecord_episode"), ("q", "stop_recording")], +) +def test_init_terminal_key_routing(monkeypatch, key, flag): + """Arrows and their letter equivalents drive the same events (terminal backend).""" + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) + listener, events = init_keyboard_listener() + listener._on_key(key) + assert events[flag] is True + + +# --- Shared factory + pynput key resolver ----------------------------------- +def test_resolve_pynput_key_char_fallback(): + """Unmapped keys fall back to ``.char`` (and yield None when there is none).""" + assert ki._resolve_pynput_key(type("K", (), {"char": "s"})()) == "s" + assert ki._resolve_pynput_key(type("K", (), {"char": None})()) is None + assert ki._resolve_pynput_key(type("K", (), {"char": ""})()) is None # empty char -> no key + + +def test_create_key_listener_routes_to_dispatch(monkeypatch): + """The terminal backend forwards canonical key names straight to ``dispatch``.""" + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=True) + monkeypatch.setattr(TerminalKeyListener, "start", lambda self: None) + seen = [] + listener = create_key_listener(seen.append, controls_help="save='s'") + assert isinstance(listener, TerminalKeyListener) + listener._on_key("space") + assert seen == ["space"] + + +def test_create_key_listener_none_without_tty(monkeypatch): + monkeypatch.setattr(ki, "pynput_can_capture", lambda: False) + _set_tty(monkeypatch, is_tty=False) + assert create_key_listener(lambda name: None) is None diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index c171763c2..e3705409b 100644 --- a/tests/utils/test_train_utils.py +++ b/tests/utils/test_train_utils.py @@ -136,3 +136,18 @@ def test_save_load_training_state(tmp_path, optimizer, scheduler): assert loaded_step == 10 assert loaded_optimizer is optimizer assert loaded_scheduler is scheduler + + +def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler): + # FSDP loads optimizer separately (after accelerator.prepare) + # load_training_state(load_optimizer=False) must restore step + scheduler but leave the + # optimizer untouched and never touch the on-disk optimizer state. + save_training_state(tmp_path, 10, optimizer, scheduler) + with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state: + loaded_step, loaded_optimizer, loaded_scheduler = load_training_state( + tmp_path, optimizer, scheduler, load_optimizer=False + ) + mock_load_optimizer_state.assert_not_called() + assert loaded_step == 10 + assert loaded_optimizer is optimizer + assert loaded_scheduler is scheduler diff --git a/uv.lock b/uv.lock index 5c55f8a3d..e24b1d884 100644 --- a/uv.lock +++ b/uv.lock @@ -768,34 +768,46 @@ wheels = [ [[package]] name = "cmeel-tinyxml2" -version = "11.0.0" +version = "10.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9b/96/4311533fee0a364bb605b585762f04c249f47857b33548a8ea837a7eb860/cmeel_tinyxml2-11.0.0.tar.gz", hash = "sha256:85d9c7680b3369af4c6b40a0dce70bbd84aa67832755622e57eb260cd95abe40", size = 645900, upload-time = "2026-05-21T11:49:32.652Z" } +sdist = { url = "https://files.pythonhosted.org/packages/28/9f/030eca702c485f7a641f975f167fa93164911b3329f005fb0730ff5e793f/cmeel_tinyxml2-10.0.0.tar.gz", hash = "sha256:00252aefc1c94a55b89f25ad08ee79fda2da8d1d94703e051598ddb52a9088fe", size = 645297, upload-time = "2025-02-06T10:29:00.106Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/79/f0/90c1640c53b623359d75ab1c70bdf19dc0afe82722bc5df57d09f8eaf83a/cmeel_tinyxml2-11.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:b0bd974e549b8c444626671a8e645897603ebf5225734cbe04a9dd3461477754", size = 111719, upload-time = "2026-05-21T11:49:25.999Z" }, - { url = "https://files.pythonhosted.org/packages/56/40/166447150a31bc3b794ffb493d5a634f67ffbc75dd8b4c46373701b7ef15/cmeel_tinyxml2-11.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:9a1406f408262c37ae7c4566b1d67801c4b10c4980903fb1ef0ba45fa4407072", size = 109146, upload-time = "2026-05-21T11:49:27.829Z" }, - { url = "https://files.pythonhosted.org/packages/d8/ca/3cc665afe2d76999f15454bb3b2f7c05f0088ad7de35718648291a536fd9/cmeel_tinyxml2-11.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:6f830007917c3e36f26b27d170ce84a619a62f46104d3cce435dff0125dd665f", size = 157109, upload-time = "2026-05-21T11:49:29.358Z" }, - { url = "https://files.pythonhosted.org/packages/87/4e/dcc0d9756d93be734d824e2a570cc9ac68909a1d7d3b6fc87c2fb32726c0/cmeel_tinyxml2-11.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:18674156bd41f3993dc1d5199da04fa496674358daa6588090fb9f86c71917b0", size = 148825, upload-time = "2026-05-21T11:49:31.035Z" }, + { url = "https://files.pythonhosted.org/packages/fe/5d/bc3a932eb7996a0a789979426a9bb8a3948bf57f3f17bab87dddbef62433/cmeel_tinyxml2-10.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:924499bb1b60b9a17bd001d12a9af88ddbee4ca888638ae684ba7f0f3ce49e87", size = 111913, upload-time = "2025-02-06T10:28:45.723Z" }, + { url = "https://files.pythonhosted.org/packages/92/bf/67d11e123313c034712896e94038291fe506bb099bdb75a136392002ffd0/cmeel_tinyxml2-10.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:26a1eb30c2a00bfc172e89ed015a18b8efb2b383546252ca8859574aed684686", size = 109487, upload-time = "2025-02-06T10:28:47.546Z" }, + { url = "https://files.pythonhosted.org/packages/ca/48/d8c81ce19b4b278ed0e8f81f93ae8670209bf3a9ac20141b9c386bb40cc7/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_17_i686.whl", hash = "sha256:53d86e02864c712f51f9a9adfcd8b6046b2ed51d44a0c34a8438d93b72b48325", size = 160118, upload-time = "2025-02-06T10:28:49.627Z" }, + { url = "https://files.pythonhosted.org/packages/87/4e/62193e27c9581f8ba7aeaeca7805632a64f2f4a824b1db37ad02ee953e8a/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:74112e2e9473afbf6ee2d25c9942553e9f6a40465e714533db72db48bc7658e1", size = 158477, upload-time = "2025-02-06T10:28:51.667Z" }, + { url = "https://files.pythonhosted.org/packages/14/f9/d0420c39e9ade99beeec61cd3abc68880fe6e14d85e9df292af8fabe65c8/cmeel_tinyxml2-10.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:ecd6e99caa2a06ac0d4b333b740c20fca526d0ca426f99eb5c0a0039117afdb6", size = 147025, upload-time = "2025-02-06T10:28:53.944Z" }, + { url = "https://files.pythonhosted.org/packages/66/9e/df63147fc162ab487217fa5596778ab7a81a82d9b3ce4236fd3a1e48cecb/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:30993fffb7032a45d5d3b1e5670cb879dad667a13144cd68c8f4e0371a8a3d2e", size = 150958, upload-time = "2025-02-06T10:28:55.301Z" }, + { url = "https://files.pythonhosted.org/packages/0e/a8/b03567275fd83f5af33ddb61de942689dec72c5b21bec01e6a5b11101aa5/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_i686.whl", hash = "sha256:8c09ede51784af54211a6225884dc7ddbb02ea1681656d173060c7ad2a5b9a3c", size = 160300, upload-time = "2025-02-06T10:28:57.189Z" }, + { url = "https://files.pythonhosted.org/packages/3a/ec/2781635b66c1059ca1243ae0f5a0410e171a5d8b8a71be3e34cb172f9f2d/cmeel_tinyxml2-10.0.0-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:3bd511d6d0758224efdebc23d3ead6e94f0755b04141ebf7d5493377829e8332", size = 149184, upload-time = "2025-02-06T10:28:58.734Z" }, ] [[package]] name = "cmeel-urdfdom" -version = "6.0.0" +version = "4.0.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cmeel" }, { name = "cmeel-console-bridge" }, { name = "cmeel-tinyxml2" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/21/75/4e8aff079e98582aeeb8e752805081da0c2dea405e79bafeefb555defe9f/cmeel_urdfdom-6.0.0.tar.gz", hash = "sha256:65c0fdc6021300fc55b2d0c03ab64dedc328034a74e40498e671bc894bb1dcf7", size = 303688, upload-time = "2026-05-21T12:08:56.663Z" } +sdist = { url = "https://files.pythonhosted.org/packages/31/09/be81a5e7db56f34b6ccdbe7afe855c95a18c8439e173519e0146e9276a8c/cmeel_urdfdom-4.0.1.tar.gz", hash = "sha256:2e3f41e8483889e195b574acb326a4464cf11a3c0a8724031ac28bcda2223efc", size = 291511, upload-time = "2025-02-12T12:07:09.699Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/40/51ba667135f01631179eee1614557193f8453740f248302d1b8b7f9f693e/cmeel_urdfdom-6.0.0-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:53d55cebb137a6e4dac6c16fa53f2dc2b7b9b5cda644bd1637a5bb849cd96e52", size = 381501, upload-time = "2026-05-21T12:08:48.758Z" }, - { url = "https://files.pythonhosted.org/packages/3c/d1/2b49a8c940fa75abc13df9842c14e577e6a82d5854b6d52597ce3bb04894/cmeel_urdfdom-6.0.0-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:0ef424735bd30f4afa4d1b4ddca9b297498c43005ddd775c080e55f62e9e0466", size = 377159, upload-time = "2026-05-21T12:08:50.485Z" }, - { url = "https://files.pythonhosted.org/packages/db/ac/0efde3a48220b55707bafb6d2e2dcca562f99dcd5c2c15311f7696eeacce/cmeel_urdfdom-6.0.0-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:0436f5230f1484c8e583284ef48c7291b230ada3dc5fb2937941f582e72409ec", size = 506000, upload-time = "2026-05-21T12:08:52.273Z" }, - { url = "https://files.pythonhosted.org/packages/32/d4/dfd617e598100e4e53ae3d228a968facff80bae53038fb18e2dccb1ab03a/cmeel_urdfdom-6.0.0-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:7ab1be680a8ec866d5422c617b641d1f0e38774061df28b8b426fb26edce6337", size = 530049, upload-time = "2026-05-21T12:08:54.224Z" }, + { url = "https://files.pythonhosted.org/packages/be/d0/20147dd6bb723afc44a58d89ea624df2bad1bed7b898a2df112aaca4a479/cmeel_urdfdom-4.0.1-0-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:2fe56939c6b47f6ec57021aac154123da47ecdcd79a217f3a5e3c4b705a07dee", size = 300860, upload-time = "2025-02-12T12:06:58.536Z" }, + { url = "https://files.pythonhosted.org/packages/8e/98/f832bca347e2d987c6b0ebb6930caf7b2c402535324aeed466b6aa2c4513/cmeel_urdfdom-4.0.1-0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:00a0aba78b68c428b27abeed1db58d73e65319ed966911a0e97b37367442e756", size = 300616, upload-time = "2025-02-12T12:07:00.556Z" }, + { url = "https://files.pythonhosted.org/packages/cf/10/bf5765b6f388037cff166a754a0958ac2fee34ca3c0975ef64d0324e4647/cmeel_urdfdom-4.0.1-0-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:a701a8f9671331f11b18ecf37a6537db546a21e6a0e5d0ff53341fea0693ed7f", size = 385951, upload-time = "2025-02-12T12:07:02.556Z" }, + { url = "https://files.pythonhosted.org/packages/c3/82/cb3f8f587d293a17bdbea15b50cdaa4a1e28e04583eb4cb4821685b89466/cmeel_urdfdom-4.0.1-0-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:12e39fc388c077d79fc9b3841d3d972a1da90b90de754d3363194c1540e18abf", size = 399619, upload-time = "2025-02-12T12:07:04.388Z" }, + { url = "https://files.pythonhosted.org/packages/24/77/322d7ac92c692d8dfaeda9de2d937087d15e2b564dc457d656e5fde3991d/cmeel_urdfdom-4.0.1-0-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:c4a83925df1d5923c4485c3eb2b80b3a61b14f119ab724fb5bd04cec494690ee", size = 373969, upload-time = "2025-02-12T12:07:06.222Z" }, + { url = "https://files.pythonhosted.org/packages/9f/63/bdc6b55cc8bd99bb9dce6be801b30feffaa1c3841ecb7f4fe4d137424518/cmeel_urdfdom-4.0.1-0-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:4c4f44270971b3d05c45a4e21b1fb2df7e05a750363ae918f59532bff0bfe0e1", size = 388237, upload-time = "2025-02-12T12:07:08.326Z" }, + { url = "https://files.pythonhosted.org/packages/1d/2d/8463fc23230612daf4da1e31d3229f47708381f3ae4d1500f0f007ac0f92/cmeel_urdfdom-4.0.1-1-py3-none-macosx_10_9_x86_64.whl", hash = "sha256:f7535158f45992eb2ba79e90d9db1bf9adc3846d9c7ed3e7a8c1c4d5343afa37", size = 301006, upload-time = "2025-02-13T11:42:08.8Z" }, + { url = "https://files.pythonhosted.org/packages/0f/d5/c8cdf500e49300d85624cbc3ef804107ddcdc9c541b1d3f726bfb58a9fc1/cmeel_urdfdom-4.0.1-1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:fef2a01a00d61d41b3d35dd4958bba973e9025c26eea1d3c9880932f4dba89a5", size = 300758, upload-time = "2025-02-13T11:42:10.449Z" }, + { url = "https://files.pythonhosted.org/packages/cf/b3/2f7bac1544113a7f8e0f6d8b1fab5e75c6a3d27ffbb584b03267251b2165/cmeel_urdfdom-4.0.1-1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7a52eb36950ce982014d99a55717ca29985da056e3705f20746f15d3244c1f7a", size = 386043, upload-time = "2025-02-13T11:42:11.923Z" }, + { url = "https://files.pythonhosted.org/packages/86/03/8bdeb36ba6a3e8125d523ecfc010403049e463fe589f9896858d4bdcaf1e/cmeel_urdfdom-4.0.1-1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:9f3b9c80b10d7246821ff61c2573f799e3da23d483e6f7367ddcad8a48baf58f", size = 399719, upload-time = "2025-02-13T11:42:14.325Z" }, + { url = "https://files.pythonhosted.org/packages/3f/ed/43f99e7512460294cd8acc5753ba25f8a20bdf28d62e143eaf3ec7a28bb6/cmeel_urdfdom-4.0.1-1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:2de69f47e8312cc09157624802d5bdaad6406443f863fb4b9ec62a19b4de3c72", size = 374073, upload-time = "2025-02-13T11:42:17.907Z" }, + { url = "https://files.pythonhosted.org/packages/17/c6/2e9bde6d7c02c1cf203ea896f8ce1afd441412f09b44830f1ee4a96d77de/cmeel_urdfdom-4.0.1-1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7708c1402de450fbeab21f7ca264a9a4676ed4c1cdf8d84d840bc5d057aac920", size = 388337, upload-time = "2025-02-13T11:42:19.657Z" }, ] [[package]] @@ -1147,7 +1159,7 @@ name = "decord" version = "0.6.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "numpy", marker = "(platform_machine != 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, + { name = "numpy", marker = "(platform_machine != 'arm64' and platform_machine != 's390x' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'linux')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/11/79/936af42edf90a7bd4e41a6cac89c913d4b47fa48a26b042d5129a9242ee3/decord-0.6.0-py3-none-manylinux2010_x86_64.whl", hash = "sha256:51997f20be8958e23b7c4061ba45d0efcd86bffd5fe81c695d0befee0d442976", size = 13602299, upload-time = "2021-06-14T21:30:55.486Z" }, @@ -2788,6 +2800,8 @@ accelerate-dep = [ all = [ { name = "accelerate" }, { name = "av" }, + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "contourpy" }, { name = "datasets" }, { name = "debugpy" }, @@ -2971,6 +2985,8 @@ hardware = [ ] hilserl = [ { name = "av" }, + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "datasets" }, { name = "grpcio" }, { name = "gym-hil" }, @@ -2993,6 +3009,8 @@ intelrealsense = [ { name = "pyrealsense2-macosx", marker = "sys_platform == 'darwin'" }, ] kinematics = [ + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "placo" }, ] lekiwi = [ @@ -3066,6 +3084,8 @@ pi = [ { name = "transformers" }, ] placo-dep = [ + { name = "cmeel-tinyxml2" }, + { name = "cmeel-urdfdom" }, { name = "placo" }, ] pusht = [ @@ -3186,6 +3206,8 @@ requires-dist = [ { name = "accelerate", marker = "extra == 'accelerate-dep'", specifier = ">=1.14.0,<2.0.0" }, { name = "av", marker = "extra == 'av-dep'", specifier = ">=15.0.0,<16.0.0" }, { name = "cmake", specifier = ">=3.29.0.1,<4.2.0" }, + { name = "cmeel-tinyxml2", marker = "extra == 'placo-dep'", specifier = "<11" }, + { name = "cmeel-urdfdom", marker = "extra == 'placo-dep'", specifier = ">=4,<5" }, { name = "contourpy", marker = "extra == 'matplotlib-dep'", specifier = ">=1.3.0,<2.0.0" }, { name = "datasets", marker = "extra == 'dataset'", specifier = ">=4.7.0,<5.0.0" }, { name = "debugpy", marker = "extra == 'dev'", specifier = ">=1.8.1,<1.9.0" },