update image

Add GaussianNoise, MotionBlur, JPEGCompression, GaussianPatchBrightness,
RandomShadow, CoarseDropout, GammaCorrection, and PlanckianJitter.

Each transform addresses a real-world failure mode not covered by the
existing 6 defaults (sensor noise, motion blur, compression artifacts,
uneven lighting, cast shadows, partial occlusion, exposure variation,
color temperature shift).

All transforms are pure PyTorch, follow the make_params/transform
pattern, and integrate with ImageTransformConfig via a registry.

README.md

@@ -101,13 +101,13 @@ lerobot-train \
   --dataset.repo_id=lerobot/aloha_mobile_cabinet
 ```
 
-| Category                   | Models                                                                                                                                                                                                                                                                                                                                                                                                                       |
-| -------------------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| **Imitation Learning**     | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md)                                                                                                                                                                                                      |
-| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon)                                                                                                                                                                                                                                                                                                                  |
-| **VLAs Models**            | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [Pi052](./docs/source/pi052.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx), [EVO1](./docs/source/evo1.mdx) |
-| **World Models**           | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx)                                                                                                                                                                                                                                                                                                     |
-| **Reward Models**          | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx)                                                                                                                                                                                                                                                                                                           |
+| Category                   | Models                                                                                                                                                                                                                                                                                                                                                                                     |
+| -------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
+| **Imitation Learning**     | [ACT](./docs/source/policy_act_README.md), [Diffusion](./docs/source/policy_diffusion_README.md), [VQ-BeT](./docs/source/policy_vqbet_README.md), [Multitask DiT Policy](./docs/source/policy_multi_task_dit_README.md)                                                                                                                                                                    |
+| **Reinforcement Learning** | [HIL-SERL](./docs/source/hilserl.mdx), [TDMPC](./docs/source/policy_tdmpc_README.md) & QC-FQL (coming soon)                                                                                                                                                                                                                                                                                |
+| **VLAs Models**            | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx), [EVO1](./docs/source/evo1.mdx) |
+| **World Models**           | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx)                                                                                                                                                                                                                                                                   |
+| **Reward Models**          | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx)                                                                                                                                                                                                                                                                         |
 
 Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub.
 

docker/Dockerfile.internal

@@ -68,17 +68,16 @@ ENV HOME=/home/user_lerobot \
 # issues with MuJoCo and OpenGL drivers.
 RUN uv venv --python python${PYTHON_VERSION}
 
-# Install Python dependencies for caching
+# Install third-party dependencies separately for layer caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-COPY --chown=user_lerobot:user_lerobot src/ src/
-
-RUN uv sync --locked --extra all --no-cache
+RUN uv sync --locked --extra all --no-install-project --no-cache
 
 RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
 
-# Copy the rest of the application source code
+# Copy the application source code and install the local project
 # Make sure to have the git-LFS files for testing
 COPY --chown=user_lerobot:user_lerobot . .
+RUN uv sync --locked --extra all --no-cache
 
 # Set the default command
 CMD ["/bin/bash"]

docker/Dockerfile.user

@@ -60,15 +60,14 @@ ENV HOME=/home/user_lerobot \
 # run other Python projects in the same container without dependency conflicts.
 RUN uv venv
 
-# Install Python dependencies for caching
+# Install third-party dependencies separately for layer caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-COPY --chown=user_lerobot:user_lerobot src/ src/
+RUN uv sync --locked --extra all --no-install-project --no-cache
 
-RUN uv sync --locked --extra all --no-cache
-
-# Copy the rest of the application code
+# Copy the application code and install the local project
 # Make sure to have the git-LFS files for testing
 COPY --chown=user_lerobot:user_lerobot . .
+RUN uv sync --locked --extra all --no-cache
 
 # Set the default command
 CMD ["/bin/bash"]

docs/source/_toctree.yml

@@ -63,8 +63,6 @@
     title: π₀-FAST (Pi0Fast)
   - local: pi05
     title: π₀.₅ (Pi05)
-  - local: pi052
-    title: π₀.₅ with language supervision (Pi052)
   - local: molmoact2
     title: MolmoAct2
   - local: vla_jepa

docs/source/bring_your_own_policies.mdx

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

docs/source/cameras.mdx

@@ -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>
 

docs/source/inference.mdx

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

docs/source/language_and_recipes.mdx

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

docs/source/pi052.mdx

@@ -1,274 +0,0 @@
-# π₀.₅ with language supervision (Pi052)
-
-Pi052 extends [Pi05](./pi05) with a trainable PaliGemma language head and a
-runtime that alternates language generation with action generation. A single
-checkpoint can predict a low-level subtask, optionally update memory or answer
-visual questions, and condition its flow-matching action expert on that text.
-
-Use Pi05 when you only need task-conditioned actions. Use Pi052 when the policy
-must generate or consume intermediate language during a rollout.
-
-## How Pi052 differs from Pi05
-
-| Capability          | Pi05                                                   | Pi052                                                                             |
-| ------------------- | ------------------------------------------------------ | --------------------------------------------------------------------------------- |
-| Action model        | PaliGemma vision-language prefix + Gemma action expert | Same base architecture                                                            |
-| Language head       | Not trained for runtime generation                     | Re-enabled and trained with text cross-entropy                                    |
-| Action conditioning | Episode task                                           | Active low-level subtask plus normalized robot state                              |
-| Training targets    | Flow-matching actions                                  | Flow actions, recipe-selected text, and optional FAST action tokens               |
-| Dataset requirement | Standard images, state, actions, and task              | The same fields plus language annotations for every language capability you train |
-| Rollout             | Direct task-to-action policy                           | Hierarchical task → subtask → action loop, with optional memory and VQA           |
-
-Pi052 can initialize from a Pi05 checkpoint. The policy architecture remains
-compatible, while Pi052 builds its own processors so recipe labels and FAST
-labels are not silently replaced by the Pi05 processor stack.
-
-## Install
-
-Install LeRobot with the PI dependencies:
-
-```bash
-git clone https://github.com/huggingface/lerobot.git
-cd lerobot
-python -m venv .venv
-source .venv/bin/activate
-pip install -e ".[pi]"
-```
-
-The `pi` extra includes the PaliGemma/FAST dependencies. Install
-`liger-kernel` for the supported fused training kernels; optional FlashRT
-backends also require the Hugging Face `kernels` package and a supported CUDA
-GPU.
-
-## Prepare language-annotated data
-
-Pi052 does not infer supervised subtasks from a normal LeRobot dataset during
-training. The dataset must contain the language targets used by the selected
-recipe in the optional `language_persistent` and `language_events` columns.
-
-At minimum, annotate a continuous `subtask` timeline so each training frame has
-an active low-level instruction. Add `memory`, VQA, interjections, and speech
-annotations only if the recipe trains those capabilities.
-
-The provided recipes are:
-
-| Recipe                                | Required annotations                                                    | Trains                                                             |
-| ------------------------------------- | ----------------------------------------------------------------------- | ------------------------------------------------------------------ |
-| `recipes/subtask.yaml`                | `subtask`                                                               | Subtask prediction and subtask-conditioned actions                 |
-| `recipes/subtask_joint.yaml`          | `subtask`                                                               | Paper-style joint sequence: subtask text and actions in one sample |
-| `recipes/subtask_mem.yaml`            | `subtask`, `memory`                                                     | Subtasks, actions, and memory updates                              |
-| `recipes/subtask_mem_vqa_speech.yaml` | `subtask`, `memory`, `vqa`; interjection/speech rows for those branches | Subtasks, actions, memory, VQA, and spoken replies                 |
-
-The blend recipes factorize training into separate high-level (task → subtask)
-and low-level (subtask → actions) samples, matching how inference decomposes
-π(a|o, subtask)·π(subtask|o, task). `recipes/subtask_joint.yaml` instead uses
-the π0.5 paper's single-sequence layout — the supervised subtask span is
-attended causally and conditions the FAST and flow losses in the same forward.
-Checkpoints trained with the joint recipe must set
-`--policy.joint_subtask_conditioning=true` at inference so the flow prefix
-rebuilds the same layout (task turn with state, then the generated subtask as a
-causal assistant turn); leave it `false` for the blend recipes.
-
-Use `lerobot-annotate` to generate these columns. The repository includes a
-Hugging Face Jobs launcher that you can edit for your source and destination
-datasets. For a local annotation run, first install
-`pip install -e ".[annotations]"`:
-
-```bash
-HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
-```
-
-Before a long training run, inspect several episodes and verify that subtasks
-are temporally correct and cover the full demonstration. See
-[Annotation Pipeline](./annotation_pipeline) for generation and validation, and
-[Language Columns and Recipes](./language_and_recipes) for the schema and
-recipe resolver.
-
-<Tip>
-  If a dataset has no language columns, recipe rendering becomes a no-op and
-  Pi052 falls back to the plain Pi05 prompt path. This is useful for
-  compatibility but does not train the language planner.
-</Tip>
-
-## Train Pi052
-
-This example initializes Pi052 from the native Pi052 initialization checkpoint
-and trains the default subtask-and-memory recipe:
-
-```bash
-lerobot-train \
-    --dataset.repo_id=${HF_USER}/my_language_annotated_dataset \
-    --policy.type=pi052 \
-    --policy.pretrained_path=lerobot/pi052_base \
-    --policy.recipe_path=recipes/subtask_mem.yaml \
-    --policy.dtype=bfloat16 \
-    --policy.device=cuda \
-    --policy.freeze_vision_encoder=false \
-    --policy.gradient_checkpointing=true \
-    --batch_size=8 \
-    --steps=30000 \
-    --output_dir=outputs/pi052 \
-    --job_name=pi052 \
-    --wandb.enable=true
-```
-
-For subtask-only data, change the recipe to `recipes/subtask.yaml` and disable
-memory during rollout. Start with a small run and confirm that W&B examples show
-the expected prompt, text target, and action endpoints before scaling up.
-
-### Main training controls
-
-| Option                              |                    Default | Purpose                                                             |
-| ----------------------------------- | -------------------------: | ------------------------------------------------------------------- |
-| `policy.recipe_path`                | `recipes/subtask_mem.yaml` | Selects the language/action objective mixture                       |
-| `policy.text_loss_weight`           |                      `1.0` | Language-head cross-entropy weight; `0` disables text training      |
-| `policy.flow_loss_weight`           |                     `10.0` | Continuous action flow-loss weight                                  |
-| `policy.enable_fast_action_loss`    |                     `true` | Adds discrete FAST action-token supervision                         |
-| `policy.fast_action_loss_weight`    |                      `1.0` | FAST cross-entropy weight                                           |
-| `policy.knowledge_insulation`       |                     `true` | Blocks action-loss gradients through the VLM K/V path               |
-| `policy.flow_num_repeats`           |                        `5` | Reuses one VLM prefix for independent denoising targets             |
-| `policy.lm_head_lr_scale`           |                      `1.0` | Scales language-head learning rate; `1.0` uses the base rate        |
-| `policy.fast_skip_tokens`           |                     `1152` | FAST id offset; skips `<seg>`+`<loc>` so VQA and FAST never collide |
-| `policy.joint_subtask_conditioning` |                    `false` | Rebuilds the joint-sequence prefix at inference (see recipes)       |
-
-`fast_skip_tokens=1152` places FAST codes below PaliGemma's `<loc>` range.
-openpi's pi0-FAST convention is `128` (FAST occupies the `<loc>` ids); use that
-value only to stay weight-compatible with checkpoints trained that way, and
-avoid combining it with the VQA recipe, whose `<loc>` targets would share
-embedding rows with FAST codes.
-
-The loss weights are starting points, not dataset-independent constants. Track
-flow loss and text/FAST losses separately, and inspect generated subtasks rather
-than selecting a checkpoint from total loss alone.
-
-### Dataset-specific FAST tokenizer
-
-The universal FAST tokenizer works out of the box. For a large or
-embodiment-specific dataset, Pi052 can fit and cache a tokenizer on normalized
-actions before training:
-
-```bash
-lerobot-train \
-    ... \
-    --policy.auto_fit_fast_tokenizer=true \
-    --policy.fast_tokenizer_fit_samples=4096
-```
-
-The fit runs once per dataset/tokenizer configuration. Keep
-`auto_fit_fast_tokenizer=false` when you do not want the extra preprocessing
-pass.
-
-## Training performance
-
-Pi052 uses optimized training paths by default:
-
-- batches repeated flow targets and suffix projections instead of replaying
-  small operations in Python;
-- caches constant action masks and computes RoPE positions once per forward;
-- selects the text/FAST cross-entropy implementation from target shape and
-  sparsity;
-- skips the mathematically dead VLM/vision backward on knowledge-insulated,
-  flow-only batches;
-- uses native non-reentrant SigLIP layer checkpointing when gradient
-  checkpointing is enabled; and
-- retains the Liger RoPE/GeGLU kernels while avoiding the slower LayerNorm
-  patch at SigLIP shapes.
-
-Optional training backends are disabled by default:
-
-| Option                                 | When to try it                                                                                    |
-| -------------------------------------- | ------------------------------------------------------------------------------------------------- |
-| `policy.use_flashrt_adarms=true`       | Fused adaptive RMSNorm and gated residuals on supported CUDA GPUs                                 |
-| `policy.use_compiled_text_ce=true`     | Compiled materialized-logit CE buckets                                                            |
-| `policy.use_compiled_vision=true`      | Compiled vision only when the vision pass has no gradients                                        |
-| `policy.use_flex_attention=true`       | Profiled CUDA setups with knowledge insulation and `flow_num_repeats > 1`; otherwise SDPA is used |
-| `policy.use_manual_attention=true`     | Explicitly profiled shapes where materialized attention is faster                                 |
-| `policy.manual_attention_scope=action` | Restricts manual attention to action queries                                                      |
-
-Do not enable every backend blindly. Flex and manual attention are mutually
-exclusive, and attention/AdaRMS alternatives require knowledge insulation.
-The benchmark-best configuration used compiled text CE and FlashRT AdaRMS,
-with Flex/manual attention and compiled vision disabled.
-
-### Reported training benchmarks
-
-These benchmarks measure complete optimizer steps with three real camera
-inputs, BF16 transformer/action execution, FP32 vision, fused AdamW, and no
-video decoding or network I/O. Results vary with GPU, batch shape, annotation
-mixture, and checkpointing:
-
-| Workload                   |     RTX PRO 6000 Blackwell |                 A100 80 GB |
-| -------------------------- | -------------------------: | -------------------------: |
-| Full flow + text, batch 1  | 4.75× vs checkpointing off | 3.33× vs checkpointing off |
-| Full flow + text, batch 8  | 2.16× vs checkpointing off | 1.66× vs checkpointing off |
-| Full flow + text, batch 64 |  1.24× vs checkpointing on |  1.15× vs checkpointing on |
-| Flow-only, batch 1         | 3.70× vs checkpointing off | 3.58× vs checkpointing off |
-| Flow-only, batch 64        |  3.76× vs checkpointing on |  3.61× vs checkpointing on |
-
-On those 80 GB GPUs, full training was fastest without gradient checkpointing
-through batch 8, then required checkpointing at batch 16 and above. Treat that
-as a tuning rule to test on your hardware, not a universal threshold. Flow-only
-means both text and FAST supervision are disabled; it is useful for action-only
-ablation or post-training but does not learn the language runtime.
-
-## Inference performance
-
-Pi052 has two inference loops, and both avoid repeatedly encoding the expensive
-multimodal prefix:
-
-1. **Action denoising** encodes the image/language prefix once, reuses its KV
-   cache across flow steps, precomputes the timestep schedule on-device, and
-   crops temporary suffix K/V instead of cloning the prefix cache.
-2. **Language decoding** uses autoregressive KV caching, so each new token only
-   processes the sampled token against cached image/language keys instead of
-   rerunning the full prefix.
-
-The runtime also runs language and actions at different rates. Increase
-`--subtask_chunks_per_gen` when a subtask remains valid across several action
-chunks, lower `--high_level_hz`, or use `--direct_subtask` to bypass language
-generation entirely. These settings reduce compute but also slow replanning.
-
-`--fp8` enables the optional FlashRT inference MLP swap on supported CUDA GPUs.
-It calibrates on the first observation and falls back to BF16 when unavailable;
-because FP8 can change outputs slightly, validate task success before using it
-for production rollouts.
-
-## Run a checkpoint
-
-RoboCasa:
-
-```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 view. Without
-`--sim.direct_subtask`, Pi052 generates the low-level subtask; with it, each
-prompt becomes the action policy's subtask directly.
-
-The same runtime supports real robots. See [Interactive language
-control](./inference#interactive-language-control) for the real-arm command,
-safety behavior, and runtime controls.
-
-## Troubleshooting
-
-- **No text loss or generated subtasks:** confirm the selected recipe can bind
-  the annotations on sampled frames and that `policy.text_loss_weight > 0`.
-- **Subtasks look plausible but actions fail:** verify subtask boundaries,
-  normalized state/action statistics, and that low-level recipe samples are
-  present.
-- **Text collapses to repeated or location tokens:** inspect text-target
-  coverage, language-head learning rate, and the balance between flow, FAST,
-  and text losses.
-- **Out of memory:** reduce batch size first, then enable gradient
-  checkpointing. Do not enable compiled or alternative attention backends
-  without profiling their memory on your camera count.
-- **Slow rollout:** separate action latency from language latency, then tune
-  `--subtask_chunks_per_gen`, `--high_level_hz`, and the number of flow
-  inference steps.

docs/source/pi0fast.mdx

@@ -109,21 +109,15 @@ lerobot-train \
 
 ### Key Training Parameters
 
-| Parameter                               | Description                                        | Default                         |
-| --------------------------------------- | -------------------------------------------------- | ------------------------------- |
-| `--policy.gradient_checkpointing=true`  | Reduces memory usage significantly during training | `false`                         |
-| `--policy.dtype=bfloat16`               | Use mixed precision training for efficiency        | `float32`                       |
-| `--policy.chunk_size`                   | Number of action steps to predict (action horizon) | `50`                            |
-| `--policy.n_action_steps`               | Number of decoded action steps to execute          | `50`                            |
-| `--policy.max_action_tokens`            | Maximum number of FAST tokens per action chunk     | `256`                           |
-| `--policy.action_tokenizer_name`        | FAST tokenizer to use                              | `lerobot/fast-action-tokenizer` |
-| `--policy.auto_fit_fast_tokenizer=true` | Fit and cache a tokenizer for the training dataset | `false`                         |
-| `--policy.compile_model=true`           | Enable torch.compile for faster training           | `false`                         |
-
-Set `--policy.auto_fit_fast_tokenizer=true` to sample action chunks from the
-training dataset and cache a fitted tokenizer under
-`~/.cache/lerobot/fast_tokenizers`. This also works when fine-tuning with
-`--policy.path`; leave it disabled to retain the checkpoint's tokenizer.
+| Parameter                              | Description                                        | Default                         |
+| -------------------------------------- | -------------------------------------------------- | ------------------------------- |
+| `--policy.gradient_checkpointing=true` | Reduces memory usage significantly during training | `false`                         |
+| `--policy.dtype=bfloat16`              | Use mixed precision training for efficiency        | `float32`                       |
+| `--policy.chunk_size`                  | Number of action steps to predict (action horizon) | `50`                            |
+| `--policy.n_action_steps`              | Number of action steps to execute                  | `50`                            |
+| `--policy.max_action_tokens`           | Maximum number of FAST tokens per action chunk     | `256`                           |
+| `--policy.action_tokenizer_name`       | FAST tokenizer to use                              | `lerobot/fast-action-tokenizer` |
+| `--policy.compile_model=true`          | Enable torch.compile for faster training           | `false`                         |
 
 ## Inference
 
@@ -157,9 +151,6 @@ actions = policy.predict_action_chunk(batch)
 
 The model takes images, text instructions, and robot state as input, and outputs discrete FAST tokens that are decoded back to continuous actions.
 
-PI0-FAST always decodes a complete `chunk_size` action chunk. `n_action_steps` controls only
-how many actions from that chunk are executed before the policy predicts again.
-
 ## Configuration Options
 
 | Parameter            | Description                                     | Default    |

pyproject.toml

@@ -150,7 +150,6 @@ pygame-dep = ["pygame>=2.5.1,<2.7.0"]
 # There is no cmeel-urdfdom 5.x; <5 selects the 4.x ABI the placo/pin wheels are built against.
 placo-dep = ["placo>=0.9.6,<0.9.16", "cmeel-urdfdom>=4,<5", "cmeel-tinyxml2<11"]
 transformers-dep = ["transformers>=5.4.0,<5.6.0"]
-sentencepiece-dep = ["sentencepiece>=0.2.0,<0.3.0"] # FAST action tokenizer backend (pi052, pi0_fast)
 grpcio-dep = ["grpcio>=1.73.1,<2.0.0", "protobuf>=6.31.1,<8.0.0"]
 accelerate-dep = ["accelerate>=1.14.0,<2.0.0"]
 can-dep = ["python-can>=4.2.0,<5.0.0"]
@@ -213,7 +212,7 @@ wallx = [
     "torchdiffeq>=0.2.4,<0.3.0",
     "lerobot[qwen-vl-utils-dep]",
 ]
-pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]", "lerobot[sentencepiece-dep]"]
+pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
 molmoact2 = ["lerobot[transformers-dep]", "lerobot[peft-dep]", "lerobot[scipy-dep]"]
 smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "lerobot[accelerate-dep]"]
 multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
@@ -495,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.*"

src/lerobot/cameras/opencv/camera_opencv.py

@@ -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, 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
+        self.capture_width: int | None = None
+        self.capture_height: int | None = None
+        self._reset_connection_settings()
 
     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,17 +172,25 @@ class OpenCVCamera(Camera):
                 f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
             )
 
-        self._configure_capture_settings()
-        self._start_read_thread()
+        try:
+            self._configure_capture_settings()
+            self._start_read_thread()
 
-        if warmup and self.warmup_s > 0:
-            start_time = time.time()
-            while time.time() - start_time < self.warmup_s:
-                self.async_read(timeout_ms=self.warmup_s * 1000)
-                time.sleep(0.1)
-            with self.frame_lock:
-                if self.latest_frame is None:
-                    raise ConnectionError(f"{self} failed to capture frames during warmup.")
+            if warmup and self.warmup_s > 0:
+                start_time = time.time()
+                while time.time() - start_time < self.warmup_s:
+                    self.async_read(timeout_ms=self.warmup_s * 1000)
+                    time.sleep(0.1)
+                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,32 +328,36 @@ class OpenCVCamera(Camera):
 
         for target in targets_to_scan:
             camera = cv2.VideoCapture(target)
-            if camera.isOpened():
-                default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
-                default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
-                default_fps = camera.get(cv2.CAP_PROP_FPS)
-                default_format = camera.get(cv2.CAP_PROP_FORMAT)
+            try:
+                if camera.isOpened():
+                    default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
+                    default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
+                    default_fps = camera.get(cv2.CAP_PROP_FPS)
+                    default_format = camera.get(cv2.CAP_PROP_FORMAT)
 
-                # 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)])
+                    # 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)]
+                    )
 
-                camera_info = {
-                    "name": f"OpenCV Camera @ {target}",
-                    "type": "OpenCV",
-                    "id": target,
-                    "backend_api": camera.getBackendName(),
-                    "default_stream_profile": {
-                        "format": default_format,
-                        "fourcc": default_fourcc,
-                        "width": default_width,
-                        "height": default_height,
-                        "fps": default_fps,
-                    },
-                }
+                    camera_info = {
+                        "name": f"OpenCV Camera @ {target}",
+                        "type": "OpenCV",
+                        "id": target,
+                        "backend_api": camera.getBackendName(),
+                        "default_stream_profile": {
+                            "format": default_format,
+                            "fourcc": default_fourcc,
+                            "width": default_width,
+                            "height": default_height,
+                            "fps": default_fps,
+                        },
+                    }
 
-                found_cameras_info.append(camera_info)
+                    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._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()
+        self._cleanup_resources()
 
         logger.info(f"{self} disconnected.")

src/lerobot/cameras/realsense/camera_realsense.py

@@ -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, 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
+        self.capture_width: int | None = None
+        self.capture_height: int | None = None
+        self._reset_connection_settings()
 
     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,22 +206,31 @@ class RealSenseCamera(Camera):
                 f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
             ) from e
 
-        self._configure_capture_settings()
-        self._start_read_thread()
+        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.
-        self.warmup_s = max(self.warmup_s, 1)
+            # 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.
+            self.warmup_s = max(self.warmup_s, 1)
 
-        warmup_read = self.async_read if self.use_rgb else self.async_read_depth
-        start_time = time.time()
-        while time.time() - start_time < self.warmup_s:
-            warmup_read(timeout_ms=self.warmup_s * 1000)
-            time.sleep(0.1)
-        with self.frame_lock:
-            if (self.use_rgb and self.latest_color_frame is None) or (
-                self.use_depth and self.latest_depth_frame is None
-            ):
-                raise ConnectionError(f"{self} failed to capture frames during warmup.")
+            warmup_read = self.async_read if self.use_rgb else self.async_read_depth
+            start_time = time.time()
+            while time.time() - start_time < self.warmup_s:
+                warmup_read(timeout_ms=self.warmup_s * 1000)
+                time.sleep(0.1)
+            with self.frame_lock:
+                if (self.use_rgb and self.latest_color_frame is None) or (
+                    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._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()
-
+        self._cleanup_resources()
         logger.info(f"{self} disconnected.")

src/lerobot/cameras/realsense/configuration_realsense.py

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

src/lerobot/configs/dataset.py

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

src/lerobot/configs/default.py

@@ -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

src/lerobot/configs/recipe.py

@@ -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)
-        has_low_level = any(turn.stream == "low_level" for turn in self.messages)
-        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)."
-            )
+        if not any(turn.target for turn in self.messages):
+            raise ValueError("Message recipes must contain at least one target turn.")
 
     def _validate_blend_recipe(self) -> None:
         """Ensure each blend component is a non-empty, weighted message recipe."""

src/lerobot/configs/recipes/subtask.yaml

@@ -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}

src/lerobot/configs/recipes/subtask_joint.yaml

@@ -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}

src/lerobot/configs/recipes/subtask_mem.yaml

@@ -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}

src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml

@@ -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}

src/lerobot/datasets/dataset_metadata.py

@@ -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.episodes = load_episodes(self.root)
+        self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
+        self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
         self.stats = load_stats(self.root)
 
     def ensure_readable(self) -> None:

src/lerobot/datasets/dataset_reader.py

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

src/lerobot/datasets/factory.py

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

src/lerobot/datasets/language_render.py

@@ -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.
-    has_low_level = any(stream == "low_level" for stream in streams)
-    if not target_indices and not has_low_level:
+    if not target_indices:
         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 and not any(s == "low_level" for s in streams):
-        raise ValueError("Rendered samples must contain a target message or a low_level-stream message.")
+    if not target_indices:
+        raise ValueError("Rendered samples must contain at least one target 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(

src/lerobot/envs/configs.py

@@ -560,13 +560,7 @@ class RoboCasaEnv(EnvConfig):
             kwargs["split"] = self.split
         return kwargs
 
-    def create_envs(
-        self,
-        n_envs: int,
-        use_async_envs: bool = False,
-        terminate_on_success: bool = True,
-        horizon: int | None = None,
-    ):
+    def create_envs(self, n_envs: int, use_async_envs: bool = False):
         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,
         )
 
 

src/lerobot/envs/libero.py

@@ -384,7 +384,12 @@ class LiberoEnv(gym.Env):
 
     def close(self):
         if self._env is not None:
-            self._env.close()
+            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(

src/lerobot/envs/robocasa.py

@@ -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)
-ACTION_DIM = 12  # ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
+OBS_STATE_DIM = 16  # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
+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):
-        ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
+    Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
     """
     return {
-        "action.end_effector_position": flat_action[0:3],
-        "action.end_effector_rotation": flat_action[3:6],
-        "action.gripper_close": flat_action[6:7],
-        "action.base_motion": flat_action[7:11],
-        "action.control_mode": flat_action[11:12],
+        "action.base_motion": flat_action[0:4],
+        "action.control_mode": flat_action[4:5],
+        "action.end_effector_position": flat_action[5:8],
+        "action.end_effector_rotation": flat_action[8:11],
+        "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:

src/lerobot/envs/robotwin.py

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

src/lerobot/envs/vlabench.py

@@ -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
 

src/lerobot/motors/dynamixel/tables.py

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

src/lerobot/optim/optimizers.py

@@ -104,8 +104,6 @@ class AdamWConfig(OptimizerConfig):
     eps: float = 1e-8
     weight_decay: float = 1e-2
     grad_clip_norm: float = 10.0
-    foreach: bool | None = None
-    fused: bool | None = None
 
     def build(self, params: OptimizerParams) -> torch.optim.Optimizer:
         kwargs = asdict(self)

src/lerobot/policies/__init__.py

@@ -28,7 +28,6 @@ from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig as M
 from .pi0.configuration_pi0 import PI0Config as PI0Config
 from .pi0_fast.configuration_pi0_fast import PI0FastConfig as PI0FastConfig
 from .pi05.configuration_pi05 import PI05Config as PI05Config
-from .pi052.configuration_pi052 import PI052Config as PI052Config
 from .pretrained import PreTrainedPolicy as PreTrainedPolicy
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
@@ -57,7 +56,6 @@ __all__ = [
     "PI0Config",
     "PI0FastConfig",
     "PI05Config",
-    "PI052Config",
     "SmolVLAConfig",
     "TDMPCConfig",
     "VLAJEPAConfig",

src/lerobot/policies/factory.py

@@ -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
@@ -137,12 +144,6 @@ class ProcessorConfigKwargs(TypedDict, total=False):
     preprocessor_overrides: dict[str, Any] | None
     postprocessor_overrides: dict[str, Any] | None
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None
-    # Dataset source used by policies that optionally fit processor artifacts.
-    dataset_repo_id: str | None
-    dataset_root: str | None
-    dataset_revision: str | None
-    dataset_episodes: list[int] | None
-    dataset_exclude_episodes: list[int] | None
     dataset_meta: Any | None
 
 
@@ -177,10 +178,6 @@ def make_pre_post_processors(
         ValueError: If no processor factory exists for the given policy configuration type.
     """
     if pretrained_path:
-        # Register the PI052-only stateful tokenizer step before deserializing its pipeline.
-        if policy_cfg.type == "pi052":
-            from .pi052 import processor_pi052 as _processor_pi052  # noqa: F401
-
         if isinstance(policy_cfg, GrootConfig):
             from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
 
@@ -200,29 +197,12 @@ def make_pre_post_processors(
                 ),
             )
 
-        preprocessor_overrides = dict(kwargs.get("preprocessor_overrides") or {})
-        if policy_cfg.type == "pi0_fast" and getattr(policy_cfg, "auto_fit_fast_tokenizer", False):
-            from .pi052.fit_fast_tokenizer import resolve_fast_tokenizer
-
-            fitted_tokenizer = resolve_fast_tokenizer(
-                policy_cfg,
-                kwargs.get("dataset_repo_id"),
-                kwargs.get("dataset_root"),
-                kwargs.get("dataset_stats"),
-                kwargs.get("dataset_revision"),
-                kwargs.get("dataset_episodes"),
-                kwargs.get("dataset_exclude_episodes"),
-            )
-            tokenizer_overrides = dict(preprocessor_overrides.get("action_tokenizer_processor") or {})
-            tokenizer_overrides["action_tokenizer_name"] = fitted_tokenizer
-            preprocessor_overrides["action_tokenizer_processor"] = tokenizer_overrides
-
         preprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_model_name_or_path=pretrained_path,
             config_filename=kwargs.get(
                 "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
             ),
-            overrides=preprocessor_overrides,
+            overrides=kwargs.get("preprocessor_overrides", {}),
             to_transition=batch_to_transition,
             to_output=transition_to_batch,
             revision=pretrained_revision,
@@ -254,11 +234,6 @@ def make_pre_post_processors(
         config=policy_cfg,
         dataset_stats=kwargs.get("dataset_stats"),
         dataset_meta=kwargs.get("dataset_meta"),
-        dataset_repo_id=kwargs.get("dataset_repo_id"),
-        dataset_root=kwargs.get("dataset_root"),
-        dataset_revision=kwargs.get("dataset_revision"),
-        episodes=kwargs.get("dataset_episodes"),
-        exclude_episodes=kwargs.get("dataset_exclude_episodes"),
     )
 
 
@@ -366,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"]:
@@ -382,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:
@@ -456,7 +439,6 @@ def _make_processors_from_policy_config(
     config: PreTrainedConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
     dataset_meta: Any | None = None,
-    **optional_kwargs: Any,
 ) -> tuple[Any, Any]:
     """Create pre- and post-processors from a policy configuration using dynamic imports.
 
@@ -492,9 +474,7 @@ def _make_processors_from_policy_config(
     function = getattr(module, function_name, None)
     if function is None:
         raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
-    parameters = inspect.signature(function).parameters
     call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
-    if "dataset_meta" in parameters:
+    if "dataset_meta" in inspect.signature(function).parameters:
         call_kwargs["dataset_meta"] = dataset_meta
-    call_kwargs.update({name: value for name, value in optional_kwargs.items() if name in parameters})
     return function(config, **call_kwargs)

src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py

@@ -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

src/lerobot/policies/molmoact2/modeling_molmoact2.py

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

src/lerobot/policies/pi05/modeling_pi05.py

@@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
 
 import torch
 import torch.nn.functional as F  # noqa: N812
-from safetensors.torch import load_file
 from torch import Tensor, nn
 
 from lerobot.utils.import_utils import _transformers_available, require_package
@@ -31,7 +30,6 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 if TYPE_CHECKING or _transformers_available:
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
-    from transformers.utils import cached_file
 
     from ..pi_gemma import (
         PaliGemmaForConditionalGenerationWithPiGemma,
@@ -46,21 +44,20 @@ else:
     _gated_residual = None
     layernorm_forward = None
     PaliGemmaForConditionalGenerationWithPiGemma = None
-    cached_file = None
 from lerobot.configs import PreTrainedConfig
 from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
     OBS_LANGUAGE_TOKENS,
-    OPENPI_ATTENTION_MASK_VALUE,
 )
 
-from ..common.flow_matching import sample_noise, sample_time_beta
+from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
 from ..common.vla_utils import (
     clone_past_key_values,
     create_sinusoidal_pos_embedding,
     make_att_2d_masks,
     pad_vector,
+    prepare_attention_masks_4d,
     resize_with_pad_torch,
 )
 from ..pretrained import PreTrainedPolicy, T
@@ -74,9 +71,6 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
-_SAFETENSORS_FILE = "model.safetensors"
-
-
 # Define the complete layer computation function for gradient checkpointing
 def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
     query_states = []
@@ -407,12 +401,6 @@ class PaliGemmaWithExpertModel(
 class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
     """Core PI05 PyTorch model."""
 
-    use_hf_vision_checkpointing_api = False
-    checkpoint_vision_embeddings = True
-    use_typed_attention_masks = False
-    use_on_device_suffix_mask = False
-    precompute_denoise_times = False
-
     def __init__(self, config: PI05Config, rtc_processor: RTCProcessor | None = None):
         super().__init__()
         self.config = config
@@ -456,11 +444,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         """Enable gradient checkpointing for memory optimization."""
         self.gradient_checkpointing_enabled = True
         self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
-        vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
-        if self.use_hf_vision_checkpointing_api:
-            vision_tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
-        else:
-            vision_tower.gradient_checkpointing = True
+        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = True
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
         logging.info("Enabled gradient checkpointing for PI05Pytorch model")
 
@@ -468,11 +452,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         """Disable gradient checkpointing."""
         self.gradient_checkpointing_enabled = False
         self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = False
-        vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
-        if self.use_hf_vision_checkpointing_api:
-            vision_tower.gradient_checkpointing_disable()
-        else:
-            vision_tower.gradient_checkpointing = False
+        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = False
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
         logging.info("Disabled gradient checkpointing for PI05Pytorch model")
 
@@ -487,14 +467,6 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
-    def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
-        """Helper method to prepare 4D attention masks for transformer."""
-        att_2d_masks_4d = att_2d_masks[:, None, :, :]
-        result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
-        if dtype is not None:
-            result = result.to(dtype=dtype)
-        return result
-
     def sample_noise(self, shape, device):
         return sample_noise(shape, device)
 
@@ -516,16 +488,13 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         pad_masks = []
         att_masks = []
 
-        if self.checkpoint_vision_embeddings:
+        # Process images
+        for img, img_mask in zip(images, img_masks, strict=True):
 
-            def embed_image(img):
-                return self._apply_checkpoint(self.paligemma_with_expert.embed_image, img)
+            def image_embed_func(img):
+                return self.paligemma_with_expert.embed_image(img)
 
-            img_embs = [embed_image(img) for img in images]
-        else:
-            img_embs = [self.paligemma_with_expert.embed_image(img) for img in images]
-
-        for img_emb, img_mask in zip(img_embs, img_masks, strict=True):
+            img_emb = self._apply_checkpoint(image_embed_func, img)
             bsize, num_img_embs = img_emb.shape[:2]
 
             embs.append(img_emb)
@@ -587,15 +556,8 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         # Set attention masks so that image, language and state inputs do not attend to action tokens
         att_masks += [1] + ([0] * (self.config.chunk_size - 1))
-
-        if self.use_on_device_suffix_mask:
-            n = len(att_masks)
-            att_masks = torch.zeros(n, dtype=action_emb.dtype, device=action_emb.device)
-            att_masks[0] = 1
-            att_masks = att_masks[None, :].expand(bsize, n)
-        else:
-            att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
-            att_masks = att_masks[None, :].expand(bsize, len(att_masks))
+        att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
+        att_masks = att_masks[None, :].expand(bsize, len(att_masks))
 
         return action_emb, pad_masks, att_masks, adarms_cond
 
@@ -621,7 +583,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
         position_ids = torch.cumsum(pad_masks, dim=1) - 1
 
-        att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)
+        att_2d_masks_4d = prepare_attention_masks_4d(att_2d_masks)
 
         def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
             (_, suffix_out), _ = self.paligemma_with_expert.forward(
@@ -679,8 +641,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
         prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
 
-        mask_dtype = prefix_embs.dtype if self.use_typed_attention_masks else None
-        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks, dtype=mask_dtype)
+        prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
         self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
 
         _, past_key_values = self.paligemma_with_expert.forward(
@@ -691,52 +652,21 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             use_cache=True,
         )
 
-        dt = -1.0 / num_steps
-
-        times = None
-        if self.precompute_denoise_times:
-            times = torch.tensor(
-                [1.0 + step * dt for step in range(num_steps)], dtype=torch.float32, device=device
-            )
-
-        x_t = noise
-        for step in range(num_steps):
-            time = 1.0 + step * dt
-            if times is None:
-                time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
-            else:
-                time_tensor = times[step].expand(bsize)
-
-            def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
-                return self.denoise_step(
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    x_t=input_x_t,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled():
-                inference_delay = kwargs.get("inference_delay")
-                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
-                execution_horizon = kwargs.get("execution_horizon")
-
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=prev_chunk_left_over,
-                    inference_delay=inference_delay,
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-
-            x_t = x_t + dt * v_t
-
-            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
-        return x_t
+        return euler_integrate(
+            lambda input_x_t, current_timestep: self.denoise_step(
+                prefix_pad_masks=prefix_pad_masks,
+                past_key_values=past_key_values,
+                x_t=input_x_t,
+                timestep=current_timestep,
+            ),
+            noise,
+            num_steps,
+            rtc_processor=self.rtc_processor,
+            rtc_enabled=self._rtc_enabled(),
+            inference_delay=kwargs.get("inference_delay"),
+            prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
+            execution_horizon=kwargs.get("execution_horizon"),
+        )
 
     def denoise_step(
         self,
@@ -759,7 +689,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
         position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
 
-        full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
+        full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks)
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
         past_key_values = clone_past_key_values(past_key_values)
@@ -783,10 +713,6 @@ class PI05Policy(PreTrainedPolicy):
 
     config_class = PI05Config
     name = "pi05"
-    model_class = PI05Pytorch
-    eval_after_pretrained_load = False
-    show_openpi_disclaimer = True
-    use_native_pretrained_loader = False
 
     def __init__(
         self,
@@ -804,7 +730,7 @@ class PI05Policy(PreTrainedPolicy):
 
         # Initialize the core PI05 model
         self.init_rtc_processor()
-        self.model = self.model_class(config, rtc_processor=self.rtc_processor)
+        self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
 
         # Enable gradient checkpointing if requested
         if config.gradient_checkpointing:
@@ -830,31 +756,16 @@ class PI05Policy(PreTrainedPolicy):
         strict: bool = True,
         **kwargs,
     ) -> T:
-        """Load a native LeRobot checkpoint or convert the PI05 base checkpoint."""
-        if cls.use_native_pretrained_loader:
-            return super().from_pretrained(
-                pretrained_name_or_path,
-                config=config,
-                force_download=force_download,
-                resume_download=resume_download,
-                proxies=proxies,
-                token=token,
-                cache_dir=cache_dir,
-                local_files_only=local_files_only,
-                revision=revision,
-                strict=strict,
-                **kwargs,
-            )
-
-        if cls.show_openpi_disclaimer:
-            print(
-                "The PI05 model is a direct port of the OpenPI implementation. \n"
-                "This implementation follows the original OpenPI structure for compatibility. \n"
-                "Original implementation: https://github.com/Physical-Intelligence/openpi"
-            )
+        """Override the from_pretrained method to handle key remapping and display important disclaimer."""
+        print(
+            "The PI05 model is a direct port of the OpenPI implementation. \n"
+            "This implementation follows the original OpenPI structure for compatibility. \n"
+            "Original implementation: https://github.com/Physical-Intelligence/openpi"
+        )
         if pretrained_name_or_path is None:
             raise ValueError("pretrained_name_or_path is required")
 
+        # Use provided config if available, otherwise create default config
         if config is None:
             config = PreTrainedConfig.from_pretrained(
                 pretrained_name_or_path=pretrained_name_or_path,
@@ -868,41 +779,85 @@ class PI05Policy(PreTrainedPolicy):
                 **kwargs,
             )
 
+        # Initialize model without loading weights
+        # Check if dataset_stats were provided in kwargs
         model = cls(config, **kwargs)
-        model_id = str(pretrained_name_or_path)
-        resolved_file = cached_file(
-            model_id,
-            _SAFETENSORS_FILE,
-            _raise_exceptions_for_missing_entries=False,
-            force_download=force_download,
-            resume_download=resume_download,
-            proxies=proxies,
-            token=token,
-            cache_dir=cache_dir,
-            local_files_only=local_files_only,
-            revision=revision,
-        )
-        if resolved_file is None:
-            raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
 
-        fixed_state_dict = model._fix_pytorch_state_dict_keys(load_file(resolved_file), model.config)
-        remapped_state_dict = {
-            key if key.startswith("model.") else f"model.{key}": value
-            for key, value in fixed_state_dict.items()
-        }
-        remapped_state_dict = model._prepare_pretrained_state_dict(remapped_state_dict)
-        missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
-        if missing_keys:
-            logging.warning("Missing %s checkpoint keys: %s", cls.name, missing_keys)
-        if unexpected_keys:
-            logging.warning("Unexpected %s checkpoint keys: %s", cls.name, unexpected_keys)
-        if model.eval_after_pretrained_load:
-            model.eval()
+        # Load state dict (expects keys with "model." prefix)
+        try:
+            print(f"Loading model from: {pretrained_name_or_path}")
+            try:
+                from transformers.utils import cached_file
+
+                resolved_file = cached_file(
+                    pretrained_name_or_path,
+                    "model.safetensors",
+                    cache_dir=kwargs.get("cache_dir"),
+                    force_download=kwargs.get("force_download", False),
+                    resume_download=kwargs.get("resume_download"),
+                    proxies=kwargs.get("proxies"),
+                    token=kwargs.get("token"),
+                    revision=kwargs.get("revision"),
+                    local_files_only=kwargs.get("local_files_only", False),
+                )
+                from safetensors.torch import load_file
+
+                original_state_dict = load_file(resolved_file)
+                print("✓ Loaded state dict from model.safetensors")
+            except Exception as e:
+                print(f"Could not load state dict from remote files: {e}")
+                print("Returning model without loading pretrained weights")
+                return model
+
+            # First, fix any key differences (see openpi model.py, _fix_pytorch_state_dict_keys)
+            fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)
+
+            # Then add "model." prefix for all keys that don't already have it
+            remapped_state_dict = {}
+            remap_count = 0
+
+            for key, value in fixed_state_dict.items():
+                if not key.startswith("model."):
+                    new_key = f"model.{key}"
+                    remapped_state_dict[new_key] = value
+                    remap_count += 1
+                else:
+                    remapped_state_dict[key] = value
+
+            if remap_count > 0:
+                print(f"Remapped {remap_count} state dict keys")
+
+            # Load the remapped state dict into the model
+            missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
+
+            if missing_keys:
+                print(f"Missing keys when loading state dict: {len(missing_keys)} keys")
+                if len(missing_keys) <= 5:
+                    for key in missing_keys:
+                        print(f"  - {key}")
+                else:
+                    for key in missing_keys[:5]:
+                        print(f"  - {key}")
+                    print(f"  ... and {len(missing_keys) - 5} more")
+
+            if unexpected_keys:
+                print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
+                if len(unexpected_keys) <= 5:
+                    for key in unexpected_keys:
+                        print(f"  - {key}")
+                else:
+                    for key in unexpected_keys[:5]:
+                        print(f"  - {key}")
+                    print(f"  ... and {len(unexpected_keys) - 5} more")
+
+            if not missing_keys and not unexpected_keys:
+                print("All keys loaded successfully!")
+
+        except Exception as e:
+            print(f"Warning: Could not load state dict: {e}")
+
         return model
 
-    def _prepare_pretrained_state_dict(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
-        return state_dict
-
     def _fix_pytorch_state_dict_keys(
         self, state_dict, model_config
     ):  # see openpi `BaseModelConfig, _fix_pytorch_state_dict_keys`
@@ -1073,16 +1028,12 @@ class PI05Policy(PreTrainedPolicy):
 
         # Action queue logic for n_action_steps > 1
         if len(self._action_queue) == 0:
-            action_batch = self._prepare_action_batch(batch)
-            actions = self.predict_action_chunk(action_batch)[:, : self.config.n_action_steps]
+            actions = self.predict_action_chunk(batch)[:, : self.config.n_action_steps]
             # Transpose to get shape (n_action_steps, batch_size, action_dim)
             self._action_queue.extend(actions.transpose(0, 1))
 
         return self._action_queue.popleft()
 
-    def _prepare_action_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        return batch
-
     @torch.no_grad()
     def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
         """Predict a chunk of actions given environment observations."""

src/lerobot/policies/pi052/__init__.py

@@ -1,19 +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.
-
-"""PI052 configuration; model and processors are imported lazily by their factories."""
-
-from .configuration_pi052 import PI052Config
-
-__all__ = ["PI052Config"]

src/lerobot/policies/pi052/configuration_pi052.py

@@ -1,170 +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.
-
-"""PI0.5 with hierarchical text generation and flow-matched actions."""
-
-from dataclasses import dataclass
-
-from lerobot.configs import PreTrainedConfig
-from lerobot.optim.optimizers import AdamWConfig
-
-from ..pi05.configuration_pi05 import PI05Config
-
-
-@PreTrainedConfig.register_subclass("pi052")
-@dataclass
-class PI052Config(PI05Config):
-    """PI0.5 with recipe-driven text and action supervision."""
-
-    # Recipe / language stack ---------------------------------------------
-    recipe_path: str | None = "recipes/subtask_mem.yaml"
-    """Recipe path, or ``None`` for the plain PI0.5 prompt."""
-
-    apply_chat_template: bool = False
-    """Apply the tokenizer's chat template."""
-
-    # Balance frequent recipe text supervision against the paper's α=10 flow weight.
-    text_loss_weight: float = 1.0
-    """Text cross-entropy weight; ``0`` disables it."""
-
-    flow_loss_weight: float = 10.0
-    """Flow-matching loss weight."""
-
-    # Backbone training ---------------------------------------------------
-    unfreeze_lm_head: bool = True
-    """Train PaliGemma's language head."""
-
-    # Optional context dropout improves tolerance to missing or stale language state.
-    plan_dropout_prob: float = 0.0
-    memory_dropout_prob: float = 0.0
-    subtask_dropout_prob: float = 0.0
-
-    # FAST adds discrete-action CE to the text and flow objectives from paper §III.B-C.
-    enable_fast_action_loss: bool = True
-    """Add FAST action-token cross-entropy."""
-
-    action_tokenizer_name: str = "physical-intelligence/fast"
-    """FAST tokenizer identifier."""
-
-    max_action_tokens: int = 256
-    """Maximum FAST tokens per action chunk."""
-
-    fast_skip_tokens: int = 1152
-    """Reserved vocabulary IDs skipped by FAST token mapping."""
-
-    fast_action_loss_weight: float = 1.0
-    """FAST action-token loss weight."""
-
-    subtask_replan_steps: int = 0
-    """Steps between subtask generations; non-positive replans every chunk."""
-
-    joint_subtask_conditioning: bool = False
-    """Condition actions on the task and generated subtask."""
-
-    auto_fit_fast_tokenizer: bool = False
-    """Fit and cache a dataset-specific FAST tokenizer."""
-
-    fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
-    """Cache directory for fitted FAST tokenizers."""
-
-    fast_tokenizer_fit_samples: int = 1024
-    """Action chunks sampled for tokenizer fitting."""
-
-    fast_tokenizer_validation_samples: int = 256
-    """Held-out chunks used for tokenizer validation."""
-
-    fast_tokenizer_max_reconstruction_rmse: float = 0.10
-    """Maximum validation reconstruction RMSE."""
-
-    fast_tokenizer_max_dim_rmse: float = 0.20
-    """Maximum per-dimension validation RMSE."""
-
-    # Knowledge insulation detaches VLM K/V from action-loss gradients (paper §III.B).
-    knowledge_insulation: bool = True
-    """Detach VLM keys and values from action-loss gradients."""
-
-    # Optional training backends. Defaults preserve the eager/SDPA path.
-    use_flashrt_adarms: bool = False
-    """Use FlashRT adaptive RMSNorm kernels."""
-
-    use_compiled_text_ce: bool = False
-    """Compile text and FAST cross-entropy."""
-
-    use_compiled_vision: bool = False
-    """Compile the SigLIP vision tower."""
-
-    use_flex_attention: bool = False
-    """Use FlexAttention for knowledge insulation."""
-
-    use_manual_attention: bool = False
-    """Use manual attention for profiled KI shapes."""
-
-    manual_attention_scope: str = "all"
-    """Manual-attention scope: ``all`` or ``action``."""
-
-    # Scale language-head updates relative to the base optimizer schedule.
-    lm_head_lr_scale: float = 1.0
-
-    # Scale backbone and action-expert optimizer groups independently.
-    backbone_lr_scale: float = 1.0
-    action_expert_lr_scale: float = 1.0
-
-    # Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow.
-    flow_num_repeats: int = 5
-
-    # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
-    text_ce_z_loss_weight: float = 1e-4
-
-    use_flashrt_fp8_mlp: bool = False
-    """Use calibrated FlashRT FP8 MLP kernels."""
-
-    # Keep serialized PI052 AdamW options local because PI05Config lacks them.
-    optimizer_foreach: bool | None = False
-    optimizer_fused: bool | None = True
-
-    def get_optimizer_preset(self) -> AdamWConfig:
-        return AdamWConfig(
-            lr=self.optimizer_lr,
-            betas=self.optimizer_betas,
-            eps=self.optimizer_eps,
-            weight_decay=self.optimizer_weight_decay,
-            grad_clip_norm=self.optimizer_grad_clip_norm,
-            foreach=self.optimizer_foreach,
-            fused=self.optimizer_fused,
-        )
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        if self.enable_fast_action_loss and not self.recipe_path:
-            raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
-        if self.text_loss_weight > 0 and self.unfreeze_lm_head:
-            self.train_expert_only = False
-        if self.flow_num_repeats < 1:
-            raise ValueError(f"flow_num_repeats must be >= 1, got {self.flow_num_repeats}")
-        if self.fast_tokenizer_validation_samples < 1:
-            raise ValueError("fast_tokenizer_validation_samples must be >= 1")
-        if self.fast_tokenizer_max_reconstruction_rmse <= 0 or self.fast_tokenizer_max_dim_rmse <= 0:
-            raise ValueError("FAST tokenizer reconstruction thresholds must be positive")
-        if self.manual_attention_scope not in {"all", "action"}:
-            raise ValueError(
-                f"manual_attention_scope must be 'all' or 'action', got {self.manual_attention_scope!r}"
-            )
-        if self.use_flex_attention and self.use_manual_attention:
-            raise ValueError("use_flex_attention and use_manual_attention are mutually exclusive")
-        if self.use_flex_attention and self.flow_num_repeats == 1:
-            raise ValueError("use_flex_attention requires flow_num_repeats > 1")
-        if not self.knowledge_insulation and (
-            self.use_flex_attention or self.use_manual_attention or self.use_flashrt_adarms
-        ):
-            raise ValueError("KI attention and AdaRMS optimizations require knowledge_insulation=True")

src/lerobot/policies/pi052/fit_fast_tokenizer.py

@@ -1,522 +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.
-
-"""Fit and cache a FAST tokenizer for a dataset's action distribution.
-
-Training invokes this automatically when FAST loss and automatic fitting are enabled.
-"""
-
-from __future__ import annotations
-
-import hashlib
-import json
-import logging
-import os
-import shutil
-import time
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-logger = logging.getLogger(__name__)
-
-# ``ProcessorMixin.save_pretrained`` writes this shared cache sentinel.
-_CACHE_SENTINEL = "processor_config.json"
-
-
-def _is_global_leader() -> bool:
-    return int(os.environ.get("RANK", "0")) == 0
-
-
-def _jsonable(value: Any) -> Any:
-    if hasattr(value, "detach"):
-        value = value.detach().cpu().numpy()
-    if isinstance(value, np.ndarray):
-        return value.tolist()
-    if isinstance(value, dict):
-        return {key: _jsonable(item) for key, item in sorted(value.items())}
-    if isinstance(value, (list, tuple)):
-        return [_jsonable(item) for item in value]
-    return value
-
-
-def _dataset_signature(
-    dataset_repo_id: str,
-    base_tokenizer_name: str,
-    n_samples: int,
-    chunk_size: int,
-    normalization_mode: str,
-    dataset_revision: str | None = None,
-    episodes: list[int] | None = None,
-    exclude_episodes: list[int] | None = None,
-    action_stats: dict | None = None,
-    use_relative_actions: bool = False,
-    relative_action_mask: list[bool] | None = None,
-    validation_samples: int = 256,
-    max_reconstruction_rmse: float = 0.10,
-    max_dim_rmse: float = 0.20,
-) -> str:
-    """Hash every input that changes the fitted action distribution."""
-    payload = {
-        "dataset_repo_id": dataset_repo_id,
-        "dataset_revision": dataset_revision,
-        "base_tokenizer_name": base_tokenizer_name,
-        "n_samples": n_samples,
-        "chunk_size": chunk_size,
-        "normalization_mode": normalization_mode,
-        "episodes": episodes,
-        "exclude_episodes": exclude_episodes,
-        "action_stats": action_stats,
-        "use_relative_actions": use_relative_actions,
-        "relative_action_mask": relative_action_mask,
-        "validation_samples": validation_samples,
-        "max_reconstruction_rmse": max_reconstruction_rmse,
-        "max_dim_rmse": max_dim_rmse,
-    }
-    encoded = json.dumps(_jsonable(payload), sort_keys=True, separators=(",", ":")).encode()
-    return hashlib.sha256(encoded).hexdigest()[:16]
-
-
-def _select_episode_indices(
-    available_episodes: list[int],
-    episodes: list[int] | None,
-    exclude_episodes: list[int] | None,
-) -> list[int]:
-    allowed = set(episodes) if episodes is not None else set(available_episodes)
-    excluded = set(exclude_episodes or [])
-    return [episode for episode in available_episodes if episode in allowed and episode not in excluded]
-
-
-def _apply_relative_actions(
-    actions: np.ndarray,
-    states: np.ndarray,
-    relative_action_mask: list[bool] | None,
-) -> np.ndarray:
-    """Match RelativeActionsProcessorStep before tokenizer fitting."""
-    action_dim = actions.shape[-1]
-    mask = list(relative_action_mask) if relative_action_mask is not None else [True] * action_dim
-    if len(mask) < action_dim:
-        mask.extend([True] * (action_dim - len(mask)))
-    mask_array = np.asarray(mask[:action_dim], dtype=np.float32)
-    relative = actions.copy()
-    relative -= states[:, None, :action_dim] * mask_array
-    return relative
-
-
-def _normalize_actions(
-    actions: np.ndarray,
-    normalization_mode: str,
-    action_stats: dict | None = None,
-) -> np.ndarray:
-    """Match the action normalization applied by the training preprocessor."""
-    mode = getattr(normalization_mode, "value", normalization_mode).upper()
-    flat = actions.reshape(-1, actions.shape[-1])
-    stats = action_stats or {}
-
-    def stat(name: str, fallback) -> np.ndarray:
-        value = stats.get(name)
-        if value is None:
-            value = fallback()
-        if hasattr(value, "detach"):
-            value = value.detach().cpu().numpy()
-        return np.asarray(value, dtype=np.float32)
-
-    if mode == "IDENTITY":
-        return actions
-    if mode == "MEAN_STD":
-        mean = stat("mean", lambda: flat.mean(axis=0))
-        std = stat("std", lambda: flat.std(axis=0))
-        return ((actions - mean) / np.where(std == 0, 1e-8, std)).astype(np.float32)
-    if mode in {"QUANTILES", "QUANTILE10"}:
-        low_name, high_name, low_q, high_q = (
-            ("q01", "q99", 0.01, 0.99) if mode == "QUANTILES" else ("q10", "q90", 0.10, 0.90)
-        )
-        low = stat(low_name, lambda: np.quantile(flat, low_q, axis=0))
-        high = stat(high_name, lambda: np.quantile(flat, high_q, axis=0))
-    elif mode == "MIN_MAX":
-        low = stat("min", lambda: flat.min(axis=0))
-        high = stat("max", lambda: flat.max(axis=0))
-    else:
-        raise ValueError(f"Unsupported FAST tokenizer normalization mode: {mode}")
-
-    return (2.0 * (actions - low) / np.where(high == low, 1e-8, high - low) - 1.0).astype(np.float32)
-
-
-def _validate_fast_reconstruction(
-    tokenizer: Any,
-    actions: np.ndarray,
-    max_reconstruction_rmse: float,
-    max_dim_rmse: float,
-) -> tuple[dict[str, Any], np.ndarray]:
-    """Decode held-out chunks and reject tokenizers with excessive quantization error."""
-    decoded = np.asarray(tokenizer.decode(tokenizer(actions)), dtype=np.float32)
-    if decoded.shape != actions.shape:
-        raise RuntimeError(
-            f"FAST tokenizer reconstruction shape mismatch: expected {actions.shape}, got {decoded.shape}."
-        )
-    if not np.isfinite(decoded).all():
-        raise RuntimeError("FAST tokenizer reconstruction contains non-finite values.")
-
-    squared_error = np.square(decoded - actions)
-    rmse = float(np.sqrt(squared_error.mean()))
-    dim_rmse = np.sqrt(squared_error.mean(axis=(0, 1)))
-    nonconstant_dims = np.ptp(actions, axis=(0, 1)) > 1e-8
-    max_observed_dim_rmse = float(dim_rmse[nonconstant_dims].max(initial=0.0))
-    report = {
-        "num_validation_chunks": int(actions.shape[0]),
-        "reconstruction_rmse": rmse,
-        "max_dim_rmse": max_observed_dim_rmse,
-        "dim_rmse": dim_rmse.tolist(),
-        "max_reconstruction_rmse": max_reconstruction_rmse,
-        "max_allowed_dim_rmse": max_dim_rmse,
-    }
-    if rmse > max_reconstruction_rmse or max_observed_dim_rmse > max_dim_rmse:
-        raise RuntimeError(
-            "FAST tokenizer reconstruction error exceeds the configured limit: "
-            f"rmse={rmse:.4f} (max {max_reconstruction_rmse:.4f}), "
-            f"max_dim_rmse={max_observed_dim_rmse:.4f} (max {max_dim_rmse:.4f})."
-        )
-    return report, decoded
-
-
-def _load_fast_fitter(base_tokenizer_name: str) -> Any:
-    """Load FAST's fitting implementation without requiring its universal BPE weights."""
-    from transformers import AutoProcessor  # noqa: PLC0415
-
-    try:
-        return AutoProcessor.from_pretrained(base_tokenizer_name, trust_remote_code=True)
-    except ValueError as error:
-        if base_tokenizer_name != "physical-intelligence/fast":
-            raise
-        logger.warning(
-            "Could not load the universal FAST tokenizer backend; loading its fitting class directly: %s",
-            error,
-        )
-        from transformers.dynamic_module_utils import get_class_from_dynamic_module  # noqa: PLC0415
-
-        return get_class_from_dynamic_module(
-            "processing_action_tokenizer.UniversalActionProcessor",
-            base_tokenizer_name,
-        )
-
-
-def fit_fast_tokenizer(
-    *,
-    dataset_repo_id: str,
-    cache_dir: str | Path,
-    base_tokenizer_name: str = "physical-intelligence/fast",
-    n_samples: int = 1024,
-    chunk_size: int = 50,
-    seed: int = 42,
-    dataset_root: str | Path | None = None,
-    dataset_revision: str | None = None,
-    episodes: list[int] | None = None,
-    exclude_episodes: list[int] | None = None,
-    normalization_mode: str = "QUANTILES",
-    action_stats: dict | None = None,
-    use_relative_actions: bool = False,
-    relative_action_mask: list[bool] | None = None,
-    validation_samples: int = 256,
-    max_reconstruction_rmse: float = 0.10,
-    max_dim_rmse: float = 0.20,
-) -> str:
-    """Fit a FAST tokenizer on a LeRobot dataset's action distribution.
-
-    Args:
-        dataset_repo_id: HF Hub repo id of the LeRobotDataset to fit on.
-        cache_dir: Directory under which to save (and look up) fitted
-            tokenizers. The actual save path is
-            ``{cache_dir}/{signature}``.
-        base_tokenizer_name: HF identifier for the base FAST tokenizer
-            to finetune from. ``physical-intelligence/fast`` is the
-            universal one.
-        n_samples: Number of action chunks to sample for the fit. The
-            FAST paper uses a few thousand; ``1024`` is a good default
-            for medium datasets.
-        chunk_size: Length of each action chunk (matches
-            ``policy.chunk_size``). The FAST tokenizer is fit on
-            sequences of this length.
-        seed: RNG seed for sample selection.
-
-    Returns:
-        The local path to the fitted tokenizer. Passed directly to
-        ``--policy.action_tokenizer_name`` for the training run.
-
-    Raises:
-        ImportError: If the ``transformers`` library doesn't expose
-            ``AutoProcessor`` or the FAST tokenizer doesn't have a
-            ``.fit()`` method (then you're on an older FAST snapshot —
-            update to the current published model).
-        FileNotFoundError: If the dataset can't be loaded.
-    """
-    cache_dir = Path(cache_dir)
-    normalization_mode = getattr(normalization_mode, "value", normalization_mode).upper()
-    sig = _dataset_signature(
-        dataset_repo_id,
-        base_tokenizer_name,
-        n_samples,
-        chunk_size,
-        normalization_mode,
-        dataset_revision,
-        episodes,
-        exclude_episodes,
-        action_stats,
-        use_relative_actions,
-        relative_action_mask,
-        validation_samples,
-        max_reconstruction_rmse,
-        max_dim_rmse,
-    )
-    out_dir = cache_dir / sig
-
-    if out_dir.exists() and (out_dir / _CACHE_SENTINEL).exists():
-        logger.info(
-            "FAST tokenizer cache hit: %s — re-using fitted tokenizer for dataset=%s base=%s n_samples=%d",
-            out_dir,
-            dataset_repo_id,
-            base_tokenizer_name,
-            n_samples,
-        )
-        return str(out_dir)
-
-    # One global rank populates the shared cache; every other rank waits for the atomic publish.
-    is_leader = _is_global_leader()
-    if not is_leader:
-        timeout_s = 1800.0  # 30 min — covers ~1024-sample fits on cold caches
-        start = time.monotonic()
-        while not (out_dir / _CACHE_SENTINEL).exists():
-            if time.monotonic() - start > timeout_s:
-                raise RuntimeError(
-                    f"FAST tokenizer fit: non-leader rank timed out after "
-                    f"{timeout_s:.0f}s waiting for {out_dir / _CACHE_SENTINEL}. "
-                    "Leader rank likely crashed during the fit."
-                )
-            time.sleep(2.0)
-        logger.info("FAST tokenizer ready (leader populated cache): %s", out_dir)
-        return str(out_dir)
-
-    logger.info(
-        "FAST tokenizer cache miss — fitting on dataset=%s base=%s n_samples=%d chunk_size=%d → %s",
-        dataset_repo_id,
-        base_tokenizer_name,
-        n_samples,
-        chunk_size,
-        out_dir,
-    )
-
-    # Read action columns directly to avoid video decoding and bound memory to sampled episodes.
-    rng = np.random.default_rng(seed)
-    actions_buf: list[np.ndarray] = []
-
-    # Read v3 parquet shards directly to avoid split lookup failures and repeated metadata parsing.
-    import pyarrow as _pa  # noqa: PLC0415
-    import pyarrow.parquet as _pq  # noqa: PLC0415
-
-    if dataset_root is not None:
-        snap = Path(dataset_root)
-    else:
-        from huggingface_hub import snapshot_download  # noqa: PLC0415
-
-        snap = Path(
-            snapshot_download(repo_id=dataset_repo_id, repo_type="dataset", revision=dataset_revision)
-        )
-    data_files = sorted((snap / "data").glob("chunk-*/file-*.parquet"))
-    if not data_files:
-        raise RuntimeError(f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}.")
-
-    columns = ["episode_index", "action"]
-    if use_relative_actions:
-        columns.append("observation.state")
-    tables = [_pq.read_table(f, columns=columns) for f in data_files]
-    table = _pa.concat_tables(tables)
-    eps = table["episode_index"].to_numpy()
-    acts_col = table["action"]
-    # Normalize Arrow action representations into an (N, D) array.
-    try:
-        acts = np.stack(acts_col.to_numpy(zero_copy_only=False)).astype(np.float32)
-    except Exception:  # noqa: BLE001
-        # Fallback path for nested-list types: flatten via to_pylist().
-        acts = np.asarray(acts_col.to_pylist(), dtype=np.float32)
-    if acts.ndim != 2:
-        raise RuntimeError(f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}.")
-    states = None
-    if use_relative_actions:
-        try:
-            states = np.stack(table["observation.state"].to_numpy(zero_copy_only=False)).astype(np.float32)
-        except Exception:  # noqa: BLE001
-            states = np.asarray(table["observation.state"].to_pylist(), dtype=np.float32)
-        if states.ndim != 2:
-            raise RuntimeError(
-                f"FAST fit: expected ``observation.state`` rows to be 1-D vectors; got {states.shape}."
-            )
-
-    # Sort once because episode order is only guaranteed within each shard.
-    order = np.argsort(eps, kind="stable")
-    eps_sorted = eps[order]
-    boundaries = np.searchsorted(eps_sorted, np.arange(int(eps_sorted.max()) + 2))
-    ep_to_slice: dict[int, tuple[int, int]] = {
-        int(ep): (int(boundaries[ep]), int(boundaries[ep + 1]))
-        for ep in range(len(boundaries) - 1)
-        if boundaries[ep] < boundaries[ep + 1]
-    }
-    num_episodes = len(ep_to_slice)
-    # ``acts`` is in original (un-sorted-by-episode) row order; reorder
-    # so per-episode slices are contiguous.
-    acts = acts[order]
-    if states is not None:
-        states = states[order]
-
-    ep_indices = _select_episode_indices(list(ep_to_slice), episodes, exclude_episodes)
-    if not ep_indices:
-        raise RuntimeError("FAST fit: episode selection is empty after applying exclusions.")
-    total_samples = n_samples + validation_samples
-    samples_per_episode = max(1, (total_samples + len(ep_indices) - 1) // len(ep_indices))
-    collected = 0
-    eps_visited = 0
-    short_episodes = 0
-    states_buf: list[np.ndarray] = []
-    for ep_idx in rng.permutation(ep_indices):
-        if collected >= total_samples:
-            break
-        start, stop = ep_to_slice[int(ep_idx)]
-        ep_actions = acts[start:stop]
-        if ep_actions.shape[0] < chunk_size:
-            short_episodes += 1
-            continue
-        starts = rng.integers(0, ep_actions.shape[0] - chunk_size + 1, size=samples_per_episode)
-        for s in starts:
-            actions_buf.append(ep_actions[int(s) : int(s) + chunk_size])
-            if states is not None:
-                states_buf.append(states[start + int(s)])
-            collected += 1
-            if collected >= total_samples:
-                break
-        eps_visited += 1
-
-    if not actions_buf:
-        raise RuntimeError(
-            f"FAST fit collected zero action chunks from {dataset_repo_id!r}: "
-            f"all {num_episodes} episodes were shorter than chunk_size="
-            f"{chunk_size} ({short_episodes} too short) or had an unreadable "
-            "``action`` column. Lower ``chunk_size`` to match your episode "
-            "lengths."
-        )
-
-    actions = np.stack(actions_buf, axis=0).astype(np.float32)  # (N, H, D)
-    if states is not None:
-        actions = _apply_relative_actions(actions, np.stack(states_buf), relative_action_mask)
-    logger.info(
-        "FAST fit: collected %d chunks of shape %s from %d episodes",
-        actions.shape[0],
-        actions.shape[1:],
-        eps_visited,
-    )
-
-    actions = _normalize_actions(actions, normalization_mode, action_stats)
-
-    base = _load_fast_fitter(base_tokenizer_name)
-    if not hasattr(base, "fit"):
-        raise ImportError(
-            f"Base FAST tokenizer {base_tokenizer_name!r} has no ``.fit()`` "
-            "method — your transformers / model snapshot is too old. Update "
-            "to the current ``physical-intelligence/fast`` revision."
-        )
-
-    if actions.shape[0] < total_samples:
-        raise RuntimeError(
-            f"FAST fit collected {actions.shape[0]} chunks, but {total_samples} are required "
-            f"for {n_samples} fit and {validation_samples} validation chunks."
-        )
-    fit_actions = actions[:n_samples]
-    validation_actions = actions[n_samples:total_samples]
-    fitted = base.fit(fit_actions)
-    validation_report, decoded_actions = _validate_fast_reconstruction(
-        fitted,
-        validation_actions,
-        max_reconstruction_rmse,
-        max_dim_rmse,
-    )
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    staging_dir = cache_dir / f".{sig}.tmp-{os.getpid()}"
-    shutil.rmtree(staging_dir, ignore_errors=True)
-    fitted.save_pretrained(str(staging_dir))
-    (staging_dir / "reconstruction_validation.json").write_text(
-        json.dumps(validation_report, indent=2) + "\n"
-    )
-    np.savez_compressed(
-        staging_dir / "reconstruction_examples.npz",
-        original=validation_actions[:8],
-        decoded=decoded_actions[:8],
-    )
-    if out_dir.exists():
-        shutil.rmtree(out_dir)
-    staging_dir.replace(out_dir)
-    logger.info("FAST fit: saved fitted tokenizer to %s", out_dir)
-    return str(out_dir)
-
-
-def resolve_fast_tokenizer(
-    config: Any,
-    dataset_repo_id: str | None,
-    dataset_root: str | Path | None = None,
-    dataset_stats: dict | None = None,
-    dataset_revision: str | None = None,
-    episodes: list[int] | None = None,
-    exclude_episodes: list[int] | None = None,
-) -> str:
-    """Return the configured tokenizer, fitting a cached dataset-specific one when requested."""
-    if not getattr(config, "auto_fit_fast_tokenizer", False) or dataset_repo_id is None:
-        return config.action_tokenizer_name
-
-    relative_action_mask = None
-    if getattr(config, "use_relative_actions", False):
-        action_names = getattr(config, "action_feature_names", None)
-        exclude_tokens = [
-            str(name).lower() for name in getattr(config, "relative_exclude_joints", []) if name
-        ]
-        if action_names is not None and exclude_tokens:
-            relative_action_mask = [
-                not any(token == str(name).lower() or token in str(name).lower() for token in exclude_tokens)
-                for name in action_names
-            ]
-
-    fit_kwargs = {
-        "dataset_repo_id": dataset_repo_id,
-        "cache_dir": Path(config.fast_tokenizer_cache_dir).expanduser(),
-        "base_tokenizer_name": config.action_tokenizer_name,
-        "n_samples": config.fast_tokenizer_fit_samples,
-        "chunk_size": config.chunk_size,
-        "dataset_root": dataset_root,
-        "dataset_revision": dataset_revision,
-        "episodes": episodes,
-        "exclude_episodes": exclude_episodes,
-        "normalization_mode": config.normalization_mapping.get("ACTION", "QUANTILES"),
-        "action_stats": (dataset_stats or {}).get("action"),
-        "use_relative_actions": getattr(config, "use_relative_actions", False),
-        "relative_action_mask": relative_action_mask,
-    }
-    validation_fields = {
-        "validation_samples": "fast_tokenizer_validation_samples",
-        "max_reconstruction_rmse": "fast_tokenizer_max_reconstruction_rmse",
-        "max_dim_rmse": "fast_tokenizer_max_dim_rmse",
-    }
-    fit_kwargs.update(
-        {
-            argument: getattr(config, attribute)
-            for argument, attribute in validation_fields.items()
-            if hasattr(config, attribute)
-        }
-    )
-    return fit_fast_tokenizer(**fit_kwargs)

src/lerobot/policies/pi052/flashrt_fp8.py

@@ -1,263 +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.
-
-"""Optional FlashRT FP8 MLP kernels with one-pass calibration and BF16 fallback."""
-
-from __future__ import annotations
-
-import logging
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F  # noqa: N812
-
-logger = logging.getLogger(__name__)
-
-_FP8_MAX = 448.0
-
-
-def _roundtrip_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
-    """Quantize->dequantize an activation through FP8 E4M3 at ``scale`` (f32)."""
-    q = torch.clamp(x.float() / scale.float(), -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
-    return q.float() * scale.float()
-
-
-_SWIGLU_REPO = "flashrt/flashrt-fp8-swiglu-ffn"
-_GELU_REPO = "flashrt/flashrt-fp8-ffn"
-_GEMM_REPO = "flashrt/flashrt-gemm-epilogues"
-
-
-def _get_kernel(repo: str):
-    """Load a cached FlashRT Hub package."""
-    from kernels import get_kernel
-
-    return get_kernel(repo, version=1)
-
-
-def _quantize_fp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
-    scale = max(weight.detach().float().abs().max().item(), 1e-12) / _FP8_MAX
-    fp8 = torch.clamp(weight.float() / scale, -_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn)
-    return fp8.contiguous(), torch.tensor([scale], dtype=torch.float32)
-
-
-def _static_scale(amax: float, safety: float) -> torch.Tensor:
-    return torch.tensor([max(amax, 1e-12) / _FP8_MAX * safety], dtype=torch.float32)
-
-
-class _FlashRTGeGLU(nn.Module):
-    """FP8 Gemma GeGLU MLP."""
-
-    def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety, fuse_weight=None):
-        super().__init__()
-        self.ffn_ops = ffn_ops
-        self.quant_ops = quant_ops
-        self.in_features = mlp.gate_proj.weight.shape[1]
-        device = mlp.gate_proj.weight.device
-        gate_up = torch.cat([mlp.gate_proj.weight, mlp.up_proj.weight], dim=0).float()
-        # Fold fixed RMSNorm weights into GEMM; adaptive norms use identity scaling.
-        if fuse_weight is not None:
-            f = 1.0 + fuse_weight.detach().float()
-            gate_up = gate_up * f[None, :]
-            channel_scale = (1.0 / f).to(torch.bfloat16)
-        else:
-            channel_scale = torch.ones(self.in_features, dtype=torch.bfloat16)
-        gate_up_fp8, gate_up_scale = _quantize_fp8(gate_up)
-        down_fp8, down_scale = _quantize_fp8(mlp.down_proj.weight)
-        self.register_buffer("gate_up_fp8", gate_up_fp8.to(device))
-        self.register_buffer("down_fp8", down_fp8.to(device))
-        self.register_buffer("gate_up_scale", gate_up_scale.to(device))
-        self.register_buffer("down_scale", down_scale.to(device))
-        self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device))
-        self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device))
-        self.register_buffer("channel_scale", channel_scale.to(device))
-        self.safety = safety
-        self.calibrating = False
-        self._ia = 0.0
-        self._ha = 0.0
-
-    def _calibrate_step(self, x):
-        # Track input and hidden maxima on live FP8-propagated activations.
-        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
-        xq = flat.float() * self.channel_scale.float()
-        self._ia = max(self._ia, xq.abs().max().item())
-        self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device))
-        xdq = _roundtrip_fp8(xq, self.input_scale)
-        wdq = self.gate_up_fp8.float() * self.gate_up_scale.float()
-        gate, up = (xdq @ wdq.t()).chunk(2, dim=-1)
-        hidden = F.gelu(gate, approximate="tanh") * up
-        self._ha = max(self._ha, hidden.abs().max().item())
-        self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device))
-
-    def forward(self, x):
-        if self.calibrating:
-            self._calibrate_step(x)
-        shape = x.shape
-        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
-        x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(
-            flat, self.channel_scale, self.input_scale
-        )
-        out = self.ffn_ops.fp8_geglu_mlp_bf16(
-            x_fp8,
-            self.gate_up_fp8,
-            self.down_fp8,
-            self.input_scale,
-            self.gate_up_scale,
-            self.hidden_scale,
-            self.down_scale,
-        )
-        return out.reshape(shape)
-
-
-class _FlashRTGeluMLP(nn.Module):
-    """FP8 SigLIP GELU MLP."""
-
-    def __init__(self, mlp, in_amax, hid_amax, ffn_ops, quant_ops, safety):
-        super().__init__()
-        self.ffn_ops = ffn_ops
-        self.quant_ops = quant_ops
-        self.in_features = mlp.fc1.weight.shape[1]
-        self.out_features = mlp.fc2.weight.shape[0]
-        device = mlp.fc1.weight.device
-        up_fp8, up_scale = _quantize_fp8(mlp.fc1.weight)
-        down_fp8, down_scale = _quantize_fp8(mlp.fc2.weight)
-        self.register_buffer("up_fp8", up_fp8.to(device))
-        self.register_buffer("down_fp8", down_fp8.to(device))
-        self.register_buffer("up_scale", up_scale.to(device))
-        self.register_buffer("down_scale", down_scale.to(device))
-        self.register_buffer("up_bias", mlp.fc1.bias.detach().to(torch.bfloat16))
-        self.register_buffer("down_bias", mlp.fc2.bias.detach().to(torch.bfloat16))
-        self.register_buffer("input_scale", _static_scale(in_amax, safety).to(device))
-        self.register_buffer("hidden_scale", _static_scale(hid_amax, safety).to(device))
-        self.register_buffer(
-            "channel_scale", torch.ones(self.in_features, device=device, dtype=torch.bfloat16)
-        )
-        self.safety = safety
-        self.calibrating = False
-        self._ia = 0.0
-        self._ha = 0.0
-
-    def _calibrate_step(self, x):
-        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
-        self._ia = max(self._ia, flat.float().abs().max().item())
-        self.input_scale.copy_(_static_scale(self._ia, self.safety).to(self.input_scale.device))
-        xdq = _roundtrip_fp8(flat.float(), self.input_scale)
-        hid = (xdq @ (self.up_fp8.float() * self.up_scale.float()).t()) + self.up_bias.float()
-        hid = F.gelu(hid, approximate="tanh")
-        self._ha = max(self._ha, hid.abs().max().item())
-        self.hidden_scale.copy_(_static_scale(self._ha, self.safety).to(self.hidden_scale.device))
-
-    def forward(self, x):
-        if self.calibrating:
-            self._calibrate_step(x)
-        shape = x.shape
-        dtype = x.dtype
-        flat = x.reshape(-1, self.in_features).to(torch.bfloat16)
-        x_fp8 = self.quant_ops.channel_scale_quantize_fp8_static_bf16(
-            flat, self.channel_scale, self.input_scale
-        )
-        out = self.ffn_ops.fp8_gelu_mlp_bf16(
-            x_fp8,
-            self.up_fp8,
-            self.up_bias,
-            self.down_fp8,
-            self.down_bias,
-            self.input_scale,
-            self.up_scale,
-            self.hidden_scale,
-            self.down_scale,
-        )
-        return out.reshape(*shape[:-1], self.out_features).to(dtype)
-
-
-def _siglip_mlps(model) -> list:
-    tower = model.paligemma_with_expert.paligemma.model.vision_tower
-    return [m for _, m in tower.named_modules() if type(m).__name__ == "SiglipMLP"]
-
-
-def _run_forward(policy, batches) -> None:
-    """Run eager action prediction so calibration reaches Python module forwards."""
-    model = policy.model
-    saved = {name: vars(model).pop(name) for name in ("sample_actions", "forward") if name in vars(model)}
-    with torch.inference_mode():
-        for batch in batches:
-            policy.predict_action_chunk(
-                {k: (v.clone() if torch.is_tensor(v) else v) for k, v in batch.items()}
-            )
-    torch.cuda.synchronize()
-    vars(model).update(saved)
-
-
-def _fixed_norm_weight(norm):
-    """Return a fixed RMSNorm fold weight, or ``None`` for adaptive norms."""
-    return norm.weight if getattr(norm, "dense", None) is None else None
-
-
-def _fp8_supported(device) -> bool:
-    """Return whether the device supports FP8 E4M3 tensor cores (CUDA SM >= 8.9)."""
-    if device.type != "cuda" or not torch.cuda.is_available():
-        return False
-    major, minor = torch.cuda.get_device_capability(device)
-    return (major, minor) >= (8, 9)
-
-
-def apply_fp8_mlp(policy, batch, *, safety: float = 1.05) -> bool:
-    """Replace Gemma and SigLIP MLPs with FlashRT FP8 kernels calibrated on the supplied batch.
-
-    Returns ``False`` without modifying BF16 execution when FP8 or its kernels are unavailable.
-    """
-    device = next(policy.parameters()).device
-    if not _fp8_supported(device):
-        logger.warning(
-            "PI052: device %s has no FP8 (E4M3) support (needs CUDA SM>=8.9); keeping BF16.",
-            device,
-        )
-        return False
-    batches = batch if isinstance(batch, (list, tuple)) else [batch]
-    try:
-        ffn_ops = _get_kernel(_SWIGLU_REPO)
-        gelu_ops = _get_kernel(_GELU_REPO)
-        quant_ops = _get_kernel(_GEMM_REPO)
-    except Exception as exc:  # noqa: BLE001
-        logger.warning("PI052: FlashRT FP8 kernels unavailable (%s); keeping BF16.", exc)
-        return False
-
-    model = policy.model
-    calibrating = []
-
-    gemma_layers = list(model.paligemma_with_expert.gemma_expert.model.layers) + list(
-        model.paligemma_with_expert.paligemma.model.language_model.layers
-    )
-    for layer in gemma_layers:
-        fw = _fixed_norm_weight(layer.post_attention_layernorm)
-        layer.mlp = _FlashRTGeGLU(layer.mlp, 1.0, 1.0, ffn_ops, quant_ops, safety, fuse_weight=fw).to(device)
-        calibrating.append(layer.mlp)
-
-    siglip = _siglip_mlps(model)
-    for mlp_parent in model.paligemma_with_expert.paligemma.model.vision_tower.vision_model.encoder.layers:
-        mlp_parent.mlp = _FlashRTGeluMLP(mlp_parent.mlp, 1.0, 1.0, gelu_ops, quant_ops, safety).to(device)
-        calibrating.append(mlp_parent.mlp)
-
-    # Calibrate every swapped module in one FP8-propagated forward.
-    for m in calibrating:
-        m.calibrating = True
-    _run_forward(policy, batches)
-    for m in calibrating:
-        m.calibrating = False
-
-    logger.info(
-        "PI052: FlashRT FP8 enabled (%d Gemma + %d SigLIP MLPs).",
-        len(gemma_layers),
-        len(siglip),
-    )
-    return True

src/lerobot/policies/pi052/inference/__init__.py

@@ -1,19 +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.
-
-"""PI052 adapter for the policy-agnostic language runtime."""
-
-from .pi052_adapter import PI052PolicyAdapter
-
-__all__ = ["PI052PolicyAdapter"]

src/lerobot/policies/pi052/inference/pi052_adapter.py

@@ -1,254 +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.
-
-"""PI052 actions and text generation for the generic language runtime."""
-
-from __future__ import annotations
-
-import logging
-from typing import Any
-
-from lerobot.runtime import RuntimeState
-from lerobot.runtime.adapter import BaseLanguageAdapter
-
-logger = logging.getLogger(__name__)
-
-_LOC_TOKENIZER_CACHE: dict[str, Any] = {}
-
-
-class PI052PolicyAdapter(BaseLanguageAdapter):
-    """Runtime bridge for PI052 policies."""
-
-    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
-        import torch  # noqa: PLC0415
-
-        from lerobot.utils.constants import (  # noqa: PLC0415
-            OBS_LANGUAGE_ATTENTION_MASK,
-            OBS_LANGUAGE_TOKENS,
-            OBS_STATE,
-        )
-
-        subtask = state.language_context.get("subtask") or state.task or ""
-        # Match the training prompt by conditioning on both subtask and discretized state.
-        state_str = None
-        obs_state = observation.get(OBS_STATE)
-        if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
-            from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str  # noqa: PLC0415
-
-            state_row = obs_state[0] if obs_state.ndim > 1 else obs_state
-            state_str = discretize_state_str(state_row)
-
-        batch = dict(observation)
-        if getattr(self.policy.config, "joint_subtask_conditioning", False):
-            # Joint sequences keep the task turn (with state) and render the
-            # subtask as a causal assistant turn, exactly as trained.
-            from transformers import AutoTokenizer  # noqa: PLC0415
-
-            from lerobot.policies.pi052.text_processor_pi052 import (  # noqa: PLC0415
-                encode_prompt_with_targets,
-                register_paligemma_loc_tokens,
-            )
-            from lerobot.utils.constants import OBS_LANGUAGE_CAUSAL_MARKS  # noqa: PLC0415
-
-            task = state.task or ""
-            task_content = task if state_str is None else f"{task}, State: {state_str};"
-            tok_name = getattr(self.policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
-            tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens)
-            ids, attn, marks = encode_prompt_with_targets(
-                tokenizer,
-                [
-                    {"role": "user", "content": task_content},
-                    {"role": "assistant", "content": subtask},
-                ],
-                target_indices=[1],
-            )
-            device = getattr(self.policy.config, "device", None)
-            if device is not None:
-                ids, attn, marks = ids.to(device), attn.to(device), marks.to(device)
-            batch[OBS_LANGUAGE_TOKENS] = ids
-            batch[OBS_LANGUAGE_ATTENTION_MASK] = attn
-            batch[OBS_LANGUAGE_CAUSAL_MARKS] = marks
-        else:
-            content = subtask if state_str is None else f"{subtask}, State: {state_str};"
-            text_batch = _build_text_batch(
-                self.policy,
-                [{"role": "user", "content": content}],
-                add_generation_prompt=False,
-            )
-            batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"]
-            batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"]
-        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=user_text)
-        if kind == "subtask" and getattr(self.policy.config, "joint_subtask_conditioning", False):
-            # Joint samples carry state on the task turn, so the subtask must be
-            # generated from the same state-bearing prompt.
-            import torch  # noqa: PLC0415
-
-            from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str  # noqa: PLC0415
-            from lerobot.utils.constants import OBS_STATE  # noqa: PLC0415
-
-            obs_state = (observation or {}).get(OBS_STATE)
-            if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
-                state_row = obs_state[0] if obs_state.ndim > 1 else obs_state
-                for m in reversed(messages):
-                    if m.get("role") == "user":
-                        m["content"] = f"{m.get('content', '')}, State: {discretize_state_str(state_row)};"
-                        break
-        return _generate_with_policy(
-            self.policy,
-            messages,
-            observation=observation,
-            state=state,
-            label=f"{kind} gen",
-            min_new_tokens=self.gen.min_new_tokens,
-            temperature=self.gen.temperature,
-            top_p=self.gen.top_p,
-            suppress_loc_tokens=True,  # all runtime text is prose; never emit <loc>
-        )
-
-    def build_messages(
-        self,
-        kind: str,
-        state: RuntimeState,
-        *,
-        user_text: str | None = None,
-    ) -> list[dict[str, Any]]:
-        if kind in ("subtask", "plan"):
-            return [{"role": "user", "content": state.task or ""}]
-        if kind == "memory":
-            messages = [{"role": "user", "content": state.task or ""}]
-            if state.language_context.get("memory"):
-                messages.append(
-                    {"role": "assistant", "content": f"Previous memory: {state.language_context['memory']}"}
-                )
-            if state.extra.get("prior_subtask"):
-                messages.append(
-                    {"role": "user", "content": f"Completed subtask: {state.extra['prior_subtask']}"}
-                )
-            return messages
-        if kind == "interjection":
-            messages = [{"role": "user", "content": state.task or ""}]
-            if state.language_context.get("plan"):
-                messages.append(
-                    {"role": "assistant", "content": f"Previous plan:\n{state.language_context['plan']}"}
-                )
-            if user_text:
-                messages.append({"role": "user", "content": user_text})
-            return messages
-        raise ValueError(f"Unknown PI052 text kind: {kind}")
-
-
-def _get_loc_tokenizer(tok_name: str, auto_tokenizer_cls: Any, register_loc_fn: Any) -> Any:
-    tokenizer = _LOC_TOKENIZER_CACHE.get(tok_name)
-    if tokenizer is None:
-        tokenizer = register_loc_fn(auto_tokenizer_cls.from_pretrained(tok_name))
-        _LOC_TOKENIZER_CACHE[tok_name] = tokenizer
-    return tokenizer
-
-
-def _build_text_batch(
-    policy: Any,
-    prompt_messages: list[dict[str, Any]],
-    *,
-    add_generation_prompt: bool = True,
-) -> dict[str, Any]:
-    import torch  # noqa: PLC0415
-    from transformers import AutoTokenizer  # noqa: PLC0415
-
-    from lerobot.policies.pi052.text_processor_pi052 import (  # noqa: PLC0415
-        _flatten_say_tool_calls,
-        _format_messages,
-        _strip_blocks,
-        register_paligemma_loc_tokens,
-    )
-
-    tok_name = getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
-    tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens)
-
-    messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in prompt_messages]
-    prompt, _spans = _format_messages(messages)
-    if add_generation_prompt:
-        # No trailing space: SentencePiece folds it into the first target token
-        # ("▁move"), so a space-suffixed prefill ends in a lone "▁" the model
-        # never saw at this position during training.
-        prompt = prompt + "Assistant:"
-
-    encoded = tokenizer(prompt, return_tensors="pt")
-    ids = encoded["input_ids"]
-    attn = encoded.get("attention_mask")
-    if attn is None and tokenizer.pad_token_id is not None:
-        attn = ids != tokenizer.pad_token_id
-    if attn is not None and hasattr(attn, "dtype") and attn.dtype != torch.bool:
-        attn = attn.bool()
-
-    device = getattr(getattr(policy, "config", None), "device", None)
-    if device is not None:
-        try:
-            ids = ids.to(device)
-            if attn is not None and hasattr(attn, "to"):
-                attn = attn.to(device)
-        except Exception as exc:  # noqa: BLE001
-            logger.debug("could not move pi052 lang tokens to %s: %s", device, exc)
-    return {"lang_tokens": ids, "lang_masks": attn, "tokenizer": tokenizer}
-
-
-def _generate_with_policy(
-    policy: Any,
-    messages: list[dict[str, Any]],
-    *,
-    observation: dict[str, Any] | None = None,
-    state: RuntimeState | None = None,
-    label: str = "select_message",
-    min_new_tokens: int = 0,
-    temperature: float = 0.0,
-    top_p: float = 1.0,
-    suppress_loc_tokens: bool = False,
-) -> str:
-    if not hasattr(policy, "select_message"):
-        if state is not None:
-            state.log(f"  [warn] policy has no select_message — skipping {label}")
-        return ""
-    text_batch = _build_text_batch(policy, messages)
-    try:
-        from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS  # noqa: PLC0415
-
-        batch: dict[str, Any] = {
-            OBS_LANGUAGE_TOKENS: text_batch["lang_tokens"],
-            OBS_LANGUAGE_ATTENTION_MASK: text_batch["lang_masks"],
-        }
-        if observation:
-            for k, v in observation.items():
-                if isinstance(k, str) and k.startswith("observation.") and k not in batch:
-                    batch[k] = v
-        return policy.select_message(
-            batch,
-            tokenizer=text_batch["tokenizer"],
-            min_new_tokens=min_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            suppress_loc_tokens=suppress_loc_tokens,
-        )
-    except Exception as exc:  # noqa: BLE001
-        logger.warning("%s failed: %s", label, exc, exc_info=logger.isEnabledFor(logging.DEBUG))
-        if state is not None:
-            state.log(f"  [warn] {label} failed: {type(exc).__name__}: {exc}")
-        return ""

src/lerobot/policies/pi052/processor_pi052.py

@@ -1,164 +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.
-
-"""PI052 processor factory with optional recipe rendering and text tokenization.
-
-Without a recipe it delegates to the standard PI0.5 pipeline.
-"""
-
-from __future__ import annotations
-
-from pathlib import Path
-from typing import Any
-
-import torch
-
-from lerobot.configs.recipe import TrainingRecipe
-from lerobot.processor import (
-    AbsoluteActionsProcessorStep,
-    ActionTokenizerProcessorStep,
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
-    PolicyAction,
-    PolicyProcessorPipeline,
-    RelativeActionsProcessorStep,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
-)
-
-# Import directly to keep optional language dependencies out of ``lerobot.processor``.
-from lerobot.processor.render_messages_processor import RenderMessagesStep
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-
-from ..pi05.processor_pi05 import make_pi05_pre_post_processors
-from .configuration_pi052 import PI052Config
-from .text_processor_pi052 import PI052TextTokenizerStep
-
-
-def make_pi052_pre_post_processors(
-    config: PI052Config,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    dataset_repo_id: str | None = None,
-    dataset_root: str | None = None,
-    dataset_revision: str | None = None,
-    episodes: list[int] | None = None,
-    exclude_episodes: list[int] | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """Build PI0.5-v2's pre/post-processor pipelines.
-
-    Falls through to π0.5's stock pipeline when ``recipe_path`` is unset.
-    """
-    if not config.recipe_path:
-        if getattr(config, "enable_fast_action_loss", False):
-            raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
-        return make_pi05_pre_post_processors(config, dataset_stats=dataset_stats)
-
-    recipe = _load_recipe(config.recipe_path)
-
-    relative_step = RelativeActionsProcessorStep(
-        enabled=config.use_relative_actions,
-        exclude_joints=getattr(config, "relative_exclude_joints", []),
-        action_names=getattr(config, "action_feature_names", None),
-    )
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        relative_step,
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        RenderMessagesStep(recipe=recipe),
-        PI052TextTokenizerStep(
-            tokenizer_name="google/paligemma-3b-pt-224",
-            max_length=config.tokenizer_max_length,
-            plan_dropout_prob=getattr(config, "plan_dropout_prob", 0.0),
-            memory_dropout_prob=getattr(config, "memory_dropout_prob", 0.0),
-            subtask_dropout_prob=getattr(config, "subtask_dropout_prob", 0.0),
-        ),
-    ]
-
-    # Add FAST action-token supervision only when explicitly enabled.
-    if getattr(config, "enable_fast_action_loss", False):
-        from .fit_fast_tokenizer import resolve_fast_tokenizer  # noqa: PLC0415
-
-        input_steps.append(
-            ActionTokenizerProcessorStep(
-                action_tokenizer_name=resolve_fast_tokenizer(
-                    config,
-                    dataset_repo_id,
-                    dataset_root,
-                    dataset_stats,
-                    dataset_revision,
-                    episodes,
-                    exclude_episodes,
-                ),
-                max_action_tokens=config.max_action_tokens,
-                fast_skip_tokens=config.fast_skip_tokens,
-                paligemma_tokenizer_name="google/paligemma-3b-pt-224",
-                allow_truncation=False,
-            )
-        )
-
-    input_steps.append(DeviceProcessorStep(device=config.device))
-
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features,
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        AbsoluteActionsProcessorStep(
-            enabled=config.use_relative_actions,
-            relative_step=relative_step,
-        ),
-        DeviceProcessorStep(device="cpu"),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )
-
-
-def _load_recipe(path_str: str) -> TrainingRecipe:
-    """Resolve ``path_str`` to a ``TrainingRecipe``.
-
-    Accepts an absolute path or a path relative to
-    ``src/lerobot/configs/``.
-    """
-    p = Path(path_str)
-    if not p.is_absolute() and not p.exists():
-        from lerobot.configs import recipe as _recipe_module  # noqa: PLC0415
-
-        configs_dir = Path(_recipe_module.__file__).resolve().parent
-        candidate = configs_dir / path_str
-        if candidate.exists():
-            p = candidate
-    return TrainingRecipe.from_yaml(p)

src/lerobot/policies/pi052/text_processor_pi052.py

@@ -1,521 +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.
-
-"""Tokenize PI052 messages and build text/action supervision masks."""
-
-from __future__ import annotations
-
-import json
-import logging
-from dataclasses import dataclass
-from typing import Any
-
-import numpy as np
-import torch
-from torch import Tensor
-
-from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.processor.pipeline import ProcessorStep, ProcessorStepRegistry
-from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
-
-logger = logging.getLogger(__name__)
-
-
-def discretize_state_str(state_row: Any) -> str:
-    """Format one normalized state row with PI0.5's 256-bin convention."""
-    arr = state_row.detach().cpu().numpy() if hasattr(state_row, "detach") else np.asarray(state_row)
-    disc = np.digitize(arr, bins=np.linspace(-1, 1, 256 + 1)[:-1]) - 1
-    return " ".join(str(int(x)) for x in disc.reshape(-1).tolist())
-
-
-def _state_row_at(state_all: Any, pos: int) -> Any:
-    """Select the per-sample state row from a (possibly batched) state tensor."""
-    if state_all is None:
-        return None
-    if hasattr(state_all, "ndim") and state_all.ndim >= 2:
-        return state_all[pos]
-    return state_all
-
-
-def _content_to_text(content: Any) -> str:
-    """Collapse a message's ``content`` (string or multimodal blocks) to text."""
-    if isinstance(content, str):
-        return content
-    if isinstance(content, list):
-        parts = [
-            b["text"]
-            for b in content
-            if isinstance(b, dict) and b.get("type") == "text" and isinstance(b.get("text"), str)
-        ]
-        return "\n".join(parts)
-    return ""
-
-
-def _flatten_say_tool_calls(message: dict[str, Any]) -> dict[str, Any]:
-    """Move ``say`` tool calls into text markers that PaliGemma can learn."""
-    tool_calls = message.get("tool_calls")
-    if not tool_calls:
-        return message
-    say_texts: list[str] = []
-    for call in tool_calls:
-        if not isinstance(call, dict):
-            continue
-        fn = call.get("function") or {}
-        if fn.get("name") != "say":
-            continue
-        args = fn.get("arguments")
-        if isinstance(args, str):
-            try:
-                import json  # noqa: PLC0415
-
-                args = json.loads(args)
-            except (ValueError, TypeError):
-                args = {}
-        text = args.get("text", "") if isinstance(args, dict) else ""
-        if text:
-            say_texts.append(str(text))
-    new = dict(message)
-    new.pop("tool_calls", None)
-    if not say_texts:
-        return new
-    base = _content_to_text(new.get("content")).strip()
-    marker = "".join(f"<say>{t}</say>" for t in say_texts)
-    new["content"] = f"{base}\n{marker}" if base else marker
-    return new
-
-
-def _strip_blocks(message: dict[str, Any]) -> dict[str, Any]:
-    """Flatten text blocks and drop image blocks handled by observation inputs."""
-    new = dict(message)
-    new.pop("stream", None)
-    new.pop("target", None)
-    content = new.get("content")
-    if content is None:
-        new["content"] = ""
-    elif isinstance(content, str):
-        pass
-    elif isinstance(content, list):
-        parts: list[str] = []
-        for block in content:
-            if not isinstance(block, dict):
-                continue
-            if block.get("type") == "text":
-                t = block.get("text", "")
-                if isinstance(t, str):
-                    parts.append(t)
-        new["content"] = "\n".join(parts)
-    else:
-        new["content"] = str(content)
-    return new
-
-
-def _is_batched_messages(messages: Any) -> bool:
-    return isinstance(messages, list) and bool(messages) and isinstance(messages[0], list)
-
-
-def _sample_indices(value: Any, batch_size: int) -> list[int | None]:
-    if value is None:
-        return [None] * batch_size
-    if isinstance(value, torch.Tensor):
-        if value.numel() == 1:
-            return [int(value.item())] * batch_size
-        values = value.reshape(-1).tolist()
-        return [int(v) for v in values[:batch_size]]
-    if isinstance(value, (list, tuple)):
-        if len(value) == 1:
-            return _sample_indices(value[0], batch_size)
-        return [int(v.item() if hasattr(v, "item") else v) for v in value[:batch_size]]
-    return [int(value)] * batch_size
-
-
-_VQA_COORD_SCALE = 1000.0
-
-
-def register_paligemma_loc_tokens(tokenizer: Any) -> Any:
-    """Register PaliGemma's reserved ``<locDDDD>`` strings as single tokens.
-
-    Without registration, the stock tokenizer splits each location into generic text pieces.
-    """
-    if "<loc0000>" in getattr(tokenizer, "added_tokens_encoder", {}):
-        return tokenizer
-    tokenizer.add_tokens([f"<loc{i:04d}>" for i in range(1024)])
-    return tokenizer
-
-
-def _loc_token(coord: float, scale: float = _VQA_COORD_SCALE) -> str:
-    """PaliGemma ``<locNNNN>`` for a coord on a ``[0, scale]`` axis."""
-    idx = round(float(coord) / scale * 1023) if scale > 0 else 0
-    return f"<loc{max(0, min(1023, idx)):04d}>"
-
-
-def _vqa_answer_to_loc(answer: dict[str, Any]) -> str | None:
-    """Convert normalized bbox/keypoint answers to label-first PaliGemma locations.
-
-    Label-first targets prevent location tokens from dominating every assistant turn; non-spatial answers return ``None``.
-    """
-    point = answer.get("point")
-    if isinstance(point, list | tuple) and len(point) == 2 and "point_format" in answer:
-        try:
-            x, y = float(point[0]), float(point[1])
-        except (TypeError, ValueError):
-            return None
-        label = str(answer.get("label", "")).strip()
-        if not label:
-            return None
-        return f"{label} {_loc_token(y)}{_loc_token(x)}"
-
-    detections = answer.get("detections")
-    if isinstance(detections, list) and detections:
-        parts: list[str] = []
-        for det in detections:
-            if not isinstance(det, dict):
-                continue
-            box = det.get("bbox")
-            if not (isinstance(box, list | tuple) and len(box) == 4):
-                continue
-            try:
-                x1, y1, x2, y2 = (float(v) for v in box)
-            except (TypeError, ValueError):
-                continue
-            label = str(det.get("label", "")).strip()
-            if not label:
-                continue
-            toks = f"{_loc_token(y1)}{_loc_token(x1)}{_loc_token(y2)}{_loc_token(x2)}"
-            parts.append(f"{label} {toks}")
-        return " ; ".join(parts) if parts else None
-    return None
-
-
-def _messages_vqa_to_loc(
-    messages: list[dict[str, Any]],
-    target_indices: list[int],
-) -> list[dict[str, Any]]:
-    """Rewrite spatial VQA target JSON as camera-independent ``<loc>`` text."""
-    if not target_indices:
-        return messages
-    out = list(messages)
-    for idx in target_indices:
-        if not (0 <= idx < len(out)):
-            continue
-        content = out[idx].get("content")
-        if not isinstance(content, str) or not content.strip():
-            continue
-        try:
-            answer = json.loads(content)
-        except (ValueError, TypeError):
-            continue
-        if not isinstance(answer, dict):
-            continue
-        loc_text = _vqa_answer_to_loc(answer)
-        if loc_text is not None:
-            out[idx] = {**out[idx], "content": loc_text}
-    return out
-
-
-def _format_messages(
-    messages: list[dict[str, Any]],
-    target_indices: list[int] | None = None,
-    eos_token: str | None = None,
-) -> tuple[str, list[tuple[int, int]]]:
-    """Build the flat PI0.5 prompt and each message's payload span.
-
-    Supervised targets include EOS so generation learns when to stop.
-    """
-    targets = set(target_indices or [])
-    parts: list[str] = []
-    spans: list[tuple[int, int]] = []
-    cursor = 0
-    for i, m in enumerate(messages):
-        role = m.get("role", "user")
-        content = m.get("content", "") or ""
-        header = f"{role.capitalize()}: "
-        body = content + eos_token if (eos_token and i in targets) else content
-        full = header + body + "\n"
-        start = cursor + len(header)
-        end = start + len(body)
-        parts.append(full)
-        spans.append((start, end))
-        cursor += len(full)
-    return "".join(parts), spans
-
-
-def encode_prompt_with_targets(
-    tokenizer: Any, messages: list[dict[str, Any]], target_indices: list[int]
-) -> tuple[Tensor, Tensor, Tensor]:
-    """Tokenize a flat prompt and mark the token positions of target spans.
-
-    Inference-side twin of ``PI052TextTokenizerStep._encode_messages``: same
-    serialization (role headers, target EOS) and the same offset-overlap span
-    arithmetic, but unpadded and returning a boolean target mask instead of
-    labels. Used to rebuild joint-sequence prompts whose target spans must be
-    attended causally, matching ``_mark_target_span_causal`` at train time.
-
-    Returns ``(input_ids, attention_mask, target_marks)``, each ``(1, L)``.
-    """
-    prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
-    encoded = tokenizer(prompt, return_tensors="pt", return_offsets_mapping=True)
-    input_ids = encoded["input_ids"][0]
-    attention_mask = encoded.get("attention_mask")
-    if attention_mask is None:
-        attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
-    else:
-        attention_mask = attention_mask[0].bool()
-    offsets = encoded["offset_mapping"][0]
-
-    marks = torch.zeros_like(input_ids, dtype=torch.bool)
-    for idx in target_indices:
-        if idx >= len(spans):
-            continue
-        char_start, char_end = spans[idx]
-        for token_pos in range(input_ids.shape[0]):
-            if not attention_mask[token_pos]:
-                continue
-            tok_start, tok_end = int(offsets[token_pos, 0]), int(offsets[token_pos, 1])
-            if tok_end <= char_start or tok_start >= char_end:
-                continue
-            marks[token_pos] = True
-    return input_ids.unsqueeze(0), attention_mask.unsqueeze(0), marks.unsqueeze(0)
-
-
-@dataclass
-@ProcessorStepRegistry.register(name="pi052_text_tokenizer")
-class PI052TextTokenizerStep(ProcessorStep):
-    """Convert flat role-delimited messages into tokens and supervision masks."""
-
-    tokenizer_name: str = "google/paligemma-3b-pt-224"
-    max_length: int = 200
-    padding: str = "max_length"
-    padding_side: str = "right"
-    plan_dropout_prob: float = 0.0
-    memory_dropout_prob: float = 0.0
-    subtask_dropout_prob: float = 0.0
-    interjection_dropout_prob: float = 0.0
-    dropout_seed: int | None = None
-
-    def __post_init__(self) -> None:
-        self._tokenizer: Any = None
-
-    def get_config(self) -> dict[str, Any]:
-        return {
-            "tokenizer_name": self.tokenizer_name,
-            "max_length": self.max_length,
-            "padding": self.padding,
-            "padding_side": self.padding_side,
-            "plan_dropout_prob": self.plan_dropout_prob,
-            "memory_dropout_prob": self.memory_dropout_prob,
-            "subtask_dropout_prob": self.subtask_dropout_prob,
-            "interjection_dropout_prob": self.interjection_dropout_prob,
-            "dropout_seed": self.dropout_seed,
-        }
-
-    def _ensure_tokenizer(self) -> Any:
-        if self._tokenizer is not None:
-            return self._tokenizer
-        from transformers import AutoTokenizer  # noqa: PLC0415
-
-        self._tokenizer = register_paligemma_loc_tokens(AutoTokenizer.from_pretrained(self.tokenizer_name))
-        return self._tokenizer
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition | None:
-        transition = transition.copy()
-        complementary = transition.get(TransitionKey.COMPLEMENTARY_DATA, {}) or {}
-        messages = complementary.get("messages") or []
-
-        if not messages:
-            return transition
-
-        tokenizer = self._ensure_tokenizer()
-        state_all = (transition.get(TransitionKey.OBSERVATION) or {}).get(OBS_STATE)
-        if _is_batched_messages(messages):
-            indices_iter = _sample_indices(complementary.get("index"), len(messages))
-            encoded = [
-                self._encode_messages(
-                    tokenizer,
-                    msg,
-                    list(streams),
-                    list(tgt_indices),
-                    complementary,
-                    sample_idx=int(s_idx) if s_idx is not None else None,
-                    state_row=_state_row_at(state_all, pos),
-                )
-                for pos, (msg, streams, tgt_indices, s_idx) in enumerate(
-                    zip(
-                        messages,
-                        complementary.get("message_streams") or [[] for _ in messages],
-                        complementary.get("target_message_indices") or [[] for _ in messages],
-                        indices_iter,
-                        strict=False,
-                    )
-                )
-            ]
-        else:
-            sample_idx = _sample_indices(complementary.get("index"), 1)[0]
-            encoded = [
-                self._encode_messages(
-                    tokenizer,
-                    messages,
-                    list(complementary.get("message_streams") or []),
-                    list(complementary.get("target_message_indices") or []),
-                    complementary,
-                    sample_idx=sample_idx,
-                    state_row=_state_row_at(state_all, 0),
-                )
-            ]
-
-        obs = dict(transition.get(TransitionKey.OBSERVATION) or {})
-        obs[OBS_LANGUAGE_TOKENS] = torch.stack([ids for ids, _, _, _, _ in encoded])
-        obs[OBS_LANGUAGE_ATTENTION_MASK] = torch.stack([attn for _, attn, _, _, _ in encoded])
-        transition[TransitionKey.OBSERVATION] = obs
-
-        transition[TransitionKey.COMPLEMENTARY_DATA] = {
-            **complementary,
-            "text_labels": torch.stack([labels for _, _, labels, _, _ in encoded]),
-            "predict_actions": torch.stack([pred for _, _, _, pred, _ in encoded]),
-        }
-        return transition
-
-    def _encode_messages(
-        self,
-        tokenizer: Any,
-        messages: list[dict[str, Any]],
-        message_streams: list[str | None],
-        target_indices: list[int],
-        complementary: dict[str, Any],
-        sample_idx: int | None = None,
-        state_row: Any = None,
-    ) -> tuple[Tensor, Tensor, Tensor, Tensor, str]:
-        if (
-            self.plan_dropout_prob
-            or self.memory_dropout_prob
-            or self.subtask_dropout_prob
-            or self.interjection_dropout_prob
-        ):
-            messages, target_indices = self._apply_prompt_dropout(
-                messages,
-                target_indices,
-                complementary,
-                sample_idx=sample_idx,
-            )
-
-        messages = _messages_vqa_to_loc(messages, target_indices)
-
-        messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in messages]
-        # Only low-level prompts carry PI0.5-style proprioception.
-        if state_row is not None and any(s == "low_level" for s in message_streams):
-            state_str = discretize_state_str(state_row)
-            for m in reversed(messages):
-                if m.get("role") == "user":
-                    base = _content_to_text(m.get("content", ""))
-                    m["content"] = f"{base}, State: {state_str};"
-                    break
-        prompt, spans = _format_messages(messages, target_indices, getattr(tokenizer, "eos_token", None))
-
-        encoded = tokenizer(
-            prompt,
-            max_length=self.max_length,
-            padding=self.padding,
-            truncation=True,
-            return_tensors="pt",
-            return_offsets_mapping=True,
-            padding_side=self.padding_side,
-        )
-
-        input_ids = encoded["input_ids"][0]
-        attention_mask = encoded["attention_mask"][0].bool()
-        offsets = encoded["offset_mapping"][0]
-
-        labels = torch.full_like(input_ids, fill_value=-100)
-        for idx in target_indices:
-            if idx >= len(spans):
-                continue
-            char_start, char_end = spans[idx]
-            for token_pos in range(input_ids.shape[0]):
-                if not attention_mask[token_pos]:
-                    continue
-                tok_start, tok_end = int(offsets[token_pos, 0]), int(offsets[token_pos, 1])
-                if tok_end <= char_start or tok_start >= char_end:
-                    continue
-                labels[token_pos] = input_ids[token_pos]
-
-        predict_actions = torch.tensor(
-            bool(any(s == "low_level" for s in message_streams)),
-            dtype=torch.bool,
-        )
-        return input_ids, attention_mask, labels, predict_actions, prompt
-
-    def _apply_prompt_dropout(
-        self,
-        messages: list[dict[str, Any]],
-        target_indices: list[int],
-        complementary: dict[str, Any],
-        sample_idx: int | None = None,
-    ) -> tuple[list[dict[str, Any]], list[int]]:
-        """Drop sampled context messages and remap the retained target positions."""
-        import random  # noqa: PLC0415
-
-        seed = self.dropout_seed
-        if seed is None:
-            seed_src = sample_idx if sample_idx is not None else complementary.get("index", 0)
-            try:
-                if hasattr(seed_src, "item"):
-                    seed_src = seed_src.item()
-                seed = int(seed_src)
-            except (TypeError, ValueError):
-                seed = 0
-        rng = random.Random(seed)
-
-        keep_indices: list[int] = []
-        for idx, msg in enumerate(messages):
-            if idx in target_indices:
-                keep_indices.append(idx)
-                continue
-            kind = _classify_for_dropout(msg)
-            prob = {
-                "plan": self.plan_dropout_prob,
-                "memory": self.memory_dropout_prob,
-                "subtask": self.subtask_dropout_prob,
-                "interjection": self.interjection_dropout_prob,
-            }.get(kind, 0.0)
-            if prob > 0.0 and rng.random() < prob:
-                continue
-            keep_indices.append(idx)
-
-        new_messages = [messages[i] for i in keep_indices]
-        old_to_new = {old: new for new, old in enumerate(keep_indices)}
-        new_targets = [old_to_new[t] for t in target_indices if t in old_to_new]
-        return new_messages, new_targets
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-def _classify_for_dropout(message: dict[str, Any]) -> str | None:
-    """Classify context from its rendered text prefix."""
-    content = message.get("content")
-    if isinstance(content, list):
-        text_parts = [b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"]
-        content = " ".join(text_parts)
-    elif content is None or not isinstance(content, str):
-        return None
-    s = content.strip()
-    if s.startswith("Plan:") or s.startswith("Previous plan"):
-        return "plan"
-    if s.startswith("Memory:") or s.startswith("Previous memory"):
-        return "memory"
-    if s.startswith("Current subtask") or s.startswith("Completed subtask"):
-        return "subtask"
-    return None

src/lerobot/policies/pi0_fast/configuration_pi0_fast.py

@@ -61,21 +61,21 @@ class PI0FastConfig(PreTrainedConfig):
     tokenizer_max_length: int = 200  # see openpi `__post_init__`
     text_tokenizer_name: str = "google/paligemma-3b-pt-224"
     action_tokenizer_name: str = "lerobot/fast-action-tokenizer"
-    auto_fit_fast_tokenizer: bool = False
-    fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
-    fast_tokenizer_fit_samples: int = 1024
     temperature: float = 0.0
     max_decoding_steps: int = 256
     fast_skip_tokens: int = 128
 
+    # Whether to validate that decoded action tokens start with "Action: " prefix
+    validate_action_token_prefix: bool = True
+
     # Whether to use KV cache for faster autoregressive decoding
     use_kv_cache: bool = True
 
     normalization_mapping: dict[str, NormalizationMode] = field(
         default_factory=lambda: {
             "VISUAL": NormalizationMode.IDENTITY,
-            "STATE": NormalizationMode.QUANTILES,
-            "ACTION": NormalizationMode.QUANTILES,
+            "STATE": NormalizationMode.MEAN_STD,  # Pi0Fast uses quantiles for state
+            "ACTION": NormalizationMode.MEAN_STD,  # Pi0Fast uses quantiles for action
         }
     )
 

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -24,7 +24,13 @@ import numpy as np
 import torch
 from torch import Tensor, nn
 
-from lerobot.utils.import_utils import _transformers_available, require_package
+from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
+
+# Conditional import for type checking and lazy loading
+if TYPE_CHECKING or _scipy_available:
+    from scipy.fftpack import idct
+else:
+    idct = None
 
 if TYPE_CHECKING or _transformers_available:
     from transformers import AutoProcessor, AutoTokenizer
@@ -60,32 +66,6 @@ class ActionSelectKwargs(TypedDict, total=False):
     temperature: float | None
 
 
-def _gather_last_valid_language_hidden(
-    hidden_states: Tensor,
-    language_masks: Tensor,
-    image_token_count: int,
-) -> Tensor:
-    """Gather each sample's last non-padding language hidden state."""
-    last_language_indices = image_token_count + language_masks.long().sum(dim=1) - 1
-    if torch.any(last_language_indices < image_token_count):
-        raise ValueError("PI0-FAST requires at least one valid language token per sample")
-    batch_indices = torch.arange(hidden_states.shape[0], device=hidden_states.device)
-    return hidden_states[batch_indices, last_language_indices]
-
-
-def _reduce_fast_token_loss(token_loss: Tensor, token_mask: Tensor) -> Tensor:
-    """Give every sample equal weight regardless of its FAST token count."""
-    sample_loss = (token_loss * token_mask).sum(dim=1) / token_mask.sum(dim=1).clamp(min=1)
-    return sample_loss.mean()
-
-
-def _sample_next_token(logits: Tensor, temperature: float) -> Tensor:
-    if temperature > 0:
-        probabilities = torch.softmax(logits / temperature, dim=-1)
-        return torch.multinomial(probabilities, num_samples=1)
-    return torch.argmax(logits, dim=-1, keepdim=True)
-
-
 class GemmaConfig:  # see openpi `gemma.py: Config`
     """Configuration for Gemma model variants."""
 
@@ -260,6 +240,7 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         # Compile model if requested
         if config.compile_model:
             torch.set_float32_matmul_precision("high")
+            self.sample_actions_fast = torch.compile(self.sample_actions_fast, mode=config.compile_mode)
             self.forward = torch.compile(self.forward, mode=config.compile_mode)
 
     def gradient_checkpointing_enable(self):
@@ -486,12 +467,18 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         # only compute logits for the positions that predict FAST tokens
         lm_head = self.paligemma_with_expert.paligemma.lm_head
 
-        # The last valid prompt token predicts "Action:", then each FAST token predicts the next one.
-        fast_hidden = prefix_out[:, -num_fast_embs:, :]
-        last_language_hidden = _gather_last_valid_language_hidden(prefix_out, masks, total_t_images)
-        prediction_hidden = torch.cat([last_language_hidden[:, None], fast_hidden[:, :-1]], dim=1)
-        fast_logits_for_pred = lm_head(prediction_hidden)
-        fast_targets = fast_action_tokens
+        # Targets are the FAST action tokens
+        fast_targets = fast_action_tokens  # (B, num_fast_embs)
+
+        # extract logits for FAST token prediction
+        fast_hidden = prefix_out[:, -fast_targets.shape[1] :, :]
+        fast_logits_for_pred = lm_head(fast_hidden)  # (B, num_fast_embs, gemma_vocab_size)
+
+        # Shift left for next-step prediction and shift target
+        # logits[:, i] predicts targets[:, i+1]
+        fast_logits_for_pred = fast_logits_for_pred[:, :-1, :]  # shift logits left
+        fast_targets = fast_targets[:, 1:]  # shift targets right
+        fast_action_masks = fast_action_masks[:, 1:]  # shift masks to match targets
 
         # compute cross-entropy loss
         loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
@@ -501,7 +488,9 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         fast_loss_per_token = loss_fct(fast_logits_flat, fast_targets_flat)
         fast_loss_per_token = fast_loss_per_token.reshape(fast_targets.shape)
 
-        fast_loss = _reduce_fast_token_loss(fast_loss_per_token, fast_action_masks.float())
+        # apply mask and compute mean loss
+        masked_fast_loss = fast_loss_per_token * fast_action_masks.float()
+        fast_loss = masked_fast_loss.sum() / fast_action_masks.sum().clamp(min=1)
 
         return {
             "ce_loss": fast_loss,
@@ -530,7 +519,15 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         device = tokens.device
         lm_head = self.paligemma_with_expert.paligemma.lm_head
 
-        # 1. Initial embedding: the prompt's existing BOS is the only BOS in the sequence.
+        # add bos token after tokens
+        bos_token = torch.full(
+            (bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
+        )
+        tokens = torch.cat([tokens, bos_token], dim=1)
+        masks = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
+
+        # 1. Initial Embedding (matches training prefix)
+        # prefix_embs will include [Images, Language Prompt, BOS]
         prefix_embs, prefix_pad_masks, prefix_att_masks, total_t_images, _ = self.embed_prefix_fast(
             images, img_masks, tokens, masks, fast_action_tokens=None, fast_action_masks=None
         )
@@ -542,8 +539,6 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
             prefix_embs = prefix_embs.to(dtype=torch.bfloat16)
 
         generated_action_tokens = torch.zeros((bsize, max_decoding_steps), dtype=torch.long, device=device)
-        eos_token_id = self._paligemma_tokenizer.eos_token_id
-        finished = torch.zeros(bsize, dtype=torch.bool, device=device)
 
         # 2. Decoding Loop (each step re-computes full sequence)
         for t in range(max_decoding_steps):
@@ -561,24 +556,16 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
                 adarms_cond=[None, None],
             )
 
-            if t == 0:
-                prediction_hidden = _gather_last_valid_language_hidden(prefix_out, masks, total_t_images)
-            else:
-                prediction_hidden = prefix_out[:, -1]
-            next_token = _sample_next_token(lm_head(prediction_hidden), temperature)
+            # predict next token from the very last sequence position
+            last_logits = lm_head(prefix_out[:, -1:, :])  # (B, 1, vocab_size)
 
-            active = ~finished
-            generated_action_tokens[:, t] = torch.where(
-                active, next_token.squeeze(-1), torch.zeros_like(next_token.squeeze(-1))
-            )
-            finished |= active & next_token.squeeze(-1).eq(eos_token_id)
-            if finished.all():
-                break
-            next_token = torch.where(
-                finished[:, None],
-                torch.full_like(next_token, eos_token_id),
-                next_token,
-            )
+            if temperature > 0:
+                probs = torch.softmax(last_logits[:, -1] / temperature, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(last_logits[:, -1], dim=-1, keepdim=True)
+
+            generated_action_tokens[:, t] = next_token.squeeze(-1)
 
             # 3. Update sequence for next iteration (unless it's the last step)
             if t < max_decoding_steps - 1:
@@ -625,14 +612,20 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         device = tokens.device
         lm_head = self.paligemma_with_expert.paligemma.lm_head
 
-        generated_action_tokens = torch.zeros((bsize, max_decoding_steps), dtype=torch.long, device=device)
-        if max_decoding_steps == 0:
-            return generated_action_tokens
-
         # --- 1. PREFILL PHASE ---
+        # Process Images + Text Prompt + BOS token once to populate the KV cache.
+
+        # Add BOS token to the prompt
+        bos_token = torch.full(
+            (bsize, 1), self._paligemma_tokenizer.bos_token_id, dtype=torch.long, device=device
+        )
+        tokens_in = torch.cat([tokens, bos_token], dim=1)
+        masks_in = torch.cat([masks, torch.ones((bsize, 1), dtype=torch.bool, device=device)], dim=1)
+
+        # Embed prefix [Images, Language, BOS]
         # fast_action_tokens=None means we are just embedding the condition (images+text)
         prefix_embs, prefix_pad_masks, prefix_att_masks, total_t_images, _ = self.embed_prefix_fast(
-            images, img_masks, tokens, masks, fast_action_tokens=None, fast_action_masks=None
+            images, img_masks, tokens_in, masks_in, fast_action_tokens=None, fast_action_masks=None
         )
 
         # Ensure correct precision (bfloat16/float32)
@@ -659,18 +652,17 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
             adarms_cond=[None, None],
         )
 
-        prediction_hidden = _gather_last_valid_language_hidden(prefix_out, masks, total_t_images)
-        next_token = _sample_next_token(lm_head(prediction_hidden), temperature)
+        # Sample the first action token from the last logit of the prefix
+        last_logits = lm_head(prefix_out[:, -1:, :])  # (B, 1, V)
+        if temperature > 0:
+            probs = torch.softmax(last_logits[:, -1] / temperature, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1)
+        else:
+            next_token = torch.argmax(last_logits[:, -1], dim=-1, keepdim=True)
+
+        # Initialize storage for generated tokens
+        generated_action_tokens = torch.zeros((bsize, max_decoding_steps), dtype=torch.long, device=device)
         generated_action_tokens[:, 0] = next_token.squeeze(-1)
-        eos_token_id = self._paligemma_tokenizer.eos_token_id
-        finished = next_token.squeeze(-1).eq(eos_token_id)
-        if finished.all():
-            return generated_action_tokens
-        next_token = torch.where(
-            finished[:, None],
-            torch.full_like(next_token, eos_token_id),
-            next_token,
-        )
 
         # Track valid tokens mask (0 for pad, 1 for valid)
         # We need this to tell the new token what it can attend to (images + text + past actions)
@@ -711,19 +703,15 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
                 adarms_cond=[None, None],
             )
 
-            next_token = _sample_next_token(lm_head(step_out[:, -1]), temperature)
-            active = ~finished
-            generated_action_tokens[:, t] = torch.where(
-                active, next_token.squeeze(-1), torch.zeros_like(next_token.squeeze(-1))
-            )
-            finished |= active & next_token.squeeze(-1).eq(eos_token_id)
-            if finished.all():
-                break
-            next_token = torch.where(
-                finished[:, None],
-                torch.full_like(next_token, eos_token_id),
-                next_token,
-            )
+            # Sample next token
+            last_logits = lm_head(step_out[:, -1:, :])
+            if temperature > 0:
+                probs = torch.softmax(last_logits[:, -1] / temperature, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(last_logits[:, -1], dim=-1, keepdim=True)
+
+            generated_action_tokens[:, t] = next_token.squeeze(-1)
 
         return generated_action_tokens
 
@@ -1036,7 +1024,7 @@ class PI0FastPolicy(PreTrainedPolicy):
         return self._paligemma_tokenizer.vocab_size - 1 - self.config.fast_skip_tokens - tokens
 
     def decode_actions_with_fast(
-        self, token_ids: list[Tensor], time_horizon: int, action_dim: int
+        self, token_ids: list[int], time_horizon: int, action_dim: int, relaxed_decoding: bool = True
     ) -> np.ndarray:
         """
         Decodes action token IDs back to continuous action values using the FAST tokenizer.
@@ -1045,6 +1033,8 @@ class PI0FastPolicy(PreTrainedPolicy):
             token_ids: List of token IDs to decode.
             time_horizon: The number of timesteps for actions.
             action_dim: The dimensionality of each action.
+            relaxed_decoding: Whether to use relaxed decoding (allows partial sequences).
+
         Returns:
             A numpy array representing the decoded actions.
         """
@@ -1052,23 +1042,40 @@ class PI0FastPolicy(PreTrainedPolicy):
 
         for token in token_ids:
             try:
-                expected_shape = (time_horizon, action_dim)
-                decoded_action = np.asarray(
-                    self.action_tokenizer.decode(
-                        [token.tolist()], time_horizon=time_horizon, action_dim=action_dim
-                    )[0],
-                    dtype=np.float32,
+                decoded_tokens = self.action_tokenizer.bpe_tokenizer.decode(token)
+                decoded_dct_coeff = np.array(list(map(ord, decoded_tokens))) + self.action_tokenizer.min_token
+
+                if relaxed_decoding:
+                    # expected sequence length
+                    expected_seq_len = time_horizon * action_dim
+                    diff = expected_seq_len - decoded_dct_coeff.shape[0]
+
+                    # apply truncation if too long
+                    if diff < 0:
+                        decoded_dct_coeff = decoded_dct_coeff[:expected_seq_len]  # truncate on the right
+
+                    # apply padding if too short
+                    elif diff > 0:
+                        decoded_dct_coeff = np.pad(
+                            decoded_dct_coeff, (0, diff), mode="constant", constant_values=0
+                        )
+
+                decoded_dct_coeff = decoded_dct_coeff.reshape(-1, action_dim)
+                assert decoded_dct_coeff.shape == (
+                    time_horizon,
+                    action_dim,
+                ), (
+                    f"Decoded DCT coefficients have shape {decoded_dct_coeff.shape}, expected ({time_horizon}, {action_dim})"
                 )
-                if decoded_action.shape != expected_shape:
-                    raise ValueError(
-                        f"decoded action shape {decoded_action.shape} does not match {expected_shape}"
-                    )
 
             except Exception as e:
-                logging.warning("Invalid FAST action sequence; returning a zero action chunk: %s", e)
-                decoded_action = np.zeros((time_horizon, action_dim))
+                logging.warning(f"Error decoding tokens: {e}")
+                logging.warning(f"Tokens: {token}")
+                decoded_dct_coeff = np.zeros((time_horizon, action_dim))
 
-            decoded_actions.append(decoded_action)
+            decoded_actions.append(
+                idct(decoded_dct_coeff / self.action_tokenizer.scale, axis=0, norm="ortho")
+            )
 
         return np.stack(decoded_actions)
 
@@ -1098,28 +1105,53 @@ class PI0FastPolicy(PreTrainedPolicy):
         if single_sample:
             tokens = tokens.unsqueeze(0)
 
-        action_tokens = []
-        for token_sequence in tokens:
-            try:
-                token_ids = token_sequence.tolist()
-                eos_token_id = self._paligemma_tokenizer.eos_token_id
-                if eos_token_id in token_ids:
-                    token_ids = token_ids[: token_ids.index(eos_token_id) + 1]
-                decoded_text = self._paligemma_tokenizer.decode(token_ids)
-                if not decoded_text.startswith("Action: ") or "|" not in decoded_text:
-                    raise ValueError(f"expected 'Action: <codes>|', got {decoded_text!r}")
-                action_text = decoded_text.removeprefix("Action: ").split("|", maxsplit=1)[0]
-                raw_action_tokens = torch.tensor(
-                    self._paligemma_tokenizer.encode(action_text, add_special_tokens=False),
-                    dtype=torch.long,
-                    device=tokens.device,
+        # Convert token IDs to token strings
+        decoded_tokens = [self._paligemma_tokenizer.convert_ids_to_tokens(seq.tolist()) for seq in tokens]
+        # Get the token sequence for "Action: " to remove it
+        action_prefix_ids = self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False)
+        action_prefix_tokens = self._paligemma_tokenizer.convert_ids_to_tokens(action_prefix_ids)
+        action_prefix_len = len(action_prefix_tokens)
+
+        # Clean tokens by removing everything after the first "|" (end-of-action marker)
+        # and removing all occurrences of "Action: " token sequence
+        # assert that beginning contain "Action: "
+        if self.config.validate_action_token_prefix:
+            for token_seq in decoded_tokens:
+                assert len(token_seq) >= 2 and token_seq[0] == "Action" and token_seq[1] == ":", (
+                    f"Token sequence does not start with ['Action', ':']: {token_seq}"
                 )
-                if raw_action_tokens.numel() == 0:
-                    raise ValueError("empty FAST action payload")
-                action_tokens.append(self._paligemma_tokens_to_act_tokens(raw_action_tokens))
-            except Exception as e:
-                logging.warning("Invalid generated PI0-FAST text; returning zeros for this sample: %s", e)
-                action_tokens.append(torch.empty(0, dtype=torch.long, device=tokens.device))
+
+        cleaned_tokens = []
+        for token_seq in decoded_tokens:
+            # Remove everything after "|"
+            if "|" in token_seq:
+                token_seq = token_seq[: token_seq.index("|")]
+
+            # Remove all occurrences of "Action: " token sequence
+            i = 0
+            while i <= len(token_seq) - action_prefix_len:
+                if token_seq[i : i + action_prefix_len] == action_prefix_tokens:
+                    # Found a match, remove it
+                    token_seq = token_seq[:i] + token_seq[i + action_prefix_len :]
+                else:
+                    i += 1
+
+            cleaned_tokens.append(token_seq)
+
+        # Convert token strings back to IDs
+        raw_action_tokens = [
+            torch.tensor(
+                self._paligemma_tokenizer.convert_tokens_to_ids(token_seq),
+                dtype=torch.long,
+                device=tokens.device,
+            )
+            for token_seq in cleaned_tokens
+        ]
+
+        # Convert PaliGemma tokens to action tokens
+        action_tokens = [
+            self._paligemma_tokens_to_act_tokens(raw_action_token) for raw_action_token in raw_action_tokens
+        ]
 
         # Decode action tokens to continuous actions
         actions = self.decode_actions_with_fast(
@@ -1188,7 +1220,7 @@ class PI0FastPolicy(PreTrainedPolicy):
             )
 
         # Detokenize action tokens to continuous actions
-        action_horizon = self.config.chunk_size
+        action_horizon = self.config.n_action_steps
         action_dim = self.config.output_features[ACTION].shape[0]
 
         continuous_actions = self.detokenize_actions(

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -70,7 +70,7 @@ class Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(ProcessorStep):
 
         full_prompts = []
         for i, task in enumerate(tasks):
-            cleaned_text = task.strip().replace("_", " ").replace("\n", " ").lower()
+            cleaned_text = task.strip().replace("_", " ").replace("\n", " ")
             state_str = " ".join(map(str, discretized_states[i]))
             full_prompt = f"Task: {cleaned_text}, State: {state_str};\n"
             full_prompts.append(full_prompt)
@@ -92,11 +92,6 @@ class Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(ProcessorStep):
 def make_pi0_fast_pre_post_processors(
     config: PI0FastConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    dataset_repo_id: str | None = None,
-    dataset_root: str | None = None,
-    dataset_revision: str | None = None,
-    episodes: list[int] | None = None,
-    exclude_episodes: list[int] | None = None,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -141,18 +136,6 @@ def make_pi0_fast_pre_post_processors(
     # state from the observation but does not change it. NormalizerProcessorStep still runs
     # before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
     # continues to receive normalized state in [-1, 1] as expected.
-    from ..pi052.fit_fast_tokenizer import resolve_fast_tokenizer  # noqa: PLC0415
-
-    action_tokenizer_path = resolve_fast_tokenizer(
-        config,
-        dataset_repo_id,
-        dataset_root,
-        dataset_stats,
-        dataset_revision,
-        episodes,
-        exclude_episodes,
-    )
-
     input_steps: list[ProcessorStep] = [
         steps.rename_observations,  # To mimic the same processor as pretrained one
         steps.add_batch_dim,
@@ -166,11 +149,10 @@ def make_pi0_fast_pre_post_processors(
             padding="max_length",
         ),
         ActionTokenizerProcessorStep(
-            action_tokenizer_name=action_tokenizer_path,
+            action_tokenizer_name=config.action_tokenizer_name,
             max_action_tokens=config.max_action_tokens,
             fast_skip_tokens=config.fast_skip_tokens,
             paligemma_tokenizer_name=config.text_tokenizer_name,
-            prepend_bos=False,
         ),
         steps.to_device,
     ]

src/lerobot/policies/pi_gemma.py

@@ -18,7 +18,6 @@ from typing import TYPE_CHECKING
 
 import torch
 from torch import nn
-from torch.nn import functional as F  # noqa: N812
 
 from lerobot.utils.import_utils import _transformers_available
 
@@ -122,10 +121,7 @@ class PiGemmaRMSNorm(nn.Module):
         if cond.shape[-1] != self.cond_dim:
             raise ValueError(f"Expected cond dim {self.cond_dim}, got {cond.shape[-1]}")
         modulation = self.dense(cond)
-        # Per-sample cond (B, cond_dim) → broadcast over the sequence. A
-        # per-token cond (B, T, cond_dim) is already aligned with x and must
-        # not be unsqueezed (used by pi052's amortized K_repeat path).
-        if len(x.shape) == 3 and modulation.dim() == 2:
+        if len(x.shape) == 3:
             modulation = modulation.unsqueeze(1)
         scale, shift, gate = modulation.chunk(3, dim=-1)
         normed = normed * (1 + scale.float()) + shift.float()
@@ -279,8 +275,6 @@ class PiGemmaModel(GemmaModel):  # type: ignore[misc]
         # Convert to bfloat16 if the first layer uses bfloat16
         if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
             hidden_states = hidden_states.to(torch.bfloat16)
-        if causal_mask is not None and torch.is_floating_point(causal_mask):
-            causal_mask = causal_mask.to(dtype=hidden_states.dtype)
 
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -373,45 +367,3 @@ __all__ = [
     "PaliGemmaModelWithPiGemma",
     "PaliGemmaForConditionalGenerationWithPiGemma",
 ]
-
-
-# PI0.5 / PI052 dual-expert backbone: generic PaliGemma + Gemma action-expert
-# transformer machinery used by the pi052 policy. GemmaVariantConfig is openpi's
-# width/depth variant config (renamed from GemmaConfig to avoid clashing with
-# transformers' GemmaConfig).
-
-
-def sdpa_attention_forward(
-    module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: torch.Tensor | None,
-    scaling: float,
-    dropout: float = 0.0,
-):
-    """Drop-in for ``modeling_gemma.eager_attention_forward`` using
-    ``torch.nn.functional.scaled_dot_product_attention``.
-
-    PyTorch SDPA picks the memory-efficient kernel for arbitrary additive
-    bias masks (the FA backend only accepts causal/sliding-window). On
-    H100 that is ~1.3-1.7x faster and uses ~30-40% less attention memory
-    than the eager softmax(QK^T)+matmul path. Mirrors eager's signature
-    and output shape (``(B, Lq, H, D)``) so call sites are unchanged.
-    """
-    n_rep = module.num_key_value_groups
-    if n_rep > 1:
-        key = key.repeat_interleave(n_rep, dim=1)
-        value = value.repeat_interleave(n_rep, dim=1)
-    if attention_mask is not None and attention_mask.dtype != query.dtype:
-        attention_mask = attention_mask.to(dtype=query.dtype)
-    attn_output = F.scaled_dot_product_attention(
-        query,
-        key,
-        value,
-        attn_mask=attention_mask,
-        dropout_p=dropout if module.training else 0.0,
-        is_causal=False,
-        scale=scaling,
-    )
-    return attn_output.transpose(1, 2).contiguous(), None

src/lerobot/policies/pretrained.py

@@ -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,
@@ -338,7 +346,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             "smolvla": "lerobot/smolvla_base",
             "pi0": "lerobot/pi0_base",
             "pi05": "lerobot/pi05_base",
-            "pi052": "lerobot/pi052_base",
             "pi0_fast": "lerobot/pi0fast-base",
             "xvla": "lerobot/xvla-base",
         }
@@ -385,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:
@@ -456,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()
 
@@ -481,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"
@@ -508,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")

src/lerobot/processor/batch_processor.py

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

src/lerobot/processor/delta_action_processor.py

@@ -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,)
             )

src/lerobot/processor/pipeline.py

@@ -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,
-            overrides or {},
-            model_id,
-            base_path,
-            config_filename,
-            hub_download_kwargs,
-            is_local_source,
+            loaded_config, overrides or {}, model_id, base_path, 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_entry,
-                model_id,
-                base_path,
-                config_filename,
-                hub_download_kwargs,
-                is_local_source,
+                step_instance, step_entry, model_id, base_path, 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,
             )

src/lerobot/processor/render_messages_processor.py

@@ -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)
+            return transition
 
         timestamp = complementary_data.get("timestamp")
         if timestamp is None:
@@ -78,147 +67,18 @@ class RenderMessagesStep(ProcessorStep):
             dataset_ctx=self.dataset_ctx,
         )
         if rendered is None:
-            rendered = _fallback_low_level_render(complementary_data.get("task"))
-            if rendered is None:
-                return 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": [],
-    }

src/lerobot/processor/tokenizer_processor.py

@@ -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,8 +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
-    prepend_bos: 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)
@@ -416,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.
-        tokens, mask, code_mask = self._tokenize_action(action)
+        # Tokenize and get both tokens and mask
+        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
@@ -435,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.
 
@@ -464,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)
@@ -482,79 +476,65 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             if tokens.dim() > 1:
                 tokens = tokens.flatten()
 
-            action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens)
-            prompt_tokens = torch.tensor(
-                self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
-                device=action.device,
+            bos_id = self._paligemma_tokenizer.bos_token_id
+            # add bos
+            tokens = torch.cat(
+                [
+                    torch.tensor([bos_id], device=action.device),
+                    torch.tensor(
+                        self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
+                        device=action.device,
+                    ),
+                    self._act_tokens_to_paligemma_tokens(tokens),
+                    torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device),
+                ]
             )
-            end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device)
-
-            token_parts = []
-            if self.prepend_bos:
-                token_parts.append(
-                    torch.tensor([self._paligemma_tokenizer.bos_token_id], device=action.device)
-                )
-            code_start = sum(len(part) for part in token_parts) + len(prompt_tokens)
-            code_end = code_start + len(action_code_tokens)
-            tokens = torch.cat([*token_parts, prompt_tokens, action_code_tokens, end_tokens])
-            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]:
@@ -570,10 +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,
-            "prepend_bos": self.prepend_bos,
         }
 
         # Only save tokenizer_name if it was used to create the tokenizer
@@ -582,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]]:

src/lerobot/rl/actor.py

@@ -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
-
-        mp.set_start_method("spawn")
+        ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
         display_pid = True
 
     # Create logs directory to ensure it exists

src/lerobot/rl/learner.py

@@ -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
-
-        mp.set_start_method("spawn")
+        ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
 
     # Use the job_name from the config
     train(

src/lerobot/robots/so_follower/config_so_follower.py

@@ -46,6 +46,12 @@ class SOFollowerConfig:
     position_i_coefficient: int = 0
     position_d_coefficient: int = 32
 
+    # Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
+    # return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
+    # at once, which otherwise aborts the control loop. Retries are immediate (no sleep) and only happen on
+    # failure, so the steady-state read cost is unchanged.
+    num_read_retries: int = 2
+
 
 @RobotConfig.register_subclass("so101_follower")
 @RobotConfig.register_subclass("so100_follower")

src/lerobot/robots/so_follower/so_follower.py

@@ -180,7 +180,7 @@ class SOFollower(Robot):
     def get_observation(self) -> RobotObservation:
         # Read arm position
         start = time.perf_counter()
-        obs_dict = self.bus.sync_read("Present_Position")
+        obs_dict = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
         obs_dict = {f"{motor}.pos": val for motor, val in obs_dict.items()}
         dt_ms = (time.perf_counter() - start) * 1e3
         logger.debug(f"{self} read state: {dt_ms:.1f}ms")
@@ -221,7 +221,7 @@ class SOFollower(Robot):
         # Cap goal position when too far away from present position.
         # /!\ Slower fps expected due to reading from the follower.
         if self.config.max_relative_target is not None:
-            present_pos = self.bus.sync_read("Present_Position")
+            present_pos = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
             goal_present_pos = {key: (g_pos, present_pos[key]) for key, g_pos in goal_pos.items()}
             goal_pos = ensure_safe_goal_position(goal_present_pos, self.config.max_relative_target)
 

src/lerobot/robots/utils.py

@@ -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)}"
         )

src/lerobot/rollout/context.py

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

src/lerobot/runtime/__init__.py

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

src/lerobot/runtime/adapter.py

@@ -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

src/lerobot/runtime/language_runtime.py

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

src/lerobot/runtime/registry.py

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

src/lerobot/runtime/sim_robocasa.py

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

src/lerobot/scripts/augment_dataset_quantile_stats.py

@@ -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] = {}
-    for idx in range(start_idx, end_idx):
-        item = dataset[idx]
-        for key, value in item.items():
-            if key not in dataset.features:
-                continue
+    episode_len = end_idx - start_idx
 
-            if key not in collected_data:
-                collected_data[key] = []
-            collected_data[key].append(value)
+    # Images/video are the memory hog, so sub-sample those frames per episode;
+    # numeric columns are cheap, so read them in full (exact).
+    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,
     )
 
 

src/lerobot/scripts/convert_dataset_v21_to_v30.py

@@ -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*"],

src/lerobot/scripts/lerobot_annotate.py

@@ -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)",
-            flush=True,
+            f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
         )
 
 

src/lerobot/scripts/lerobot_dataset_viz.py

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

src/lerobot/scripts/lerobot_eval.py

@@ -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:
-    """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])
+logger = logging.getLogger(__name__)
 
 
 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:
-            from peft import PeftModel
-
-            if not isinstance(policy, PeftModel):
-                raise exc
-        except ImportError:
-            raise exc from None
+        if not _peft_available:
+            raise exc
+        require_package("peft", extra="peft")
+        if not isinstance(policy, PeftModel):
+            raise exc
 
     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])
-        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))
+            ep_frames.append(np.stack(env.call("render")[:n_to_render_now]))
 
     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:")
-        print(info["overall"])
+        logger.info("Overall Aggregated Metrics:")
+        logger.info(info["overall"])
 
         # Print per-suite stats
         for task_group, task_group_info in info.items():
-            print(f"\nAggregated Metrics for {task_group}:")
-            print(task_group_info)
+            logger.info(f"\nAggregated Metrics for {task_group}:")
+            logger.info(task_group_info)
     # Close all vec envs
     close_envs(envs)
 

src/lerobot/scripts/lerobot_find_cameras.py

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

src/lerobot/scripts/lerobot_rollout.py

@@ -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 = {
-    "--language",
-    "--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.
-    """
+def main():
+    """CLI entry point for ``lerobot-rollout``."""
     register_third_party_plugins()
-    cli_args = list(sys.argv[1:] if argv is None else argv)
-    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
+    rollout()
 
 
 if __name__ == "__main__":

src/lerobot/scripts/lerobot_train.py

@@ -20,11 +20,10 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
-import os
 import sys
 import time
-from contextlib import nullcontext
-from datetime import timedelta
+from collections.abc import Iterator
+from contextlib import contextmanager, nullcontext
 from pprint import pformat
 from typing import TYPE_CHECKING, Any
 
@@ -59,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 (
@@ -70,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
@@ -93,7 +111,6 @@ def update_policy(
     lr_scheduler=None,
     lock=None,
     sample_weighter=None,
-    log_metrics: bool = True,
 ) -> tuple[MetricsTracker, dict | None]:
     """
     Performs a single training step to update the policy's weights.
@@ -111,7 +128,6 @@ def update_policy(
         lr_scheduler: An optional learning rate scheduler.
         lock: An optional lock for thread-safe optimizer updates.
         sample_weighter: Optional SampleWeighter instance for per-sample loss weighting.
-        log_metrics: Whether to synchronize and record GPU metrics this step.
 
     Returns:
         A tuple containing:
@@ -179,20 +195,12 @@ def update_policy(
     if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"):
         accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
 
+    train_metrics.loss = loss.item()
+    train_metrics.grad_norm = grad_norm.item()
     train_metrics.lr = optimizer.param_groups[0]["lr"]
+    train_metrics.update_s = time.perf_counter() - start_time
     if torch.cuda.is_available():
         train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
-    train_metrics.accumulate_tensor("loss", loss)
-    train_metrics.accumulate_tensor("grad_norm", grad_norm)
-    train_metrics.update_s = time.perf_counter() - start_time
-    # Synchronize accumulated GPU metrics only when logging.
-    if log_metrics:
-        train_metrics.materialize_tensors()
-        # Materialize detached loss components during the same logging synchronization.
-        if output_dict:
-            output_dict = {
-                k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items()
-            }
     # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
     if output_dict:
         train_metrics.update_metrics(output_dict)
@@ -219,11 +227,9 @@ 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, InitProcessGroupKwargs
+    from accelerate.utils import DistributedDataParallelKwargs, DistributedType
 
     cfg.validate()
 
@@ -232,16 +238,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes
     # We set find_unused_parameters=True to handle models with conditional computation
     if accelerator is None:
-        # Static graphs restore DDP overlap when conditional parameter usage is stable.
-        # Environment flags retain the existing defaults.
-        ddp_find_unused = os.environ.get("LEROBOT_DDP_FIND_UNUSED", "1") == "1"
-        ddp_static_graph = os.environ.get("LEROBOT_DDP_STATIC_GRAPH", "0") == "1"
-        ddp_kwargs = DistributedDataParallelKwargs(
-            find_unused_parameters=ddp_find_unused and not ddp_static_graph,
-            static_graph=ddp_static_graph,
-        )
-        # Allow rank 0 enough time to index large datasets before other ranks leave the barrier.
-        ipg_kwargs = InitProcessGroupKwargs(timeout=timedelta(hours=2))
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
         # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
         # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
         force_cpu = cfg.trainable_config.device == "cpu"
@@ -251,7 +248,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         accelerator = Accelerator(
             step_scheduler_with_optimizer=False,
             mixed_precision=mixed_precision,
-            kwargs_handlers=[ddp_kwargs, ipg_kwargs],
+            kwargs_handlers=[ddp_kwargs],
             cpu=force_cpu,
         )
 
@@ -298,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")
@@ -333,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.")
@@ -347,14 +336,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     active_cfg = cfg.trainable_config
     processor_pretrained_path = active_cfg.pretrained_path
-    # A weight checkpoint may contain PI05 or differently configured PI052 processors.
-    if cfg.policy.type == "pi052" and processor_pretrained_path is not None and not cfg.resume:
-        logging.warning(
-            "pi052 is loading pretrained weights from %s, but building processors from the current "
-            "pi052 config so recipe text labels and FAST action labels are generated.",
-            processor_pretrained_path,
-        )
-        processor_pretrained_path = None
 
     processor_kwargs = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
@@ -363,13 +344,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
-    if cfg.policy.type in {"pi0_fast", "pi052"}:
-        processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
-        processor_kwargs["dataset_revision"] = cfg.dataset.revision
-        processor_kwargs["dataset_episodes"] = cfg.dataset.episodes
-        processor_kwargs["dataset_exclude_episodes"] = cfg.dataset.exclude_episodes
-        processor_kwargs["dataset_root"] = cfg.dataset.root
-
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
         preprocessor_overrides = {
             "device_processor": {"device": device.type},
@@ -466,17 +440,13 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         # same permutation. accelerate then shards it disjointly across ranks via BatchSamplerShard
         # without needing a `generator` attribute to synchronize an RNG, and resume is sample-exact.
         shuffle = False
-        from_indices = dataset.meta.episodes["dataset_from_index"]
-        to_indices = dataset.meta.episodes["dataset_to_index"]
-        seed = cfg.seed if cfg.seed is not None else 0
-
         sampler = EpisodeAwareSampler(
-            from_indices,
-            to_indices,
+            dataset.meta.episodes["dataset_from_index"],
+            dataset.meta.episodes["dataset_to_index"],
             episode_indices_to_use=dataset.episodes,
             drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
             shuffle=True,
-            seed=seed,
+            seed=cfg.seed if cfg.seed is not None else 0,
             absolute_to_relative_idx=dataset.absolute_to_relative_idx,
         )
         if cfg.resume and step > 0:
@@ -623,10 +593,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        # Synchronize GPU metrics only for updates that will be logged.
-        log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0
-
-        train_tracker, output_dict = update_policy(
+        train_tracker, _ = update_policy(
             train_tracker,
             policy,
             batch,
@@ -635,7 +602,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             accelerator=accelerator,
             lr_scheduler=lr_scheduler,
             sample_weighter=sample_weighter,
-            log_metrics=log_metrics,
         )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
@@ -736,11 +702,10 @@ 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}")
-                eval_target_policy = accelerator.unwrap_model(policy)
-                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=eval_target_policy,
+                        policy=accelerator.unwrap_model(policy),
                         env_preprocessor=env_preprocessor,
                         env_postprocessor=env_postprocessor,
                         preprocessor=preprocessor,
@@ -785,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:

src/lerobot/scripts/lerobot_train_tokenizer.py

@@ -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:
-        print(f"Error processing episode {ep_idx}: {e}")
-        import traceback
-
-        traceback.print_exc()
+    except Exception:
+        logger.exception("Error processing episode %s", ep_idx)
         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...")
-    print(f"Action chunk shape: {action_chunks.shape}")
-    print(f"Vocab size: {vocab_size}")
-    print(f"DCT scale: {scale}")
+    logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
+    logger.info(f"Action chunk shape: {action_chunks.shape}")
+    logger.info(f"Vocab size: {vocab_size}")
+    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:")
-    print(f"  Average compression ratio: {stats['compression_ratio']:.2f}x")
-    print(f"  Mean token length: {stats['mean_token_length']:.1f}")
-    print(f"  P99 token length: {stats['p99_token_length']:.0f}")
-    print(f"  Min token length: {stats['min_token_length']:.0f}")
-    print(f"  Max token length: {stats['max_token_length']:.0f}")
+    logger.info("Compression Statistics:")
+    logger.info(f"  Average compression ratio: {stats['compression_ratio']:.2f}x")
+    logger.info(f"  Mean token length: {stats['mean_token_length']:.1f}")
+    logger.info(f"  P99 token length: {stats['p99_token_length']:.0f}")
+    logger.info(f"  Min token length: {stats['min_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}")
-    print(f"State key: {cfg.state_key}")
+    logger.info(f"Action horizon: {cfg.action_horizon}")
+    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:")
-    print(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("\nBefore normalization - overall stats:")
+    logger.info(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(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:")
-            print(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("\nAfter normalization - overall stats:")
+            logger.info(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(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}")
-    print(f"Metadata: {json.dumps(metadata, indent=2)}")
+    logger.info(f"\nSaved FAST tokenizer to {output_path}")
+    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}")
-        print(f"   Private: {cfg.hub_private}")
+        logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
+        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}")
-            print("   Make sure you're logged in with `huggingface-cli login`")
+            logger.error(f"Error pushing to hub: {e}")
+            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()
 
 

src/lerobot/teleoperators/so_leader/config_so_leader.py

@@ -29,6 +29,12 @@ class SOLeaderConfig:
     # Whether to use degrees for angles
     use_degrees: bool = True
 
+    # Number of extra attempts when a `sync_read` of the motors fails. Feetech buses can occasionally
+    # return a corrupted status packet ("Incorrect status packet!"), especially when several joints move
+    # at once, which otherwise aborts the teleoperation loop. Retries are immediate (no sleep) and only
+    # happen on failure, so the steady-state read cost is unchanged.
+    num_read_retries: int = 2
+
 
 @TeleoperatorConfig.register_subclass("so101_leader")
 @TeleoperatorConfig.register_subclass("so100_leader")

src/lerobot/teleoperators/so_leader/so_leader.py

@@ -145,7 +145,7 @@ class SOLeader(Teleoperator):
     @check_if_not_connected
     def get_action(self) -> dict[str, float]:
         start = time.perf_counter()
-        action = self.bus.sync_read("Present_Position")
+        action = self.bus.sync_read("Present_Position", num_retry=self.config.num_read_retries)
         action = {f"{motor}.pos": val for motor, val in action.items()}
         dt_ms = (time.perf_counter() - start) * 1e3
         logger.debug(f"{self} read action: {dt_ms:.1f}ms")

src/lerobot/transforms/__init__.py

@@ -13,18 +13,34 @@
 # limitations under the License.
 
 from .transforms import (
+    CoarseDropout,
+    GammaCorrection,
+    GaussianNoise,
+    GaussianPatchBrightness,
     ImageTransformConfig,
     ImageTransforms,
     ImageTransformsConfig,
+    JPEGCompression,
+    MotionBlur,
+    PlanckianJitter,
+    RandomShadow,
     RandomSubsetApply,
     SharpnessJitter,
     make_transform_from_config,
 )
 
 __all__ = [
+    "CoarseDropout",
+    "GammaCorrection",
+    "GaussianNoise",
+    "GaussianPatchBrightness",
     "ImageTransformConfig",
     "ImageTransforms",
     "ImageTransformsConfig",
+    "JPEGCompression",
+    "MotionBlur",
+    "PlanckianJitter",
+    "RandomShadow",
     "RandomSubsetApply",
     "SharpnessJitter",
     "make_transform_from_config",

src/lerobot/transforms/transforms.py

@@ -14,11 +14,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import collections
+import math
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 from typing import Any
 
 import torch
+from torchvision.io import decode_image, encode_jpeg
 from torchvision.transforms import v2
 from torchvision.transforms.v2 import (
     Transform,
@@ -41,7 +43,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 +52,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 +71,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 +121,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.")
@@ -144,6 +146,471 @@ class SharpnessJitter(Transform):
         return self._call_kernel(F.adjust_sharpness, inpt, sharpness_factor=sharpness_factor)
 
 
+class GaussianNoise(Transform):
+    """Add Gaussian noise to simulate camera sensor noise.
+
+    Models readout noise from ADC quantization, which increases in low-light conditions.
+    Common in real-robot setups where wrist cameras operate in suboptimal lighting.
+
+    Args:
+        std: Range (min, max) for noise standard deviation in pixel-value scale (0-255).
+    """
+
+    def __init__(self, std: float | Sequence[float] = (5.0, 25.0)) -> None:
+        super().__init__()
+        if isinstance(std, (int, float)):
+            self.std = (0.0, float(std))
+        elif isinstance(std, Sequence) and len(std) == 2:
+            self.std = (float(std[0]), float(std[1]))
+        else:
+            raise TypeError("std must be a number or a sequence with length 2.")
+        if not 0.0 <= self.std[0] <= self.std[1]:
+            raise ValueError(f"std must satisfy 0 <= min <= max, but got {self.std}.")
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        return {
+            "std": torch.empty(1).uniform_(self.std[0], self.std[1]).item(),
+            "seed": torch.randint(0, torch.iinfo(torch.int64).max, ()).item(),
+        }
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if isinstance(inpt, torch.Tensor) and inpt.is_floating_point():
+            generator = torch.Generator(device=inpt.device).manual_seed(params["seed"])
+            noise = torch.randn(inpt.shape, device=inpt.device, dtype=inpt.dtype, generator=generator)
+            return (inpt + noise * (params["std"] / 255.0)).clamp(0.0, 1.0)
+        return inpt
+
+
+class MotionBlur(Transform):
+    """Apply directional motion blur to simulate fast robot or object movement.
+
+    Generates a 1D averaging kernel along a random direction, applied via depthwise convolution.
+
+    Args:
+        kernel_size: An odd kernel size or a range containing at least one odd kernel size.
+    """
+
+    def __init__(self, kernel_size: int | Sequence[int] = (3, 11)) -> None:
+        super().__init__()
+        if isinstance(kernel_size, int):
+            self.kernel_size = (kernel_size, kernel_size)
+        elif isinstance(kernel_size, Sequence) and len(kernel_size) == 2:
+            self.kernel_size = (int(kernel_size[0]), int(kernel_size[1]))
+        else:
+            raise TypeError("kernel_size must be an int or a sequence with length 2.")
+        if not 1 <= self.kernel_size[0] <= self.kernel_size[1]:
+            raise ValueError(f"kernel_size must satisfy 1 <= min <= max, but got {self.kernel_size}.")
+        self._first_odd_kernel_size = self.kernel_size[0] + (self.kernel_size[0] + 1) % 2
+        if self._first_odd_kernel_size > self.kernel_size[1]:
+            raise ValueError(f"kernel_size range must contain an odd value, but got {self.kernel_size}.")
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        num_odd_sizes = (self.kernel_size[1] - self._first_odd_kernel_size) // 2 + 1
+        size_index = int(torch.randint(0, num_odd_sizes, ()).item())
+        ks = self._first_odd_kernel_size + 2 * size_index
+        angle = torch.empty(1).uniform_(0, 360).item()
+        return {"kernel_size": ks, "angle": angle}
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
+            return inpt
+        if inpt.ndim < 3:
+            raise ValueError(f"MotionBlur expects [..., C, H, W] input, but got shape {inpt.shape}.")
+
+        kernel_size = params["kernel_size"]
+        radius = kernel_size // 2
+        angle = math.radians(params["angle"])
+        positions = torch.linspace(-radius, radius, kernel_size, device=inpt.device)
+        x_coords = (positions * math.cos(angle)).round().to(torch.long) + radius
+        y_coords = (positions * math.sin(angle)).round().to(torch.long) + radius
+        kernel = torch.zeros((kernel_size, kernel_size), device=inpt.device, dtype=inpt.dtype)
+        kernel[y_coords, x_coords] = 1
+        kernel /= kernel.sum()
+
+        channels, height, width = inpt.shape[-3:]
+        flat_input = inpt.reshape(-1, channels, height, width)
+        depthwise_kernel = kernel.expand(channels, 1, kernel_size, kernel_size)
+        padded = torch.nn.functional.pad(flat_input, (radius,) * 4, mode="replicate")
+        output = torch.nn.functional.conv2d(padded, depthwise_kernel, groups=channels)
+        return output.reshape(inpt.shape).clamp(0.0, 1.0)
+
+
+class JPEGCompression(Transform):
+    """Simulate JPEG compression artifacts (block artifacts, color banding).
+
+    Models quality degradation from video compression in network-streamed camera feeds.
+
+    Args:
+        quality: Range (min, max) for JPEG quality factor (lower = more artifacts).
+    """
+
+    def __init__(self, quality: int | Sequence[int] = (15, 75)) -> None:
+        super().__init__()
+        if isinstance(quality, int):
+            self.quality = (quality, quality)
+        elif isinstance(quality, Sequence) and len(quality) == 2:
+            self.quality = (int(quality[0]), int(quality[1]))
+        else:
+            raise TypeError("quality must be an int or a sequence with length 2.")
+        if not 1 <= self.quality[0] <= self.quality[1] <= 100:
+            raise ValueError(f"quality must satisfy 1 <= min <= max <= 100, but got {self.quality}.")
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        return {"quality": int(torch.randint(self.quality[0], self.quality[1] + 1, (1,)).item())}
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
+            return inpt
+        if inpt.ndim < 3:
+            raise ValueError(f"JPEGCompression expects [..., C, H, W] input, but got shape {inpt.shape}.")
+
+        channels, height, width = inpt.shape[-3:]
+        if channels not in (1, 3):
+            raise ValueError(f"JPEGCompression expects 1 or 3 channels, but got {channels}.")
+
+        flat_input = inpt.reshape(-1, channels, height, width)
+        flat_uint8 = (flat_input.clamp(0.0, 1.0) * 255).round().to(torch.uint8).cpu()
+        decoded_frames = [
+            decode_image(encode_jpeg(frame, quality=params["quality"])) for frame in flat_uint8.unbind()
+        ]
+        output = torch.stack(decoded_frames).to(device=inpt.device, dtype=inpt.dtype) / 255.0
+        return output.reshape(inpt.shape)
+
+
+class GaussianPatchBrightness(Transform):
+    """Apply spatially-varying brightness with Gaussian patches.
+
+    Simulates uneven overhead lighting, spotlights, and shadow patches commonly
+    encountered in real robot workspaces with multiple light sources.
+
+    Args:
+        num_patches: Range (min, max) for number of brightness patches.
+        sigma_range: Range for Gaussian sigma as fraction of image size.
+        factor_range: Range for brightness factor (< 1 darkens, > 1 brightens).
+    """
+
+    def __init__(
+        self,
+        num_patches: int | Sequence[int] = (1, 4),
+        sigma_range: Sequence[float] = (0.05, 0.25),
+        factor_range: Sequence[float] = (0.4, 1.6),
+    ) -> None:
+        super().__init__()
+        if isinstance(num_patches, int):
+            self.num_patches = (num_patches, num_patches)
+        elif isinstance(num_patches, Sequence) and len(num_patches) == 2:
+            self.num_patches = (int(num_patches[0]), int(num_patches[1]))
+        else:
+            raise TypeError("num_patches must be an int or a sequence with length 2.")
+        if not 1 <= self.num_patches[0] <= self.num_patches[1]:
+            raise ValueError(f"num_patches must satisfy 1 <= min <= max, but got {self.num_patches}.")
+        if not isinstance(sigma_range, Sequence) or len(sigma_range) != 2:
+            raise TypeError("sigma_range must be a sequence with length 2.")
+        self.sigma_range = (float(sigma_range[0]), float(sigma_range[1]))
+        if not 0.0 < self.sigma_range[0] <= self.sigma_range[1]:
+            raise ValueError(f"sigma_range must satisfy 0 < min <= max, but got {self.sigma_range}.")
+        if not isinstance(factor_range, Sequence) or len(factor_range) != 2:
+            raise TypeError("factor_range must be a sequence with length 2.")
+        self.factor_range = (float(factor_range[0]), float(factor_range[1]))
+        if not 0.0 <= self.factor_range[0] <= self.factor_range[1]:
+            raise ValueError(f"factor_range must satisfy 0 <= min <= max, but got {self.factor_range}.")
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        n = int(torch.randint(self.num_patches[0], self.num_patches[1] + 1, (1,)).item())
+        return {
+            "centers": torch.rand(n, 2).tolist(),
+            "sigmas": torch.empty(n).uniform_(self.sigma_range[0], self.sigma_range[1]).tolist(),
+            "factors": torch.empty(n).uniform_(self.factor_range[0], self.factor_range[1]).tolist(),
+        }
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
+            return inpt
+        h, w = inpt.shape[-2:]
+        mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype)
+        grid_y = torch.linspace(0, 1, h, device=inpt.device, dtype=inpt.dtype)
+        grid_x = torch.linspace(0, 1, w, device=inpt.device, dtype=inpt.dtype)
+        yy, xx = torch.meshgrid(grid_y, grid_x, indexing="ij")
+        for (cy, cx), sigma, factor in zip(
+            params["centers"], params["sigmas"], params["factors"], strict=True
+        ):
+            gauss = torch.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * sigma**2))
+            mask = mask * (1.0 + (factor - 1.0) * gauss)
+        broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
+        return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
+
+
+class RandomShadow(Transform):
+    """Add random vertical band shadow with smooth edges.
+
+    Simulates cast shadows from objects or people near the robot workspace.
+    Symmetric: randomly brightens or darkens to prevent BatchNorm stats shift.
+
+    Args:
+        opacity: Range (min, max) for shadow/highlight opacity.
+    """
+
+    def __init__(self, opacity: float | Sequence[float] = (0.3, 0.6)) -> None:
+        super().__init__()
+        if isinstance(opacity, (int, float)):
+            self.opacity = (float(opacity), float(opacity))
+        elif isinstance(opacity, Sequence) and len(opacity) == 2:
+            self.opacity = (float(opacity[0]), float(opacity[1]))
+        else:
+            raise TypeError("opacity must be a number or a sequence with length 2.")
+        if not 0.0 <= self.opacity[0] <= self.opacity[1] <= 1.0:
+            raise ValueError(f"opacity must satisfy 0 <= min <= max <= 1, but got {self.opacity}.")
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        return {
+            "opacity": torch.empty(1).uniform_(self.opacity[0], self.opacity[1]).item(),
+            "start": torch.rand(1).item(),
+            "width": torch.empty(1).uniform_(1 / 3, 2 / 3).item(),
+            "direction": -1.0 if torch.rand(1).item() < 0.5 else 1.0,
+        }
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
+            return inpt
+        if inpt.ndim < 3:
+            raise ValueError(f"RandomShadow expects [..., C, H, W] input, but got shape {inpt.shape}.")
+
+        h, w = inpt.shape[-2:]
+        band_width = max(1, min(w, round(params["width"] * w)))
+        x_start = round(params["start"] * (w - band_width))
+        x_end = x_start + band_width
+        mask = torch.ones(h, w, device=inpt.device, dtype=inpt.dtype)
+        mask[:, x_start:x_end] = 1.0 + params["direction"] * params["opacity"]
+
+        smoothing_size = min(8, h, w)
+        if smoothing_size > 1:
+            batched_mask = mask[None, None]
+            small = torch.nn.functional.avg_pool2d(batched_mask, smoothing_size, stride=smoothing_size)
+            mask = torch.nn.functional.interpolate(small, size=(h, w), mode="bilinear", align_corners=False)[
+                0, 0
+            ]
+
+        broadcast_shape = (1,) * (inpt.ndim - 2) + (h, w)
+        return (inpt * mask.reshape(broadcast_shape)).clamp(0.0, 1.0)
+
+
+class CoarseDropout(Transform):
+    """Drop random rectangular patches to simulate partial occlusion.
+
+    Models objects, hands, or cables passing through the camera field of view
+    during robot manipulation.
+
+    Args:
+        max_holes: Maximum number of rectangular patches to drop.
+        max_height_frac: Maximum patch height as fraction of image height.
+        max_width_frac: Maximum patch width as fraction of image width.
+        fill_value: Value to fill dropped regions with.
+    """
+
+    def __init__(
+        self,
+        max_holes: int = 8,
+        max_height_frac: float = 0.07,
+        max_width_frac: float = 0.07,
+        fill_value: float = 0.0,
+    ) -> None:
+        super().__init__()
+        if not isinstance(max_holes, int):
+            raise TypeError("max_holes must be an int.")
+        if max_holes < 1:
+            raise ValueError(f"max_holes must be at least 1, but got {max_holes}.")
+        if not 0.0 < max_height_frac <= 1.0:
+            raise ValueError(f"max_height_frac must be in (0, 1], but got {max_height_frac}.")
+        if not 0.0 < max_width_frac <= 1.0:
+            raise ValueError(f"max_width_frac must be in (0, 1], but got {max_width_frac}.")
+        if not 0.0 <= fill_value <= 1.0:
+            raise ValueError(f"fill_value must be in [0, 1], but got {fill_value}.")
+        self.max_holes = max_holes
+        self.max_height_frac = max_height_frac
+        self.max_width_frac = max_width_frac
+        self.fill_value = fill_value
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        n = int(torch.randint(1, self.max_holes + 1, (1,)).item())
+        sizes = torch.rand(n, 2)
+        sizes[:, 0] *= self.max_height_frac
+        sizes[:, 1] *= self.max_width_frac
+        return {"sizes": sizes.tolist(), "positions": torch.rand(n, 2).tolist()}
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
+            return inpt
+        if inpt.ndim < 3:
+            raise ValueError(f"CoarseDropout expects [..., C, H, W] input, but got shape {inpt.shape}.")
+
+        h, w = inpt.shape[-2:]
+        result = inpt.clone()
+        for (height_frac, width_frac), (y_frac, x_frac) in zip(
+            params["sizes"], params["positions"], strict=True
+        ):
+            hole_h = max(1, min(h, round(height_frac * h)))
+            hole_w = max(1, min(w, round(width_frac * w)))
+            y = round(y_frac * (h - hole_h))
+            x = round(x_frac * (w - hole_w))
+            result[..., y : y + hole_h, x : x + hole_w] = self.fill_value
+        return result
+
+
+class GammaCorrection(Transform):
+    """Apply random gamma correction to simulate exposure variation.
+
+    Models different camera auto-exposure settings and sensor response curves.
+    Uses log-symmetric sampling so brightening and darkening are equally likely,
+    preventing BatchNorm statistics shift.
+
+    Args:
+        gamma: Range (min, max) for gamma value. Values < 1 brighten, > 1 darken.
+    """
+
+    def __init__(self, gamma: float | Sequence[float] = (0.5, 2.0)) -> None:
+        super().__init__()
+        if isinstance(gamma, (int, float)):
+            gamma = float(gamma)
+            if gamma <= 0:
+                raise ValueError(f"gamma must be positive, but got {gamma}.")
+            self.gamma = (min(gamma, 1.0 / gamma), max(gamma, 1.0 / gamma))
+        elif isinstance(gamma, Sequence) and len(gamma) == 2:
+            self.gamma = (float(gamma[0]), float(gamma[1]))
+        else:
+            raise TypeError("gamma must be a number or a sequence with length 2.")
+        if not 0.0 < self.gamma[0] <= self.gamma[1]:
+            raise ValueError(f"gamma must satisfy 0 < min <= max, but got {self.gamma}.")
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        log_lo = math.log(self.gamma[0])
+        log_hi = math.log(self.gamma[1])
+        gamma = math.exp(torch.empty(1).uniform_(log_lo, log_hi).item())
+        return {"gamma": gamma}
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if isinstance(inpt, torch.Tensor) and inpt.is_floating_point():
+            return inpt.pow(params["gamma"]).clamp(0.0, 1.0)
+        return inpt
+
+
+# From the paper authors' MIT-licensed reference implementation:
+# https://github.com/TheZino/PlanckianJitter
+_PLANCKIAN_BLACKBODY_COEFFICIENTS = (
+    (0.6743, 0.4029, 0.0013),
+    (0.6281, 0.4241, 0.1665),
+    (0.5919, 0.4372, 0.2513),
+    (0.5623, 0.4457, 0.3154),
+    (0.5376, 0.4515, 0.3672),
+    (0.5163, 0.4555, 0.4103),
+    (0.4979, 0.4584, 0.4468),
+    (0.4816, 0.4604, 0.4782),
+    (0.4672, 0.4619, 0.5053),
+    (0.4542, 0.4630, 0.5289),
+    (0.4426, 0.4638, 0.5497),
+    (0.4320, 0.4644, 0.5681),
+    (0.4223, 0.4648, 0.5844),
+    (0.4135, 0.4651, 0.5990),
+    (0.4054, 0.4653, 0.6121),
+    (0.3980, 0.4654, 0.6239),
+    (0.3911, 0.4655, 0.6346),
+    (0.3847, 0.4656, 0.6444),
+    (0.3787, 0.4656, 0.6532),
+    (0.3732, 0.4656, 0.6613),
+    (0.3680, 0.4655, 0.6688),
+    (0.3632, 0.4655, 0.6756),
+    (0.3586, 0.4655, 0.6820),
+    (0.3544, 0.4654, 0.6878),
+    (0.3503, 0.4653, 0.6933),
+)
+_PLANCKIAN_MIN_TEMPERATURE = 3_000
+_PLANCKIAN_MAX_TEMPERATURE = 15_000
+_PLANCKIAN_TEMPERATURE_STEP = 500
+
+
+class PlanckianJitter(Transform):
+    """Simulate color temperature shift along the Planckian locus.
+
+    Samples one black-body temperature and applies the corresponding correlated red
+    and blue channel scaling while preserving the green channel. Coefficients between
+    the tabulated 500 K intervals are linearly interpolated.
+
+    Reference: Zini et al., "Planckian Jitter", CVPR 2022 Workshop.
+
+    Args:
+        temperature: A fixed color temperature or range in Kelvin. Supported values
+            are between 3000 K and 15000 K.
+    """
+
+    def __init__(self, temperature: int | Sequence[int] = (3_000, 15_000)) -> None:
+        super().__init__()
+        if isinstance(temperature, int):
+            self.temperature = (temperature, temperature)
+        elif isinstance(temperature, Sequence) and len(temperature) == 2:
+            self.temperature = (int(temperature[0]), int(temperature[1]))
+        else:
+            raise TypeError("temperature must be an int or a sequence with length 2.")
+        if not (
+            _PLANCKIAN_MIN_TEMPERATURE
+            <= self.temperature[0]
+            <= self.temperature[1]
+            <= _PLANCKIAN_MAX_TEMPERATURE
+        ):
+            raise ValueError(
+                "temperature must satisfy "
+                f"{_PLANCKIAN_MIN_TEMPERATURE} <= min <= max <= {_PLANCKIAN_MAX_TEMPERATURE}, "
+                f"but got {self.temperature}."
+            )
+
+    def make_params(self, flat_inputs: list[Any]) -> dict[str, Any]:
+        temperature = int(torch.randint(self.temperature[0], self.temperature[1] + 1, ()).item())
+        return {"temperature": temperature}
+
+    def transform(self, inpt: Any, params: dict[str, Any]) -> Any:
+        if not isinstance(inpt, torch.Tensor) or not inpt.is_floating_point():
+            return inpt
+        if inpt.ndim < 3 or inpt.shape[-3] != 3:
+            raise ValueError(f"PlanckianJitter expects [..., 3, H, W] input, but got shape {inpt.shape}.")
+
+        table_position = (params["temperature"] - _PLANCKIAN_MIN_TEMPERATURE) / _PLANCKIAN_TEMPERATURE_STEP
+        left_index = math.floor(table_position)
+        right_index = min(left_index + 1, len(_PLANCKIAN_BLACKBODY_COEFFICIENTS) - 1)
+        interpolation_weight = table_position - left_index
+
+        left = torch.tensor(
+            _PLANCKIAN_BLACKBODY_COEFFICIENTS[left_index],
+            device=inpt.device,
+            dtype=inpt.dtype,
+        )
+        right = torch.tensor(
+            _PLANCKIAN_BLACKBODY_COEFFICIENTS[right_index],
+            device=inpt.device,
+            dtype=inpt.dtype,
+        )
+        coefficients = torch.lerp(left, right, interpolation_weight)
+        scale = torch.stack(
+            (
+                coefficients[0] / coefficients[1],
+                coefficients.new_tensor(1.0),
+                coefficients[2] / coefficients[1],
+            )
+        )
+        broadcast_shape = (1,) * (inpt.ndim - 3) + (3, 1, 1)
+        return (inpt * scale.reshape(broadcast_shape)).clamp(0.0, 1.0)
+
+
+_CUSTOM_TRANSFORMS: dict[str, type[Transform]] = {
+    "SharpnessJitter": SharpnessJitter,
+    "GaussianNoise": GaussianNoise,
+    "MotionBlur": MotionBlur,
+    "JPEGCompression": JPEGCompression,
+    "GaussianPatchBrightness": GaussianPatchBrightness,
+    "RandomShadow": RandomShadow,
+    "CoarseDropout": CoarseDropout,
+    "GammaCorrection": GammaCorrection,
+    "PlanckianJitter": PlanckianJitter,
+}
+
+
 @dataclass
 class ImageTransformConfig:
     """
@@ -215,17 +682,18 @@ class ImageTransformsConfig:
     )
 
 
-def make_transform_from_config(cfg: ImageTransformConfig):
-    if cfg.type == "SharpnessJitter":
-        return SharpnessJitter(**cfg.kwargs)
+def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
+    if cfg.type in _CUSTOM_TRANSFORMS:
+        return _CUSTOM_TRANSFORMS[cfg.type](**cfg.kwargs)
 
     transform_cls = getattr(v2, cfg.type, None)
     if isinstance(transform_cls, type) and issubclass(transform_cls, Transform):
         return transform_cls(**cfg.kwargs)
 
+    valid_custom = ", ".join(sorted(_CUSTOM_TRANSFORMS.keys()))
     raise ValueError(
         f"Transform '{cfg.type}' is not valid. It must be a class in "
-        f"torchvision.transforms.v2 or 'SharpnessJitter'."
+        f"torchvision.transforms.v2 or one of: {valid_custom}."
     )
 
 
@@ -236,8 +704,8 @@ class ImageTransforms(Transform):
         super().__init__()
         self._cfg = cfg
 
-        self.weights = []
-        self.transforms = {}
+        self.weights: list[float] = []
+        self.transforms: dict[str, Transform] = {}
         for tf_name, tf_cfg in cfg.tfs.items():
             if tf_cfg.weight <= 0.0:
                 continue

src/lerobot/utils/collate.py

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

src/lerobot/utils/constants.py

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

src/lerobot/utils/io_utils.py

@@ -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)
-    try:
+    from .import_utils import require_package
+
+    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)
-    finally:
-        writer.close()
+            if height != orig_height or width != orig_width:
+                frame_array = frame_array[:height, :width]
+            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:

src/lerobot/utils/logging_utils.py

@@ -105,8 +105,6 @@ class MetricsTracker:
         "epochs",
         "accelerator",
         "_caller_metrics",
-        "_tensor_sums",
-        "_tensor_counts",
     ]
 
     def __init__(
@@ -135,8 +133,6 @@ class MetricsTracker:
         # Meter names the caller registered up front. update_metrics() leaves these untouched, so a
         # policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter.
         self._caller_metrics: set[str] = set(self.metrics)
-        self._tensor_sums: dict[str, torch.Tensor] = {}
-        self._tensor_counts: dict[str, int] = {}
 
     def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
         if name in self.__dict__:
@@ -164,22 +160,6 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
 
-    def accumulate_tensor(self, name: str, value: torch.Tensor) -> None:
-        """Accumulate a detached metric on-device until the next logging step."""
-        if name not in self.metrics:
-            raise KeyError(f"Unknown metric {name!r}.")
-        value = value.detach()
-        self._tensor_sums[name] = self._tensor_sums.get(name, torch.zeros_like(value)) + value
-        self._tensor_counts[name] = self._tensor_counts.get(name, 0) + 1
-
-    def materialize_tensors(self) -> None:
-        """Transfer pending tensor averages to their meters with one sync per metric."""
-        for name, total in self._tensor_sums.items():
-            count = self._tensor_counts[name]
-            self.metrics[name].update((total / count).item(), n=count)
-        self._tensor_sums.clear()
-        self._tensor_counts.clear()
-
     def update_metrics(self, values: dict[str, Any]) -> None:
         """Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
 
@@ -187,7 +167,7 @@ class MetricsTracker:
         Caller-registered metrics (those passed to the constructor) are never overridden.
         """
         for name, value in values.items():
-            if isinstance(value, bool) or not isinstance(value, int | float):
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
                 continue
             if name in self._caller_metrics:
                 continue
@@ -255,5 +235,3 @@ class MetricsTracker:
         """Resets average meters."""
         for m in self.metrics.values():
             m.reset()
-        self._tensor_sums.clear()
-        self._tensor_counts.clear()

src/lerobot/utils/process.py

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

src/lerobot/utils/rerun_visualization.py

@@ -38,10 +38,7 @@ def _is_scalar(x):
 
 
 def init_rerun(
-    session_name: str = "lerobot_control_loop",
-    ip: str | None = None,
-    port: int | None = None,
-    web_port: int | None = None,
+    session_name: str = "lerobot_control_loop", ip: str | None = None, 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)
 

src/lerobot/utils/utils.py

@@ -133,10 +133,13 @@ def say(text: str, blocking: bool = False):
     else:
         raise RuntimeError("Unsupported operating system for text-to-speech.")
 
-    if blocking:
-        subprocess.run(cmd, check=True)
-    else:
-        subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
+    try:
+        if blocking:
+            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):

src/lerobot/utils/video_annotation.py

@@ -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

tests/cameras/test_opencv.py

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

tests/cameras/test_realsense.py

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

tests/configs/test_recipe.py

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

tests/datasets/test_augment_quantile_stats.py

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

tests/datasets/test_image_transforms.py

@@ -28,9 +28,17 @@ from lerobot.scripts.lerobot_imgtransform_viz import (
     save_each_transform,
 )
 from lerobot.transforms import (
+    CoarseDropout,
+    GammaCorrection,
+    GaussianNoise,
+    GaussianPatchBrightness,
     ImageTransformConfig,
     ImageTransforms,
     ImageTransformsConfig,
+    JPEGCompression,
+    MotionBlur,
+    PlanckianJitter,
+    RandomShadow,
     RandomSubsetApply,
     SharpnessJitter,
     make_transform_from_config,
@@ -455,3 +463,153 @@ def test_save_each_transform(img_tensor_factory, tmp_path):
             assert (transform_dir / file_name).exists(), (
                 f"{file_name} was not found in {transform} directory."
             )
+
+
+# --- Tests for robotics-relevant augmentations ---
+
+ROBOTICS_TRANSFORMS = [
+    ("GaussianNoise", GaussianNoise, {"std": (5.0, 25.0)}),
+    ("MotionBlur", MotionBlur, {"kernel_size": (3, 11)}),
+    ("JPEGCompression", JPEGCompression, {"quality": (15, 75)}),
+    ("GaussianPatchBrightness", GaussianPatchBrightness, {}),
+    ("RandomShadow", RandomShadow, {"opacity": (0.3, 0.6)}),
+    ("CoarseDropout", CoarseDropout, {"max_holes": 8}),
+    ("GammaCorrection", GammaCorrection, {"gamma": (0.5, 2.0)}),
+    ("PlanckianJitter", PlanckianJitter, {"temperature": (3_000, 15_000)}),
+]
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_shape_preserved(name, cls, kwargs, img_tensor_factory):
+    img = img_tensor_factory()
+    tf = cls(**kwargs)
+    out = tf(img)
+    assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_output_range(name, cls, kwargs, img_tensor_factory):
+    img = img_tensor_factory()
+    tf = cls(**kwargs)
+    out = tf(img)
+    assert out.min() >= -0.01, f"{name} min below range: {out.min():.4f}"
+    assert out.max() <= 1.01, f"{name} max above range: {out.max():.4f}"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_float_output(name, cls, kwargs, img_tensor_factory):
+    img = img_tensor_factory()
+    tf = cls(**kwargs)
+    out = tf(img)
+    assert out.is_floating_point(), f"{name} output dtype={out.dtype}"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_non_float_passthrough(name, cls, kwargs):
+    int_img = torch.randint(0, 255, (3, 32, 32), dtype=torch.uint8)
+    tf = cls(**kwargs)
+    out = tf(int_img)
+    assert torch.equal(out, int_img), f"{name} modified non-float input"
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+def test_robotics_transform_via_config(name, cls, kwargs):
+    cfg = ImageTransformConfig(type=name, kwargs=kwargs)
+    tf = make_transform_from_config(cfg)
+    assert isinstance(tf, cls), f"Config produced {type(tf)}, expected {cls}"
+
+
+def test_make_transform_error_message_includes_custom():
+    """Error message should list all registered custom transforms."""
+    with pytest.raises(ValueError, match="GaussianNoise"):
+        make_transform_from_config(ImageTransformConfig(type="NonExistent"))
+
+
+@pytest.mark.parametrize("name,cls,kwargs", ROBOTICS_TRANSFORMS, ids=[t[0] for t in ROBOTICS_TRANSFORMS])
+@pytest.mark.parametrize("shape", [(4, 3, 32, 32), (2, 4, 3, 16, 16)])
+def test_robotics_transform_supports_temporal_batches(name, cls, kwargs, shape):
+    img = torch.rand(shape)
+    out = cls(**kwargs)(img)
+    assert out.shape == img.shape, f"{name} changed shape: {img.shape} -> {out.shape}"
+    assert out.min() >= 0
+    assert out.max() <= 1
+
+
+@pytest.mark.parametrize(
+    "cls,kwargs",
+    [
+        (GaussianNoise, {"std": (25.0, 25.0)}),
+        (MotionBlur, {"kernel_size": 5}),
+        (JPEGCompression, {"quality": 10}),
+        (
+            GaussianPatchBrightness,
+            {"num_patches": 1, "sigma_range": (0.2, 0.2), "factor_range": (0.5, 0.5)},
+        ),
+        (RandomShadow, {"opacity": 0.5}),
+        (CoarseDropout, {"max_holes": 1, "fill_value": 0.0}),
+        (GammaCorrection, {"gamma": (2.0, 2.0)}),
+        (PlanckianJitter, {"temperature": 3_000}),
+    ],
+)
+def test_robotics_transform_is_not_silent_noop(cls, kwargs):
+    img = torch.rand(3, 32, 32)
+    out = cls(**kwargs)(img)
+    assert not torch.equal(out, img)
+
+
+@pytest.mark.parametrize(
+    "transform",
+    [
+        GaussianNoise(std=25),
+        RandomShadow(opacity=0.5),
+        CoarseDropout(max_holes=4),
+    ],
+)
+def test_robotics_transform_random_params_are_reused(transform):
+    img = torch.rand(3, 32, 32)
+    params = transform.make_params([img])
+    torch.testing.assert_close(transform.transform(img, params), transform.transform(img, params))
+
+
+def test_motion_blur_kernel_size_stays_in_configured_range():
+    transform = MotionBlur(kernel_size=(4, 10))
+    sampled_sizes = {transform.make_params([])["kernel_size"] for _ in range(100)}
+    assert sampled_sizes <= {5, 7, 9}
+    assert sampled_sizes
+
+
+def test_gamma_correction_scalar_below_one_defines_symmetric_range():
+    transform = GammaCorrection(gamma=0.5)
+    assert transform.gamma == (0.5, 2.0)
+    assert transform(torch.rand(3, 8, 8)).shape == (3, 8, 8)
+
+
+def test_planckian_jitter_uses_correlated_temperature_coefficients():
+    img = torch.full((2, 3, 8, 8), 0.25)
+    out = PlanckianJitter(temperature=3_000)(img)
+    torch.testing.assert_close(out[:, 1], img[:, 1])
+    assert torch.all(out[:, 0] > out[:, 1])
+    assert torch.all(out[:, 2] < out[:, 1])
+
+
+def test_random_shadow_supports_small_images():
+    img = torch.rand(3, 7, 7)
+    assert RandomShadow()(img).shape == img.shape
+
+
+@pytest.mark.parametrize(
+    "cls,kwargs",
+    [
+        (GaussianNoise, {"std": (-1.0, 1.0)}),
+        (MotionBlur, {"kernel_size": 4}),
+        (JPEGCompression, {"quality": (0, 75)}),
+        (GaussianPatchBrightness, {"sigma_range": (0.0, 0.25)}),
+        (RandomShadow, {"opacity": (0.3, 1.1)}),
+        (CoarseDropout, {"max_holes": 0}),
+        (GammaCorrection, {"gamma": 0.0}),
+        (PlanckianJitter, {"temperature": (2_000, 6_500)}),
+    ],
+)
+def test_robotics_transform_rejects_invalid_config(cls, kwargs):
+    with pytest.raises(ValueError):
+        cls(**kwargs)

tests/datasets/test_language_render.py

@@ -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

tests/datasets/test_lerobot_dataset.py

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

tests/datasets/test_sampler.py

@@ -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
 
 

tests/envs/test_robocasa_ordering.py

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

tests/motors/test_feetech.py

@@ -294,6 +294,19 @@ def test__sync_read(addr, length, ids_values, mock_motors, dummy_motors):
     assert read_values == ids_values
 
 
+def test__sync_read_retries_after_transient_failure(mock_motors, dummy_motors):
+    addr, length, ids_values = (10, 4, {1: 1337})
+    stub = mock_motors.build_sync_read_stub(addr, length, ids_values, num_invalid_try=1)
+    bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
+    bus.connect(handshake=False)
+
+    read_values, read_comm = bus._sync_read(addr, length, list(ids_values), num_retry=1)
+
+    assert read_comm == scs.COMM_SUCCESS
+    assert read_values == ids_values
+    assert mock_motors.stubs[stub].calls == 2
+
+
 @pytest.mark.parametrize("raise_on_error", (True, False))
 def test__sync_read_comm(raise_on_error, mock_motors, dummy_motors):
     addr, length, ids_values = (10, 4, {1: 1337})