fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames (#3749 )

* fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames Per-episode stats previously materialized every frame (and decoded up to 16 episodes in parallel), so peak memory scaled with episode length and OOM'd on large datasets (#2889). Numeric features are now read in full from the table (exact), while only image/video frames are sub-sampled per episode using the existing sample_indices heuristic. Worker count is configurable via LEROBOT_STATS_MAX_WORKERS; --no-sampling restores exact behavior. * Update tests/datasets/test_augment_quantile_stats.py Co-authored-by: Haoming Song <1847575517@qq.com> Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> --------- Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Haoming Song <1847575517@qq.com>
fix(eval): prevent eval_policy crash when start_seed is None and num_envs>1 (#4203 )
2026-07-29 12:39:41 +00:00 · 2026-07-29 12:32:03 +02:00 · 2026-07-28 18:41:28 +02:00 · 2026-07-28 17:54:58 +02:00 · 2026-07-28 16:45:48 +02:00 · 2026-07-28 16:45:32 +02:00
78 changed files with 1282 additions and 4095 deletions
@@ -191,162 +191,6 @@ def make_my_policy_pre_post_processors(
 ---
 ## Adding high- and low-level language control
 The policy API above is sufficient for training and standard evaluation. To use a language-conditioned policy with interactive `lerobot-rollout`, also register a runtime adapter. The adapter keeps policy-specific prompting and tokenization out of the generic control loop.
 The runtime supports two policy shapes:
 | Policy shape     | Behavior                                                                | Adapter                                        |
 | ---------------- | ----------------------------------------------------------------------- | ---------------------------------------------- |
 | Low-level / flat | The operator's task or subtask directly conditions action prediction.   | Reuse `DirectTaskPolicyAdapter`.               |
 | High + low level | The policy generates subtasks or memory, then conditions actions on it. | Subclass `BaseLanguageAdapter`, as PI052 does. |
 During a rollout, `RuntimeState` stores the high-level task and the active language context:
 ```text
 task ──> adapter.generate_text("subtask", ...) ──> state.language_context["subtask"]
                                                       │
 observation ──> processors ──> adapter.select_action() ─┴─> action chunk ──> robot
 ```
 The generic runtime handles generation frequency, pause/resume, prompt replacement, action queues, and dispatch. The adapter only translates between that runtime contract and your policy.
 ### Low-level policies
 If your policy already consumes the live task through its normal preprocessor and implements `predict_action_chunk`, register the shared direct adapter. PI0.5 and MolmoAct2 use this path:
 ```python
 # src/lerobot/runtime/registry.py
 _ADAPTERS = {
    # ...
    "my_policy": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
 }
 ```
 Run it with direct-subtask mode so the operator supplies the instruction used by the action policy:
 ```bash
 lerobot-rollout \
    --language \
    --policy.path=user/my_policy_checkpoint \
    --robot.type=so101_follower \
    --robot.port=/dev/ttyACM0 \
    --direct_subtask
 ```
 The rollout context builds the observation batch with the current instruction before `DirectTaskPolicyAdapter` calls `policy.predict_action_chunk(observation)`. No text-generation method is required.
 ### Hierarchical policies
 For a policy that generates language and actions, subclass [`BaseLanguageAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/runtime/adapter.py) and implement two methods:
 - `generate_text(kind, observation, state, user_text=None) -> str` generates a `subtask`, `memory`, or interjection response.
 - `select_action(observation, state)` builds the low-level prompt from the active context and returns an action chunk.
 This abbreviated adapter follows [`PI052PolicyAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py):
 ```python
 # inference/my_policy_adapter.py
 from typing import Any
 from lerobot.runtime import RuntimeState
 from lerobot.runtime.adapter import BaseLanguageAdapter
 from lerobot.utils.constants import (
    OBS_LANGUAGE_ATTENTION_MASK,
    OBS_LANGUAGE_TOKENS,
 )
 class MyPolicyAdapter(BaseLanguageAdapter):
    def select_action(self, observation: dict[str, Any], state: RuntimeState):
        instruction = state.language_context.get("subtask") or state.task or ""
        tokens, attention_mask = tokenize_instruction(instruction)
        batch = dict(observation)
        batch[OBS_LANGUAGE_TOKENS] = tokens
        batch[OBS_LANGUAGE_ATTENTION_MASK] = attention_mask
        return self.policy.predict_action_chunk(batch)
    def generate_text(
        self,
        kind: str,
        observation: dict[str, Any] | None,
        state: RuntimeState,
        user_text: str | None = None,
    ) -> str:
        messages = self.build_messages(kind, state, user_text)
        batch, tokenizer = tokenize_messages(messages, observation)
        return self.policy.select_message(
            batch,
            tokenizer=tokenizer,
            min_new_tokens=self.gen.min_new_tokens,
            temperature=self.gen.temperature,
            top_p=self.gen.top_p,
        )
    def build_messages(
        self, kind: str, state: RuntimeState, user_text: str | None
    ) -> list[dict[str, str]]:
        if kind == "subtask":
            return [{"role": "user", "content": state.task or ""}]
        if kind == "memory":
            return [
                {"role": "user", "content": state.task or ""},
                {
                    "role": "user",
                    "content": f"Completed subtask: {state.extra.get('prior_subtask', '')}",
                },
            ]
        if kind == "interjection":
            return [
                {"role": "user", "content": state.task or ""},
                {"role": "user", "content": user_text or ""},
            ]
        raise ValueError(f"Unsupported text kind: {kind}")
 ```
 `tokenize_instruction` and `tokenize_messages` are policy-specific helpers. They must reproduce the prompt format used during training; PI052, for example, adds the discretized robot state to its low-level subtask prompt and uses the same PaliGemma formatting for `select_message`.
 `BaseLanguageAdapter` provides the default hierarchy: regenerate a subtask at action-chunk boundaries, update memory when the subtask changes, and handle user interjections. Override `_regenerate_context` only if your policy uses a different hierarchy.
 Register the adapter with a lazy import so importing LeRobot does not load the model or its optional dependencies:
 ```python
 # src/lerobot/runtime/registry.py
 _ADAPTERS = {
    # ...
    "my_policy": "lerobot.policies.my_policy.inference.my_policy_adapter:MyPolicyAdapter",
 }
 ```
 The key must match the policy's registered type. Once registered, the same checkpoint works through the shared entry point:
 ```bash
 lerobot-rollout \
    --language \
    --policy.path=user/my_hierarchical_checkpoint \
    --robot.type=so101_follower \
    --robot.port=/dev/ttyACM0 \
    --task="put the cup in the sink"
 ```
 For RoboCasa-compatible policies, replace the robot arguments with `--sim --sim.task=<task>`. Without `--direct_subtask`, the adapter generates the low-level subtask; with it, the operator bypasses high-level generation and supplies each subtask.
 ### Keep training and deployment aligned
 The adapter is intentionally small, but its prompts are part of the model contract:
 - Use the same tokenizer, role formatting, special tokens, image ordering, and state encoding as training.
 - Condition `select_action` on `state.language_context["subtask"]`, falling back to `state.task` for direct or not-yet-generated prompts.
 - Return a full action chunk from `select_action`; the runtime handles control-rate dispatch.
 - Keep optional model dependencies inside lazy imports.
 - Test adapter selection, generated-message routing, action-batch construction, and direct-subtask behavior with a lightweight fake policy.
 PI052 is the complete in-tree reference: its [processor](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/processor_pi052.py) renders the training recipe, its [policy](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/modeling_pi052.py) exposes text and action generation, and its [adapter](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py) reconstructs those same prompts at deployment.
 ---
 ## Path A: Out-of-tree plugin
 The fastest way to ship a policy: package it as a standalone Python distribution and install it alongside LeRobot. No PR required, you own the release cycle, and you can publish to PyPI under your own namespace.
@@ -136,6 +136,10 @@ config = RealSenseCameraConfig(
    height=480,
    color_mode=ColorMode.RGB,
    use_depth=True,
    # Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
    exposure=120,
    gain=64,
    white_balance=4600,
    rotation=Cv2Rotation.NO_ROTATION
 )
@@ -154,6 +158,15 @@ finally:
 ```
 <!-- prettier-ignore-end -->
 Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
 supported ranges vary by camera model; an invalid value raises an error at connection time that
 includes the range reported by the sensor. Requesting an unsupported control also raises an error.
 Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
 require `use_rgb=True`.
 On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
 exposure or gain also affects the depth stream.
 </hfoption>
 </hfoptions>
@@ -1,6 +1,6 @@
 # Policy Deployment (lerobot-rollout)
-`lerobot-rollout` is the single CLI for deploying trained policies on real robots or in an interactive simulator. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording, language-driven control, and human-in-the-loop data collection.
+`lerobot-rollout` is the single CLI for deploying trained policies on real robots. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection.
 ## Quick Start
@@ -197,52 +197,6 @@ Teleop is optional — if omitted the robot holds its position during the reset
 ---
 ## Interactive language control
 Language-conditioned policies can expose a high-level text head in addition to
 their action head. Add `--language` to open-prompt one of these policies on a
 real robot. Language-only flags such as `--direct_subtask` select this mode
 automatically.
 MolmoAct2 has no high-level planner, so use direct-subtask mode and type each
 next low-level instruction yourself:
 ```bash
 lerobot-rollout \
    --policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \
    --policy.device=cuda \
    --robot.type=so101_follower \
    --robot.port=/dev/ttyACM1 \
    --robot.cameras='{"cam0":{"type":"opencv","index_or_path":"/dev/video0","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200},"cam1":{"type":"opencv","index_or_path":"/dev/video2","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200}}' \
    --direct_subtask \
    --robot.max_relative_target='{"shoulder_pan":5,"shoulder_lift":5,"elbow_flex":5,"wrist_flex":5,"wrist_roll":5,"gripper":5}'
 ```
 The robot starts paused. Type a subtask, then use `/resume` and `/pause` to
 control action dispatch. Check the workspace and motion limits before resuming.
 Without `--direct_subtask`, a policy such as PI052 generates its active subtask
 from the high-level `--task` itself.
 RoboCasa uses the same runtime and processor path. `--sim` selects it
 automatically, so no robot configuration is needed:
 ```bash
 MUJOCO_GL=egl lerobot-rollout \
    --policy.path=lerobot/pi052_robocasa \
    --sim --sim.task=CloseFridge --sim.split=pretrain \
    --task="close the fridge" \
    --disable_memory \
    --sim.render_size=384 \
    --sim.views=robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right \
    --mode=action --ctrl_hz=20
 ```
 Open `http://localhost:8010` for the live simulator view. Add
 `--sim.direct_subtask` to bypass the language planner and make each typed prompt
 the action policy's current subtask.
 ---
 ## Inference Backends
 Select a backend with `--inference.type=<name>`. All strategies work with both backends.
@@ -141,17 +141,6 @@ sample["target_message_indices"]
 The renderer does not apply a tokenizer chat template. Policy processors decide how to serialize the messages for their backbone, which keeps the same dataset usable across SmolVLA, Pi0.5, and any future VLM that expects OpenAI-style chat messages.
 ## Blends
 Blend recipes select one weighted sub-recipe deterministically from the sample index.
 `recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses.
 A message recipe with a supervised assistant turn on the `low_level` stream trains
 the π0.5 paper's joint sequence instead of a blend: the target span gets text CE
 while also conditioning the action losses in the same forward.
 `recipes/subtask_joint.yaml` is the provided example; pair it with
 `--policy.joint_subtask_conditioning=true` at inference.
 ## Graceful absence
 If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.
@@ -494,6 +494,19 @@ ignore_errors = true
 module = "lerobot.envs.*"
 ignore_errors = false
 [[tool.mypy.overrides]]
 module = "lerobot.annotations.*"
 ignore_errors = false
 disallow_untyped_defs = true
 disallow_incomplete_defs = true
 check_untyped_defs = true
 [[tool.mypy.overrides]]
 module = "lerobot.transforms.*"
 ignore_errors = false
 disallow_untyped_defs = true
 disallow_incomplete_defs = true
 check_untyped_defs = true
 # [[tool.mypy.overrides]]
 # module = "lerobot.utils.*"
@@ -120,14 +120,22 @@ class OpenCVCamera(Camera):
        self.rotation: int | None = get_cv2_rotation(config.rotation)
        self.backend: int = config.backend
-        if self.height and self.width:
+        self.capture_width: int | None = None
-            self.capture_width, self.capture_height = self.width, self.height
+        self.capture_height: int | None = None
-            if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
+        self._reset_connection_settings()
                self.capture_width, self.capture_height = self.height, self.width
    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.index_or_path})"
    def _reset_connection_settings(self) -> None:
        """Restore settings that may have been auto-detected during a failed connection."""
        self.fps = self.config.fps
        self.width = self.config.width
        self.height = self.config.height
        self.capture_width, self.capture_height = self.width, self.height
        if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
            self.capture_width, self.capture_height = self.height, self.width
    @property
    def is_connected(self) -> bool:
        """Checks if the camera is currently connected and opened."""
@@ -164,6 +172,7 @@ class OpenCVCamera(Camera):
                f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
            )
        try:
            self._configure_capture_settings()
            self._start_read_thread()
@@ -175,6 +184,13 @@ class OpenCVCamera(Camera):
                with self.frame_lock:
                    if self.latest_frame is None:
                        raise ConnectionError(f"{self} failed to capture frames during warmup.")
        except BaseException:
            try:
                self._cleanup_resources()
            except Exception:
                logger.exception(f"Failed to fully clean up {self} after connect() failed.")
            self._reset_connection_settings()
            raise
        logger.info(f"{self} connected.")
@@ -312,6 +328,7 @@ class OpenCVCamera(Camera):
        for target in targets_to_scan:
            camera = cv2.VideoCapture(target)
            try:
                if camera.isOpened():
                    default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
                    default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -321,7 +338,9 @@ class OpenCVCamera(Camera):
                    # Get FOURCC code and convert to string
                    default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
                    default_fourcc_code_int = int(default_fourcc_code)
-                default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
+                    default_fourcc = "".join(
                        [chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
                    )
                    camera_info = {
                        "name": f"OpenCV Camera @ {target}",
@@ -338,6 +357,7 @@ class OpenCVCamera(Camera):
                    }
                    found_cameras_info.append(camera_info)
            finally:
                camera.release()
        return found_cameras_info
@@ -496,6 +516,26 @@ class OpenCVCamera(Camera):
            self.latest_timestamp = None
            self.new_frame_event.clear()
    def _cleanup_resources(self) -> None:
        """Stop background reads and release the capture, including after partial setup."""
        read_thread = self.thread
        videocapture = self.videocapture
        try:
            self._stop_read_thread()
        finally:
            self.videocapture = None
            try:
                if videocapture is not None:
                    videocapture.release()
            finally:
                # Releasing the device may unblock a hardware read that outlived
                # the first bounded join in _stop_read_thread().
                if read_thread is not None and read_thread.is_alive():
                    read_thread.join(timeout=2.0)
                    if read_thread.is_alive():  # pragma: no cover
                        logger.warning(f"{self} read thread remained alive after releasing the capture.")
    @check_if_not_connected
    def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
        """
@@ -586,16 +626,6 @@ class OpenCVCamera(Camera):
        if not self.is_connected and self.thread is None:
            raise DeviceNotConnectedError(f"{self} not connected.")
-        if self.thread is not None:
+        self._cleanup_resources()
            self._stop_read_thread()
        if self.videocapture is not None:
            self.videocapture.release()
            self.videocapture = None
        with self.frame_lock:
            self.latest_frame = None
            self.latest_timestamp = None
            self.new_frame_event.clear()
        logger.info(f"{self} disconnected.")
@@ -121,6 +121,9 @@ class RealSenseCamera(Camera):
        self.config = config
        self.width: int | None = config.width
        self.height: int | None = config.height
        if config.serial_number_or_name.isdigit():
            self.serial_number = config.serial_number_or_name
        else:
@@ -131,6 +134,9 @@ class RealSenseCamera(Camera):
        self.use_rgb = config.use_rgb
        self.use_depth = config.use_depth
        self.warmup_s = config.warmup_s
        self.exposure: int | None = config.exposure
        self.gain: int | None = config.gain
        self.white_balance: int | None = config.white_balance
        self.rs_pipeline: rs.pipeline | None = None
        self.rs_profile: rs.pipeline_profile | None = None
@@ -145,14 +151,23 @@ class RealSenseCamera(Camera):
        self.rotation: int | None = get_cv2_rotation(config.rotation)
-        if self.height and self.width:
+        self.capture_width: int | None = None
-            self.capture_width, self.capture_height = self.width, self.height
+        self.capture_height: int | None = None
-            if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
+        self._reset_connection_settings()
                self.capture_width, self.capture_height = self.height, self.width
    def __str__(self) -> str:
        return f"{self.__class__.__name__}({self.serial_number})"
    def _reset_connection_settings(self) -> None:
        """Restore settings that may have been auto-detected during a failed connection."""
        self.fps = self.config.fps
        self.width = self.config.width
        self.height = self.config.height
        self.warmup_s = self.config.warmup_s
        self.capture_width, self.capture_height = self.width, self.height
        if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
            self.capture_width, self.capture_height = self.height, self.width
    @property
    def is_connected(self) -> bool:
        """Checks if the camera pipeline is started and streams are active."""
@@ -172,7 +187,8 @@ class RealSenseCamera(Camera):
        Raises:
            DeviceAlreadyConnectedError: If the camera is already connected.
-            ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
+            ValueError: If the configuration is invalid, a requested sensor option is unsupported,
                or a requested sensor value is invalid.
            ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
            RuntimeError: If the pipeline starts but fails to apply requested settings.
        """
@@ -190,7 +206,9 @@ class RealSenseCamera(Camera):
                f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
            ) from e
        try:
            self._configure_capture_settings()
            self._configure_sensor_options()
            self._start_read_thread()
            # NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
@@ -206,6 +224,13 @@ class RealSenseCamera(Camera):
                    self.use_depth and self.latest_depth_frame is None
                ):
                    raise ConnectionError(f"{self} failed to capture frames during warmup.")
        except BaseException:
            try:
                self._cleanup_resources()
            except Exception:
                logger.exception(f"Failed to fully clean up {self} after connect() failed.")
            self._reset_connection_settings()
            raise
        logger.info(f"{self} connected.")
@@ -339,6 +364,111 @@ class RealSenseCamera(Camera):
        self.new_frame_event.clear()
        return self._async_read(timeout_ms=10000, read_depth=read_depth)
    def _get_color_sensor(self) -> "rs.sensor":
        """Returns the sensor that controls the color stream.
        Most RealSense cameras expose "RGB Camera" for color. The D405 has no
        separate RGB module — its color stream comes from "Stereo Module".
        We try RGB Camera first, then fall back to Stereo Module.
        """
        if self.rs_profile is None:
            raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
        device = self.rs_profile.get_device()
        sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
        for name in ("RGB Camera", "Stereo Module"):
            if name in sensors:
                return sensors[name]
        available = list(sensors.keys())
        raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
    def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
        """Sets a sensor option, re-raising range errors with actionable diagnostics."""
        try:
            sensor.set_option(option, value)
        except Exception as e:
            range_info = ""
            try:
                option_range = sensor.get_option_range(option)
                range_info = (
                    f" (supported range: min={option_range.min}, max={option_range.max}, "
                    f"step={option_range.step}, default={option_range.default})"
                )
            except Exception:
                range_info = " (option range unavailable)"
            raise ValueError(
                f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
            ) from e
    def _configure_sensor_options(self) -> None:
        """Applies manual sensor options (exposure, gain, white balance) to the color sensor.
        When exposure or gain is set, auto-exposure is disabled first. When white_balance
        is set, auto white balance is disabled first. An omitted option is left unchanged,
        and configuration is skipped entirely if all options are omitted.
        Raises:
            ValueError: If the sensor does not support a requested option or a requested
                value is invalid. Invalid-value errors include the option name, requested
                value, and supported range when available.
        """
        if self.exposure is None and self.gain is None and self.white_balance is None:
            return
        color_sensor = self._get_color_sensor()
        requested_options = (
            (rs.option.exposure, self.exposure, "exposure"),
            (rs.option.gain, self.gain, "gain"),
            (rs.option.white_balance, self.white_balance, "white balance"),
        )
        unsupported_options = [
            label
            for option, value, label in requested_options
            if value is not None and not color_sensor.supports(option)
        ]
        if unsupported_options:
            raise ValueError(
                f"{self}: color sensor does not support requested manual options: {unsupported_options}."
            )
        manual_exposure_requested = self.exposure is not None or self.gain is not None
        if manual_exposure_requested:
            if color_sensor.supports(rs.option.enable_auto_exposure):
                self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
                logger.info(f"{self} auto-exposure disabled.")
            else:
                logger.warning(
                    f"{self} sensor does not support disabling auto-exposure; "
                    "applying manual exposure/gain directly."
                )
        if self.exposure is not None:
            self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
            logger.info(f"{self} exposure set to {self.exposure}.")
        if self.gain is not None:
            self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
            logger.info(f"{self} gain set to {self.gain}.")
        if self.white_balance is not None:
            if color_sensor.supports(rs.option.enable_auto_white_balance):
                self._set_sensor_option(
                    color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
                )
                logger.info(f"{self} auto white balance disabled.")
            else:
                logger.warning(
                    f"{self} sensor does not support disabling auto white balance; "
                    "applying manual white balance directly."
                )
            self._set_sensor_option(
                color_sensor, rs.option.white_balance, self.white_balance, "white balance"
            )
            logger.info(f"{self} white balance set to {self.white_balance}.")
    @check_if_not_connected
    def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
        """
@@ -541,6 +671,27 @@ class RealSenseCamera(Camera):
            self.latest_timestamp = None
            self.new_frame_event.clear()
    def _cleanup_resources(self) -> None:
        """Stop background reads and stop the pipeline, including after partial setup."""
        read_thread = self.thread
        rs_pipeline = self.rs_pipeline
        try:
            self._stop_read_thread()
        finally:
            self.rs_pipeline = None
            self.rs_profile = None
            try:
                if rs_pipeline is not None:
                    rs_pipeline.stop()
            finally:
                # Stopping the pipeline may unblock a hardware read that outlived
                # the first bounded join in _stop_read_thread().
                if read_thread is not None and read_thread.is_alive():
                    read_thread.join(timeout=2.0)
                    if read_thread.is_alive():  # pragma: no cover
                        logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
    def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
        """Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
        if self.thread is None or not self.thread.is_alive():
@@ -684,18 +835,5 @@ class RealSenseCamera(Camera):
                f"Attempted to disconnect {self}, but it appears already disconnected."
            )
-        if self.thread is not None:
+        self._cleanup_resources()
            self._stop_read_thread()
        if self.rs_pipeline is not None:
            self.rs_pipeline.stop()
            self.rs_pipeline = None
            self.rs_profile = None
        with self.frame_lock:
            self.latest_color_frame = None
            self.latest_depth_frame = None
            self.latest_timestamp = None
            self.new_frame_event.clear()
        logger.info(f"{self} disconnected.")
@@ -46,6 +46,17 @@ class RealSenseCameraConfig(CameraConfig):
        use_depth: Whether to enable depth stream. Defaults to False.
        rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
        warmup_s: Time reading frames before returning from connect (in seconds)
        exposure: Manual exposure value for the color sensor. When set, auto-exposure is
            disabled and this fixed value is used. Valid ranges are camera-model specific
            and reported if the value is rejected. Defaults to None (leave unchanged).
        gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
            and this fixed gain is used, which also freezes exposure at its current value
            when no exposure is configured. Valid ranges are camera-model specific and
            reported if the value is rejected. Defaults to None (leave unchanged).
        white_balance: Manual white balance value for the color sensor. When set, auto
            white balance is disabled and this fixed value is used. Valid ranges are
            camera-model specific and reported if the value is rejected. Defaults to None
            (leave unchanged).
    Note:
        - Either name or serial_number must be specified.
@@ -61,6 +72,9 @@ class RealSenseCameraConfig(CameraConfig):
    use_depth: bool = False
    rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
    warmup_s: int = 1
    exposure: int | None = None
    gain: int | None = None
    white_balance: int | None = None
    def __post_init__(self) -> None:
        self.color_mode = ColorMode(self.color_mode)
@@ -69,6 +83,18 @@ class RealSenseCameraConfig(CameraConfig):
        if not self.use_rgb and not self.use_depth:
            raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
        manual_color_options = {
            "exposure": self.exposure,
            "gain": self.gain,
            "white_balance": self.white_balance,
        }
        configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
        if configured_color_options and not self.use_rgb:
            raise ValueError(
                "Manual color sensor options require `use_rgb=True`. "
                f"Configured options: {configured_color_options}."
            )
        values = (self.fps, self.width, self.height)
        if any(v is not None for v in values) and any(v is None for v in values):
            raise ValueError(
@@ -71,13 +71,19 @@ class DatasetRecordConfig:
    # Number of threads per encoder instance. None = auto (codec default).
    # Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
    encoder_threads: int | None = None
    # Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
    # (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
    no_stamp: bool = False
    def stamp_repo_id(self) -> None:
        """Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
        Must be called explicitly at dataset *creation* time — not on resume,
        where the existing ``repo_id`` (already stamped) must be preserved.
        No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
        """
        if self.no_stamp:
            return
        if self.repo_id:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            self.repo_id = f"{self.repo_id}_{timestamp}"
@@ -33,8 +33,6 @@ class DatasetConfig:
    # looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
    root: str | None = None
    episodes: list[int] | None = None
    # Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`.
    exclude_episodes: list[int] | None = None
    image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
    revision: str | None = None
    use_imagenet_stats: bool = True
@@ -64,10 +62,6 @@ class DatasetConfig:
            if len(self.episodes) != len(set(self.episodes)):
                duplicates = sorted({ep for ep in self.episodes if self.episodes.count(ep) > 1})
                raise ValueError(f"Episode indices contain duplicates: {duplicates}")
        if self.exclude_episodes is not None and any(ep < 0 for ep in self.exclude_episodes):
            raise ValueError(
                f"exclude_episodes must be non-negative, got: {[ep for ep in self.exclude_episodes if ep < 0]}"
            )
@dataclass
@@ -78,7 +78,7 @@ class MessageTurn:
            raise ValueError(f"Unsupported message stream: {self.stream!r}")
        if self.content is None and self.tool_calls_from is None:
            raise ValueError("MessageTurn.content is required unless tool_calls_from is set.")
-        if self.content is not None and not isinstance(self.content, str | list):
+        if self.content is not None and not isinstance(self.content, (str, list)):
            raise TypeError("MessageTurn.content must be a string, a list of HF-style blocks, or None.")
        if isinstance(self.content, list):
            for block in self.content:
@@ -147,7 +147,7 @@ class TrainingRecipe:
        return cls.from_dict(data)
    def _validate_message_recipe(self) -> None:
-        """Validate bindings and require text or low-level action supervision."""
+        """Ensure every templated binding is known and at least one turn is a target."""
        assert self.messages is not None
        known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
@@ -156,14 +156,8 @@ class TrainingRecipe:
            if missing:
                raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}")
-        has_target = any(turn.target for turn in self.messages)
+        if not any(turn.target for turn in self.messages):
-        has_low_level = any(turn.stream == "low_level" for turn in self.messages)
+            raise ValueError("Message recipes must contain at least one target turn.")
        if not (has_target or has_low_level):
            raise ValueError(
                "Message recipes must contain at least one supervised turn — "
                "either ``target: true`` (text CE) or ``stream: low_level`` "
                "(flow/action loss)."
            )
    def _validate_blend_recipe(self) -> None:
        """Ensure each blend component is a non-empty, weighted message recipe."""
@@ -1,16 +0,0 @@
 # Predicts subtasks from tasks and trains subtask-conditioned action flow without memory or plans.
 # Requires `subtask` annotations; samples with missing `if_present` bindings do not render.
 blend:
  high_level_subtask:
    weight: 0.30
    messages:
      - {role: user, content: "${task}", stream: high_level}
      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
  low_level_execution:
    weight: 0.70
    messages:
      # The low-level stream trains action flow on the generated or annotated subtask.
      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
@@ -1,13 +0,0 @@
 # Paper-style joint sequence (pi0.5 §IV-B): one sample supervises the subtask
 # text with CE and, because the assistant turn is part of the prefix, conditions
 # the FAST and flow action losses on the same annotated subtask in one forward.
 # The supervised span is attended causally; the action losses see task + subtask.
 #
 # Pair with `--policy.joint_subtask_conditioning=true` at inference so the flow
 # prefix reproduces this layout (task turn with state + causal generated subtask).
 # Samples without a `subtask` annotation fall back to a plain task-prompt
 # low-level sample via `if_present`.
 messages:
  - {role: user, content: "${task}", stream: low_level}
  - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask}
@@ -1,30 +0,0 @@
 # Trains subtask prediction, subtask-conditioned action flow, and memory updates without plans.
 # Requires `subtask` and `memory`; missing `if_present` bindings skip the affected sub-recipe.
 blend:
  high_level_subtask:
    weight: 0.25
    messages:
      - {role: user, content: "${task}", stream: high_level}
      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
  low_level_execution:
    weight: 0.60
    messages:
      # The low-level stream trains action flow on the generated or annotated subtask.
      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
  memory_update:
    # `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
    # Inference controls update timing through `subtask_change` events.
    weight: 0.15
    bindings:
      prior_memory: "nth_prev(style=memory, offset=1)"
      current_memory: "active_at(t, style=memory)"
      completed_subtask: "nth_prev(style=subtask, offset=1)"
    messages:
      - {role: user, content: "${task}", stream: high_level}
      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
@@ -1,70 +0,0 @@
 # Adds memory, spoken interjection responses, and camera-grounded VQA to subtask/action training.
 # Missing optional annotations skip only their sub-recipe; `say` tool calls tokenize as `<say>...</say>`.
 blend:
  high_level_subtask:
    weight: 0.25
    messages:
      - {role: user, content: "${task}", stream: high_level}
      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
  low_level_execution:
    weight: 0.40
    messages:
      # The low-level stream trains action flow on the generated or annotated subtask.
      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
  memory_update:
    # `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
    # Inference controls update timing through `subtask_change` events.
    weight: 0.10
    bindings:
      prior_memory: "nth_prev(style=memory, offset=1)"
      current_memory: "active_at(t, style=memory)"
      completed_subtask: "nth_prev(style=subtask, offset=1)"
    messages:
      - {role: user, content: "${task}", stream: high_level}
      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
  user_interjection_response:
    weight: 0.10
    bindings:
      interjection: "emitted_at(t, style=interjection)"
      speech: "emitted_at(t, role=assistant, tool_name=say)"
    messages:
      - {role: user, content: "${task}", stream: high_level}
      - {role: user, content: "${interjection}", stream: high_level, if_present: interjection}
      # The assistant target is a `say` tool call flattened to a `<say>...</say>` marker.
      - {role: assistant, stream: high_level, target: true, if_present: speech, tool_calls_from: speech}
  # Each camera uses a separate VQA sub-recipe for view-specific binding.
  ask_vqa_top:
    weight: 0.075
    bindings:
      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)"
      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.front)"
    messages:
      - role: user
        stream: high_level
        if_present: vqa_query
        content:
          - {type: image, feature: observation.images.front}
          - {type: text, text: "${vqa_query}"}
      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
  ask_vqa_wrist:
    weight: 0.075
    bindings:
      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
    messages:
      - role: user
        stream: high_level
        if_present: vqa_query
        content:
          - {type: image, feature: observation.images.wrist}
          - {type: text, text: "${vqa_query}"}
      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
    def _load_metadata(self):
        self.info = load_info(self.root)
        check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
-        self.tasks = load_tasks(self.root)
+        self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
-        self.episodes = load_episodes(self.root)
+        self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
        self.stats = load_stats(self.root)
    def ensure_readable(self) -> None:
@@ -163,40 +163,10 @@ class DatasetReader:
    def _load_hf_dataset(self) -> datasets.Dataset:
        """hf_dataset contains all the observations, states, actions, rewards, etc."""
        features = get_hf_features_from_features(self._meta.features)
        # Annotated datasets may have language columns absent from metadata.
        # Extend the schema before the strict Parquet cast.
        features = self._extend_features_with_language_columns(features)
        hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes)
        hf_dataset.set_transform(hf_transform_to_torch)
        return hf_dataset
    def _extend_features_with_language_columns(self, features: datasets.Features) -> datasets.Features:
        """Register language columns found in Parquet but missing from metadata."""
        # Leave empty datasets to fail through the normal loading path.
        try:
            sample = next((self.root / "data").glob("*/*.parquet"))
        except StopIteration:
            return features
        from pyarrow import parquet as _pq  # noqa: PLC0415
        schema_names = set(_pq.read_schema(sample).names)
        from .language import (  # noqa: PLC0415
            LANGUAGE_EVENTS,
            LANGUAGE_PERSISTENT,
            language_events_column_feature,
            language_persistent_column_feature,
        )
        extra: dict[str, object] = {}
        if LANGUAGE_PERSISTENT in schema_names and LANGUAGE_PERSISTENT not in features:
            extra[LANGUAGE_PERSISTENT] = language_persistent_column_feature()
        if LANGUAGE_EVENTS in schema_names and LANGUAGE_EVENTS not in features:
            extra[LANGUAGE_EVENTS] = language_events_column_feature()
        if not extra:
            return features
        return datasets.Features({**features, **extra})
    def _check_cached_episodes_sufficient(self) -> bool:
        """Check if the cached dataset contains all requested episodes and their video files."""
        if self.hf_dataset is None or len(self.hf_dataset) == 0:
@@ -66,17 +66,6 @@ def resolve_delta_timestamps(
    return delta_timestamps
 def _resolve_episodes(
    episodes: list[int] | None, exclude_episodes: list[int] | None, total_episodes: int
 ) -> list[int] | None:
    """Apply an episode exclusion list on top of an optional allowlist."""
    if not exclude_episodes:
        return episodes
    base = episodes if episodes is not None else list(range(total_episodes))
    excluded = set(exclude_episodes)
    return [episode for episode in base if episode not in excluded]
 def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
    """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
@@ -98,14 +87,11 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
            cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
        )
        delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
        episodes = _resolve_episodes(
            cfg.dataset.episodes, cfg.dataset.exclude_episodes, ds_meta.total_episodes
        )
        if not cfg.dataset.streaming:
            dataset = LeRobotDataset(
                cfg.dataset.repo_id,
                root=cfg.dataset.root,
-                episodes=episodes,
+                episodes=cfg.dataset.episodes,
                delta_timestamps=delta_timestamps,
                image_transforms=image_transforms,
                revision=cfg.dataset.revision,
@@ -118,7 +104,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
            dataset = StreamingLeRobotDataset(
                cfg.dataset.repo_id,
                root=cfg.dataset.root,
-                episodes=episodes,
+                episodes=cfg.dataset.episodes,
                delta_timestamps=delta_timestamps,
                image_transforms=image_transforms,
                revision=cfg.dataset.revision,
@@ -162,28 +162,14 @@ def render_sample(
    task: str | None = None,
    dataset_ctx: Any | None = None,
 ) -> RenderedMessages | None:
-    """Resolve one sample's bindings and render its message recipe.
+    """Render the chat-style messages for a single dataset sample.
-    Returns ``None`` when no text or low-level action supervision applies.
+    Resolves the recipe's bindings against ``persistent`` and ``events`` rows
    at frame timestamp ``t``, then expands the recipe's message templates.
    Returns ``None`` if the resolved sample contains no target message.
    """
    persistent_rows = _normalize_rows(persistent or [])
    event_rows = _normalize_rows(events or [])
    # Route sparse VQA frames to a matching view-specific component before weighted selection.
    # This avoids dropping annotated frames or selecting VQA without annotations.
    if recipe.blend is not None:
        vqa_rendered = _render_vqa_if_present(
            recipe,
            persistent=persistent_rows,
            events=event_rows,
            t=t,
            sample_idx=sample_idx,
            task=task,
            dataset_ctx=dataset_ctx,
        )
        if vqa_rendered is not None:
            return vqa_rendered
    selected_recipe = _select_recipe(recipe, sample_idx)
    bindings = _resolve_bindings(
        selected_recipe,
@@ -197,55 +183,6 @@ def render_sample(
    return _render_message_recipe(selected_recipe, bindings)
 def _render_vqa_if_present(
    recipe: TrainingRecipe,
    *,
    persistent: Sequence[LanguageRow],
    events: Sequence[LanguageRow],
    t: float,
    sample_idx: int,
    task: str | None,
    dataset_ctx: Any | None,
 ) -> RenderedMessages | None:
    """Render a matching VQA component, or return ``None`` for normal selection.
    Multiple matching views are selected deterministically by relative weight.
    """
    assert recipe.blend is not None
    renderable: list[tuple[float, RenderedMessages]] = []
    for name, component in recipe.blend.items():
        if not name.startswith("ask_vqa"):
            continue
        bindings = _resolve_bindings(
            component,
            persistent=persistent,
            events=events,
            t=t,
            sample_idx=sample_idx,
            task=task,
            dataset_ctx=dataset_ctx,
        )
        rendered = _render_message_recipe(component, bindings)
        if rendered is not None:
            renderable.append((float(component.weight or 0.0), rendered))
    if not renderable:
        return None
    if len(renderable) == 1:
        return renderable[0][1]
    # Choose among matching cameras by relative weight, or uniformly when all weights are zero.
    total = sum(w for w, _ in renderable) or float(len(renderable))
    digest = hashlib.blake2b(f"vqa:{sample_idx}".encode(), digest_size=8).digest()
    draw = int.from_bytes(digest, "big") / 2**64 * total
    cumulative = 0.0
    for w, rendered in renderable:
        cumulative += w or (total / len(renderable))
        if draw < cumulative:
            return rendered
    return renderable[-1][1]
 def _select_recipe(recipe: TrainingRecipe, sample_idx: int) -> TrainingRecipe:
    """Pick a deterministic blend component for ``sample_idx`` (or return ``recipe``)."""
    if recipe.blend is None:
@@ -409,9 +346,7 @@ def _render_message_recipe(
        if turn.target:
            target_indices.append(message_idx)
-    # Keep samples with either text targets or low-level action supervision.
+    if not target_indices:
    has_low_level = any(stream == "low_level" for stream in streams)
    if not target_indices and not has_low_level:
        return None
    rendered = {
@@ -468,12 +403,14 @@ def _validate_rendered(rendered: RenderedMessages) -> None:
    if len(streams) != len(messages):
        raise ValueError("message_streams must be aligned with messages.")
-    # Require text or low-level action supervision.
+    if not target_indices:
-    if not target_indices and not any(s == "low_level" for s in streams):
+        raise ValueError("Rendered samples must contain at least one target message.")
        raise ValueError("Rendered samples must contain a target message or a low_level-stream message.")
    for idx in target_indices:
        if idx < 0 or idx >= len(messages):
            raise ValueError(f"Target message index {idx} is out of bounds.")
    # ``stream`` is enforced non-None at MessageTurn construction time
    # (see ``MessageTurn.__post_init__``), so a missing stream here would
    # mean the dataclass invariant was bypassed; no need to re-check.
 def _nth_relative(
@@ -560,13 +560,7 @@ class RoboCasaEnv(EnvConfig):
            kwargs["split"] = self.split
        return kwargs
-    def create_envs(
+    def create_envs(self, n_envs: int, use_async_envs: bool = False):
        self,
        n_envs: int,
        use_async_envs: bool = False,
        terminate_on_success: bool = True,
        horizon: int | None = None,
    ):
        from .robocasa import create_robocasa_envs
        if self.task is None:
@@ -580,8 +574,6 @@ class RoboCasaEnv(EnvConfig):
            env_cls=env_cls,
            episode_length=self.episode_length,
            obj_registries=tuple(self.obj_registries),
            terminate_on_success=terminate_on_success,
            horizon=horizon,
        )
@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
    def close(self):
        if self._env is not None:
            try:
                self._env.close()
            finally:
                # LIBERO deletes its inner env on close, so this wrapper must
                # be recreated before the next reset.
                self._env = None
 def _make_env_fns(
@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
 # Dimensions for the flat action/state vectors used by the LeRobot wrapper.
 # These correspond to the PandaOmron robot in RoboCasa365.
-OBS_STATE_DIM = 16  # ee_pos_rel(3) + ee_quat_rel(4) + base_pos(3) + base_quat(4) + gripper_qpos(2)
+OBS_STATE_DIM = 16  # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
-ACTION_DIM = 12  # ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
+ACTION_DIM = 12  # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
 ACTION_LOW = -1.0
 ACTION_HIGH = 1.0
@@ -101,15 +101,14 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
 def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
    """Split a flat (12,) action vector into a RoboCasa action dict.
-    Layout (openpi / robocasa.utils.env_utils.convert_action order):
+    Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
        ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
    """
    return {
-        "action.end_effector_position": flat_action[0:3],
+        "action.base_motion": flat_action[0:4],
-        "action.end_effector_rotation": flat_action[3:6],
+        "action.control_mode": flat_action[4:5],
-        "action.gripper_close": flat_action[6:7],
+        "action.end_effector_position": flat_action[5:8],
-        "action.base_motion": flat_action[7:11],
+        "action.end_effector_rotation": flat_action[8:11],
-        "action.control_mode": flat_action[11:12],
+        "action.gripper_close": flat_action[11:12],
    }
@@ -137,16 +136,9 @@ class RoboCasaEnv(gym.Env):
        episode_length: int | None = None,
        obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
        episode_index: int = 0,
        terminate_on_success: bool = True,
        horizon: int | None = None,
    ):
        super().__init__()
        self.task = task
        # When False, a task-success does NOT end/reset the episode — used by the
        # interactive sim so one kitchen persists across sequential prompts.
        self.terminate_on_success = terminate_on_success
        # Underlying robosuite horizon (steps before truncation). None -> default.
        self.horizon = horizon
        self.obs_type = obs_type
        self.render_mode = render_mode
        self.observation_width = observation_width
@@ -218,16 +210,12 @@ class RoboCasaEnv(gym.Env):
        # (only None/"all"/"pretrain"/"target" are valid). Always pass a
        # valid value so we don't hit that default. Extra kwargs are
        # forwarded to the underlying kitchen env via create_env/robosuite.make.
        extra_kwargs: dict[str, Any] = {}
        if self.horizon is not None:
            extra_kwargs["horizon"] = int(self.horizon)
        self._env = RoboCasaGymEnv(
            env_name=self.task,
            camera_widths=self.observation_width,
            camera_heights=self.observation_height,
            split=self.split if self.split is not None else "all",
            obj_registries=self.obj_registries,
            **extra_kwargs,
        )
        ep_meta = self._env.env.get_ep_meta()
@@ -242,14 +230,12 @@ class RoboCasaEnv(gym.Env):
            return {"pixels": images}
        # `state.*` keys come from PandaOmronKeyConverter inside the wrapper.
        # openpi state order: ee first, then base, then gripper (matches the
        # openpi robocasa pipeline / examples/robocasa/main.py state layout).
        agent_pos = np.concatenate(
            [
                raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
                raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
                raw_obs.get("state.base_position", np.zeros(3)),
                raw_obs.get("state.base_rotation", np.zeros(4)),
                raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
                raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
                raw_obs.get("state.gripper_qpos", np.zeros(2)),
            ],
            axis=-1,
@@ -294,7 +280,7 @@ class RoboCasaEnv(gym.Env):
        raw_obs, reward, done, truncated, info = self._env.step(action_dict)
        is_success = bool(info.get("success", False))
-        terminated = done or (is_success and self.terminate_on_success)
+        terminated = done or is_success
        info.update({"task": self.task, "done": done, "is_success": is_success})
        observation = self._format_raw_obs(raw_obs)
@@ -327,8 +313,6 @@ def _make_env_fns(
    split: str | None,
    episode_length: int | None,
    obj_registries: Sequence[str],
    terminate_on_success: bool = True,
    horizon: int | None = None,
 ) -> list[Callable[[], RoboCasaEnv]]:
    """Build n_envs factory callables for a single task.
@@ -351,8 +335,6 @@ def _make_env_fns(
            episode_length=episode_length,
            obj_registries=obj_registries,
            episode_index=episode_index,
            terminate_on_success=terminate_on_success,
            horizon=horizon,
        )
    return [partial(_make_env, i) for i in range(n_envs)]
@@ -366,8 +348,6 @@ def create_robocasa_envs(
    env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
    episode_length: int | None = None,
    obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
    terminate_on_success: bool = True,
    horizon: int | None = None,
 ) -> dict[str, dict[int, Any]]:
    """Create vectorized RoboCasa365 environments with a consistent return shape.
@@ -429,8 +409,6 @@ def create_robocasa_envs(
            split=split,
            episode_length=episode_length,
            obj_registries=obj_registries,
            terminate_on_success=terminate_on_success,
            horizon=horizon,
        )
        if is_async:
@@ -384,7 +384,9 @@ class RoboTwinEnv(gym.Env):
        self._env: Any | None = None  # deferred — created on first reset() inside worker
        self._step_count: int = 0
-        self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
+        self._black_frame: np.ndarray = np.zeros(
            (self.observation_height, self.observation_width, 3), dtype=np.uint8
        )
        image_spaces = {
            cam: spaces.Box(
@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
        if action.shape[0] != 7:
            # Unknown layout — fall back to zero-pad so the sim doesn't crash.
-            padded = np.zeros(ctrl_dim, dtype=np.float64)
+            padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
            padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
            return padded
@@ -122,6 +122,9 @@ MODEL_ENCODING_TABLE = {
    "xm430-w350": X_SERIES_ENCODINGS_TABLE,
    "xm540-w270": X_SERIES_ENCODINGS_TABLE,
    "xc430-w150": X_SERIES_ENCODINGS_TABLE,
    "xh540-w150": X_SERIES_ENCODINGS_TABLE,
    "xc330-t288": X_SERIES_ENCODINGS_TABLE,
    "xc330-t181": X_SERIES_ENCODINGS_TABLE,
 }
 # {model: model_resolution}
@@ -134,6 +137,9 @@ MODEL_RESOLUTION = {
    "xm430-w350": 4096,
    "xm540-w270": 4096,
    "xc430-w150": 4096,
    "xh540-w150": 4096,
    "xc330-t288": 4096,
    "xc330-t181": 4096,
 }
 # {model: model_number}
@@ -145,6 +151,9 @@ MODEL_NUMBER_TABLE = {
    "xm430-w350": 1020,
    "xm540-w270": 1120,
    "xc430-w150": 1070,
    "xh540-w150": 1110,
    "xc330-t288": 1220,
    "xc330-t181": 1210,
 }
 # {model: available_operating_modes}
@@ -156,6 +165,9 @@ MODEL_OPERATING_MODES = {
    "xm430-w350": [0, 1, 3, 4, 5, 16],
    "xm540-w270": [0, 1, 3, 4, 5, 16],
    "xc430-w150": [1, 3, 4, 16],
    "xh540-w150": [0, 1, 3, 4, 5, 16],
    "xc330-t288": [0, 1, 3, 4, 5, 16],
    "xc330-t181": [0, 1, 3, 4, 5, 16],
 }
 MODEL_CONTROL_TABLE = {
@@ -166,6 +178,9 @@ MODEL_CONTROL_TABLE = {
    "xm430-w350": X_SERIES_CONTROL_TABLE,
    "xm540-w270": X_SERIES_CONTROL_TABLE,
    "xc430-w150": X_SERIES_CONTROL_TABLE,
    "xh540-w150": X_SERIES_CONTROL_TABLE,
    "xc330-t288": X_SERIES_CONTROL_TABLE,
    "xc330-t181": X_SERIES_CONTROL_TABLE,
 }
 MODEL_BAUDRATE_TABLE = {
@@ -176,6 +191,9 @@ MODEL_BAUDRATE_TABLE = {
    "xm430-w350": X_SERIES_BAUDRATE_TABLE,
    "xm540-w270": X_SERIES_BAUDRATE_TABLE,
    "xc430-w150": X_SERIES_BAUDRATE_TABLE,
    "xh540-w150": X_SERIES_BAUDRATE_TABLE,
    "xc330-t288": X_SERIES_BAUDRATE_TABLE,
    "xc330-t181": X_SERIES_BAUDRATE_TABLE,
 }
 AVAILABLE_BAUDRATES = [
@@ -44,12 +44,19 @@ from lerobot.utils.constants import (
    POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
 from lerobot.utils.feature_utils import dataset_to_policy_features
 from lerobot.utils.import_utils import _peft_available, require_package
 from .evo1.configuration_evo1 import Evo1Config
 from .groot.configuration_groot import GrootConfig
 from .pretrained import PreTrainedPolicy
 from .utils import validate_visual_features_consistency
 if TYPE_CHECKING or _peft_available:
    from peft import PeftConfig, PeftModel
 else:
    PeftConfig = None
    PeftModel = None
 def _reconnect_relative_absolute_steps(
    preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -334,12 +341,15 @@ def make_policy(
        # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
        # of the adapter and the adapter's config contains the path to the base policy. So we need the
        # adapter config first, then load the correct policy and then apply PEFT.
-        from peft import PeftConfig, PeftModel
+        require_package("peft", extra="peft")
        logging.info("Loading policy's PEFT adapter.")
        peft_pretrained_path = str(cfg.pretrained_path)
-        peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
+        peft_config = PeftConfig.from_pretrained(
            peft_pretrained_path,
            revision=cfg.pretrained_revision,
        )
        kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
        if not kwargs["pretrained_name_or_path"]:
@@ -350,9 +360,14 @@ def make_policy(
                "the adapter was trained."
            )
        kwargs["revision"] = peft_config.revision
        policy = policy_cls.from_pretrained(**kwargs)
        policy = PeftModel.from_pretrained(
-            policy, peft_pretrained_path, config=peft_config, is_trainable=True
+            policy,
            peft_pretrained_path,
            config=peft_config,
            revision=cfg.pretrained_revision,
            is_trainable=True,
        )
    else:
@@ -37,13 +37,19 @@ def is_image_feature(key: str) -> bool:
@dataclass
 class ConcurrencyConfig:
    """Configuration for the concurrency of the actor and learner.
    Possible values are:
    - "threads": Use threads for the actor and learner.
    - "processes": Use processes for the actor and learner.
    ``multiprocessing_context`` selects the process-wide start method when
    processes are used. Set it to ``None`` to preserve Python's default or a
    method already selected by the embedding application.
    """
    actor: str = "threads"
    learner: str = "threads"
    multiprocessing_context: str | None = "spawn"
@dataclass
@@ -43,11 +43,22 @@ from torch.distributions import Beta
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION
-from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
+from lerobot.utils.import_utils import (
    _peft_available,
    _scipy_available,
    _transformers_available,
    require_package,
 )
 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_molmoact2 import MolmoAct2Config
 if TYPE_CHECKING or _peft_available:
    from peft import LoraConfig, get_peft_model
 else:
    LoraConfig = None
    get_peft_model = None
 logger = logging.getLogger(__name__)
@@ -1731,13 +1742,11 @@ class MolmoAct2Policy(PreTrainedPolicy):
    def _build_inner_lora_config(self):
        require_package("peft", extra="molmoact2")
        from peft import LoraConfig
        return LoraConfig(**self._get_inner_peft_targets())
    def _apply_lora_adapters(self) -> None:
        require_package("peft", extra="molmoact2")
        from peft import get_peft_model
        peft_config = self._build_inner_lora_config()
        self._validate_peft_config(peft_config)
@@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
 from lerobot.utils.import_utils import _peft_available, require_package
 from .utils import log_model_loading_keys
-T = TypeVar("T", bound="PreTrainedPolicy")
+if TYPE_CHECKING or _peft_available:
    from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
 else:
    PEFT_TYPE_TO_CONFIG_MAPPING = None
    PeftType = None
    get_peft_model = None
 if TYPE_CHECKING:
    from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
 T = TypeVar("T", bound="PreTrainedPolicy")
 def _build_card_context(
    cfg: TrainPipelineConfig | None,
@@ -384,7 +392,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
            peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
                These are merged with policy defaults to build the final config.
        """
-        from peft import get_peft_model
+        require_package("peft", extra="peft")
        # If user provided a complete config, use it directly (with overrides)
        if peft_config is not None:
@@ -455,7 +463,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        Returns:
            Preprocessed dict with renamed keys and init_type mapped to method-specific key.
        """
-        from peft import PeftType
+        require_package("peft", extra="peft")
        cli_overrides = cli_overrides.copy()
@@ -480,7 +488,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
    def _build_peft_config(self, cli_overrides: dict):
        """Build a PEFT config from policy defaults and CLI overrides."""
-        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
+        require_package("peft", extra="peft")
        # Determine PEFT method type (default to LORA)
        method_type_str = cli_overrides.get("method_type") or "lora"
@@ -507,7 +515,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
    def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
        """Apply CLI overrides to an existing PEFT config."""
-        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
+        require_package("peft", extra="peft")
        # Get method type from existing config or CLI override
        method_type_str = cli_overrides.get("method_type")
@@ -175,6 +175,9 @@ class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
            if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
                complementary_data["task_index"] = task_index_value.unsqueeze(0)
        complementary_data.pop("language_persistent", None)
        complementary_data.pop("language_events", None)
        if "messages" in complementary_data:
            messages = complementary_data["messages"]
            if isinstance(messages, list) and (not messages or isinstance(messages[0], dict)):
@@ -132,10 +132,20 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
    def transform_features(
        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        for axis in ["x", "y", "z", "gripper"]:
+        for axis in ["x", "y", "z"]:
            features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
        features[PipelineFeatureType.ACTION].pop("gripper", None)
-        for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
+        for feat in [
            "enabled",
            "target_x",
            "target_y",
            "target_z",
            "target_wx",
            "target_wy",
            "target_wz",
            "gripper_vel",
        ]:
            features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
                type=FeatureType.ACTION, shape=(1,)
            )
@@ -41,7 +41,7 @@ from pathlib import Path
 from typing import Any, TypedDict, TypeVar, cast
 import torch
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
 from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -205,10 +205,6 @@ class ProcessorStep(ABC):
        """
        return None
    def save_artifacts(self, save_directory: Path) -> dict[str, str]:
        """Save non-tensor assets and map constructor arguments to relative paths."""
        return {}
    def reset(self) -> None:
        """Resets the internal state of the processor step, if any."""
        return None
@@ -553,22 +549,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
        pipeline_config = self.get_config()
        pipeline_state_dict = self.state_dict()
        for processor_step, step_entry in zip(self.steps, pipeline_config["steps"], strict=True):
            artifacts = processor_step.save_artifacts(save_directory)
            if artifacts:
                for config_key, relative_path in artifacts.items():
                    artifact_path = Path(relative_path)
                    if artifact_path.is_absolute() or ".." in artifact_path.parts:
                        raise ValueError(
                            f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
                        )
                    if not (save_directory / artifact_path).exists():
                        raise FileNotFoundError(
                            f"Processor step did not save declared artifact '{relative_path}'"
                        )
                    step_entry["config"][config_key] = artifact_path.as_posix()
                step_entry["artifacts"] = artifacts
        for state_key, step_state_dict in pipeline_state_dict.items():
            state_filename = f"{state_key}.safetensors"
            save_file(step_state_dict, save_directory / state_filename)
@@ -753,13 +733,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
        # 3. Build steps with overrides
        steps, validated_overrides = cls._build_steps_with_overrides(
-            loaded_config,
+            loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source
            overrides or {},
            model_id,
            base_path,
            config_filename,
            hub_download_kwargs,
            is_local_source,
        )
        # 4. Validate that all overrides were used
@@ -948,7 +922,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
        overrides: dict[str, Any],
        model_id: str,
        base_path: Path | None,
        config_filename: str,
        hub_download_kwargs: dict[str, Any],
        is_local_source: bool = False,
    ) -> tuple[list[ProcessorStep], set[str]]:
@@ -1003,68 +976,15 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
            ImportError: If a step class cannot be imported or found in registry
            ValueError: If a step cannot be instantiated with its configuration
        """
        loaded_config = deepcopy(loaded_config)
        cls._resolve_artifact_paths(
            loaded_config,
            model_id,
            base_path,
            config_filename,
            hub_download_kwargs,
        )
        steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
        for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
            cls._load_step_state(
-                step_instance,
+                step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source
                step_entry,
                model_id,
                base_path,
                config_filename,
                hub_download_kwargs,
                is_local_source,
            )
        return steps, remaining_override_keys
    @classmethod
    def _resolve_artifact_paths(
        cls,
        loaded_config: dict[str, Any],
        model_id: str,
        base_path: Path | None,
        config_filename: str,
        hub_download_kwargs: dict[str, Any],
    ) -> None:
        """Resolve declared relative processor artifacts before step construction."""
        is_local = Path(model_id).is_dir() or Path(model_id).is_file()
        for step_entry in loaded_config["steps"]:
            artifacts = step_entry.get("artifacts", {})
            for config_key, relative_path in artifacts.items():
                artifact_path = Path(relative_path)
                if artifact_path.is_absolute() or ".." in artifact_path.parts:
                    raise ValueError(
                        f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
                    )
                resolved_path = base_path / artifact_path if base_path is not None else artifact_path
                if not resolved_path.exists() and not is_local:
                    repository_path = Path(config_filename).parent / artifact_path
                    snapshot_download(
                        repo_id=model_id,
                        repo_type="model",
                        allow_patterns=f"{repository_path.as_posix()}/**",
                        **hub_download_kwargs,
                    )
                if not resolved_path.exists():
                    step_name = step_entry.get("registry_name", step_entry.get("class", "unknown"))
                    raise FileNotFoundError(
                        f"Missing processor artifact '{relative_path}' for step '{step_name}' "
                        f"next to '{config_filename}'. Checkpoint artifacts are incomplete."
                    )
                step_entry["config"][config_key] = str(resolved_path)
    @classmethod
    def _build_steps_from_config(
        cls,
@@ -1224,7 +1144,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
        step_entry: dict[str, Any],
        model_id: str,
        base_path: Path | None,
        config_filename: str,
        hub_download_kwargs: dict[str, Any],
        is_local_source: bool = False,
    ) -> None:
@@ -1290,7 +1209,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
            # Download from Hub
            state_path = hf_hub_download(
                repo_id=model_id,
-                filename=(Path(config_filename).parent / state_filename).as_posix(),
+                filename=state_filename,
                repo_type="model",
                **hub_download_kwargs,
            )
@@ -16,7 +16,7 @@
 from __future__ import annotations
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 from typing import Any
 from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -32,18 +32,17 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry
@dataclass
@ProcessorStepRegistry.register(name="render_messages_processor")
 class RenderMessagesStep(ProcessorStep):
-    """Render language columns into recipe-defined messages and supervision metadata."""
+    """Processor step that turns raw language columns into rendered chat messages.
    Reads ``language_persistent`` and ``language_events`` from the transition's
    complementary data, renders them through ``recipe`` at the sample timestamp,
    and replaces the raw columns with the resulting ``messages`` /
    ``message_streams`` / ``target_message_indices`` keys.
    """
    recipe: TrainingRecipe
    dataset_ctx: Any | None = None
    def __post_init__(self) -> None:
        if isinstance(self.recipe, dict):
            self.recipe = TrainingRecipe.from_dict(self.recipe)
    def get_config(self) -> dict[str, Any]:
        return {"recipe": asdict(self.recipe)}
    def __call__(self, transition: EnvTransition) -> EnvTransition | None:
        """Render messages for a single transition; return ``None`` to drop it."""
        complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
@@ -51,17 +50,7 @@ class RenderMessagesStep(ProcessorStep):
        events = complementary_data.get(LANGUAGE_EVENTS) or []
        if not persistent and not events:
            rendered = _fallback_low_level_render(complementary_data.get("task"))
            if rendered is None:
            return transition
            new_transition = transition.copy()
            new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
            new_complementary_data.update(rendered)
            new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
            return new_transition
        if _is_batched_language(persistent) or _is_batched_language(events):
            return self._call_batch(transition, complementary_data, persistent, events)
        timestamp = complementary_data.get("timestamp")
        if timestamp is None:
@@ -77,148 +66,19 @@ class RenderMessagesStep(ProcessorStep):
            task=complementary_data.get("task"),
            dataset_ctx=self.dataset_ctx,
        )
        if rendered is None:
            rendered = _fallback_low_level_render(complementary_data.get("task"))
        if rendered is None:
            return None
        new_transition = transition.copy()
-        new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
+        new_complementary_data = dict(complementary_data)
        new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
        new_complementary_data.pop(LANGUAGE_EVENTS, None)
        new_complementary_data.update(rendered)
        new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
        return new_transition
    def _call_batch(
        self,
        transition: EnvTransition,
        complementary_data: dict[str, Any],
        persistent_batch: list,
        events_batch: list,
    ) -> EnvTransition | None:
        timestamp = complementary_data.get("timestamp")
        if timestamp is None:
            raise KeyError("RenderMessagesStep requires sample timestamp in complementary data.")
        batch_size = max(len(persistent_batch), len(events_batch))
        messages: list[list[dict[str, Any]]] = []
        message_streams: list[list[str | None]] = []
        target_message_indices: list[list[int]] = []
        keep_indices: list[int] = []
        for i in range(batch_size):
            rendered = render_sample(
                recipe=self.recipe,
                persistent=persistent_batch[i] if i < len(persistent_batch) else [],
                events=events_batch[i] if i < len(events_batch) else [],
                t=_batch_value(timestamp, i),
                sample_idx=int(_batch_value(complementary_data.get("index", 0), i)),
                task=_batch_value(complementary_data.get("task"), i),
                dataset_ctx=self.dataset_ctx,
            )
            if rendered is None:
                rendered = _fallback_low_level_render(_batch_value(complementary_data.get("task"), i))
                if rendered is None:
                    continue
            keep_indices.append(i)
            messages.append(rendered["messages"])
            message_streams.append(rendered["message_streams"])
            target_message_indices.append(rendered["target_message_indices"])
        if not messages:
            return None
        new_transition = (
            _select_batch_indices(transition, keep_indices)
            if len(keep_indices) != batch_size
            else transition.copy()
        )
        new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
        new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
        new_complementary_data.pop(LANGUAGE_EVENTS, None)
        new_complementary_data["messages"] = messages
        new_complementary_data["message_streams"] = message_streams
        new_complementary_data["target_message_indices"] = target_message_indices
        new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
        return new_transition
    def transform_features(
        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
        """Pass features through unchanged; rendering only touches complementary data."""
        return features
 def _scalar(value: Any) -> float | int:
    """Unwrap a tensor/array/single-element list into a Python scalar."""
    if hasattr(value, "item"):
        return value.item()
    if isinstance(value, list):
        if len(value) != 1:
            raise ValueError(f"Expected a scalar, got list of length {len(value)}: {value!r}")
        return _scalar(value[0])
    return value
 def _is_batched_language(value: Any) -> bool:
    return isinstance(value, list) and bool(value) and isinstance(value[0], list)
 def _batch_value(value: Any, index: int) -> Any:
    if value is None:
        return None
    if isinstance(value, list):
        return value[index]
    if hasattr(value, "ndim") and value.ndim > 0:
        return _scalar(value[index])
    return _scalar(value)
 def _select_batch_indices(transition: EnvTransition, indices: list[int]) -> EnvTransition:
    selected = transition.copy()
    for key in (TransitionKey.OBSERVATION, TransitionKey.COMPLEMENTARY_DATA):
        data = selected.get(key)
        if isinstance(data, dict):
            selected[key] = {k: _select_value(v, indices) for k, v in data.items()}
    action = selected.get(TransitionKey.ACTION)
    if action is not None:
        selected[TransitionKey.ACTION] = _select_value(action, indices)
    return selected
 def _select_value(value: Any, indices: list[int]) -> Any:
    if isinstance(value, list) and len(value) >= len(indices):
        return [value[i] for i in indices]
    if hasattr(value, "index_select") and hasattr(value, "new_tensor") and getattr(value, "ndim", 0) > 0:
        return value.index_select(0, value.new_tensor(indices).long())
    return value
 def _fallback_low_level_render(task: Any) -> dict[str, Any] | None:
    """Keep action-only samples trainable when no recipe branch matches."""
    if hasattr(task, "item"):
        task = task.item()
    if isinstance(task, list):
        messages = []
        message_streams = []
        target_message_indices = []
        for t in task:
            rendered = _fallback_low_level_render(t)
            if rendered is None:
                return None
            messages.append(rendered["messages"])
            message_streams.append(rendered["message_streams"])
            target_message_indices.append(rendered["target_message_indices"])
        return {
            "messages": messages,
            "message_streams": message_streams,
            "target_message_indices": target_message_indices,
        }
    if not isinstance(task, str) or not task:
        return None
    return {
        "messages": [{"role": "user", "content": task}],
        "message_streams": ["low_level"],
        "target_message_indices": [],
    }
@@ -25,7 +25,6 @@ from __future__ import annotations
 import logging
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 import torch
@@ -33,7 +32,6 @@ import torch
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.types import EnvTransition, RobotObservation, TransitionKey
 from lerobot.utils.constants import (
    ACTION_CODE_TOKEN_MASK,
    ACTION_TOKEN_MASK,
    ACTION_TOKENS,
    OBS_LANGUAGE_ATTENTION_MASK,
@@ -138,7 +136,7 @@ class TokenizerProcessorStep(ObservationProcessorStep):
        # Standardize to a list of strings for the tokenizer
        if isinstance(task, str):
            return [task]
-        elif isinstance(task, list | tuple) and all(isinstance(t, str) for t in task):
+        elif isinstance(task, (list, tuple)) and all(isinstance(t, str) for t in task):
            return list(task)
        return None
@@ -351,7 +349,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
    max_action_tokens: int = 256
    fast_skip_tokens: int = 128
    paligemma_tokenizer_name: str = "google/paligemma-3b-pt-224"
    allow_truncation: bool = True
    # Internal tokenizer instance (not part of the config)
    action_tokenizer: Any = field(default=None, init=False, repr=False)
    _paligemma_tokenizer: Any = field(default=None, init=False, repr=False)
@@ -415,15 +412,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
            # During inference, no action is available, skip tokenization
            return new_transition
-        # Tokenize and get masks for the full formatted sequence and the discrete action codes.
+        # Tokenize and get both tokens and mask
-        tokens, mask, code_mask = self._tokenize_action(action)
+        tokens, mask = self._tokenize_action(action)
        # Store mask in complementary data
        complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
        if complementary_data is None:
            complementary_data = {}
        complementary_data[ACTION_TOKEN_MASK] = mask
        complementary_data[ACTION_CODE_TOKEN_MASK] = code_mask
        complementary_data[ACTION_TOKENS] = tokens
        new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
        return new_transition
@@ -434,7 +430,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
        """
        return self._paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens - tokens
-    def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Tokenizes the action tensor and creates a mask.
@@ -463,7 +459,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
        # The fast tokenizer expects action data and returns token IDs
        tokens_list = []
        masks_list = []
        code_masks_list = []
        for i in range(batch_size):
            # Tokenize single action (move to CPU first as tokenizer uses scipy which requires numpy)
@@ -481,82 +476,65 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
            if tokens.dim() > 1:
                tokens = tokens.flatten()
            action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens)
            bos_id = self._paligemma_tokenizer.bos_token_id
-            prompt_tokens = torch.tensor(
+            # add bos
                self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
                device=action.device,
            )
            end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device)
            code_start = 1 + len(prompt_tokens)
            code_end = code_start + len(action_code_tokens)
            tokens = torch.cat(
                [
                    torch.tensor([bos_id], device=action.device),
-                    prompt_tokens,
+                    torch.tensor(
-                    action_code_tokens,
+                        self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
-                    end_tokens,
+                        device=action.device,
                    ),
                    self._act_tokens_to_paligemma_tokens(tokens),
                    torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device),
                ]
            )
            code_mask = torch.zeros(len(tokens), dtype=torch.bool, device=action.device)
            code_mask[code_start:code_end] = True
            # Truncate or pad to max_action_tokens
            if len(tokens) > self.max_action_tokens:
                if not self.allow_truncation:
                    raise ValueError(
                        f"FAST action sequence has {len(tokens)} tokens, exceeding "
                        f"max_action_tokens={self.max_action_tokens}."
                    )
                logging.warning(
                    f"Token length ({len(tokens)}) exceeds max length ({self.max_action_tokens}), truncating. "
                    "Consider increasing the `max_action_tokens` in your model config if this happens frequently."
                )
                tokens = tokens[: self.max_action_tokens]
                code_mask = code_mask[: self.max_action_tokens]
                mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device)
            else:
                pad_len = self.max_action_tokens - len(tokens)
                mask = torch.cat(
                    [
                        torch.ones(len(tokens), dtype=torch.bool, device=action.device),
-                        torch.zeros(pad_len, dtype=torch.bool, device=action.device),
+                        torch.zeros(
                            self.max_action_tokens - len(tokens), dtype=torch.bool, device=action.device
                        ),
                    ]
                )
                code_mask = torch.nn.functional.pad(code_mask, (0, pad_len), value=False)
                # Pad tokens with zeros
-                tokens = torch.nn.functional.pad(tokens, (0, pad_len), value=0)
+                tokens = torch.nn.functional.pad(tokens, (0, self.max_action_tokens - len(tokens)), value=0)
            tokens_list.append(tokens)
            masks_list.append(mask)
            code_masks_list.append(code_mask)
        # Stack into batched tensors
        tokens_batch = torch.stack(tokens_list, dim=0)  # (B, max_action_tokens)
        masks_batch = torch.stack(masks_list, dim=0)  # (B, max_action_tokens)
        code_masks_batch = torch.stack(code_masks_list, dim=0)  # (B, max_action_tokens)
        # Remove batch dimension if input was single sample
        if single_sample:
            tokens_batch = tokens_batch.squeeze(0)
            masks_batch = masks_batch.squeeze(0)
            code_masks_batch = code_masks_batch.squeeze(0)
        # Move to the same device as the input
        if device is not None:
            tokens_batch = tokens_batch.to(device)
            masks_batch = masks_batch.to(device)
            code_masks_batch = code_masks_batch.to(device)
-        return tokens_batch, masks_batch, code_masks_batch
+        return tokens_batch, masks_batch
    def action(self, action: torch.Tensor) -> torch.Tensor:
        """
        This method is not used since we override __call__.
        Required by ActionProcessorStep ABC.
        """
-        tokens, _, _ = self._tokenize_action(action)
+        tokens, _ = self._tokenize_action(action)
        return tokens
    def get_config(self) -> dict[str, Any]:
@@ -572,9 +550,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
        config = {
            "trust_remote_code": self.trust_remote_code,
            "max_action_tokens": self.max_action_tokens,
            "fast_skip_tokens": self.fast_skip_tokens,
            "paligemma_tokenizer_name": self.paligemma_tokenizer_name,
            "allow_truncation": self.allow_truncation,
        }
        # Only save tokenizer_name if it was used to create the tokenizer
@@ -583,14 +558,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
        return config
    def save_artifacts(self, save_directory: Path) -> dict[str, str]:
        artifact_path = Path("action_tokenizer")
        save_pretrained = getattr(self.action_tokenizer, "save_pretrained", None)
        if save_pretrained is None:
            raise TypeError("Action tokenizer must implement save_pretrained() to save a portable pipeline.")
        save_pretrained(save_directory / artifact_path)
        return {"action_tokenizer_name": artifact_path.as_posix()}
    def transform_features(
        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
@@ -91,7 +91,7 @@ from lerobot.robots import so_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
 from lerobot.utils.device_utils import get_safe_torch_device
-from lerobot.utils.process import ProcessSignalHandler
+from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.transition import (
@@ -124,9 +124,7 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
    cfg.validate()
    display_pid = False
    if not use_threads(cfg):
-        import torch.multiprocessing as mp
+        ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
        mp.set_start_method("spawn")
        display_pid = True
    # Create logs directory to ensure it exists
@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
 )
 from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.io_utils import load_json, write_json
-from lerobot.utils.process import ProcessSignalHandler
+from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
    format_big_number,
@@ -123,9 +123,7 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
    # Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
    require_package("grpcio", extra="hilserl", import_name="grpc")
    if not use_threads(cfg):
-        import torch.multiprocessing as mp
+        ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
        mp.set_start_method("spawn")
    # Use the job_name from the config
    train(
@@ -21,8 +21,6 @@ from lerobot.utils.import_utils import make_device_from_device_class
 from .config import RobotConfig
 from .robot import Robot
 logger = logging.getLogger(__name__)
 def make_robot_from_config(config: RobotConfig) -> Robot:
    # TODO(Steven): Consider just using the make_device_from_device_class for all types
@@ -120,7 +118,7 @@ def ensure_safe_goal_position(
            }
    if warnings_dict:
-        logger.warning(
+        logging.warning(
            "Relative goal position magnitude had to be clamped to be safe.\n"
            f"{pformat(warnings_dict, indent=4)}"
        )
@@ -24,6 +24,7 @@ from __future__ import annotations
 import logging
 from dataclasses import dataclass, field
 from threading import Event
 from typing import TYPE_CHECKING
 import torch
@@ -47,6 +48,7 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
 from lerobot.robots import make_robot_from_config
 from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
 from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
 from lerobot.utils.import_utils import _peft_available, require_package
 from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
 from .inference import (
@@ -57,6 +59,12 @@ from .inference import (
 )
 from .robot_wrapper import ThreadSafeRobot
 if TYPE_CHECKING or _peft_available:
    from peft import PeftConfig, PeftModel
 else:
    PeftConfig = None
    PeftModel = None
 logger = logging.getLogger(__name__)
@@ -171,7 +179,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy
            revision=pretrained_revision,
        )
-    from peft import PeftConfig, PeftModel
+    require_package("peft", extra="peft")
    peft_path = policy_config.pretrained_path
    peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
@@ -294,12 +302,22 @@ def build_rollout_context(
    # ``observation_features`` values are either a tuple (camera shape) or the
    # ``float`` type itself used as a sentinel for scalar motor features —
    # see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
    # Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
    # scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
    # x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
    # filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
    # Pure-arm robots have no .vel state keys, so this is a no-op for them.
    observation_features_hw = {
        k: v
        for k, v in all_obs_features.items()
-        if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
+        if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
    }
-    action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
+    # Keep both joint-position (.pos) and base-velocity (.vel) action features so
    # mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
    # x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
    # a no-op for them. Without the .vel keys the base velocities are silently
    # dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
    action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
    # The action side is always needed: sync inference reads action names from
    # ``dataset_features[ACTION]`` to map policy tensors back to robot actions.
@@ -1,38 +0,0 @@
 # 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.
 """Policy-agnostic runtime for language-conditioned policies.
 Adapters registered in :mod:`lerobot.runtime.registry` are served by ``lerobot-rollout --language``.
 """
 from .adapter import BaseLanguageAdapter, GenerationConfig, LanguageDiagnostics
 from .language_runtime import (
    LanguageConditionedPolicyAdapter,
    LanguageConditionedRuntime,
    RuntimeState,
    Tick,
    TickClock,
 )
 __all__ = [
    "BaseLanguageAdapter",
    "GenerationConfig",
    "LanguageConditionedPolicyAdapter",
    "LanguageConditionedRuntime",
    "LanguageDiagnostics",
    "RuntimeState",
    "Tick",
    "TickClock",
 ]
@@ -1,165 +0,0 @@
 # 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.
 """Policy adapters for the language runtime.
 The base adapter owns generation control and diagnostics while subclasses provide policy-specific actions and text.
 """
 from __future__ import annotations
 import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
 from typing import Any
 from .language_runtime import RuntimeState
 _SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL)
@dataclass
 class GenerationConfig:
    """Text-generation settings fixed for the adapter's lifetime."""
    min_new_tokens: int = 0
    temperature: float = 0.0
    top_p: float = 1.0
    chunks_per_regen: int = 1  # regenerate the language context every N action chunks
    enable_memory: bool = True  # generate a running memory note on subtask change
    enable_subtask: bool = True  # generate the low-level subtask (off => use the given text directly)
@dataclass
 class LanguageDiagnostics:
    """Runtime-panel generation counters keyed by text kind."""
    last_raw: dict[str, str] = field(default_factory=dict)
    empty: dict[str, int] = field(default_factory=dict)
    repeat: int = 0
    def _bump(self, table: dict[str, int], kind: str) -> int:
        table[kind] = table.get(kind, 0) + 1
        return table[kind]
 class BaseLanguageAdapter(ABC):
    """Batteries-included adapter: generic high-level control, policy primitives abstract."""
    def __init__(self, policy: Any, gen: GenerationConfig | None = None) -> None:
        self.policy = policy
        self.gen = gen or GenerationConfig()
        self.diag = LanguageDiagnostics()
        self._chunks_until_regen = 0
    @abstractmethod
    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
        """Produce an action chunk from the observation + current language context."""
    @abstractmethod
    def generate_text(
        self,
        kind: str,
        observation: dict[str, Any] | None,
        state: RuntimeState,
        user_text: str | None = None,
    ) -> str:
        """Generate one text stream (``kind``) and return the decoded string."""
    def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
        """Throttled regeneration of the language context (subtask / memory / ...)."""
        if self._chunks_until_regen > 0:
            self._chunks_until_regen -= 1
            return
        self._chunks_until_regen = max(1, self.gen.chunks_per_regen) - 1
        self._regenerate_context(observation, state)
    def handle_interjection(
        self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
    ) -> None:
        """React to a mid-run user message by regenerating the plan."""
        out = self.generate_text("interjection", observation, state, user_text=user_text)
        plan = self.plan_from_text(out)
        if plan:
            state.set_context("plan", plan, label="plan")
    def plan_from_text(self, text: str) -> str:
        """Strip ``<say>`` speech markers from a generated plan."""
        plan, _speech = split_plan_and_say(text)
        return plan
    def _regenerate_context(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
        """Default hierarchy: regenerate the subtask, then memory when it changes.
        Override for a policy with a different language hierarchy.
        """
        if not self.gen.enable_subtask:
            # Preserve operator-provided subtasks in direct mode.
            return
        subtask = self._generate_filtered("subtask", observation, state)
        if subtask is None:
            return
        previous = state.language_context.get("subtask")
        if not state.set_context("subtask", subtask, label="subtask"):
            self.diag.repeat += 1
            return
        self.diag.repeat = 0
        if previous:
            state.extra["prior_subtask"] = previous
        if not self.gen.enable_memory:
            return
        memory = self._generate_filtered("memory", observation, state)
        if memory is not None:
            state.set_context("memory", memory, label="memory")
    def _generate_filtered(
        self, kind: str, observation: dict[str, Any] | None, state: RuntimeState
    ) -> str | None:
        """Generate one ``kind``, record diagnostics, and drop empty output."""
        text = self.generate_text(kind, observation, state)
        self.diag.last_raw[kind] = text or ""
        if not text:
            count = self.diag._bump(self.diag.empty, kind)
            if count == 1 or count % 5 == 0:
                state.log(f"  [info] {kind} gen returned empty (x{count})")
            return None
        return text
 class DirectTaskPolicyAdapter(BaseLanguageAdapter):
    """Adapter for flat policies whose preprocessors condition actions on the operator's task."""
    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
        return self.policy.predict_action_chunk(observation)
    def generate_text(
        self,
        kind: str,
        observation: dict[str, Any] | None,
        state: RuntimeState,
        user_text: str | None = None,
    ) -> str:
        return ""
 def split_plan_and_say(text: str) -> tuple[str, str]:
    """Split ``plan <say>speech</say>`` into ``(plan, speech)``."""
    if not text:
        return "", ""
    match = _SAY_RE.search(text)
    if not match:
        return text.strip(), ""
    speech = match.group(1).strip().strip('"').strip("'")
    plan = (text[: match.start()] + text[match.end() :]).strip()
    return plan, speech
@@ -1,349 +0,0 @@
 # 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.
 """Small reusable runtime for language-conditioned robot policies."""
 from __future__ import annotations
 import logging
 import threading
 import time
 from collections import deque
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from typing import Any, Protocol
 logger = logging.getLogger(__name__)
@dataclass
 class RuntimeState:
    """Explicit state shared by the runtime and policy adapter."""
    task: str = ""
    language_context: dict[str, str] = field(default_factory=dict)
    action_queue: deque[Any] = field(default_factory=deque)
    events: set[str] = field(default_factory=set)
    log_lines: list[str] = field(default_factory=list)
    mode: str = "action"
    stop: bool = False
    tick: Tick | None = None
    actions_dispatched: int = 0
    action_deadline: float | None = None
    extra: dict[str, Any] = field(default_factory=dict)
    revision: int = 0
    lock: Any = field(default_factory=threading.RLock, repr=False)
    def emit(self, event_name: str) -> None:
        self.events.add(event_name)
    def take_event(self, event_name: str) -> bool:
        if event_name not in self.events:
            return False
        self.events.remove(event_name)
        return True
    def log(self, line: str) -> None:
        self.log_lines.append(line)
    def set_context(self, key: str, value: str | None, *, label: str | None = None) -> bool:
        with self.lock:
            previous = self.language_context.get(key)
            if previous == value:
                return False
            if value is None:
                self.language_context.pop(key, None)
            else:
                self.language_context[key] = value
            self.revision += 1
            if label is not None and value:
                self.log(f"  {label}: {value}")
            return True
    def get(self, key: str, default: Any = None) -> Any:
        try:
            return self[key]
        except KeyError:
            return default
    def setdefault(self, key: str, default: Any = None) -> Any:
        current = self.get(key, None)
        if current is not None:
            return current
        self[key] = default
        return default
    def __getitem__(self, key: str) -> Any:
        if hasattr(self, key):
            return getattr(self, key)
        if key in self.extra:
            return self.extra[key]
        raise KeyError(key)
    def __setitem__(self, key: str, value: Any) -> None:
        with self.lock:
            if hasattr(self, key):
                if key == "mode" and self.mode != value:
                    self.revision += 1
                setattr(self, key, value)
            else:
                self.extra[key] = value
 class LanguageConditionedPolicyAdapter(Protocol):
    """Runtime policy contract, implemented directly or through ``BaseLanguageAdapter``."""
    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ...
    def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: ...
    def handle_interjection(
        self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
    ) -> None: ...
@dataclass
 class Tick:
    index: int
    monotonic_seconds: float
@dataclass
 class TickClock:
    max_rate_hz: float = 50.0
    _index: int = field(default=0, init=False)
    _last_seconds: float | None = field(default=None, init=False)
    def advance(self) -> Tick:
        period = 1.0 / max(self.max_rate_hz, 0.1)
        now = time.monotonic()
        if self._last_seconds is not None:
            sleep_for = (self._last_seconds + period) - now
            if sleep_for > 0:
                time.sleep(sleep_for)
                now = time.monotonic()
        self._last_seconds = now
        self._index += 1
        return Tick(index=self._index, monotonic_seconds=now)
@dataclass
 class _RateGate:
    hz: float
    _last_seconds: float | None = None
    def due(self, tick: Tick, *, force: bool = False) -> bool:
        if force:
            self._last_seconds = tick.monotonic_seconds
            return True
        period = 1.0 / max(self.hz, 1e-6)
        if self._last_seconds is None or tick.monotonic_seconds - self._last_seconds >= period:
            self._last_seconds = tick.monotonic_seconds
            return True
        return False
    def rearm(self) -> None:
        self._last_seconds = None
@dataclass
 class LanguageConditionedRuntime:
    """Generic tick loop for language-conditioned robot policies."""
    policy_adapter: LanguageConditionedPolicyAdapter
    observation_provider: Callable[[], dict[str, Any] | None] | None = None
    action_executor: Callable[[Any], None] | None = None
    event_collector: Callable[[RuntimeState], None] | None = None
    chunk_hz: float = 4.0
    ctrl_hz: float = 50.0
    high_level_hz: float = 1.0
    max_rate_hz: float = 50.0
    state: RuntimeState = field(default_factory=RuntimeState)
    _chunk_gate: _RateGate = field(init=False)
    _ctrl_gate: _RateGate = field(init=False)
    _language_gate: _RateGate = field(init=False)
    _stop: bool = field(default=False, init=False)
    _last_dispatch_seconds: float | None = field(default=None, init=False)
    def __post_init__(self) -> None:
        self._chunk_gate = _RateGate(self.chunk_hz)
        self._ctrl_gate = _RateGate(self.ctrl_hz)
        self._language_gate = _RateGate(self.high_level_hz)
    @property
    def policy(self) -> Any:
        return getattr(self.policy_adapter, "policy", self.policy_adapter)
    def set_task(self, task: str) -> None:
        with self.state.lock:
            if self.state.task != task:
                self.state.revision += 1
            self.state.task = task
            self.state.log(f"Task: {task}")
    def stop(self) -> None:
        self._stop = True
        self.state.stop = True
    def run(self, *, max_ticks: int | None = None) -> None:
        clock = TickClock(max_rate_hz=self.max_rate_hz)
        while not self._stop:
            tick = clock.advance()
            self._run_tick(tick)
            self._flush_logs()
            if self.state.stop:
                self._stop = True
            if max_ticks is not None and tick.index >= max_ticks:
                break
        self._on_shutdown()
    def step_once(self) -> list[str]:
        previous = self.state.tick.index if self.state.tick is not None else 0
        tick = Tick(index=previous + 1, monotonic_seconds=time.monotonic())
        self._run_tick(tick, force_rates=True)
        return list(self.state.log_lines)
    def _run_tick(self, tick: Tick, *, force_rates: bool = False) -> None:
        self.state.tick = tick
        self.state.log_lines = []
        if self.event_collector is not None:
            self.event_collector(self.state)
        self._handle_action_deadline()
        if self.state.stop:
            return
        self.maybe_update_language_state(force=force_rates)
        self.maybe_handle_user_events()
        self.maybe_enqueue_action_chunk(force=force_rates)
        self.dispatch_action(force=force_rates)
        self.state.events.clear()
    def _current_observation(self) -> dict[str, Any] | None:
        if self.observation_provider is None:
            return None
        try:
            return self.observation_provider()
        except Exception as exc:  # noqa: BLE001
            logger.debug("observation_provider failed: %s", exc)
            return None
    def maybe_update_language_state(self, *, force: bool = False) -> None:
        if self.state.mode != "action" or not self.state.task:
            return
        if self.state.action_queue:
            self._language_gate.rearm()
            return
        if self.state.tick is None or not self._language_gate.due(self.state.tick, force=force):
            return
        observation = self._current_observation()
        try:
            self.policy_adapter.update_language_state(observation, self.state)
        except Exception as exc:  # noqa: BLE001
            logger.warning("language update failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
            self.state.log(f"  [warn] language update failed: {type(exc).__name__}: {exc}")
    def maybe_handle_user_events(self) -> None:
        if self.state.take_event("user_interjection"):
            self._handle_user_interjection()
    def _handle_user_interjection(self) -> None:
        text = str(self.state.extra.get("recent_interjection") or "")
        if not text:
            return
        observation = self._current_observation()
        self.policy_adapter.handle_interjection(text, observation, self.state)
        self.state.extra["recent_interjection"] = None
    def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None:
        with self.state.lock:
            if self.state.mode != "action" or not self.state.task:
                return
            if self.state.action_queue:
                return
            if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force):
                return
            revision = self.state.revision
        observation = self._current_observation()
        if observation is None:
            return
        try:
            chunk = self.policy_adapter.select_action(observation, self.state)
        except Exception as exc:  # noqa: BLE001
            logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
            self.state.log(f"  [warn] select_action failed: {type(exc).__name__}: {exc}")
            return
        with self.state.lock:
            if (
                self.state.revision != revision
                or self.state.mode != "action"
                or self.state.stop
                or self._stop
            ):
                logger.info("Discarded an action chunk invalidated during inference.")
                return
            self._enqueue_chunk(chunk)
    def _enqueue_chunk(self, chunk: Any) -> None:
        if chunk is None:
            return
        chunk_iter = chunk[0] if getattr(chunk, "ndim", None) == 3 else chunk
        if getattr(chunk_iter, "ndim", None) == 1:
            chunk_iter = chunk_iter.unsqueeze(0)
        for step in chunk_iter:
            self.state.action_queue.append(step.unsqueeze(0) if hasattr(step, "unsqueeze") else step)
        try:
            self.state.extra["last_chunk_size"] = int(chunk_iter.shape[0])
        except Exception:  # noqa: BLE001
            self.state.extra["last_chunk_size"] = len(self.state.action_queue)
    def dispatch_action(self, *, force: bool = False) -> None:
        if self.state.mode != "action":
            self._last_dispatch_seconds = None
            return
        if self.state.tick is None or not self._ctrl_gate.due(self.state.tick, force=force):
            return
        queue = self.state.action_queue
        if not queue:
            self._last_dispatch_seconds = None
            return
        now = time.monotonic()
        if self._last_dispatch_seconds is None or self.ctrl_hz <= 0:
            n_to_pop = 1
        else:
            n_to_pop = max(1, min(len(queue), int(round((now - self._last_dispatch_seconds) * self.ctrl_hz))))
        self._last_dispatch_seconds = now
        latest = None
        for _ in range(n_to_pop):
            if not queue:
                break
            latest = queue.popleft()
            self.state.actions_dispatched += 1
        if latest is not None and self.action_executor is not None:
            self.action_executor(latest)
    def _handle_action_deadline(self) -> None:
        deadline = self.state.action_deadline
        if self.state.mode == "action" and deadline is not None and time.monotonic() >= deadline:
            self.state.mode = "paused"
            self.state.action_deadline = None
            self.state.action_queue.clear()
            self.state.log("timed action elapsed — paused")
    def _flush_logs(self) -> None:
        for line in self.state.log_lines:
            print(f"[runtime] {line}", flush=True)
    def _on_shutdown(self) -> None:
        self.state.action_queue.clear()
        print("[runtime] stopped", flush=True)
@@ -1,39 +0,0 @@
 # 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.
 """Lazy mapping from policy types to language-runtime adapters."""
 from __future__ import annotations
 import importlib
 from collections.abc import Callable
 from typing import Any
 _ADAPTERS: dict[str, str] = {
    "pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter",
    "pi05": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
    "molmoact2": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
 }
 def get_language_adapter_factory(policy_type: str) -> Callable[..., Any]:
    """Return the adapter class registered for ``policy_type``."""
    spec = _ADAPTERS.get(policy_type)
    if spec is None:
        raise ValueError(
            f"No language-runtime adapter registered for policy type {policy_type!r}. "
            f"Registered: {sorted(_ADAPTERS)}. Add an entry to lerobot.runtime.registry."
        )
    module_path, class_name = spec.split(":")
    return getattr(importlib.import_module(module_path), class_name)
@@ -1,406 +0,0 @@
 # 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.
 """RoboCasa backend for interactive language-conditioned rollouts.
 It reuses the eval observation/action pipeline while prompts control a persistent selected scene.
 """
 from __future__ import annotations
 import logging
 from collections.abc import Callable
 from datetime import datetime
 from pathlib import Path
 from typing import Any
 import numpy as np
 import torch
 from lerobot.utils.io_utils import StreamingVideoWriter
 from lerobot.utils.video_annotation import annotate_frame
 logger = logging.getLogger(__name__)
 def _short_cam_name(cam: str) -> str:
    """Human-friendly view label for a RoboCasa camera name."""
    c = cam.replace("robot0_", "")
    return {
        "agentview_left": "left",
        "agentview_right": "right",
        "eye_in_hand": "wrist",
    }.get(c, c)
 def _label_panel(img: np.ndarray, label: str) -> np.ndarray:
    """Draw a small camera-view label in the bottom-left corner of a panel."""
    try:
        import cv2  # noqa: PLC0415
    except ImportError:
        return img
    y = img.shape[0] - 6
    cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 3, cv2.LINE_AA)
    cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 0), 1, cv2.LINE_AA)
    return img
 # Two workers avoid broken single-worker EGL rendering; only env 0 is displayed.
 _SIM_N_ENVS = 2
 def create_sim_env(
    *,
    task: str,
    split: str | None,
    obj_registries: list[str],
    seed: int | None,
    render_size: int = 384,
 ) -> tuple[Any, dict]:
    """Create and reset the vectorized RoboCasa environment before CUDA initializes.
    Two workers keep EGL stable, while only env 0 is driven and displayed.
    """
    from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig  # noqa: PLC0415
    # The policy resizes inputs, so render_size only affects display quality and cost.
    env_cfg = RoboCasaEnvConfig(
        task=task,
        split=split,
        obj_registries=list(obj_registries),
        observation_height=render_size,
        observation_width=render_size,
    )
    # Keep one kitchen alive across sequential prompts.
    envs = env_cfg.create_envs(
        n_envs=_SIM_N_ENVS,
        use_async_envs=True,
        terminate_on_success=False,
        horizon=100_000,
    )
    env = envs[next(iter(envs))][0]
    logger.info("[sim] resetting RoboCasa scene task=%r split=%r (n_envs=%d)", task, split, _SIM_N_ENVS)
    seeds = None if seed is None else [seed + i for i in range(_SIM_N_ENVS)]
    obs, _ = env.reset(seed=seeds)
    return env, obs
 def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) -> Any:
    """Start an MJPEG server that shows a placeholder until ``get_frame`` returns frames."""
    import io  # noqa: PLC0415
    import threading  # noqa: PLC0415
    import time  # noqa: PLC0415
    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer  # noqa: PLC0415
    from PIL import Image  # noqa: PLC0415
    _placeholder = Image.new("RGB", (256, 256), (17, 17, 17))
    class _Handler(BaseHTTPRequestHandler):
        def log_message(self, *args):  # silence per-request logging
            pass
        def do_GET(self):  # noqa: N802
            if self.path in ("/", "/index.html"):
                self.send_response(200)
                self.send_header("Content-Type", "text/html")
                self.end_headers()
                self.wfile.write(
                    b"<html><body style='margin:0;background:#111;text-align:center'>"
                    b"<img src='/stream' style='max-width:100vw;max-height:100vh;"
                    b"image-rendering:pixelated'></body></html>"
                )
                return
            if self.path != "/stream":
                self.send_response(404)
                self.end_headers()
                return
            self.send_response(200)
            self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
            self.end_headers()
            try:
                while True:
                    frame = get_frame()
                    buf = io.BytesIO()
                    img = Image.fromarray(frame) if frame is not None else _placeholder
                    img.save(buf, format="JPEG", quality=80)
                    data = buf.getvalue()
                    self.wfile.write(
                        b"--frame\r\nContent-Type: image/jpeg\r\nContent-Length: "
                        + str(len(data)).encode()
                        + b"\r\n\r\n"
                        + data
                        + b"\r\n"
                    )
                    time.sleep(0.05)
            except (BrokenPipeError, ConnectionResetError):
                pass
    try:
        # Bind all interfaces intentionally so the viewer remains reachable
        # through the documented SSH port-forwarding workflow.
        server = ThreadingHTTPServer(("0.0.0.0", port), _Handler)  # nosec B104
    except OSError as exc:
        logger.warning("[sim] could not start live stream on port %d: %s", port, exc)
        print(f"[runtime] WARNING: live stream port {port} unavailable ({exc})", flush=True)
        return None
    threading.Thread(target=server.serve_forever, daemon=True, name="sim-mjpeg").start()
    print(
        f"[runtime] live view: http://localhost:{port}  "
        f"(over SSH: ssh -L {port}:localhost:{port} <host>) — loading until scene is ready",
        flush=True,
    )
    return server
 class RoboCasaSimBackend:
    """Expose a RoboCasa environment through the runtime observation/action contract.
    The environment must be created before the policy initializes CUDA.
    """
    def __init__(
        self,
        *,
        env: Any,
        last_obs: dict,
        task: str,
        seed: int | None,
        device: str,
        preprocessor: Any,
        postprocessor: Any,
        record: bool = True,
        output_dir: str | None = None,
        view_cams: list[str] | None = None,
    ) -> None:
        self.env = env
        self._last_obs = last_obs
        self._scene_task = task
        self._view_cams = view_cams or [
            "robot0_agentview_left",
            "robot0_eye_in_hand",
            "robot0_agentview_right",
        ]
        self.device = torch.device(device) if isinstance(device, str) else device
        self.preprocessor = preprocessor
        self.postprocessor = postprocessor
        self.seed = seed
        self.record = record
        self.output_dir = Path(output_dir) if output_dir else Path("outputs/runtime_sim")
        self._video_writer: StreamingVideoWriter | None = None
        self._video_path: Path | None = None
        self._live_counter = 0
        self._latest_frame: np.ndarray | None = None
        self._stream_server: Any = None
        self._reset_count = 0
        # Bind these after runtime construction for live annotations.
        self._task_getter: Callable[[], str | None] | None = None
        self._subtask_getter: Callable[[], str | None] | None = None
        self._memory_getter: Callable[[], str | None] | None = None
        logger.info("[sim] scene ready — task_description=%r", self._scene_description())
    def bind_runtime(self, runtime: Any) -> None:
        """Wire live task/subtask/memory getters from the runtime state."""
        self._task_getter = lambda: runtime.state.get("task")
        self._subtask_getter = lambda: runtime.state.language_context.get("subtask")
        self._memory_getter = lambda: (runtime.state.get("language_context") or {}).get("memory")
    def _scene_description(self) -> str:
        try:
            return str(self.env.get_attr("task_description")[0]) or self._scene_task
        except Exception:  # noqa: BLE001
            return self._scene_task
    def _current_task(self) -> str:
        task = self._task_getter() if self._task_getter else None
        return task or self._scene_description() or self._scene_task
    def reset_scene(self) -> None:
        """Re-roll the kitchen: reset the env to a fresh scene (new layout/style).
        Uses a new seed each call so ``/reset`` explores different kitchens.
        """
        self._reset_count += 1
        n = self.env.num_envs
        if self.seed is None:
            seeds = None
        else:
            base = self.seed + self._reset_count * 1000
            seeds = [base + i for i in range(n)]
        obs, _ = self.env.reset(seed=seeds)
        self._last_obs = obs
        logger.info("[sim] scene reset (#%d)", self._reset_count)
    def _env0_obs(self) -> dict:
        """Slice env 0 out of the batched vec-env observation (batch of 1)."""
        raw = self._last_obs or {}
        pixels = raw.get("pixels")
        out: dict[str, Any] = {}
        if isinstance(pixels, dict):
            out["pixels"] = {k: np.asarray(v)[0:1] for k, v in pixels.items()}
        agent_pos = raw.get("agent_pos")
        if agent_pos is not None:
            out["agent_pos"] = np.asarray(agent_pos)[0:1]
        return out
    def observation_provider(self) -> dict | None:
        from lerobot.envs.utils import preprocess_observation  # noqa: PLC0415
        try:
            obs = preprocess_observation(self._env0_obs())
        except Exception as exc:  # noqa: BLE001
            logger.warning("[sim] preprocess_observation failed: %s", exc)
            return None
        # The adapter later replaces this recipe input with its generated subtask.
        obs["task"] = [self._current_task()]
        if self.preprocessor is not None:
            try:
                obs = self.preprocessor(obs)
            except Exception as exc:  # noqa: BLE001
                logger.warning("[sim] preprocessor failed: %s", exc)
                return None
        return {
            k: (v.to(self.device) if isinstance(v, torch.Tensor) else v)
            for k, v in obs.items()
            if isinstance(k, str) and k.startswith("observation.")
        }
    def action_executor(self, action: Any) -> None:
        try:
            if self.postprocessor is not None:
                action = self.postprocessor(action)
            if isinstance(action, torch.Tensor):
                if action.ndim > 1 and action.shape[0] == 1:
                    action = action.squeeze(0)
                action = action.detach().to("cpu").numpy()
            # Tile env 0's action because the extra workers exist only for EGL stability.
            action_row = np.asarray(action, dtype=np.float32).reshape(-1)
            action_np = np.tile(action_row, (self.env.num_envs, 1))
            obs, _reward, terminated, truncated, _info = self.env.step(action_np)
            self._last_obs = obs
            self._capture_frame()
            # AsyncVectorEnv resets terminated sub-environments automatically.
            if bool(np.any(terminated)) or bool(np.any(truncated)):
                logger.info("[sim] episode ended — scene auto-reset")
        except Exception as exc:  # noqa: BLE001
            logger.error("[sim] env.step failed: %s", exc, exc_info=True)
    def _multiview_frame(self) -> np.ndarray | None:
        """Label and compose env 0's existing observation views without extra rendering."""
        pixels = (self._last_obs or {}).get("pixels")
        if not isinstance(pixels, dict) or not pixels:
            return None
        panels: list[np.ndarray] = []
        for cam in self._view_cams:
            v = pixels.get(cam)
            if v is None:
                continue
            img = np.asarray(v)
            if img.ndim == 4:  # (n_envs, H, W, C) -> env 0
                img = img[0]
            if img.ndim != 3 or img.shape[-1] != 3:
                continue
            panels.append(_label_panel(np.ascontiguousarray(img.astype(np.uint8)), _short_cam_name(cam)))
        if not panels:
            return None
        h = min(p.shape[0] for p in panels)
        panels = [p[:h] for p in panels]
        return np.concatenate(panels, axis=1)
    def _capture_frame(self) -> None:
        frame = self._multiview_frame()
        if frame is None:  # fallback to single env.render()
            try:
                rendered = self.env.call("render")[0]
                if isinstance(rendered, np.ndarray) and rendered.ndim == 3:
                    frame = rendered
            except Exception as exc:  # noqa: BLE001
                logger.debug("[sim] render failed: %s", exc)
        if frame is None:
            return
        subtask = self._subtask_getter() if self._subtask_getter else None
        memory = self._memory_getter() if self._memory_getter else None
        annotated = annotate_frame(
            frame,
            (("Task", self._current_task()), ("Subtask", subtask), ("Memory", memory)),
        )
        self._latest_frame = annotated  # served by the live MJPEG stream
        self._write_live_frame(annotated)
        if self.record:
            self._write_recording_frame(annotated)
    def _write_recording_frame(self, frame: np.ndarray) -> None:
        try:
            if self._video_writer is None:
                self.output_dir.mkdir(parents=True, exist_ok=True)
                stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                self._video_path = self.output_dir / f"sim_{stamp}.mp4"
                fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20))
                self._video_writer = StreamingVideoWriter(self._video_path, fps)
            self._video_writer.add_frame(frame)
        except Exception as exc:  # noqa: BLE001
            logger.warning("[sim] video encoding failed: %s", exc)
            self.record = False
    def _write_live_frame(self, frame: np.ndarray) -> None:
        """Write a rolling latest.png every few frames for live viewing over SSH.
        Open ``{output_dir}/latest.png`` in an editor/viewer and refresh to watch
        the rollout in near-real-time without a GUI window. Written atomically
        (temp + replace) so a reader never sees a half-written file.
        """
        self._live_counter += 1
        if self._live_counter % 3 != 0:
            return
        try:
            import os  # noqa: PLC0415
            from PIL import Image  # noqa: PLC0415
            self.output_dir.mkdir(parents=True, exist_ok=True)
            tmp = self.output_dir / ".latest.tmp.png"
            Image.fromarray(frame).save(tmp)
            os.replace(tmp, self.output_dir / "latest.png")
        except Exception as exc:  # noqa: BLE001
            logger.debug("[sim] live frame write failed: %s", exc)
    def _flush_video(self) -> None:
        if self._video_writer is None:
            return
        writer = self._video_writer
        self._video_writer = None
        try:
            writer.close()
            logger.info("[sim] wrote video (%d frames) to %s", writer.frames_written, self._video_path)
            print(f"[runtime] sim video saved to {self._video_path}", flush=True)
        except Exception as exc:  # noqa: BLE001
            logger.warning("[sim] video close failed: %s", exc)
    def attach_stream_server(self, server: Any) -> None:
        """Attach an already-running MJPEG server so disconnect() can stop it."""
        self._stream_server = server
    def disconnect(self) -> None:
        """Match the robot backend's cleanup contract."""
        if self._stream_server is not None:
            try:
                self._stream_server.shutdown()
            except Exception as exc:  # noqa: BLE001
                logger.debug("[sim] stream server shutdown raised %s", exc)
        self._flush_video()
        try:
            self.env.close()
        except Exception as exc:  # noqa: BLE001
            logger.debug("[sim] env.close raised %s", exc)
@@ -36,6 +36,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
 import argparse
 import concurrent.futures
 import logging
 import os
 from pathlib import Path
 import numpy as np
@@ -52,6 +53,7 @@ from lerobot.datasets import (
    get_feature_stats,
    write_stats,
 )
 from lerobot.datasets.compute_stats import sample_indices
 from lerobot.utils.utils import init_logging
@@ -77,12 +79,14 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
    return False
-def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
+def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
    """Process a single episode and return its statistics.
    Args:
        dataset: The LeRobot dataset
        episode_idx: Index of the episode to process
        use_sampling: If True, sub-sample image/video frames per episode to bound
            memory. If False, use every frame (exact, higher memory).
    Returns:
        Dictionary containing episode statistics
@@ -92,16 +96,31 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
    start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
    end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]
-    collected_data: dict[str, list] = {}
+    episode_len = end_idx - start_idx
    for idx in range(start_idx, end_idx):
        item = dataset[idx]
        for key, value in item.items():
            if key not in dataset.features:
                continue
-            if key not in collected_data:
+    # Images/video are the memory hog, so sub-sample those frames per episode;
-                collected_data[key] = []
+    # numeric columns are cheap, so read them in full (exact).
-            collected_data[key].append(value)
+    image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
    numeric_keys = [
        k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
    ]
    collected_data: dict[str, list] = {}
    # Numeric features: every frame, read directly from the underlying table.
    if numeric_keys:
        numeric_cols = dataset.hf_dataset.select_columns(numeric_keys)[start_idx:end_idx]
        for key in numeric_keys:
            collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]
    # Image/video features: decode only a sampled subset of frames.
    if image_keys:
        sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
        for offset in sampled_offsets:
            item = dataset[start_idx + offset]
            for key in image_keys:
                if key in item:
                    collected_data.setdefault(key, []).append(item[key])
    ep_stats = {}
    for key, data_list in collected_data.items():
@@ -131,11 +150,13 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int) -> dict:
    return ep_stats
-def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dict]:
+def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
    """Compute quantile statistics for all episodes in the dataset.
    Args:
        dataset: The LeRobot dataset to compute statistics for
        use_sampling: If True, sub-sample image/video frames per episode to bound
            memory. If False, use every frame (exact, higher memory).
    Returns:
        Dictionary containing aggregated statistics with quantiles
@@ -153,15 +174,15 @@ def compute_quantile_stats_for_dataset(dataset: LeRobotDataset) -> dict[str, dic
    if has_videos:
        logging.info("Dataset contains video keys - using sequential processing for thread safety")
        for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
-            ep_stats = process_single_episode(dataset, episode_idx)
+            ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
            episode_stats_list.append(ep_stats)
    else:
        logging.info("Dataset has no video keys - using parallel processing for better performance")
-        max_workers = min(dataset.num_episodes, 16)
+        max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_to_episode = {
-                executor.submit(process_single_episode, dataset, episode_idx): episode_idx
+                executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
                for episode_idx in range(dataset.num_episodes)
            }
@@ -188,6 +209,7 @@ def augment_dataset_with_quantile_stats(
    repo_id: str,
    root: str | Path | None = None,
    overwrite: bool = False,
    use_sampling: bool = True,
 ) -> None:
    """Augment a dataset with quantile statistics if they are missing.
@@ -195,6 +217,8 @@ def augment_dataset_with_quantile_stats(
        repo_id: Repository ID of the dataset
        root: Local root directory for the dataset
        overwrite: Overwrite existing quantile statistics if they already exist
        use_sampling: If True, sub-sample image/video frames per episode to bound
            memory. If False, use every frame (exact, higher memory).
    """
    logging.info(f"Loading dataset: {repo_id}")
    dataset = LeRobotDataset(
@@ -208,7 +232,7 @@ def augment_dataset_with_quantile_stats(
    logging.info("Dataset does not contain quantile statistics. Computing them now...")
-    new_stats = compute_quantile_stats_for_dataset(dataset)
+    new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
    logging.info("Updating dataset metadata with new quantile statistics")
    dataset.meta.stats = new_stats
@@ -248,6 +272,14 @@ def main():
        action="store_true",
        help="Overwrite existing quantile statistics if they already exist",
    )
    parser.add_argument(
        "--no-sampling",
        action="store_true",
        help=(
            "Compute stats over every frame (exact, higher memory). By default, "
            "image/video frames are sub-sampled per episode to bound memory."
        ),
    )
    args = parser.parse_args()
    root = Path(args.root) if args.root else None
@@ -258,6 +290,7 @@ def main():
        repo_id=args.repo_id,
        root=root,
        overwrite=args.overwrite,
        use_sampling=not args.no_sampling,
    )
@@ -94,6 +94,8 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
 from lerobot.utils.constants import HF_LEROBOT_HOME
 from lerobot.utils.utils import flatten_dict, init_logging
 logger = logging.getLogger(__name__)
 V21 = "v2.1"
 V30 = "v3.0"
@@ -476,11 +478,11 @@ def convert_dataset(
    # First check if the dataset already has a v3.0 version
    if root is None and not force_conversion:
        try:
-            print("Trying to download v3.0 version of the dataset from the hub...")
+            logger.info("Trying to download v3.0 version of the dataset from the hub...")
            snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
            return
        except Exception:
-            print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
+            logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
    # Set root based on whether local dataset path is provided
    use_local_dataset = False
@@ -488,7 +490,7 @@ def convert_dataset(
    if root.exists():
        validate_local_dataset_version(root)
        use_local_dataset = True
-        print(f"Using local dataset at {root}")
+        logger.info(f"Using local dataset at {root}")
    old_root = root.parent / f"{root.name}_old"
    new_root = root.parent / f"{root.name}_v30"
@@ -523,7 +525,7 @@ def convert_dataset(
        try:
            hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
        except (HTTPError, RevisionNotFoundError) as e:
-            print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
+            logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
            pass
        hub_api.delete_files(
            delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],
@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
    repo_id = cfg.new_repo_id or cfg.repo_id
    commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
    api = HfApi()
-    print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
+    logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
    api.create_repo(
        repo_id=repo_id,
        repo_type="dataset",
        private=cfg.push_private,
        exist_ok=True,
    )
-    print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
+    logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
    commit_info = api.upload_folder(
        folder_path=str(root),
        repo_id=repo_id,
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
        # at the source dataset; a fresh card is generated below instead.
        ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
    )
-    print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
+    logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
    dataset_info = load_info(root)
    card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
@@ -200,14 +200,13 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
        with suppress(RevisionNotFoundError):
            api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
        api.create_tag(**tag_kwargs)
-        print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
+        logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
    except Exception as exc:  # noqa: BLE001
-        print(
+        logger.warning(
            f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
            "Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
            "Run: from huggingface_hub import HfApi; "
-            f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
+            f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
            flush=True,
        )
@@ -89,6 +89,8 @@ from lerobot.datasets import LeRobotDataset
 from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
 from lerobot.utils.utils import init_logging
 logger = logging.getLogger(__name__)
 DEFAULT_FOXGLOVE_PORT = 8765
 DEFAULT_RERUN_PORT = 9090
@@ -299,7 +301,7 @@ def visualize_dataset(
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
-            print("Ctrl-C received. Exiting.")
+            logger.info("Ctrl-C received. Exiting.")
 def main():
@@ -62,7 +62,7 @@ from dataclasses import asdict
 from functools import partial
 from pathlib import Path
 from pprint import pformat
-from typing import Any, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict
 import einops
 import gymnasium as gym
@@ -87,26 +87,21 @@ from lerobot.processor import PolicyProcessorPipeline
 from lerobot.types import PolicyAction
 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.import_utils import _peft_available, register_third_party_plugins, require_package
 from lerobot.utils.io_utils import write_video
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
    init_logging,
    inside_slurm,
 )
-from lerobot.utils.video_annotation import annotate_frame
+
 if TYPE_CHECKING or _peft_available:
    from peft import PeftModel
 else:
    PeftModel = None
-def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray:
+logger = logging.getLogger(__name__)
    """Overlay the high-level task and predicted subtask onto rendered frames.
    ``frames`` is ``(n_envs, H, W, C)`` uint8. Best-effort: if OpenCV isn't
    available the frames are returned unchanged so eval never fails over a
    visualization concern.
    """
    if frames.ndim != 4 or frames.shape[-1] != 3:
        return frames
    return np.stack([annotate_frame(frame, (("Task", task), ("Subtask", subtask))) for frame in frames])
 def _env_features_to_dataset_features(env_features: dict) -> dict:
@@ -457,13 +452,11 @@ def eval_policy(
        exc = ValueError(
            f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
        )
-        try:
+        if not _peft_available:
-            from peft import PeftModel
+            raise exc
-
+        require_package("peft", extra="peft")
        if not isinstance(policy, PeftModel):
            raise exc
        except ImportError:
            raise exc from None
    start = time.time()
    # Preserve the mode for direct callers. eval_policy_all scopes the mode
@@ -490,36 +483,11 @@ def eval_policy(
            return
        n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs)
        if isinstance(env, gym.vector.SyncVectorEnv):
-            frames = np.stack([env.envs[i].render() for i in range(n_to_render_now)])  # noqa: B023
+            ep_frames.append(np.stack([env.envs[i].render() for i in range(n_to_render_now)]))  # noqa: B023
        elif hasattr(env, "call"):
            # Here we must render all frames and discard any we don't need.
            # Covers AsyncVectorEnv and _LazyAsyncVectorEnv (which wraps one).
-            frames = np.stack(env.call("render")[:n_to_render_now])
+            ep_frames.append(np.stack(env.call("render")[:n_to_render_now]))
        else:
            return
        # Overlay the high-level task and (for hierarchical policies like
        # pi052) the predicted low-level subtask onto each frame. Both are
        # best-effort: missing values just skip that line.
        try:
            tasks = list(env.call("task_description"))
        except (AttributeError, NotImplementedError):
            try:
                tasks = list(env.call("task"))
            except (AttributeError, NotImplementedError):
                tasks = None
        subtasks = getattr(policy, "last_subtasks", None)
        annotated = []
        for i in range(frames.shape[0]):
            subtask_i = subtasks[i] if subtasks is not None and i < len(subtasks) else None
            annotated.append(
                _annotate_eval_frames(
                    frames[i : i + 1],
                    tasks[i] if tasks is not None and i < len(tasks) else None,
                    subtask_i,
                )[0]
            )
        ep_frames.append(np.stack(annotated))
    if max_episodes_rendered > 0:
        video_paths: list[str] = []
@@ -596,7 +564,7 @@ def eval_policy(
        if seeds:
            all_seeds.extend(seeds)
        else:
-            all_seeds.append(None)
+            all_seeds.extend([None] * env.num_envs)
        # FIXME: episode_data is either None or it doesn't exist
        if return_episode_data:
@@ -834,13 +802,13 @@ def eval_main(cfg: EvalPipelineConfig):
            recording_repo_id=cfg.eval.recording_repo_id,
            recording_private=cfg.eval.recording_private,
        )
-        print("Overall Aggregated Metrics:")
+        logger.info("Overall Aggregated Metrics:")
-        print(info["overall"])
+        logger.info(info["overall"])
        # Print per-suite stats
        for task_group, task_group_info in info.items():
-            print(f"\nAggregated Metrics for {task_group}:")
+            logger.info(f"\nAggregated Metrics for {task_group}:")
-            print(task_group_info)
+            logger.info(task_group_info)
    # Close all vec envs
    close_envs(envs)
@@ -40,6 +40,7 @@ from PIL import Image
 from lerobot.cameras import ColorMode
 from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
 from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
 from lerobot.utils.utils import init_logging
 logger = logging.getLogger(__name__)
@@ -285,6 +286,8 @@ def save_images_from_all_cameras(
 def main():
    init_logging()
    parser = argparse.ArgumentParser(
        description="Unified camera utility script for listing cameras and capturing images."
    )
@@ -151,7 +151,6 @@ Usage examples
 """
 import logging
 import sys
 from lerobot.cameras.opencv import OpenCVCameraConfig  # noqa: F401
 from lerobot.cameras.realsense import RealSenseCameraConfig  # noqa: F401
@@ -166,6 +165,7 @@ from lerobot.robots import (  # noqa: F401
    earthrover_mini_plus,
    hope_jr,
    koch_follower,
    lekiwi,
    omx_follower,
    openarm_follower,
    reachy2,
@@ -242,69 +242,10 @@ def rollout(cfg: RolloutConfig):
    logger.info("Rollout finished")
-_LANGUAGE_RUNTIME_FLAGS = {
+def main():
-    "--language",
+    """CLI entry point for ``lerobot-rollout``."""
    "--no_robot",
    "--sim",
    "--direct_subtask",
    "--sim.direct_subtask",
    "--disable_memory",
    "--fp8",
 }
 _LANGUAGE_RUNTIME_PREFIXES = (
    "--sim.",
    "--chunk_hz",
    "--ctrl_hz",
    "--high_level_hz",
    "--subtask_chunks_per_gen",
    "--text_min_new_tokens",
    "--text_temperature",
    "--text_top_p",
 )
 def _uses_language_runtime(argv: list[str]) -> bool:
    """Return whether *argv* selects the interactive language runtime.
    ``--language`` is the explicit selector for real-robot runs whose other
    options overlap with the standard rollout CLI. Language-only options also
    select it automatically, which keeps the former language-runtime examples
    working after replacing their command name with ``lerobot-rollout``.
    """
    return any(
        arg.split("=", 1)[0] in _LANGUAGE_RUNTIME_FLAGS or arg.startswith(_LANGUAGE_RUNTIME_PREFIXES)
        for arg in argv
    )
 def main(argv: list[str] | None = None):
    """CLI entry point for ``lerobot-rollout``.
    Standard policy deployment continues through :class:`RolloutConfig`.
    Interactive language-conditioned and RoboCasa runs share this entry point
    and are selected with ``--language`` or any language-runtime-only option.
    """
    register_third_party_plugins()
-    cli_args = list(sys.argv[1:] if argv is None else argv)
+    rollout()
    if _uses_language_runtime(cli_args):
        from lerobot.runtime.cli import run as run_language_runtime
        # ``--language`` is a dispatcher flag, not part of the runtime's own
        # argparse surface. All other arguments pass through unchanged.
        runtime_args = [arg for arg in cli_args if arg != "--language"]
        return run_language_runtime(runtime_args, prog="lerobot-rollout")
    if argv is None:
        return rollout()
    # draccus reads sys.argv. Supporting an explicit argv keeps this entry
    # point easy to smoke-test and mirrors the language-runtime branch above.
    previous_argv = sys.argv
    try:
        sys.argv = [previous_argv[0], *cli_args]
        return rollout()
    finally:
        sys.argv = previous_argv
 if __name__ == "__main__":
@@ -22,7 +22,8 @@ import dataclasses
 import logging
 import sys
 import time
-from contextlib import nullcontext
+from collections.abc import Iterator
 from contextlib import contextmanager, nullcontext
 from pprint import pformat
 from typing import TYPE_CHECKING, Any
@@ -57,7 +58,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.rewards import make_reward_pre_post_processors
 from lerobot.utils.collate import lerobot_collate_fn
-from lerobot.utils.import_utils import register_third_party_plugins
+from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
@@ -68,9 +69,28 @@ from lerobot.utils.utils import (
    inside_slurm,
 )
 if TYPE_CHECKING or _peft_available:
    from peft import PeftModel
 else:
    PeftModel = None
 from .lerobot_eval import eval_policy_all
@contextmanager
 def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
    """Create evaluation environments for one run and always dispose of them."""
    envs = make_env(
        cfg.env,
        n_envs=cfg.eval.batch_size,
        use_async_envs=cfg.eval.use_async_envs,
    )
    try:
        yield envs
    finally:
        close_envs(envs)
 def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
    """Return worker-only DataLoader options, disabling them for single-process loading."""
    workers_enabled = cfg.num_workers > 0
@@ -207,8 +227,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
    if cfg.job.is_remote:
        return submit_to_hf(cfg)
    from lerobot.utils.import_utils import require_package
    require_package("accelerate", extra="training")
    from accelerate import Accelerator
    from accelerate.utils import DistributedDataParallelKwargs, DistributedType
@@ -277,14 +295,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
    if not is_main_process:
        dataset, eval_dataset = make_train_eval_datasets(cfg)
    # Create environment used for evaluating checkpoints during training on simulation data.
    # On real-world data, no need to create an environment as evaluations are done outside train.py,
    # using the eval.py instead, with gym_dora environment and dora-rs.
    eval_env = None
    if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
        logging.info("Creating env")
        eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
    if cfg.is_reward_model_training:
        if is_main_process:
            logging.info("Creating reward model")
@@ -312,7 +322,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
    if cfg.peft is not None:
        if cfg.is_reward_model_training:
            raise ValueError("PEFT is only supported for policy training. ")
-        from peft import PeftModel
+        require_package("peft", extra="peft")
        if isinstance(policy, PeftModel):
            logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
@@ -692,7 +702,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
            if is_main_process:
                step_id = get_step_identifier(step, cfg.steps)
                logging.info(f"Eval policy at step {step}")
-                with torch.no_grad(), accelerator.autocast():
+                with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
                    eval_info = eval_policy_all(
                        envs=eval_env,  # dict[suite][task_id] -> vec_env
                        policy=accelerator.unwrap_model(policy),
@@ -740,9 +750,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
    if is_main_process:
        progbar.close()
    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:
@@ -45,6 +45,7 @@ lerobot-train-tokenizer \
 """
 import json
 import logging
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -63,6 +64,9 @@ else:
 from lerobot.configs import NormalizationMode, parser
 from lerobot.datasets import LeRobotDataset
 from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.utils import init_logging
 logger = logging.getLogger(__name__)
@dataclass
@@ -274,11 +278,8 @@ def process_episode(args):
        return action_chunks
-    except Exception as e:
+    except Exception:
-        print(f"Error processing episode {ep_idx}: {e}")
+        logger.exception("Error processing episode %s", ep_idx)
        import traceback
        traceback.print_exc()
        return None
@@ -300,10 +301,10 @@ def train_fast_tokenizer(
    Returns:
        Trained FAST tokenizer
    """
-    print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
+    logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
-    print(f"Action chunk shape: {action_chunks.shape}")
+    logger.info(f"Action chunk shape: {action_chunks.shape}")
-    print(f"Vocab size: {vocab_size}")
+    logger.info(f"Vocab size: {vocab_size}")
-    print(f"DCT scale: {scale}")
+    logger.info(f"DCT scale: {scale}")
    # download the tokenizer source code (not pretrained weights)
    # we'll train a new tokenizer on our own data
@@ -314,7 +315,7 @@ def train_fast_tokenizer(
    # train the new tokenizer on our action data using .fit()
    # this trains the BPE tokenizer on DCT coefficients
-    print("Training new tokenizer (this may take a few minutes)...")
+    logger.info("Training new tokenizer (this may take a few minutes)...")
    tokenizer = base_tokenizer.fit(
        action_data_list,
        scale=scale,
@@ -322,21 +323,21 @@ def train_fast_tokenizer(
        time_horizon=action_chunks.shape[1],  # action_horizon
        action_dim=action_chunks.shape[2],  # encoded dimensions
    )
-    print("✓ Tokenizer training complete!")
+    logger.info("✓ Tokenizer training complete!")
    # validate it works
    sample_chunk = action_chunks[0]
    encoded = tokenizer(sample_chunk[None])[0]
    if isinstance(encoded, list):
        encoded = np.array(encoded)
-    print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
+    logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
    return tokenizer
 def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
    """Compute compression statistics."""
-    print("\nComputing compression statistics...")
+    logger.info("\nComputing compression statistics...")
    # sample for stats (use max 1000 chunks for speed)
    sample_size = min(1000, len(action_chunks))
@@ -366,12 +367,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
        "max_token_length": float(np.max(token_lengths)),
    }
-    print("Compression Statistics:")
+    logger.info("Compression Statistics:")
-    print(f"  Average compression ratio: {stats['compression_ratio']:.2f}x")
+    logger.info(f"  Average compression ratio: {stats['compression_ratio']:.2f}x")
-    print(f"  Mean token length: {stats['mean_token_length']:.1f}")
+    logger.info(f"  Mean token length: {stats['mean_token_length']:.1f}")
-    print(f"  P99 token length: {stats['p99_token_length']:.0f}")
+    logger.info(f"  P99 token length: {stats['p99_token_length']:.0f}")
-    print(f"  Min token length: {stats['min_token_length']:.0f}")
+    logger.info(f"  Min token length: {stats['min_token_length']:.0f}")
-    print(f"  Max token length: {stats['max_token_length']:.0f}")
+    logger.info(f"  Max token length: {stats['max_token_length']:.0f}")
    return stats
@@ -385,9 +386,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
        cfg: TokenizerTrainingConfig dataclass with all configuration parameters
    """
    # load dataset
-    print(f"Loading dataset: {cfg.repo_id}")
+    logger.info(f"Loading dataset: {cfg.repo_id}")
    dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
-    print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
+    logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
    # parse normalization mode
    try:
@@ -397,7 +398,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
            f"Invalid normalization_mode: {cfg.normalization_mode}. "
            f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
        ) from err
-    print(f"Normalization mode: {norm_mode.value}")
+    logger.info(f"Normalization mode: {norm_mode.value}")
    # parse encoded dimensions
    encoded_dim_ranges = []
@@ -406,38 +407,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
        encoded_dim_ranges.append((start, end))
    total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
-    print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
+    logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
    # parse relative dimensions
    relative_dim_list = None
    if cfg.relative_dims is not None and cfg.relative_dims.strip():
        relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
-        print(f"Relative dimensions: {relative_dim_list}")
+        logger.info(f"Relative dimensions: {relative_dim_list}")
    else:
-        print("No relative dimensions specified")
+        logger.info("No relative dimensions specified")
-    print(f"Use relative transform: {cfg.use_relative_transform}")
+    logger.info(f"Use relative transform: {cfg.use_relative_transform}")
    if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
-        print(
+        logger.warning(
            "Warning: use_relative_transform=True but no relative_dims specified. "
            "No relative transform will be applied."
        )
-    print(f"Action horizon: {cfg.action_horizon}")
+    logger.info(f"Action horizon: {cfg.action_horizon}")
-    print(f"State key: {cfg.state_key}")
+    logger.info(f"State key: {cfg.state_key}")
    # determine episodes to process
    num_episodes = dataset.num_episodes
    if cfg.max_episodes is not None:
        num_episodes = min(cfg.max_episodes, num_episodes)
-    print(f"Processing {num_episodes} episodes...")
+    logger.info(f"Processing {num_episodes} episodes...")
    # process episodes sequentially (to avoid pickling issues with dataset)
    all_chunks = []
    for ep_idx in range(num_episodes):
        if ep_idx % 10 == 0:
-            print(f"  Processing episode {ep_idx}/{num_episodes}...")
+            logger.info(f"  Processing episode {ep_idx}/{num_episodes}...")
        chunks = process_episode(
            (
@@ -455,19 +456,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
    # concatenate all chunks
    all_chunks = np.concatenate(all_chunks, axis=0)
-    print(f"Collected {len(all_chunks)} action chunks")
+    logger.info(f"Collected {len(all_chunks)} action chunks")
    # extract only encoded dimensions FIRST (before normalization)
    encoded_chunks = []
    for start, end in encoded_dim_ranges:
        encoded_chunks.append(all_chunks[:, :, start:end])
    encoded_chunks = np.concatenate(encoded_chunks, axis=-1)  # [N, H, D_encoded]
-    print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
+    logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
    # apply normalization to encoded dimensions
-    print("\nBefore normalization - overall stats:")
+    logger.info("\nBefore normalization - overall stats:")
-    print(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
+    logger.info(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
-    print(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
+    logger.info(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
    # get normalization stats from dataset
    norm_stats = dataset.meta.stats
@@ -489,9 +490,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
                    encoded_stats[stat_name] = stat_array[encoded_dim_indices]
        if encoded_stats:
-            print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
+            logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
            for stat_name, stat_values in encoded_stats.items():
-                print(
+                logger.info(
                    f"  {stat_name}: shape={stat_values.shape}, "
                    f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
                )
@@ -499,27 +500,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
            # apply normalization based on mode
            try:
                encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
-                print(f"\nApplied {norm_mode.value} normalization")
+                logger.info(f"\nApplied {norm_mode.value} normalization")
            except ValueError as e:
-                print(f"Warning: {e}. Using raw actions without normalization.")
+                logger.warning(f"Warning: {e}. Using raw actions without normalization.")
-            print("\nAfter normalization - overall stats:")
+            logger.info("\nAfter normalization - overall stats:")
-            print(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
+            logger.info(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
-            print(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
+            logger.info(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
-            print("\nPer-dimension stats (after normalization):")
+            logger.info("\nPer-dimension stats (after normalization):")
            for d in range(encoded_chunks.shape[-1]):
                dim_data = encoded_chunks[:, :, d]
-                print(
+                logger.info(
                    f"  Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
                    f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
                )
        else:
-            print("Warning: Could not extract stats for encoded dimensions, using raw actions")
+            logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
    else:
-        print("Warning: No normalization stats found in dataset, using raw actions")
+        logger.warning("Warning: No normalization stats found in dataset, using raw actions")
-    print(f"Encoded chunks shape: {encoded_chunks.shape}")
+    logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
    # train FAST tokenizer
    tokenizer = train_fast_tokenizer(
@@ -561,8 +562,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
    with open(output_path / "metadata.json", "w") as f:
        json.dump(metadata, f, indent=2)
-    print(f"\nSaved FAST tokenizer to {output_path}")
+    logger.info(f"\nSaved FAST tokenizer to {output_path}")
-    print(f"Metadata: {json.dumps(metadata, indent=2)}")
+    logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
    # push to Hugging Face Hub if requested
    if cfg.push_to_hub:
@@ -570,10 +571,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
        hub_repo_id = cfg.hub_repo_id
        if hub_repo_id is None:
            hub_repo_id = output_path.name
-            print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
+            logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
-        print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
+        logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
-        print(f"   Private: {cfg.hub_private}")
+        logger.info(f"   Private: {cfg.hub_private}")
        try:
            # use the tokenizer's push_to_hub method
@@ -593,14 +594,15 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
                commit_message="Upload tokenizer metadata",
            )
-            print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
+            logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
        except Exception as e:
-            print(f"Error pushing to hub: {e}")
+            logger.error(f"Error pushing to hub: {e}")
-            print("   Make sure you're logged in with `huggingface-cli login`")
+            logger.error("   Make sure you're logged in with `huggingface-cli login`")
 def main():
    """CLI entry point that parses arguments and runs the tokenizer training."""
    init_logging()
    train_tokenizer()
@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
    def __init__(
        self,
-        transforms: Sequence[Callable],
+        transforms: Sequence[Callable[..., Any]],
        p: list[float] | None = None,
        n_subset: int | None = None,
        random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
        if not isinstance(transforms, Sequence):
            raise TypeError("Argument transforms should be a sequence of callables")
        if p is None:
-            p = [1] * len(transforms)
+            p = [1.0] * len(transforms)
        elif len(p) != len(transforms):
            raise ValueError(
                f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
        self.n_subset = n_subset
        self.random_order = random_order
-        self.selected_transforms = None
+        self.selected_transforms: list[Callable[..., Any]] = []
    def forward(self, *inputs: Any) -> Any:
        needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
        super().__init__()
        self.sharpness = self._check_input(sharpness)
-    def _check_input(self, sharpness):
+    def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
        if isinstance(sharpness, (int | float)):
            if sharpness < 0:
                raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
    )
-def make_transform_from_config(cfg: ImageTransformConfig):
+def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
    if cfg.type == "SharpnessJitter":
        return SharpnessJitter(**cfg.kwargs)
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
        super().__init__()
        self._cfg = cfg
-        self.weights = []
+        self.weights: list[float] = []
-        self.transforms = {}
+        self.transforms: dict[str, Transform] = {}
        for tf_name, tf_cfg in cfg.tfs.items():
            if tf_cfg.weight <= 0.0:
                continue
@@ -22,7 +22,7 @@ from torch.utils.data._utils.collate import default_collate
 from lerobot.datasets.language import LANGUAGE_COLUMNS
-_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices", *LANGUAGE_COLUMNS}
+_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"}
 def lerobot_collate_fn(batch: list[dict[str, Any] | None]) -> dict[str, Any] | None:
@@ -26,7 +26,6 @@ OBS_IMAGES = OBS_IMAGE + "s"
 OBS_LANGUAGE = OBS_STR + ".language"
 OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
 OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
 OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks"
 OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
 OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
 OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
@@ -35,7 +34,6 @@ ACTION = "action"
 ACTION_PREFIX = ACTION + "."
 ACTION_TOKENS = ACTION + ".tokens"
 ACTION_TOKEN_MASK = ACTION + ".token_mask"
 ACTION_CODE_TOKEN_MASK = ACTION + ".code_token_mask"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"
@@ -23,46 +23,6 @@ logger = logging.getLogger(__name__)
 JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...]
 class StreamingVideoWriter:
    """Incrementally encode RGB frames to an MP4 without retaining them in memory."""
    def __init__(self, video_path: str | Path, fps: int) -> None:
        from .import_utils import require_package
        require_package("av", extra="av-dep")
        import av
        self._av = av
        self._container = av.open(str(video_path), mode="w")
        self._stream = self._container.add_stream("libx264", rate=fps)
        self._shape: tuple[int, int] | None = None
        self.frames_written = 0
    def add_frame(self, frame_array) -> None:
        orig_height, orig_width = frame_array.shape[:2]
        height = orig_height - orig_height % 2
        width = orig_width - orig_width % 2
        if self._shape is None:
            self._shape = (height, width)
            self._stream.width = width
            self._stream.height = height
            self._stream.pix_fmt = "yuv420p"
        elif self._shape != (height, width):
            raise ValueError(f"Video frame shape changed from {self._shape} to {(height, width)}")
        frame = self._av.VideoFrame.from_ndarray(frame_array[:height, :width], format="rgb24")
        for packet in self._stream.encode(frame):
            self._container.mux(packet)
        self.frames_written += 1
    def close(self) -> None:
        if self._container is None:
            return
        for packet in self._stream.encode():
            self._container.mux(packet)
        self._container.close()
        self._container = None
 def load_json(fpath: Path) -> Any:
    """Load data from a JSON file.
@@ -98,12 +58,36 @@ def write_video(video_path: str | Path, stacked_frames: list, fps: int) -> None:
        stacked_frames: List of HWC uint8 numpy arrays (RGB).
        fps: Frames per second for the output video.
    """
-    writer = StreamingVideoWriter(video_path, fps)
+    from .import_utils import require_package
-    try:
+
    require_package("av", extra="av-dep")
    import av
    with av.open(str(video_path), mode="w") as container:
        orig_height, orig_width = stacked_frames[0].shape[:2]
        # yuv420p requires even dimensions; crop by one pixel if needed
        height = orig_height if orig_height % 2 == 0 else orig_height - 1
        width = orig_width if orig_width % 2 == 0 else orig_width - 1
        if height != orig_height or width != orig_width:
            logger.warning(
                "Frame dimensions %dx%d are not even; cropping to %dx%d for yuv420p compatibility.",
                orig_width,
                orig_height,
                width,
                height,
            )
        stream = container.add_stream("libx264", rate=fps)
        stream.width = width
        stream.height = height
        stream.pix_fmt = "yuv420p"
        for frame_array in stacked_frames:
-            writer.add_frame(frame_array)
+            if height != orig_height or width != orig_width:
-    finally:
+                frame_array = frame_array[:height, :width]
-        writer.close()
+            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
            for packet in stream.encode(frame):
                container.mux(packet)
        for packet in stream.encode():
            container.mux(packet)
 def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:
@@ -16,11 +16,39 @@
 # limitations under the License.
 import logging
 import multiprocessing
 import os
 import signal
 import sys
 def ensure_multiprocessing_start_method(start_method: str | None) -> None:
    """Set a multiprocessing start method once, or verify the existing method matches.
    Passing ``None`` leaves Python's process-wide default untouched. This is useful
    when LeRobot is embedded in an application that owns multiprocessing setup.
    """
    if start_method is None:
        return
    available_methods = multiprocessing.get_all_start_methods()
    if start_method not in available_methods:
        raise ValueError(
            f"Multiprocessing start method must be one of {available_methods} on this platform, "
            f"got {start_method!r}."
        )
    current_method = multiprocessing.get_start_method(allow_none=True)
    if current_method is None:
        multiprocessing.set_start_method(start_method)
    elif current_method != start_method:
        raise RuntimeError(
            f"Multiprocessing start method is already {current_method!r}; cannot change it to "
            f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
            "application's existing method, or launch LeRobot in a fresh process."
        )
 class ProcessSignalHandler:
    """Utility class to attach graceful shutdown signal handlers.
@@ -38,10 +38,7 @@ def _is_scalar(x):
 def init_rerun(
-    session_name: str = "lerobot_control_loop",
+    session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
    ip: str | None = None,
    port: int | None = None,
    web_port: int | None = None,
 ) -> None:
    """
    Initializes the Rerun SDK for visualizing the control loop.
@@ -50,7 +47,6 @@ def init_rerun(
        session_name: Name of the Rerun session.
        ip: Optional IP for connecting to a Rerun server.
        port: Optional port for connecting to a Rerun server.
        web_port: Serve a headless web viewer on this port, using ``port`` for gRPC.
    """
    require_package("rerun-sdk", extra="viz", import_name="rerun")
@@ -64,10 +60,6 @@ def init_rerun(
    memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
    if ip and port:
        rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy")
    elif web_port is not None:
        grpc_port = port or 9876
        url = rr.serve_grpc(grpc_port=grpc_port)
        rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url)
    else:
        rr.spawn(memory_limit=memory_limit)
@@ -133,10 +133,13 @@ def say(text: str, blocking: bool = False):
    else:
        raise RuntimeError("Unsupported operating system for text-to-speech.")
    try:
        if blocking:
-        subprocess.run(cmd, check=True)
+            subprocess.run(cmd, check=True, timeout=5)
        else:
            subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
    except (FileNotFoundError, subprocess.TimeoutExpired) as e:
        logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
 def log_say(text: str, play_sounds: bool = True, blocking: bool = False):
@@ -1,71 +0,0 @@
 # 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.
 """Best-effort text overlays shared by evaluation and interactive rollouts."""
 from __future__ import annotations
 from collections.abc import Iterable
 import numpy as np
 def annotate_frame(frame: np.ndarray, fields: Iterable[tuple[str, str | None]]) -> np.ndarray:
    """Return an RGB frame annotated with the non-empty labeled ``fields``."""
    if frame.ndim != 3 or frame.shape[-1] != 3:
        return frame
    try:
        import cv2  # noqa: PLC0415
    except ImportError:
        return frame
    text_rows = [f"{label}: {value}" for label, value in fields if value]
    if not text_rows:
        return frame
    image = np.ascontiguousarray(frame).copy()
    font, scale, thickness, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1, 6
    max_width = image.shape[1] - 2 * margin
    lines: list[str] = []
    for text in text_rows:
        current = ""
        for word in text.split():
            candidate = f"{current} {word}".strip()
            width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
            if width > max_width and current:
                lines.append(current)
                current = word
            else:
                current = candidate
        if current:
            lines.append(current)
    line_height = 20
    header_height = min(image.shape[0], len(lines) * line_height + 6)
    backdrop = image.copy()
    cv2.rectangle(backdrop, (0, 0), (image.shape[1], header_height), (0, 0, 0), -1)
    cv2.addWeighted(backdrop, 0.55, image, 0.45, 0, dst=image)
    for index, line in enumerate(lines):
        cv2.putText(
            image,
            line,
            (margin, 18 + index * line_height),
            font,
            scale,
            (255, 255, 255),
            thickness,
            cv2.LINE_AA,
        )
    return image
@@ -20,7 +20,7 @@
 # ```
 from pathlib import Path
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 import cv2
 import numpy as np
@@ -123,6 +123,73 @@ def test_invalid_width_connect():
        camera.connect(warmup=False)
 def test_connect_cleans_up_after_settings_failure_and_allows_retry():
    config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
    camera = OpenCVCamera(config)
    opened_captures = []
    def fail_settings():
        opened_captures.append(camera.videocapture)
        raise RuntimeError("settings failed")
    with (
        patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
        pytest.raises(RuntimeError, match="settings failed"),
    ):
        camera.connect(warmup=False)
    assert camera.videocapture is None
    assert camera.thread is None
    assert not camera.is_connected
    assert opened_captures[0] is not None
    assert not opened_captures[0].isOpened()
    camera.connect(warmup=False)
    assert camera.is_connected
    camera.disconnect()
 def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
    config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
    camera = OpenCVCamera(config)
    read_threads = []
    def fail_warmup(*_args, **_kwargs):
        read_threads.append(camera.thread)
        raise TimeoutError("no frame")
    with (
        patch.object(camera, "async_read", side_effect=fail_warmup),
        pytest.raises(TimeoutError, match="no frame"),
    ):
        camera.connect()
    assert camera.videocapture is None
    assert camera.thread is None
    assert not camera.is_connected
    assert read_threads[0] is not None
    assert not read_threads[0].is_alive()
    camera.connect(warmup=False)
    assert camera.is_connected
    camera.disconnect()
 def test_find_cameras_releases_unopened_handles():
    module_path = OpenCVCamera.__module__
    unopened_capture = MagicMock()
    unopened_capture.isOpened.return_value = False
    with (
        patch(f"{module_path}.platform.system", return_value="Darwin"),
        patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
        patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
    ):
        assert OpenCVCamera.find_cameras() == []
    unopened_capture.release.assert_called_once_with()
@pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
 def test_read(index_or_path):
    config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)
@@ -20,7 +20,7 @@
 # ```
 from pathlib import Path
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 import numpy as np
 import pytest
@@ -30,6 +30,8 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
 pytest.importorskip("pyrealsense2")
 import pyrealsense2 as rs
 from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
 TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
@@ -61,6 +63,17 @@ def test_abc_implementation():
    _ = RealSenseCamera(config)
@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
 def test_manual_color_option_requires_rgb(option):
    with pytest.raises(ValueError, match="use_rgb=True"):
        RealSenseCameraConfig(
            serial_number_or_name="042",
            use_rgb=False,
            use_depth=True,
            **{option: 100},
        )
 def test_connect():
    config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
@@ -83,6 +96,27 @@ def test_connect_invalid_camera_path(patch_realsense):
        camera.connect(warmup=False)
 def test_connect_cleans_up_when_sensor_configuration_fails():
    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
    camera = RealSenseCamera(config)
    pipeline = MagicMock()
    pipeline.start.return_value = MagicMock()
    with (
        patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
        patch.object(camera, "_configure_rs_pipeline_config"),
        patch.object(camera, "_configure_capture_settings"),
        patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
        pytest.raises(ValueError, match="invalid exposure"),
    ):
        camera.connect(warmup=False)
    pipeline.stop.assert_called_once_with()
    assert camera.rs_pipeline is None
    assert camera.rs_profile is None
    assert not camera.is_connected
 def test_invalid_width_connect():
    config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
    camera = RealSenseCamera(config)
@@ -91,6 +125,33 @@ def test_invalid_width_connect():
        camera.connect(warmup=False)
 def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
    config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
    camera = RealSenseCamera(config)
    read_threads = []
    def fail_warmup(*_args, **_kwargs):
        read_threads.append(camera.thread)
        raise TimeoutError("no frame")
    with (
        patch.object(camera, "async_read", side_effect=fail_warmup),
        pytest.raises(TimeoutError, match="no frame"),
    ):
        camera.connect()
    assert camera.rs_pipeline is None
    assert camera.rs_profile is None
    assert camera.thread is None
    assert not camera.is_connected
    assert read_threads[0] is not None
    assert not read_threads[0].is_alive()
    camera.connect(warmup=False)
    assert camera.is_connected
    camera.disconnect()
 def test_read():
    config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
    with RealSenseCamera(config) as camera:
@@ -228,6 +289,203 @@ def test_read_latest_too_old():
            _ = camera.read_latest(max_age_ms=0)  # immediately too old
 def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
    """Build a fake rs.sensor that reports a name and a configurable supported-options set."""
    supported = supported_options if supported_options is not None else set()
    sensor = MagicMock()
    sensor.get_info.return_value = name
    sensor.supports.side_effect = lambda opt: opt in supported
    return sensor
 def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
    """Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
    profile = MagicMock()
    device = MagicMock()
    device.query_sensors.return_value = [sensor]
    profile.get_device.return_value = device
    camera.rs_profile = profile
 def test_get_color_sensor_prefers_rgb_camera():
    config = RealSenseCameraConfig(serial_number_or_name="042")
    camera = RealSenseCamera(config)
    rgb = _make_mock_sensor("RGB Camera")
    stereo = _make_mock_sensor("Stereo Module")
    profile = MagicMock()
    device = MagicMock()
    device.query_sensors.return_value = [stereo, rgb]
    profile.get_device.return_value = device
    camera.rs_profile = profile
    assert camera._get_color_sensor() is rgb
 def test_get_color_sensor_falls_back_to_stereo_module():
    """D405 has no separate RGB module; color comes from Stereo Module."""
    config = RealSenseCameraConfig(serial_number_or_name="042")
    camera = RealSenseCamera(config)
    stereo = _make_mock_sensor("Stereo Module")
    _attach_mock_color_sensor(camera, stereo)
    assert camera._get_color_sensor() is stereo
 def test_get_color_sensor_raises_with_available_sensors():
    config = RealSenseCameraConfig(serial_number_or_name="042")
    camera = RealSenseCamera(config)
    other = _make_mock_sensor("Motion Module")
    _attach_mock_color_sensor(camera, other)
    with pytest.raises(RuntimeError, match="Motion Module"):
        camera._get_color_sensor()
 def test_configure_sensor_options_skipped_when_none():
    config = RealSenseCameraConfig(serial_number_or_name="042")
    camera = RealSenseCamera(config)
    with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
        camera._configure_sensor_options()
        mock_get.assert_not_called()
 def test_configure_sensor_options_applies_all_values():
    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
    camera = RealSenseCamera(config)
    sensor = _make_mock_sensor(
        "RGB Camera",
        supported_options={
            rs.option.enable_auto_exposure,
            rs.option.exposure,
            rs.option.gain,
            rs.option.enable_auto_white_balance,
            rs.option.white_balance,
        },
    )
    _attach_mock_color_sensor(camera, sensor)
    camera._configure_sensor_options()
    sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
    sensor.set_option.assert_any_call(rs.option.exposure, 120)
    sensor.set_option.assert_any_call(rs.option.gain, 64)
    sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
    sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
@pytest.mark.parametrize(
    ("config_field", "option", "label"),
    [
        ("exposure", rs.option.exposure, "exposure"),
        ("gain", rs.option.gain, "gain"),
        ("white_balance", rs.option.white_balance, "white balance"),
    ],
 )
 def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
    config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
    camera = RealSenseCamera(config)
    sensor = _make_mock_sensor("RGB Camera", supported_options=set())
    _attach_mock_color_sensor(camera, sensor)
    with pytest.raises(ValueError, match=label):
        camera._configure_sensor_options()
    sensor.supports.assert_any_call(option)
    sensor.set_option.assert_not_called()
@pytest.mark.parametrize(
    ("config_field", "option", "value"),
    [
        ("exposure", rs.option.exposure, 120),
        ("gain", rs.option.gain, 64),
    ],
 )
 def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
    """white_balance=None should not touch auto white balance."""
    config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
    camera = RealSenseCamera(config)
    sensor = _make_mock_sensor(
        "RGB Camera",
        supported_options={rs.option.enable_auto_exposure, option},
    )
    _attach_mock_color_sensor(camera, sensor)
    camera._configure_sensor_options()
    calls = [call.args for call in sensor.set_option.call_args_list]
    assert (rs.option.enable_auto_exposure, 0) in calls
    assert (option, value) in calls
    for opt, _ in calls:
        assert opt != rs.option.enable_auto_white_balance
        assert opt != rs.option.white_balance
 def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
    camera = RealSenseCamera(config)
    sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
    _attach_mock_color_sensor(camera, sensor)
    with caplog.at_level("WARNING"):
        camera._configure_sensor_options()
    sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
    assert "does not support disabling auto-exposure" in caplog.text
 def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
    config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
    camera = RealSenseCamera(config)
    sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
    _attach_mock_color_sensor(camera, sensor)
    with caplog.at_level("WARNING"):
        camera._configure_sensor_options()
    sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
    assert "does not support disabling auto white balance" in caplog.text
 def test_configure_sensor_options_out_of_range_raises_value_error():
    """set_option errors should be re-raised as ValueError with range diagnostics."""
    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
    camera = RealSenseCamera(config)
    sensor = _make_mock_sensor(
        "RGB Camera",
        supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
    )
    def fake_set_option(option, value):
        if option == rs.option.exposure:
            raise RuntimeError("value out of range")
    sensor.set_option.side_effect = fake_set_option
    option_range = MagicMock(min=1, max=10000, step=1, default=156)
    sensor.get_option_range.return_value = option_range
    _attach_mock_color_sensor(camera, sensor)
    with pytest.raises(ValueError, match="exposure") as exc_info:
        camera._configure_sensor_options()
    msg = str(exc_info.value)
    assert "999999" in msg
    assert "min=1" in msg
    assert "max=10000" in msg
@pytest.mark.parametrize(
    "rotation",
    [
@@ -29,13 +29,6 @@ def test_message_recipe_validates_unknown_binding():
        )
 def test_canonical_recipe_loads():
    """The canonical PI052 blend YAML loads + validates."""
    recipe = TrainingRecipe.from_yaml(Path("src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml"))
    assert recipe.blend is not None
    assert sum(c.weight for c in recipe.blend.values()) == pytest.approx(1.0)
 def test_message_turn_requires_a_stream():
    """Every turn must declare a stream — None is rejected at construction.
@@ -0,0 +1,104 @@
 # Copyright 2025 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.
 import numpy as np
 import pytest
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 from lerobot.scripts.augment_dataset_quantile_stats import (
    compute_quantile_stats_for_dataset,
    has_quantile_stats,
 )
 def _numeric_keys(dataset):
    return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
 def _image_keys(dataset):
    return [k for k, v in dataset.features.items() if v["dtype"] in ("image", "video")]
 def test_numeric_stats_are_unaffected_by_sampling(tmp_path, lerobot_dataset_factory):
    """Sampling only touches image/video frames; numeric features are read in
    full either way, so their stats must be identical with and without sampling."""
    dataset = lerobot_dataset_factory(
        root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
    )
    exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
    sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
    numeric_keys = _numeric_keys(dataset)
    assert numeric_keys, "fixture should expose numeric features"
    for key in numeric_keys:
        if key not in exact:
            continue
        for stat in ("mean", "std", "q01", "q50", "q99"):
            if stat in exact[key]:
                np.testing.assert_allclose(
                    sampled[key][stat],
                    exact[key][stat],
                    rtol=1e-6,
                    atol=1e-6,
                    err_msg=f"numeric feature '{key}' stat '{stat}' changed under sampling",
                )
 def test_image_sampling_reduces_data_but_keeps_stats_close(tmp_path, lerobot_dataset_factory):
    """For images, sampling should reduce the number of samples considered while
    keeping the resulting statistics close to the exact ones."""
    dataset = lerobot_dataset_factory(
        root=tmp_path / "ds", total_episodes=2, total_frames=400, use_videos=False
    )
    exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
    sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
    image_keys = _image_keys(dataset)
    assert image_keys, "fixture should expose at least one image feature"
    for key in image_keys:
        # sampling actually looked at fewer pixels
        assert sampled[key]["count"][0] < exact[key]["count"][0]
        # but per-channel mean stays close
        np.testing.assert_allclose(
            sampled[key]["mean"],
            exact[key]["mean"],
            rtol=0.15,
            err_msg=f"image feature '{key}' mean drifted too far under sampling",
        )
 def test_short_episodes_use_all_frames(tmp_path, lerobot_dataset_factory):
    """With episodes shorter than the sampling floor, sampling is a no-op and
    must produce exactly the same stats as the exact path."""
    dataset = lerobot_dataset_factory(
        root=tmp_path / "ds", total_episodes=2, total_frames=40, use_videos=False
    )
    exact = compute_quantile_stats_for_dataset(dataset, use_sampling=False)
    sampled = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
    for key in _image_keys(dataset):
        assert sampled[key]["count"][0] == exact[key]["count"][0]
 def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory):
    """The computed stats should contain quantile keys for the dataset."""
    dataset = lerobot_dataset_factory(
        root=tmp_path / "ds", total_episodes=2, total_frames=200, use_videos=False
    )
    stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
    assert has_quantile_stats(stats)
@@ -343,84 +343,6 @@ def test_resolve_task_explicit_override_beats_rephrasings():
    assert rendered["messages"][0]["content"] == "explicit override wins"
 def test_flow_only_low_level_recipe_renders_without_target():
    """Regression: a flow-only ``low_level`` recipe has no ``target`` turn —
    its supervision is the action-expert flow loss, not text-CE. It must
    still render (not ``None``), otherwise every blend draw of it is dropped
    and the action expert never receives a flow loss."""
    recipe = TrainingRecipe(
        messages=[
            MessageTurn(
                role="user",
                content="${subtask}",
                stream="low_level",
                if_present="subtask",
            ),
        ],
        bindings={"subtask": "active_at(t, style=subtask)"},
    )
    rendered = render_sample(
        recipe=recipe,
        persistent=PERSISTENT,
        events=[],
        t=0.5,
        sample_idx=0,
        task="clean kitchen",
    )
    assert rendered is not None
    assert rendered["messages"] == [{"role": "user", "content": "subtask 0"}]
    assert rendered["message_streams"] == ["low_level"]
    assert rendered["target_message_indices"] == []
 def test_vqa_frame_is_consumed_over_the_weighted_blend():
    """A frame carrying a VQA annotation renders the ``ask_vqa*`` sub-recipe
    even when its blend weight is tiny — VQA annotations are sparse and must
    never be wasted on a subtask/action draw."""
    recipe = TrainingRecipe(
        blend={
            "high_level_subtask": TrainingRecipe(
                weight=0.99,
                messages=[
                    MessageTurn(role="user", content="${task}", stream="high_level"),
                    MessageTurn(role="assistant", content="a subtask", stream="high_level", target=True),
                ],
            ),
            "ask_vqa_top": TrainingRecipe(
                weight=0.01,
                bindings={
                    "vqa_query": "emitted_at(t, style=vqa, role=user, camera=observation.images.top)",
                    "vqa": "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)",
                },
                messages=[
                    MessageTurn(
                        role="user", content="${vqa_query}", stream="high_level", if_present="vqa_query"
                    ),
                    MessageTurn(
                        role="assistant",
                        content="${vqa}",
                        stream="high_level",
                        target=True,
                        if_present="vqa",
                    ),
                ],
            ),
        }
    )
    # A frame WITH a vqa event renders VQA on every sample_idx, despite the
    # ask_vqa weight being only 0.01.
    for sample_idx in range(20):
        rendered = render_sample(
            recipe=recipe, persistent=PERSISTENT, events=EVENTS_AT_1, t=1.0, sample_idx=sample_idx, task="x"
        )
        assert rendered["messages"][-1]["content"] == '{"count": 2}', sample_idx
    # A frame WITHOUT a vqa event falls back to the normal weighted blend.
    rendered = render_sample(recipe=recipe, persistent=PERSISTENT, events=[], t=1.0, sample_idx=0, task="x")
    assert rendered["messages"][-1]["content"] == "a subtask"
 def test_emitted_at_persistent_tolerates_small_timestamp_drift():
    """Persistent ``emitted_at`` should match within EMITTED_AT_TOLERANCE_S
    so callers that derive ``t`` arithmetically (``frame_idx / fps``) still
@@ -482,6 +482,20 @@ def test_add_frame_works_in_write_mode(tmp_path):
 # ── Resume mode ──────────────────────────────────────────────────────
 def test_resume_freshly_created_empty_dataset(tmp_path):
    """resume() accepts a local dataset created before any episode was recorded."""
    root = tmp_path / "resume_empty_ds"
    LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
    resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
    assert isinstance(resumed.writer, DatasetWriter)
    assert resumed.meta.total_episodes == 0
    assert resumed.meta.total_frames == 0
    assert resumed.meta.tasks is None
    assert resumed.meta.episodes is None
 def test_resume_creates_writer(tmp_path):
    """After resume(), writer is a DatasetWriter."""
    root = tmp_path / "resume_ds"
@@ -25,7 +25,7 @@ from datasets import Dataset  # noqa: E402
 from lerobot.datasets.io_utils import (
    hf_transform_to_torch,
 )
-from lerobot.datasets.sampler import EpisodeAwareSampler, compute_sampler_state
+from lerobot.datasets.sampler import EpisodeAwareSampler
 def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
@@ -154,6 +154,8 @@ def test_partial_episode_drop_warns(caplog):
 # --- seeded (seed, epoch) shuffling, resume, and state ---
 from lerobot.datasets.sampler import compute_sampler_state  # noqa: E402
 EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6])  # episodes of 2, 1 and 3 frames
@@ -1,48 +0,0 @@
 #!/usr/bin/env python
 # Copyright 2026 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import numpy as np
 from lerobot.envs.robocasa import RoboCasaEnv, convert_action
 def test_robocasa_action_uses_openpi_checkpoint_order():
    action = np.arange(12, dtype=np.float32)
    converted = convert_action(action)
    np.testing.assert_array_equal(converted["action.end_effector_position"], [0, 1, 2])
    np.testing.assert_array_equal(converted["action.end_effector_rotation"], [3, 4, 5])
    np.testing.assert_array_equal(converted["action.gripper_close"], [6])
    np.testing.assert_array_equal(converted["action.base_motion"], [7, 8, 9, 10])
    np.testing.assert_array_equal(converted["action.control_mode"], [11])
 def test_robocasa_state_uses_openpi_checkpoint_order():
    env = object.__new__(RoboCasaEnv)
    env.obs_type = "pixels_agent_pos"
    env.camera_name = []
    raw_observation = {
        "state.end_effector_position_relative": np.arange(0, 3),
        "state.end_effector_rotation_relative": np.arange(3, 7),
        "state.base_position": np.arange(7, 10),
        "state.base_rotation": np.arange(10, 14),
        "state.gripper_qpos": np.arange(14, 16),
    }
    observation = env._format_raw_obs(raw_observation)
    np.testing.assert_array_equal(observation["agent_pos"], np.arange(16, dtype=np.float32))
@@ -0,0 +1,83 @@
 # 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.
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 import torch
 import lerobot.policies.factory as policy_factory
 def test_make_policy_keeps_peft_adapter_and_base_revisions_separate(monkeypatch):
    cfg = SimpleNamespace(
        type="mock",
        device="cpu",
        pretrained_path="user/adapter",
        pretrained_revision="adapter-sha",
        use_peft=True,
        input_features={},
        output_features={},
    )
    dataset_meta = SimpleNamespace(features={}, stats={})
    base_policy = torch.nn.Linear(1, 1)
    policy_from_pretrained = MagicMock(return_value=base_policy)
    policy_class = SimpleNamespace(from_pretrained=policy_from_pretrained)
    monkeypatch.setattr(policy_factory, "get_policy_class", lambda _: policy_class)
    monkeypatch.setattr(policy_factory, "dataset_to_policy_features", lambda _: {})
    monkeypatch.setattr(policy_factory, "validate_visual_features_consistency", lambda *args: None)
    peft_config = SimpleNamespace(
        base_model_name_or_path="user/base-policy",
        revision="base-sha",
    )
    peft_config_from_pretrained = MagicMock(return_value=peft_config)
    adapted_policy = torch.nn.Linear(1, 1)
    peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
    require_package = MagicMock()
    monkeypatch.setattr(policy_factory, "require_package", require_package)
    monkeypatch.setattr(
        policy_factory,
        "PeftConfig",
        SimpleNamespace(from_pretrained=peft_config_from_pretrained),
    )
    monkeypatch.setattr(
        policy_factory,
        "PeftModel",
        SimpleNamespace(from_pretrained=peft_model_from_pretrained),
    )
    policy = policy_factory.make_policy(cfg, ds_meta=dataset_meta)
    assert policy is adapted_policy
    require_package.assert_called_once_with("peft", extra="peft")
    peft_config_from_pretrained.assert_called_once_with(
        "user/adapter",
        revision="adapter-sha",
    )
    policy_from_pretrained.assert_called_once_with(
        config=cfg,
        dataset_stats=dataset_meta.stats,
        dataset_meta=dataset_meta,
        pretrained_name_or_path="user/base-policy",
        revision="base-sha",
    )
    peft_model_from_pretrained.assert_called_once_with(
        base_policy,
        "user/adapter",
        config=peft_config,
        revision="adapter-sha",
        is_trainable=True,
    )
@@ -113,6 +113,7 @@ def test_gaussian_actor_config_default_initialization():
    # Concurrency configuration
    assert config.concurrency.actor == "threads"
    assert config.concurrency.learner == "threads"
    assert config.concurrency.multiprocessing_context == "spawn"
    assert isinstance(config.actor_network_kwargs, ActorNetworkConfig)
    assert isinstance(config.policy_kwargs, PolicyConfig)
@@ -152,6 +153,7 @@ def test_concurrency_config():
    config = ConcurrencyConfig()
    assert config.actor == "threads"
    assert config.learner == "threads"
    assert config.multiprocessing_context == "spawn"
 def test_gaussian_actor_config_custom_initialization():
@@ -12,9 +12,7 @@ from lerobot.processor.render_messages_processor import RenderMessagesStep  # no
 from lerobot.types import TransitionKey  # noqa: E402
-def test_render_messages_step_renders_task_fallback_without_language_columns():
+def test_render_messages_step_noops_without_language_columns():
    """No language columns + a task string → low-level task fallback render,
    matching what the policy sees at eval time on unannotated observations."""
    recipe = TrainingRecipe(
        messages=[
            MessageTurn(role="user", content="${task}", stream="high_level"),
@@ -23,24 +21,6 @@ def test_render_messages_step_renders_task_fallback_without_language_columns():
    )
    transition = create_transition(complementary_data={"task": "do it"})
    out = RenderMessagesStep(recipe)(transition)
    data = out[TransitionKey.COMPLEMENTARY_DATA]
    assert data["messages"] == [{"role": "user", "content": "do it"}]
    assert data["message_streams"] == ["low_level"]
    assert data["target_message_indices"] == []
    assert data["task"] == "do it"
 def test_render_messages_step_noops_without_language_columns_or_task():
    recipe = TrainingRecipe(
        messages=[
            MessageTurn(role="user", content="${task}", stream="high_level"),
            MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
        ]
    )
    transition = create_transition(complementary_data={})
    assert RenderMessagesStep(recipe)(transition) == transition
@@ -78,70 +58,3 @@ def test_render_messages_step_renders_and_drops_raw_language():
    assert data["messages"][-1]["content"] == "reach carefully"
    assert data["message_streams"] == ["high_level", "low_level"]
    assert data["target_message_indices"] == [1]
 def test_render_messages_step_falls_back_to_low_level_task_when_recipe_misses():
    recipe = TrainingRecipe(
        messages=[
            MessageTurn(
                role="assistant",
                content="${subtask}",
                stream="high_level",
                target=True,
                if_present="subtask",
            ),
        ]
    )
    transition = create_transition(
        complementary_data={
            "task": "pick the cube",
            "timestamp": torch.tensor(0.0),
            "index": torch.tensor(7),
            "language_persistent": [],
            "language_events": [{"style": "unmatched", "timestamp": 0.0}],
        }
    )
    out = RenderMessagesStep(recipe)(transition)
    data = out[TransitionKey.COMPLEMENTARY_DATA]
    assert data["messages"] == [{"role": "user", "content": "pick the cube"}]
    assert data["message_streams"] == ["low_level"]
    assert data["target_message_indices"] == []
 def test_render_messages_step_falls_back_per_sample_in_batched_language():
    recipe = TrainingRecipe(
        messages=[
            MessageTurn(
                role="assistant",
                content="${subtask}",
                stream="high_level",
                target=True,
                if_present="subtask",
            ),
        ]
    )
    transition = create_transition(
        action=torch.arange(4).reshape(2, 2),
        complementary_data={
            "task": ["pick the cube", "open the drawer"],
            "timestamp": torch.tensor([0.0, 1.0]),
            "index": torch.tensor([7, 8]),
            "language_persistent": [[], []],
            "language_events": [
                [{"style": "unmatched", "timestamp": 0.0}],
                [{"style": "unmatched", "timestamp": 1.0}],
            ],
        },
    )
    out = RenderMessagesStep(recipe)(transition)
    data = out[TransitionKey.COMPLEMENTARY_DATA]
    assert data["messages"] == [
        [{"role": "user", "content": "pick the cube"}],
        [{"role": "user", "content": "open the drawer"}],
    ]
    assert data["message_streams"] == [["low_level"], ["low_level"]]
    assert data["target_message_indices"] == [[], []]
@@ -25,7 +25,7 @@ import pytest
 import torch
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, TokenizerProcessorStep
+from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep
 from lerobot.processor.converters import create_transition, identity_transition
 from lerobot.types import TransitionKey
 from lerobot.utils.constants import (
@@ -88,46 +88,6 @@ class MockTokenizer:
        return result
 def test_action_tokenizer_config_preserves_token_mapping():
    processor = object.__new__(ActionTokenizerProcessorStep)
    processor.trust_remote_code = True
    processor.max_action_tokens = 384
    processor.fast_skip_tokens = 64
    processor.paligemma_tokenizer_name = "custom/paligemma"
    processor.allow_truncation = False
    processor.action_tokenizer_name = "custom/fast"
    processor.action_tokenizer_input_object = None
    assert processor.get_config() == {
        "trust_remote_code": True,
        "max_action_tokens": 384,
        "fast_skip_tokens": 64,
        "paligemma_tokenizer_name": "custom/paligemma",
        "allow_truncation": False,
        "action_tokenizer_name": "custom/fast",
    }
 def test_action_tokenizer_can_reject_truncated_sequences():
    processor = object.__new__(ActionTokenizerProcessorStep)
    processor.max_action_tokens = 4
    processor.fast_skip_tokens = 128
    processor.allow_truncation = False
    processor.action_tokenizer = lambda _actions: [1, 2, 3]
    processor._paligemma_tokenizer = type(
        "Tokenizer",
        (),
        {
            "vocab_size": 1000,
            "bos_token_id": 2,
            "encode": lambda _self, text, **_kwargs: [10, 11] if text == "Action: " else [12, 1],
        },
    )()
    with pytest.raises(ValueError, match="max_action_tokens=4"):
        processor._tokenize_action(torch.zeros(1, 2, 1))
@pytest.fixture
 def mock_tokenizer():
    """Provide a mock tokenizer for testing."""
@@ -1,105 +0,0 @@
 # 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.
 from lerobot.runtime import RuntimeState
 from lerobot.runtime.adapter import (
    BaseLanguageAdapter,
    DirectTaskPolicyAdapter,
    GenerationConfig,
 )
 class ScriptedAdapter(BaseLanguageAdapter):
    """Base adapter whose text generation returns queued strings per kind."""
    def __init__(self, scripts, gen=None):
        super().__init__(policy=object(), gen=gen)
        self.scripts = {k: list(v) for k, v in scripts.items()}
        self.calls = []
    def select_action(self, observation, state):
        return None
    def generate_text(self, kind, observation, state, user_text=None):
        self.calls.append(kind)
        queue = self.scripts.get(kind, [])
        return queue.pop(0) if queue else ""
 def test_cascade_sets_subtask_then_memory():
    adapter = ScriptedAdapter({"subtask": ["pick the red cup"], "memory": ["the cup is grasped"]})
    state = RuntimeState(task="clean")
    adapter.update_language_state(None, state)
    assert state.language_context["subtask"] == "pick the red cup"
    assert state.language_context["memory"] == "the cup is grasped"
    assert adapter.calls == ["subtask", "memory"]
 def test_nonempty_generation_is_used_verbatim():
    adapter = ScriptedAdapter({"subtask": [":::: ::"], "memory": ["memory"]})
    state = RuntimeState(task="clean")
    adapter.update_language_state(None, state)
    assert state.language_context["subtask"] == ":::: ::"
    assert state.language_context["memory"] == "memory"
    assert adapter.calls == ["subtask", "memory"]
 def test_throttle_regenerates_every_n_chunks():
    adapter = ScriptedAdapter(
        {
            "subtask": ["pick the first cup", "pick the second cup"],
            "memory": ["memory one two three", "memory four five six"],
        },
        gen=GenerationConfig(chunks_per_regen=2),
    )
    state = RuntimeState(task="clean")
    adapter.update_language_state(None, state)  # generates
    assert state.language_context["subtask"] == "pick the first cup"
    adapter.update_language_state(None, state)  # throttled — no generation
    assert state.language_context["subtask"] == "pick the first cup"
    adapter.update_language_state(None, state)  # generates again
    assert state.language_context["subtask"] == "pick the second cup"
 def test_handle_interjection_sets_plan_and_strips_say():
    adapter = ScriptedAdapter({"interjection": ["turn to the left now <say>heading left</say>"]})
    state = RuntimeState(task="clean")
    adapter.handle_interjection("turn", None, state)
    assert state.language_context["plan"] == "turn to the left now"
 def test_direct_task_adapter_delegates_action_chunk():
    class Policy:
        def predict_action_chunk(self, observation):
            return ("chunk", observation)
    observation = {"task": "pick up the cube"}
    adapter = DirectTaskPolicyAdapter(Policy())
    assert adapter.select_action(observation, RuntimeState()) == ("chunk", observation)
    assert adapter.generate_text("subtask", observation, RuntimeState()) == ""
 def test_flat_policy_registry_reuses_direct_task_adapter():
    from lerobot.runtime.registry import get_language_adapter_factory
    assert get_language_adapter_factory("pi05") is DirectTaskPolicyAdapter
    assert get_language_adapter_factory("molmoact2") is DirectTaskPolicyAdapter
@@ -1,75 +0,0 @@
 # 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.
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 import pytest
 import torch
 from lerobot.runtime.cli import _build_rollout_runtime_io, _parse_args
 def test_parse_args_preserves_rollout_robot_overrides():
    args = _parse_args(
        [
            "--policy.path=checkpoint",
            "--robot.type=so101_follower",
            "--robot.calibration_dir=/tmp/calibration",
        ]
    )
    assert args.robot_type == "so101_follower"
    assert "--robot.calibration_dir=/tmp/calibration" in args.raw_argv
 def test_parse_args_rejects_removed_dataset_replay_flags():
    with pytest.raises(SystemExit):
        _parse_args(["--policy.path=checkpoint", "--dataset.repo_id=dataset"])
 def test_rollout_runtime_io_uses_context_processors():
    robot = MagicMock()
    robot.robot_type = "mock_robot"
    robot.cameras = {}
    robot.get_observation.return_value = {"joint.pos": 1.5}
    ctx = SimpleNamespace(
        hardware=SimpleNamespace(robot_wrapper=robot),
        runtime=SimpleNamespace(cfg=SimpleNamespace(device="cpu")),
        processors=SimpleNamespace(
            robot_observation_processor=lambda observation: observation,
            robot_action_processor=lambda pair: pair[0],
        ),
        policy=SimpleNamespace(
            preprocessor=lambda observation: observation,
            postprocessor=lambda action: action,
        ),
        data=SimpleNamespace(
            dataset_features={
                "observation.state": {
                    "dtype": "float32",
                    "shape": (1,),
                    "names": ["joint.pos"],
                },
                "action": {"dtype": "float32", "shape": (1,), "names": ["joint.pos"]},
            }
        ),
    )
    provider, executor = _build_rollout_runtime_io(ctx, rerun_log=False, get_task=lambda: "move")
    observation = provider()
    executor(torch.tensor([[2.0]]))
    assert observation["observation.state"].shape == (1, 1)
    robot.send_action.assert_called_once_with({"joint.pos": 2.0})
@@ -1,100 +0,0 @@
 # 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.
 import threading
 import time
 from lerobot.runtime import LanguageConditionedRuntime, Tick
 class FakeAdapter:
    def __init__(self):
        self.updated = False
        self.interjections = []
    def select_action(self, observation, state):
        assert observation == {"observation.state": 1}
        assert state.task == "clean"
        return ["a0", "a1"]
    def update_language_state(self, observation, state):
        self.updated = True
        state.set_context("subtask", "pick cup", label="subtask")
    def handle_interjection(self, user_text, observation, state):
        self.interjections.append(user_text)
        state.set_context("plan", "new plan", label="plan")
 def test_runtime_tick_updates_language_enqueues_and_dispatches_action():
    adapter = FakeAdapter()
    executed = []
    runtime = LanguageConditionedRuntime(
        policy_adapter=adapter,
        observation_provider=lambda: {"observation.state": 1},
        action_executor=executed.append,
    )
    runtime.set_task("clean")
    logs = runtime.step_once()
    assert adapter.updated
    assert runtime.state.language_context["subtask"] == "pick cup"
    assert executed == ["a0"]
    assert list(runtime.state.action_queue) == ["a1"]
    assert "  subtask: pick cup" in logs
 def test_runtime_handles_user_interjection():
    adapter = FakeAdapter()
    runtime = LanguageConditionedRuntime(
        policy_adapter=adapter,
        observation_provider=lambda: {"observation.state": 1},
    )
    runtime.set_task("clean")
    runtime.state.extra["recent_interjection"] = "please say ok"
    runtime.state.emit("user_interjection")
    runtime.step_once()
    assert "please say ok" in adapter.interjections
    assert runtime.state.language_context["plan"] == "new plan"
 def test_prompt_change_discards_in_flight_action_chunk():
    started = threading.Event()
    release = threading.Event()
    class BlockingAdapter(FakeAdapter):
        def select_action(self, observation, state):
            started.set()
            assert release.wait(timeout=2)
            return ["stale"]
    runtime = LanguageConditionedRuntime(
        policy_adapter=BlockingAdapter(),
        observation_provider=lambda: {"observation.state": 1},
    )
    runtime.set_task("old task")
    runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic())
    inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True})
    inference.start()
    assert started.wait(timeout=2)
    runtime.set_task("new task")
    release.set()
    inference.join(timeout=2)
    assert not inference.is_alive()
    assert list(runtime.state.action_queue) == []
@@ -1,82 +0,0 @@
 # 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.
 import sys
 from types import SimpleNamespace
 import numpy as np
 from lerobot.runtime.sim_robocasa import RoboCasaSimBackend
 from lerobot.utils.video_annotation import annotate_frame
 def test_overlay_draws_each_label_once(monkeypatch):
    put_text_calls = []
    rectangle_calls = []
    def put_text(image, text, origin, font, scale, color, thickness, line_type):
        put_text_calls.append((text, color, thickness))
        return image
    def rectangle(image, start, end, color, thickness):
        rectangle_calls.append((start, end, color, thickness))
        return image
    def add_weighted(src1, alpha, src2, beta, gamma, *, dst):
        dst[:] = src1 * alpha + src2 * beta + gamma
        return dst
    fake_cv2 = SimpleNamespace(
        FONT_HERSHEY_SIMPLEX=0,
        LINE_AA=16,
        getTextSize=lambda text, font, scale, thickness: ((len(text) * 7, 10), 0),
        putText=put_text,
        rectangle=rectangle,
        addWeighted=add_weighted,
    )
    monkeypatch.setitem(sys.modules, "cv2", fake_cv2)
    frame = np.full((120, 480, 3), 200, dtype=np.uint8)
    annotated = annotate_frame(
        frame,
        (("Task", "close the fridge"), ("Subtask", "reach for the handle"), ("Memory", None)),
    )
    assert [call[0] for call in put_text_calls] == [
        "Task: close the fridge",
        "Subtask: reach for the handle",
    ]
    assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
    assert len(rectangle_calls) == 1
    assert not np.shares_memory(annotated, frame)
 def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch):
    backend = object.__new__(RoboCasaSimBackend)
    frame = np.full((8, 8, 3), 42, dtype=np.uint8)
    written = []
    backend.record = False
    backend.runtime_state = None
    backend._multiview_frame = lambda: frame
    backend._current_task = lambda: "task"
    backend._subtask_getter = None
    backend._memory_getter = None
    backend._latest_frame = None
    backend._write_live_frame = written.append
    monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image)
    backend._capture_frame()
    assert backend._latest_frame is frame
    assert written == [frame]
@@ -185,18 +185,25 @@ def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(m
    peft_config_from_pretrained = MagicMock(return_value=peft_config)
    adapted_policy = MagicMock()
    peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
-    monkeypatch.setitem(
+    require_package = MagicMock()
-        sys.modules,
+    monkeypatch.setattr(rollout_context, "require_package", require_package)
-        "peft",
+    monkeypatch.setattr(
-        SimpleNamespace(
+        rollout_context,
-            PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained),
+        "PeftConfig",
-            PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained),
+        SimpleNamespace(from_pretrained=peft_config_from_pretrained),
-        ),
+        raising=False,
    )
    monkeypatch.setattr(
        rollout_context,
        "PeftModel",
        SimpleNamespace(from_pretrained=peft_model_from_pretrained),
        raising=False,
    )
    policy = rollout_context._load_pretrained_policy(policy_config)
    assert policy is adapted_policy
    require_package.assert_called_once_with("peft", extra="peft")
    peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
    policy_class.from_pretrained.assert_called_once_with(
        pretrained_name_or_path="user/base-policy",
Author	SHA1	Message	Date
Predrag Cvetkovic	35339d31e5	fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames (#3749 ) * fix(datasets): bound memory of augment_dataset_quantile_stats by sampling frames Per-episode stats previously materialized every frame (and decoded up to 16 episodes in parallel), so peak memory scaled with episode length and OOM'd on large datasets (#2889). Numeric features are now read in full from the table (exact), while only image/video frames are sub-sampled per episode using the existing sample_indices heuristic. Worker count is configurable via LEROBOT_STATS_MAX_WORKERS; --no-sampling restores exact behavior. * Update tests/datasets/test_augment_quantile_stats.py Co-authored-by: Haoming Song <1847575517@qq.com> Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> --------- Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Haoming Song <1847575517@qq.com>	2026-07-29 12:32:03 +02:00
Steven Palma	f37be3edbe	fix(eval): prevent eval_policy crash when start_seed is None and num_envs>1 (#4203 ) * fix(eval): align seed list length with num_envs when unseeded eval_policy appended a single None per batch to all_seeds on the unseeded path while the reward and success lists grew by num_envs. The per-episode zip(..., strict=True) then raised ValueError for num_envs > 1. Extend all_seeds by num_envs so the lists stay aligned. * chore(tests): delete lerobot_eval test --------- Co-authored-by: Devin Lai <markauto75@gmail.com>	2026-07-28 18:41:28 +02:00
Khalil Meftah	4d076845ac	fix peft factory test mocking (#4201 )	2026-07-28 17:54:58 +02:00
Steven Palma	413972c812	fix(env): eval env lifecycle (#4194 ) Co-authored-by: itxaiohanglover <1531137510@qq.com> Co-authored-by: nickndeng <nickndeng@gmail.com> Co-authored-by: nickndeng <107904079+nickndeng@users.noreply.github.com> Co-authored-by: Pepijn <138571049+pkooij@users.noreply.github.com>	2026-07-28 16:45:48 +02:00
Steven Palma	0449aa02f6	fix(utils): handle missing/unresponsive TTS on Linux (#4199 ) * fix: handle missing/unresponsive TTS on Linux spd-say may be installed but hang indefinitely when speech-dispatcher is not running. Add a 5s timeout and catch TimeoutExpired alongside FileNotFoundError so recording continues without audio. * chore(utils): add log warning for say --------- Co-authored-by: Jiwen Cai <jiwenc@nvidia.com>	2026-07-28 16:45:32 +02:00
Alexandre Edmond	a05c0833e1	chore(mypy): cover annotations and transforms (#3860 ) Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>	2026-07-28 16:25:45 +02:00
Alexandre Edmond	7b76d94c5b	Handle resuming empty local datasets (#3859 ) Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>	2026-07-28 16:25:42 +02:00
Khalil Meftah	ec2dbc1c98	fix(policy): honor revisions when loading PEFT checkpoints (#4189 )	2026-07-28 15:41:47 +02:00
Steven Palma	d526785e47	fix(dependencies): protect peft import (#4188 ) Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>	2026-07-28 14:57:50 +02:00
Steven Palma	4af7c70664	refactor(logging): standardize logging with getLogger(__name__) in scripts (#4192 ) * refactor(logging): replace print with logger in lerobot_info * refactor(logging): replace print with logger in convert_dataset_v21_to_v30 * refactor(logging): replace print with logger in lerobot_annotate * refactor(logging): replace print with logger in lerobot_dataset_viz * refactor(logging): replace print with logger in lerobot_eval * refactor(logging): replace print with logger in lerobot_find_cameras * refactor(logging): replace print with logger in lerobot_find_joint_limits * refactor(logging): replace print with logger in lerobot_find_port * refactor(logging): replace print with logger in lerobot_imgtransform_viz * refactor(logging): replace print with logger in lerobot_setup_can * refactor(logging): replace print with logger in lerobot_teleoperate * refactor(logging): replace print with logger in lerobot_train_tokenizer * fix(logging): preserve CLI output semantics --------- Co-authored-by: ailisilob <2248345706@qq.com>	2026-07-28 14:42:38 +02:00
charlie8612	a855570097	feat(motors): add XH540-W150, XC330-T288, XC330-T181 to Dynamixel tables (#3815 ) Register three X-series Dynamixel models so they can be driven by DynamixelMotorsBus: XH540-W150 (model 1110), XC330-T288 (1220) and XC330-T181 (1210). All are standard Protocol 2.0 X-series motors that share the existing X_SERIES control/baudrate/encoding tables and 4096 resolution; only the model number and operating-mode list are model-specific. Values verified against the ROBOTIS e-manual. These motors are used by the ROBOTIS OMY-L100 arm, among others. Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>	2026-07-28 14:27:48 +02:00
Steven Palma	167e22ba51	feat(record): add `--dataset.no_stamp` to opt out of repo_id timestamping (#4190 ) * feat(record): add `--dataset.no_stamp` to opt out of repo_id timestamping stamp_repo_id() unconditionally appended a date-time tag to repo_id for every new (non-resume) dataset, so users managing their own versioned repo names (e.g. for a later lerobot-edit-dataset merge) could not opt out. Add a no_stamp field to DatasetRecordConfig and make stamp_repo_id() a no-op when it is set. The flag covers both lerobot-record and lerobot-rollout since they share this config, and no call-site changes are needed. Fixes #3722. * chore(tests): delete dataset config test --------- Co-authored-by: Philipp Sinitsin <ph.sinitsin@gmail.com>	2026-07-28 14:06:28 +02:00
Steven Palma	00c25c65c2	feat(camera): add manual exposure, gain, and white balance options for RealSense cameras (#4130 ) * feat(camera): add manual exposure, gain, and white balance options for RealSense cameras The RealSense camera integration lacked sensor-level controls, causing issues like unstable lighting from auto-exposure hunting. This adds optional `exposure`, `gain`, and `white_balance` fields to RealSenseCameraConfig that disable the corresponding auto modes and apply fixed values when set. * fix: support D405 stereo module for sensor options D405 exposes color stream via "Stereo Module", not "RGB Camera". Fall back to Stereo Module when RGB Camera is not found. * test(camera): add unit tests + range-aware errors for RealSense sensor options Address PR #3220 review: - Wrap set_option calls; re-raise ValueError with option name, value, and sensor.get_option_range() diagnostics on out-of-range values. - Add unit tests for _get_color_sensor (RGB Camera, D405 Stereo Module fallback, diagnostic error) and _configure_sensor_options (no-op, all values, unsupported warns, partial config, out-of-range raise). * fix(realsense): validate manual color controls * refactor(camera): apply feedback --------- Co-authored-by: Lev Kozlov <kozlov.l.a10@gmail.com>	2026-07-28 13:41:09 +02:00
Steven Palma	23f6d5dabd	fix(cameras): release device handle when connect() setup fails (#4187 ) Co-authored-by: Ryan Rana <39924576+RyanRana@users.noreply.github.com>	2026-07-28 13:21:06 +02:00
Xingdong Zuo	9b25b7fe0a	feat(lekiwi): support LeKiwi in the rollout/eval CLI (#3742 ) * feat(lekiwi): support LeKiwi in the rollout/eval CLI Register the lekiwi robot in lerobot_rollout.py so policies can be rolled out on a LeKiwi, and keep base-velocity (.vel) features in build_rollout_context. LeKiwi's observation.state and action are 9-dim (6 arm .pos + x/y/theta.vel) and the policy is normalized on all 9. The old filter kept only .pos features, so it fed a 6-dim vector into a 9-dim normalizer (RuntimeError, size 6 vs 9) and silently dropped the base velocities from the action, leaving the base unable to move. Keeping both .pos and .vel fixes both. Pure-arm robots have no .vel keys, so this is a no-op for them. * style: format LeKiwi rollout action features --------- Co-authored-by: Steven Palma <imstevenpmwork@ieee.org> Co-authored-by: Steven Palma <steven.palma@huggingface.co>	2026-07-28 12:52:02 +02:00
Steven Palma	c1b6ea85d6	feat(rl): add multiprocessing option to training pipeline and sets spawn as default + guard (#4140 )	2026-07-28 12:19:53 +02:00
Steven Palma	ffe25afb8f	fix(processors): wrong feature key dropped in delta-action transform_features (#4165 )	2026-07-28 11:18:23 +02:00