Merge branch 'main' into feat/unitree_g1_sonic_rebased

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)

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()

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()

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.

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)

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"}:

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:

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):

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

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

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:

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)
 

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)
 
 

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) ──────────────────────────────────
 

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
 
 

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",

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),
+    }

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.

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()

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

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

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)

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()
 

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:

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."""

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:

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:

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)
 

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):

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:

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] = []
 

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