refactor(rtc): share Pi05 training helpers

Extend trained RTC conditioning and postfix-only flow loss to plain PI05 so single-task checkpoints can use the same real-time chunking path as PI052.

Co-authored-by: Cursor <cursoragent@cursor.com>

AGENT_GUIDE.md

@@ -61,20 +61,16 @@ Full details in [`docs/source/so101.mdx`](./docs/source/so101.mdx) and [`docs/so
 **4.1 Install**
 
 ```bash
-# uv (recommended — see AGENTS.md and CLAUDE.md)
-uv sync --locked --extra feetech          # SO-100/SO-101 motor stack
-# uv sync --locked --extra all            # everything
-# uv sync --locked --extra smolvla        # add SmolVLA deps
-
-# pip (alternative, e.g. when not working from source)
-# pip install 'lerobot[feetech]'
-# pip install 'lerobot[all]'
-# pip install 'lerobot[smolvla]'
-
+pip install 'lerobot[feetech]'              # SO-100/SO-101 motor stack
+# pip install 'lerobot[all]'                # everything
+# pip install 'lerobot[aloha,pusht]'        # specific features
+# pip install 'lerobot[smolvla]'            # add SmolVLA deps
 git lfs install && git lfs pull
-hf auth login                             # required to push datasets/policies
+hf auth login                               # required to push datasets/policies
 ```
 
+Contributors can alternatively use `uv sync --locked --extra feetech` (see `AGENTS.md`).
+
 **4.2 Find USB ports** — run once per arm, unplug when prompted.
 
 ```bash

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), [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), [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)                                                                                                                                                                                                                                                                                                           |
 
 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,16 +68,17 @@ ENV HOME=/home/user_lerobot \
 # issues with MuJoCo and OpenGL drivers.
 RUN uv venv --python python${PYTHON_VERSION}
 
-# Install third-party dependencies separately for layer caching
+# Install Python dependencies for caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-RUN uv sync --locked --extra all --no-install-project --no-cache
+COPY --chown=user_lerobot:user_lerobot src/ src/
+
+RUN uv sync --locked --extra all --no-cache
 
 RUN chmod +x /lerobot/.venv/lib/python${PYTHON_VERSION}/site-packages/triton/backends/nvidia/bin/ptxas
 
-# Copy the application source code and install the local project
+# Copy the rest of the application source code
 # 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,14 +60,15 @@ ENV HOME=/home/user_lerobot \
 # run other Python projects in the same container without dependency conflicts.
 RUN uv venv
 
-# Install third-party dependencies separately for layer caching
+# Install Python dependencies for caching
 COPY --chown=user_lerobot:user_lerobot setup.py pyproject.toml uv.lock README.md MANIFEST.in ./
-RUN uv sync --locked --extra all --no-install-project --no-cache
+COPY --chown=user_lerobot:user_lerobot src/ src/
 
-# Copy the application code and install the local project
+RUN uv sync --locked --extra all --no-cache
+
+# Copy the rest of the application code
 # 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,6 +63,8 @@
     title: π₀-FAST (Pi0Fast)
   - local: pi05
     title: π₀.₅ (Pi05)
+  - local: pi052
+    title: π₀.₅ with language supervision (Pi052)
   - local: molmoact2
     title: MolmoAct2
   - local: vla_jepa

docs/source/backwardcomp.mdx

@@ -58,7 +58,7 @@ final_action = postprocessor(action)
 
 ## Hardware API redesign
 
-PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is an overview of what changed and how you can continue to work with datasets created before this pull request.
+PR [#777](https://github.com/huggingface/lerobot/pull/777) improves the LeRobot calibration but is **not backward-compatible**. Below is a overview of what changed and how you can continue to work with datasets created before this pull request.
 
 ### What changed?
 
@@ -129,8 +129,8 @@ python examples/backward_compatibility/replay.py \
 
 Policies output actions in the same format as the datasets (`torch.Tensors`). Therefore, the same transformations should be applied.
 
-To find these transformations, we recommend first replaying an episode of the dataset your policy was trained on using the section above.
-Then, add these same transformations to your inference script (shown here in the `record.py` script):
+To find these transformations, we recommend to first try and and replay an episode of the dataset your policy was trained on using the section above.
+Then, add these same transformations on your inference script (shown here in the `record.py` script):
 
 ```diff
 action_values = predict_action(

docs/source/bring_your_own_policies.mdx

@@ -191,6 +191,162 @@ 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,10 +136,6 @@ 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
 )
 
@@ -158,15 +154,6 @@ 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/env_processor.mdx

@@ -88,6 +88,20 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
 
 The same policy can work with different environment processors, and the same environment processor can work with different policies:
 
+````python
+# Use SmolVLA policy with LIBERO environment
+# Use SmolVLA policy with LIBERO environment
+libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
+    env_cfg=libero_cfg,
+    policy_cfg=smolvla_cfg,
+)
+smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
+# Or use ACT policy with the same LIBERO environment
+libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
+    env_cfg=libero_cfg,
+    policy_cfg=act_cfg,
+)
+act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
 ```python
 # Use SmolVLA policy with LIBERO environment
 libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -102,7 +116,6 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
     policy_cfg=act_cfg,
 )
 act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
-```
 
 ### 3. **Easier Experimentation**
 
@@ -132,7 +145,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
         state = torch.cat([eef_pos, eef_axisangle, eef_vel,
                           gripper_pos, gripper_vel], dim=-1)  # 14D
         return state
-```
+````
 
 ### 4. **Cleaner Environment Code**
 

docs/source/feetech.mdx

@@ -40,10 +40,10 @@ This tutorial guides you through updating the firmware of Feetech motors using t
 For each motor you want to update:
 
 1. **Select the motor** from the list by clicking on it
-2. **Click the Upgrade tab**:
-3. **Click the Online button**:
-   - If a potential firmware update is found, it will be displayed in the box
-4. **Click the Upgrade button**:
+2. **Click on Upgrade tab**:
+3. **Click on Online button**:
+   - If an potential firmware update is found, it will be displayed in the box
+4. **Click on Upgrade button**:
    - The update progress will be displayed
 
 ## Step 6: Verify Update

docs/source/hil_data_collection.mdx

@@ -59,7 +59,6 @@ The `lerobot-rollout --strategy.type=dagger` mode requires **teleoperators with
 
 - `bi_openarm_mini` - Bimanual OpenArm Mini
 - `so_leader` - SO100 / SO101 leader arm
-- `bi_so_leader` - Bimanual SO100 / SO101 leader arms
 
 > [!IMPORTANT]
 > The provided commands default to `bi_openarm_follower` + `bi_openarm_mini`.

docs/source/hope_jr.mdx

@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
 
 ### Record
 
-This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
+This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
 
 ```bash
 lerobot-record \

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. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection.
+`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.
 
 ## Quick Start
 
@@ -197,6 +197,52 @@ 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.
@@ -228,12 +274,13 @@ lerobot-rollout \
     --device=cuda
 ```
 
-| Flag                                        | Description                                                    |
-| ------------------------------------------- | -------------------------------------------------------------- |
-| `--inference.rtc.execution_horizon`         | Steps to blend with previous chunk (default: varies by policy) |
-| `--inference.rtc.max_guidance_weight`       | Consistency enforcement strength (default: varies by policy)   |
-| `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS`               |
-| `--inference.queue_threshold`               | Max queue size before backpressure (default: 30)               |
+| Flag                                        | Description                                                                     |
+| ------------------------------------------- | ------------------------------------------------------------------------------- |
+| `--inference.rtc.execution_horizon`         | Steps to blend with previous chunk (default: varies by policy)                  |
+| `--inference.rtc.mode`                      | `guided` (default) or trained-prefix `trained` for compatible Pi052 checkpoints |
+| `--inference.rtc.max_guidance_weight`       | Consistency enforcement strength (default: varies by policy)                    |
+| `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS`                                |
+| `--inference.queue_threshold`               | Backpressure threshold; trained RTC requires at least its maximum delay         |
 
 See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.
 

docs/source/integrate_hardware.mdx

@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
 - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
 
 Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
 
 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
 

docs/source/language_and_recipes.mdx

@@ -141,6 +141,17 @@ 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/lekiwi.mdx

@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
 pip install -e ".[lekiwi]"
 ```
 
-Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
+Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
 Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
 
 # Step-by-Step Assembly Instructions

docs/source/pi052.mdx

@@ -0,0 +1,334 @@
+# π₀.₅ 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.rtc_training_max_delay`     |                        `0` | Maximum clean-prefix delay; `0` disables training-time RTC          |
+| `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.
+
+### Training-time RTC
+
+Pi052 optionally supports training-time action conditioning from
+[Training-Time Action Conditioning for Efficient Real-Time Chunking](https://arxiv.org/abs/2512.05964).
+It simulates inference latency by sampling a clean action prefix for every flow
+draw, passing a per-action flow timestep to the action expert, and computing the
+flow loss only on the remaining postfix. The default value of `0` leaves the
+standard Pi052 objective unchanged.
+
+```bash
+lerobot-train \
+    --dataset.repo_id=${HF_USER}/my_language_annotated_dataset \
+    --policy.type=pi052 \
+    --policy.pretrained_path=lerobot/pi05_base \
+    --policy.recipe_path=recipes/subtask_mem.yaml \
+    --policy.rtc_training_max_delay=10 \
+    --policy.dtype=bfloat16 \
+    --policy.device=cuda \
+    --batch_size=8 \
+    --steps=30000 \
+    --output_dir=outputs/pi052_rtc \
+    --job_name=pi052_rtc
+```
+
+`rtc_training_max_delay` is measured in controller steps and must be smaller
+than `chunk_size`. Choose it to cover the largest inference latency expected at
+deployment: at 50 Hz, for example, 10 steps correspond to 200 ms. A delay of
+zero is included in the uniform sampling distribution, so the checkpoint also
+continues to receive ordinary flow-matching examples. Set rollout's
+`inference.rtc.execution_horizon` and `inference.queue_threshold` to at least
+this maximum so inference starts early enough and the previous chunk retains
+every action needed for the committed prefix.
+
+Run the resulting checkpoint with the asynchronous `lerobot-rollout` backend
+and select the trained-prefix path explicitly:
+
+```bash
+lerobot-rollout \
+    --strategy.type=base \
+    --policy.path=outputs/pi052_rtc/checkpoints/last/pretrained_model \
+    --inference.type=rtc \
+    --inference.rtc.mode=trained \
+    --inference.rtc.execution_horizon=10 \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --task="pick up the cube" \
+    --fps=50 \
+    --device=cuda
+```
+
+The rollout engine measures latency continuously, carries the still-unexecuted
+actions from the previous chunk into the next prediction, and discards the
+prefix that elapsed during inference. If the measured delay exceeds the
+checkpoint's `rtc_training_max_delay`, rollout stops with an explicit error
+instead of silently extrapolating beyond the training distribution. Use
+`--inference.rtc.mode=guided` for the original Jacobian-guided RTC path; it does
+not require a training-time RTC checkpoint but adds backward-pass work during
+denoising.
+
+### 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,15 +109,21 @@ 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 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`                         |
+| 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.
 
 ## Inference
 
@@ -151,6 +157,9 @@ 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    |
@@ -174,7 +183,7 @@ The model takes images, text instructions, and robot state as input, and outputs
 
 ## Reproducing π₀Fast results
 
-We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
 
 The finetuned model can be found here:
 

docs/source/processors_robots_teleop.mdx

@@ -22,7 +22,7 @@ With processors, you choose the learning features you want to use for your polic
 ## Three pipelines
 
 We often compose three pipelines. Depending on your setup, some can be empty if action and observation spaces already match.
-Each of these pipelines handles different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
+Each of these pipelines handle different conversions between different action and observation spaces. Below is a quick explanation of each pipeline.
 
 1. Pipeline 1: Teleop action space → dataset action space (phone pose → EE targets)
 2. Pipeline 2: Dataset action space → robot command space (EE targets → joints)
@@ -74,15 +74,15 @@ In the phone to SO-100 follower examples we use the following adapters:
 - `robot_action_to_transition`: transforms the teleop action dict to a pipeline transition.
 - `transition_to_robot_action`: transforms the pipeline transition to a robot action dict.
 - `observation_to_transition`: transforms the robot observation dict to a pipeline transition.
-- `transition_to_observation`: transforms the pipeline transition to an observation dict.
+- `transition_to_observation`: transforms the pipeline transition to a observation dict.
 
-Check out [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
+Checkout [src/lerobot/processor/converters.py](https://github.com/huggingface/lerobot/blob/main/src/lerobot/processor/converters.py) for more details.
 
 ## Dataset feature contracts
 
 Dataset features are determined by the keys saved in the dataset. Each step can declare what features it modifies in a contract called `transform_features(...)`. Once you build a processor, the processor can then aggregate all of these features with `aggregate_pipeline_dataset_features()` and merge multiple feature dicts with `combine_feature_dicts(...)`.
 
-Below is an example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
+Below is and example of how we declare features with the `transform_features` method in the phone to SO-100 follower examples:
 
 ```python
     def transform_features(

docs/source/robocasa.mdx

@@ -82,8 +82,6 @@ By default the env samples objects only from the `lightwheel` registry (what `--
 
 All eval snippets below mirror the CI command (see `.github/workflows/benchmark_tests.yml`). The `--rename_map` argument maps RoboCasa's native camera keys (`robot0_agentview_left` / `robot0_eye_in_hand` / `robot0_agentview_right`) onto the three-camera (`camera1` / `camera2` / `camera3`) input layout the released `smolvla_robocasa` policy was trained on.
 
-By default, each task uses the rollout horizon registered by RoboCasa. Set `--env.episode_length=<steps>` to apply the same explicit horizon to every selected task.
-
 ### Single-task evaluation (recommended for quick iteration)
 
 ```bash

docs/source/rtc.mdx

@@ -1,6 +1,6 @@
 # Real-Time Chunking (RTC)
 
-Real-Time Chunking (RTC) is an inference-time method that allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency.
+Real-Time Chunking (RTC) allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency. LeRobot provides the original inference-time guided mode and, for compatible Pi052 checkpoints, training-time action conditioning with cheap hard-prefix inference.
 
 These policies generate chunks of future actions (e.g., 50 steps at a time) instead of single actions.
 Because the models are large, producing each chunk takes longer than the time it takes the robot to execute it.
@@ -57,7 +57,7 @@ policy_cfg.rtc_config = RTCConfig(
 policy = PI0Policy.from_pretrained("lerobot/pi0_base", policy_cfg=policy_cfg, device="cuda")
 
 # Now use predict_action_chunk with RTC parameters
-inference_delay = 4  # How many steps of inference latency, this value should be calculated based on the inference latency of the policy
+inference_delay = 4  # How many steps of inference latency, this values should be calculated based on the inference latency of the policy
 
 # Initialize the action queue
 action_queue = ActionQueue(policy_cfg.rtc_config)
@@ -92,6 +92,15 @@ for step in range(num_steps):
 
 `RTCConfig` has the following parameters to tune:
 
+**`mode`** selects the action-prefix conditioning method:
+
+- `guided` (default) applies the original Jacobian guidance during denoising and works with ordinary flow-matching checkpoints.
+- `trained` hard-inpaints the previous chunk's prefix with per-action flow timesteps. It currently requires a Pi052 checkpoint trained with `policy.rtc_training_max_delay > 0` and avoids the guidance backward pass.
+
+For trained mode, both `execution_horizon` and the rollout backend's
+`inference.queue_threshold` must be at least the checkpoint's
+`rtc_training_max_delay`; rollout validates this before connecting the robot.
+
 **`execution_horizon`**: How many timesteps from the previous chunk to maintain consistency with. Higher values mean smoother transitions but potentially less reactivity.
 
 Typical values: 8-12 steps
@@ -100,7 +109,7 @@ Typical values: 8-12 steps
 RTCConfig(execution_horizon=10)
 ```
 
-**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is an optimal value.
+**`max_guidance_weight`**: How strongly to enforce consistency with the previous chunk. This is a hyperparameter that can be tuned to balance the smoothness of the transitions and the reactivity of the policy. For 10 steps flow matching (SmolVLA, Pi0, Pi0.5), a value of 10.0 is a optimal value.
 
 **`prefix_attention_schedule`**: How to weight consistency across the overlap region.
 
@@ -124,6 +133,10 @@ python examples/rtc/eval_dataset.py \
     --device=cuda
 ```
 
+Add `--rtc.mode=trained` when evaluating a compatible training-time RTC Pi052
+checkpoint. Unsupported policies reject trained mode instead of falling back to
+guided RTC.
+
 The script generates a visualization of the denoising process, comparing standard generation (left) with RTC (right). In the RTC plots, you can see how the first few steps (blue/purple lines) are guided to match the red ground truth trajectory (previous chunk's tail), ensuring a smooth transition between chunks.
 
 <p align="center">
@@ -141,6 +154,7 @@ lerobot-rollout \
     --strategy.type=base \
     --policy.path=${HF_USERNAME}/policy_repo_id \
     --inference.type=rtc \
+    --inference.rtc.mode=guided \
     --inference.rtc.execution_horizon=10 \
     --inference.rtc.max_guidance_weight=10.0 \
     --robot.type=so100_follower \
@@ -151,6 +165,24 @@ lerobot-rollout \
     --device=cuda
 ```
 
+For a training-time RTC Pi052 checkpoint, change the mode to `trained`. The
+checkpoint records its maximum supported delay, and rollout validates measured
+latency against it:
+
+```bash
+lerobot-rollout \
+    --strategy.type=base \
+    --policy.path=${HF_USERNAME}/pi052_training_rtc \
+    --inference.type=rtc \
+    --inference.rtc.mode=trained \
+    --inference.rtc.execution_horizon=10 \
+    --robot.type=so100_follower \
+    --robot.port=/dev/tty.usbmodem58FA0834591 \
+    --task="Move green small object into the purple platform" \
+    --duration=120 \
+    --device=cuda
+```
+
 ## How It Differs from the Async Inference in LeRobot
 
 Both RTC and [async inference](./async) improve real-time robot control, but they solve different problems.
@@ -189,3 +221,5 @@ See `examples/rtc/eval_dataset.py` for a complete example of offline RTC visuali
 - [Smooth-As-Butter Robot Policies](https://alexander-soare.github.io/robotics/2025/08/05/smooth-as-butter-robot-policies.html) - Excellent technical explanation with real robot results
 - [Physical Intelligence - Real-Time Chunking](https://www.physicalintelligence.company/research/real_time_chunking) - Original paper and research
 - [Kinetix RTC Implementation](https://github.com/Physical-Intelligence/real-time-chunking-kinetix) - Reference implementation from Physical Intelligence
+- [Training-Time Action Conditioning](https://arxiv.org/abs/2512.05964) - Efficient RTC with clean-prefix conditioning during training
+- [RLDX-1](https://github.com/RLWRLD/RLDX-1) - PyTorch reference used for the training-time RTC integration

docs/source/smolvla.mdx

@@ -93,7 +93,7 @@ lerobot-train --help
 
 ## Evaluate the finetuned model and run it in real-time
 
-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash

docs/source/so101.mdx

@@ -338,7 +338,7 @@ It is advisable to install one 3-pin cable in the motor after placing them befor
 <hfoption id="Leader">
 
 - Mount the leader holder onto the wrist and secure it with 4 M3x6mm screws.
-- Attach the handle to the leader holder using 1 M2x6mm screw.
+- Attach the handle to motor 5 using 1 M2x6mm screw.
 - Insert the gripper motor, secure it with 2 M2x6mm screws on each side, attach a motor horn using a M3x6mm horn screw.
 - Attach the follower trigger with 4 M3x6mm screws.
 

docs/source/using_dataset_tools.mdx

@@ -50,11 +50,11 @@ lerobot-edit-dataset \
 Divide a dataset into multiple subsets.
 
 ```bash
-# Split by fractions (e.g. 60% train, 20% val, 20% test)
+# Split by fractions (e.g. 80% train, 20% test, 20% val)
 lerobot-edit-dataset \
     --repo_id lerobot/pusht \
     --operation.type split \
-    --operation.splits '{"train": 0.6, "val": 0.2, "test": 0.2}'
+    --operation.splits '{"train": 0.8, "test": 0.2, "val": 0.2}'
 
 # Split by specific episode indices
 lerobot-edit-dataset \

examples/isaac_teleop_to_so101/common.py

@@ -44,7 +44,6 @@ from typing import Protocol
 
 import numpy as np
 
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -57,6 +56,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     EEBoundsAndSafety,
     InverseKinematicsEEToJoints,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import HF_LEROBOT_CALIBRATION, HF_LEROBOT_HOME, TELEOPERATORS
 from lerobot.utils.robot_utils import precise_sleep
 

examples/isaac_teleop_to_so101/isaac_teleop/teleop_so101_leader_arm.py

@@ -38,7 +38,7 @@ from typing import TYPE_CHECKING
 
 import numpy as np
 
-from lerobot.lerobot_types import RobotAction
+from lerobot.types import RobotAction
 
 from .base import _GRIPPER_MOTOR_SCALE, IsaacTeleopTeleoperator, _isaacteleop_available
 from .config_isaac_teleop import SO101LeaderArmConfig

examples/isaac_teleop_to_so101/isaac_teleop/teleop_xr_controller.py

@@ -32,7 +32,7 @@ from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
-from lerobot.lerobot_types import RobotAction
+from lerobot.types import RobotAction
 
 from .base import IsaacTeleopTeleoperator, _isaacteleop_available
 from .config_isaac_teleop import XRControllerConfig

examples/isaac_teleop_to_so101/isaac_teleop/xr_controller_processor.py

@@ -26,8 +26,8 @@ from __future__ import annotations
 from dataclasses import dataclass
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.lerobot_types import RobotAction
 from lerobot.processor import ProcessorStepRegistry, RobotActionProcessorStep
+from lerobot.types import RobotAction
 from lerobot.utils.rotation import Rotation
 
 from .base import _GRIPPER_MOTOR_SCALE

examples/phone_to_so100/evaluate.py

@@ -21,7 +21,6 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.common.control_utils import predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
@@ -39,6 +38,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener

examples/phone_to_so100/record.py

@@ -16,7 +16,6 @@
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -37,6 +36,7 @@ from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.phone import Phone, PhoneConfig
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.feature_utils import combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
 from lerobot.utils.utils import log_say

examples/phone_to_so100/replay.py

@@ -17,7 +17,6 @@
 import time
 
 from lerobot.datasets import LeRobotDataset
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -28,6 +27,7 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say

examples/phone_to_so100/rollout.py

@@ -27,7 +27,6 @@ Highlight, or DAgger via ``lerobot-rollout --strategy.type=...``.
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.configs import PreTrainedConfig
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -44,6 +43,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
 from lerobot.rollout.inference import SyncInferenceConfig
 from lerobot.rollout.strategies import BaseStrategy
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.utils import init_logging
 

examples/phone_to_so100/teleoperate.py

@@ -15,7 +15,6 @@
 
 import time
 
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -32,6 +31,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.teleoperators.phone import Phone, PhoneConfig
 from lerobot.teleoperators.phone.config_phone import PhoneOS
 from lerobot.teleoperators.phone.phone_processor import MapPhoneActionToRobotAction
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 

examples/rtc/eval_dataset.py

@@ -306,6 +306,7 @@ class RTCEvaluator:
         # Configure RTC
         rtc_config = RTCConfig(
             enabled=rtc_enabled,
+            mode=self.cfg.rtc.mode,
             execution_horizon=self.cfg.rtc.execution_horizon,
             max_guidance_weight=self.cfg.rtc.max_guidance_weight,
             prefix_attention_schedule=self.cfg.rtc.prefix_attention_schedule,

examples/so100_to_so100_EE/evaluate.py

@@ -21,7 +21,6 @@ from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.common.control_utils import predict_action
 from lerobot.configs import FeatureType, PolicyFeature
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.policies import make_pre_post_processors
 from lerobot.policies.act import ACTPolicy
@@ -39,6 +38,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     ForwardKinematicsJointsToEE,
     InverseKinematicsEEToJoints,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame, combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener

examples/so100_to_so100_EE/record.py

@@ -17,7 +17,6 @@
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.datasets import LeRobotDataset, aggregate_pipeline_dataset_features, create_initial_features
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -34,6 +33,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 )
 from lerobot.scripts.lerobot_record import record_loop
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.feature_utils import combine_feature_dicts
 from lerobot.utils.keyboard_input import init_keyboard_listener
 from lerobot.utils.utils import log_say

examples/so100_to_so100_EE/replay.py

@@ -18,7 +18,6 @@
 import time
 
 from lerobot.datasets import LeRobotDataset
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -29,6 +28,7 @@ from lerobot.robots.so_follower import SO100Follower, SO100FollowerConfig
 from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.utils import log_say

examples/so100_to_so100_EE/rollout.py

@@ -25,7 +25,6 @@ forward/inverse kinematics.
 
 from lerobot.cameras.opencv import OpenCVCameraConfig
 from lerobot.configs import PreTrainedConfig
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -42,6 +41,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
 from lerobot.rollout import BaseStrategyConfig, RolloutConfig, build_rollout_context
 from lerobot.rollout.inference import SyncInferenceConfig
 from lerobot.rollout.strategies import BaseStrategy
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.utils import init_logging
 

examples/so100_to_so100_EE/teleoperate.py

@@ -16,7 +16,6 @@
 
 import time
 
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.model.kinematics import RobotKinematics
 from lerobot.processor import (
     RobotProcessorPipeline,
@@ -31,6 +30,7 @@ from lerobot.robots.so_follower.robot_kinematic_processor import (
     InverseKinematicsEEToJoints,
 )
 from lerobot.teleoperators.so_leader import SO100Leader, SO100LeaderConfig
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.visualization_utils import init_rerun, log_rerun_data
 

pyproject.toml

@@ -67,8 +67,8 @@ dependencies = [
     "einops>=0.8.0,<0.9.0",
 
     # Config & Hub
-    "draccus>=0.11.6,<0.12.0",
-    "huggingface-hub>=1.6.0,<2.0.0",
+    "draccus==0.10.0", # TODO: Relax version constraint
+    "huggingface-hub>=1.0.0,<2.0.0",
     "requests>=2.32.0,<3.0.0",
 
     # Environments
@@ -95,7 +95,7 @@ dependencies = [
 
 # ── Feature-scoped extras ──────────────────────────────────
 dataset = [
-    "datasets>=4.8.0,<5.0.0",
+    "datasets>=4.7.0,<5.0.0",
     "pandas>=2.0.0,<3.0.0", # NOTE: Transitive dependency of datasets
     "pyarrow>=21.0.0,<30.0.0", # NOTE: Transitive dependency of datasets
     "lerobot[av-dep]",
@@ -150,6 +150,7 @@ 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"]
@@ -212,7 +213,7 @@ wallx = [
     "torchdiffeq>=0.2.4,<0.3.0",
     "lerobot[qwen-vl-utils-dep]",
 ]
-pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
+pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]", "lerobot[sentencepiece-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]"]
@@ -494,19 +495,6 @@ 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/async_inference/policy_server.py

@@ -38,7 +38,6 @@ import draccus
 import grpc
 import torch
 
-from lerobot.lerobot_types import PolicyAction
 from lerobot.policies import get_policy_class, make_pre_post_processors
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.transport import (
@@ -46,6 +45,7 @@ from lerobot.transport import (
     services_pb2_grpc,  # type: ignore
 )
 from lerobot.transport.utils import receive_bytes_in_chunks
+from lerobot.types import PolicyAction
 
 from .configs import PolicyServerConfig
 from .constants import SUPPORTED_POLICIES

src/lerobot/cameras/opencv/camera_opencv.py

@@ -120,22 +120,14 @@ class OpenCVCamera(Camera):
         self.rotation: int | None = get_cv2_rotation(config.rotation)
         self.backend: int = config.backend
 
-        self.capture_width: int | None = None
-        self.capture_height: int | None = None
-        self._reset_connection_settings()
+        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
 
     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."""
@@ -172,25 +164,17 @@ class OpenCVCamera(Camera):
                 f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
             )
 
-        try:
-            self._configure_capture_settings()
-            self._start_read_thread()
+        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.")
-        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
+        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.")
 
         logger.info(f"{self} connected.")
 
@@ -328,36 +312,32 @@ class OpenCVCamera(Camera):
 
         for target in targets_to_scan:
             camera = cv2.VideoCapture(target)
-            try:
-                if camera.isOpened():
-                    default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
-                    default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
-                    default_fps = camera.get(cv2.CAP_PROP_FPS)
-                    default_format = camera.get(cv2.CAP_PROP_FORMAT)
+            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)
-            finally:
+                found_cameras_info.append(camera_info)
                 camera.release()
 
         return found_cameras_info
@@ -516,26 +496,6 @@ 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]:
         """
@@ -626,6 +586,16 @@ class OpenCVCamera(Camera):
         if not self.is_connected and self.thread is None:
             raise DeviceNotConnectedError(f"{self} not connected.")
 
-        self._cleanup_resources()
+        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()
 
         logger.info(f"{self} disconnected.")

src/lerobot/cameras/realsense/camera_realsense.py

@@ -121,9 +121,6 @@ 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:
@@ -134,9 +131,6 @@ 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
@@ -151,23 +145,14 @@ class RealSenseCamera(Camera):
 
         self.rotation: int | None = get_cv2_rotation(config.rotation)
 
-        self.capture_width: int | None = None
-        self.capture_height: int | None = None
-        self._reset_connection_settings()
+        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
 
     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."""
@@ -187,8 +172,7 @@ class RealSenseCamera(Camera):
 
         Raises:
             DeviceAlreadyConnectedError: If the camera is already connected.
-            ValueError: If the configuration is invalid, a requested sensor option is unsupported,
-                or a requested sensor value is invalid.
+            ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
             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.
         """
@@ -206,31 +190,22 @@ class RealSenseCamera(Camera):
                 f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
             ) from e
 
-        try:
-            self._configure_capture_settings()
-            self._configure_sensor_options()
-            self._start_read_thread()
+        self._configure_capture_settings()
+        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.")
-        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
+        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.")
 
         logger.info(f"{self} connected.")
 
@@ -364,111 +339,6 @@ 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]:
         """
@@ -671,27 +541,6 @@ 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():
@@ -835,5 +684,18 @@ class RealSenseCamera(Camera):
                 f"Attempted to disconnect {self}, but it appears already disconnected."
             )
 
-        self._cleanup_resources()
+        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()
+
         logger.info(f"{self} disconnected.")

src/lerobot/cameras/realsense/configuration_realsense.py

@@ -46,17 +46,6 @@ 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.
@@ -72,9 +61,6 @@ 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)
@@ -83,18 +69,6 @@ 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/common/control_utils.py

@@ -35,9 +35,9 @@ else:
 
 if TYPE_CHECKING:
     from lerobot.datasets import LeRobotDataset
-from lerobot.lerobot_types import PolicyAction
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.robots import Robot
+from lerobot.types import PolicyAction
 
 
 def predict_action(

src/lerobot/configs/dataset.py

@@ -71,19 +71,13 @@ 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,6 +33,8 @@ 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
@@ -62,6 +64,10 @@ 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/policies.py

@@ -163,10 +163,8 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
         return None
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        # Encode against the base class so draccus includes the choice "type" key,
-        # which `from_pretrained` needs to resolve the concrete subclass.
-        with open(save_directory / CONFIG_NAME, "w") as f:
-            json.dump(draccus.encode(self, PreTrainedConfig), f, indent=4)
+        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
 
     @classmethod
     def from_pretrained(

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:
-        """Ensure every templated binding is known and at least one turn is a target."""
+        """Validate bindings and require text or low-level action supervision."""
         assert self.messages is not None
         known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
 
@@ -156,8 +156,14 @@ class TrainingRecipe:
             if missing:
                 raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}")
 
-        if not any(turn.target for turn in self.messages):
-            raise ValueError("Message recipes must contain at least one target turn.")
+        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)."
+            )
 
     def _validate_blend_recipe(self) -> None:
         """Ensure each blend component is a non-empty, weighted message recipe."""

src/lerobot/configs/recipes/subtask.yaml

@@ -0,0 +1,16 @@
+# 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

@@ -0,0 +1,13 @@
+# 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

@@ -0,0 +1,30 @@
+# 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

@@ -0,0 +1,70 @@
+# 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/configs/rewards.py

@@ -103,10 +103,8 @@ class RewardModelConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):
         pass
 
     def _save_pretrained(self, save_directory: Path) -> None:
-        # Encode against the base class so draccus includes the choice "type" key,
-        # which `from_pretrained` needs to resolve the concrete subclass.
-        with open(save_directory / CONFIG_NAME, "w") as f:
-            json.dump(draccus.encode(self, RewardModelConfig), f, indent=4)
+        with open(save_directory / CONFIG_NAME, "w") as f, draccus.config_type("json"):
+            draccus.dump(self, f, indent=4)
 
     @classmethod
     def from_pretrained(

src/lerobot/configs/train.py

@@ -194,11 +194,7 @@ class TrainPipelineConfig(HubMixin):
             )
 
         if Path(config_path).resolve().exists():
-            # `config_path` may point at the checkpoint's train_config.json or at its
-            # pretrained_model/ directory (both documented above) — resolve either to
-            # the pretrained_model/ directory.
-            config_path_obj = Path(config_path)
-            policy_dir = config_path_obj.parent if config_path_obj.is_file() else config_path_obj
+            policy_dir = Path(config_path).parent
             self.checkpoint_path = policy_dir.parent
         elif self.job.is_remote:
             return

src/lerobot/datasets/aggregate.py

@@ -19,7 +19,6 @@ import copy
 import logging
 import shutil
 from pathlib import Path
-from typing import Any, NotRequired, TypedDict
 
 import datasets
 import pandas as pd
@@ -50,32 +49,8 @@ from .utils import (
 )
 from .video_utils import concatenate_video_files, get_video_duration_in_s
 
-logger = logging.getLogger(__name__)
 
-type FeatureDict = dict[str, dict[str, Any]]
-type ChunkFile = tuple[int, int]
-
-
-class IndexState(TypedDict):
-    chunk: int
-    file: int
-    src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
-
-
-class VideoIndex(TypedDict):
-    chunk: int
-    file: int
-    latest_duration: float
-    episode_duration: float
-    src_to_offset: NotRequired[dict[ChunkFile, float]]
-    src_to_dst: NotRequired[dict[ChunkFile, ChunkFile]]
-    dst_file_durations: NotRequired[dict[ChunkFile, float]]
-
-
-type VideoIndexState = dict[str, VideoIndex]
-
-
-def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> FeatureDict:
+def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMetadata]) -> dict[str, dict]:
     """Create a merged video feature info dictionary for aggregation. The video encoder info is merged field-by-field: each key is kept only when every source agrees; otherwise that key is set to ``null`` (or ``{}`` for ``video.extra_options``) and a warning is logged.
 
     Args:
@@ -84,14 +59,14 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
     Returns:
         dict: A dictionary of merged video feature info.
     """
-    merged_info: FeatureDict = copy.deepcopy(all_metadata[0].features)
+    merged_info = copy.deepcopy(all_metadata[0].features)
     video_keys = [k for k in merged_info if merged_info[k].get("dtype") == "video"]
 
     for vk in video_keys:
         video_infos = [m.features.get(vk, {}).get("info") or {} for m in all_metadata]
         base_video_info = video_infos[0]
 
-        merged_encoder_info: dict[str, Any] = {}
+        merged_encoder_info: dict = {}
         fallback_keys: list[str] = []
         for info_key in VIDEO_ENCODER_INFO_KEYS:
             values = [info.get(info_key, None) for info in video_infos]
@@ -105,7 +80,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
                 merged_encoder_info[info_key] = {} if info_key == "video.extra_options" else None
 
         if fallback_keys:
-            logger.warning(
+            logging.warning(
                 f"Merging heterogeneous or incomplete video encoder metadata for feature {vk}. "
                 f"Setting these keys to null: {fallback_keys}.",
             )
@@ -117,7 +92,7 @@ def merge_video_feature_info_for_aggregate(all_metadata: list[LeRobotDatasetMeta
     return merged_info
 
 
-def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[int, str | None, FeatureDict]:
+def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]):
     """Validates that all dataset metadata have consistent properties.
 
     Ensures all datasets have the same fps, robot_type, and features to guarantee
@@ -154,9 +129,7 @@ def validate_all_metadata(all_metadata: list[LeRobotDatasetMetadata]) -> tuple[i
     return fps, robot_type, features
 
 
-def update_data_df(
-    df: pd.DataFrame, src_meta: LeRobotDatasetMetadata, dst_meta: LeRobotDatasetMetadata
-) -> pd.DataFrame:
+def update_data_df(df, src_meta, dst_meta):
     """Updates a data DataFrame with new indices and task mappings for aggregation.
 
     Adjusts episode indices, frame indices, and task indices to account for
@@ -181,12 +154,12 @@ def update_data_df(
 
 
 def update_meta_data(
-    df: pd.DataFrame,
-    dst_meta: LeRobotDatasetMetadata,
-    meta_idx: IndexState,
-    data_idx: IndexState,
-    videos_idx: VideoIndexState,
-) -> pd.DataFrame:
+    df,
+    dst_meta,
+    meta_idx,
+    data_idx,
+    videos_idx,
+):
     """Updates metadata DataFrame with new chunk, file, and timestamp indices.
 
     Adjusts all indices and timestamps to account for previously aggregated
@@ -316,7 +289,7 @@ def aggregate_datasets(
     chunk_size: int | None = None,
     concatenate_videos: bool = True,
     concatenate_data: bool = True,
-) -> None:
+):
     """Aggregates multiple LeRobot datasets into a single unified dataset.
 
     This is the main function that orchestrates the aggregation process by:
@@ -336,7 +309,7 @@ def aggregate_datasets(
         concatenate_videos: When False, keep one mp4 per source file instead of packing into shards.
         concatenate_data: When False, keep one parquet per source file instead of packing into shards.
     """
-    logger.info("Start aggregate_datasets")
+    logging.info("Start aggregate_datasets")
 
     if data_files_size_in_mb is None:
         data_files_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB
@@ -368,15 +341,15 @@ def aggregate_datasets(
         video_files_size_in_mb=video_files_size_in_mb,
     )
 
-    logger.info("Find all tasks")
+    logging.info("Find all tasks")
     unique_tasks = pd.concat([m.tasks for m in all_metadata]).index.unique()
     dst_meta.tasks = pd.DataFrame(
         {"task_index": range(len(unique_tasks))}, index=pd.Index(unique_tasks, name="task")
     )
 
-    meta_idx: IndexState = {"chunk": 0, "file": 0}
-    data_idx: IndexState = {"chunk": 0, "file": 0}
-    videos_idx: VideoIndexState = {
+    meta_idx = {"chunk": 0, "file": 0}
+    data_idx = {"chunk": 0, "file": 0}
+    videos_idx = {
         key: {"chunk": 0, "file": 0, "latest_duration": 0, "episode_duration": 0} for key in video_keys
     }
 
@@ -400,17 +373,12 @@ def aggregate_datasets(
         dst_meta.info.total_frames += src_meta.total_frames
 
     finalize_aggregation(dst_meta, all_metadata)
-    logger.info("Aggregation complete.")
+    logging.info("Aggregation complete.")
 
 
 def aggregate_videos(
-    src_meta: LeRobotDatasetMetadata,
-    dst_meta: LeRobotDatasetMetadata,
-    videos_idx: VideoIndexState,
-    video_files_size_in_mb: float,
-    chunk_size: int,
-    concatenate_videos: bool = True,
-) -> VideoIndexState:
+    src_meta, dst_meta, videos_idx, video_files_size_in_mb, chunk_size, concatenate_videos=True
+):
     """Aggregates video chunks from a source dataset into the destination dataset.
 
     Handles video file concatenation and rotation based on file size limits.
@@ -438,16 +406,15 @@ def aggregate_videos(
             videos_idx[key]["dst_file_durations"] = {}
 
     for key, video_idx in videos_idx.items():
-        unique_chunk_file_pairs: list[ChunkFile] = sorted(
-            {
-                (chunk, file)
-                for chunk, file in zip(
-                    src_meta.episodes[f"videos/{key}/chunk_index"],
-                    src_meta.episodes[f"videos/{key}/file_index"],
-                    strict=False,
-                )
-            }
-        )
+        unique_chunk_file_pairs = {
+            (chunk, file)
+            for chunk, file in zip(
+                src_meta.episodes[f"videos/{key}/chunk_index"],
+                src_meta.episodes[f"videos/{key}/file_index"],
+                strict=False,
+            )
+        }
+        unique_chunk_file_pairs = sorted(unique_chunk_file_pairs)
 
         chunk_idx = video_idx["chunk"]
         file_idx = video_idx["file"]
@@ -522,14 +489,7 @@ def aggregate_videos(
     return videos_idx
 
 
-def aggregate_data(
-    src_meta: LeRobotDatasetMetadata,
-    dst_meta: LeRobotDatasetMetadata,
-    data_idx: IndexState,
-    data_files_size_in_mb: float,
-    chunk_size: int,
-    concatenate_data: bool = True,
-) -> IndexState:
+def aggregate_data(src_meta, dst_meta, data_idx, data_files_size_in_mb, chunk_size, concatenate_data=True):
     """Aggregates data chunks from a source dataset into the destination dataset.
 
     Reads source data files, updates indices to match the aggregated dataset,
@@ -550,16 +510,14 @@ def aggregate_data(
     Returns:
         dict: Updated data_idx with current chunk and file indices.
     """
-    unique_chunk_file_ids: list[ChunkFile] = sorted(
-        {
-            (c, f)
-            for c, f in zip(
-                src_meta.episodes["data/chunk_index"],
-                src_meta.episodes["data/file_index"],
-                strict=False,
-            )
-        }
-    )
+    unique_chunk_file_ids = {
+        (c, f)
+        for c, f in zip(
+            src_meta.episodes["data/chunk_index"], src_meta.episodes["data/file_index"], strict=False
+        )
+    }
+
+    unique_chunk_file_ids = sorted(unique_chunk_file_ids)
     contains_images = len(dst_meta.image_keys) > 0
 
     # retrieve features schema for proper image typing in parquet
@@ -567,7 +525,7 @@ def aggregate_data(
 
     # Track source to destination file mapping for metadata update
     # This is critical for handling datasets that are already results of a merge
-    src_to_dst: dict[ChunkFile, ChunkFile] = {}
+    src_to_dst: dict[tuple[int, int], tuple[int, int]] = {}
 
     for src_chunk_idx, src_file_idx in unique_chunk_file_ids:
         src_path = src_meta.root / DEFAULT_DATA_PATH.format(
@@ -606,13 +564,7 @@ def aggregate_data(
     return data_idx
 
 
-def aggregate_metadata(
-    src_meta: LeRobotDatasetMetadata,
-    dst_meta: LeRobotDatasetMetadata,
-    meta_idx: IndexState,
-    data_idx: IndexState,
-    videos_idx: VideoIndexState,
-) -> IndexState:
+def aggregate_metadata(src_meta, dst_meta, meta_idx, data_idx, videos_idx):
     """Aggregates metadata from a source dataset into the destination dataset.
 
     Reads source metadata files, updates all indices and timestamps,
@@ -628,16 +580,16 @@ def aggregate_metadata(
     Returns:
         dict: Updated meta_idx with current chunk and file indices.
     """
-    chunk_file_ids: list[ChunkFile] = sorted(
-        {
-            (c, f)
-            for c, f in zip(
-                src_meta.episodes["meta/episodes/chunk_index"],
-                src_meta.episodes["meta/episodes/file_index"],
-                strict=False,
-            )
-        }
-    )
+    chunk_file_ids = {
+        (c, f)
+        for c, f in zip(
+            src_meta.episodes["meta/episodes/chunk_index"],
+            src_meta.episodes["meta/episodes/file_index"],
+            strict=False,
+        )
+    }
+
+    chunk_file_ids = sorted(chunk_file_ids)
     for chunk_idx, file_idx in chunk_file_ids:
         src_path = src_meta.root / DEFAULT_EPISODES_PATH.format(chunk_index=chunk_idx, file_index=file_idx)
         df = pd.read_parquet(src_path)
@@ -670,16 +622,16 @@ def aggregate_metadata(
 def append_or_create_parquet_file(
     df: pd.DataFrame,
     src_path: Path,
-    idx: IndexState,
+    idx: dict[str, int],
     max_mb: float,
     chunk_size: int,
     default_path: str,
     contains_images: bool = False,
-    aggr_root: Path | None = None,
+    aggr_root: Path = None,
     hf_features: datasets.Features | None = None,
     concatenate: bool = True,
     one_row_group_per_episode: bool = False,
-) -> tuple[IndexState, ChunkFile]:
+) -> tuple[dict[str, int], tuple[int, int]]:
     """Appends data to an existing parquet file or creates a new one based on size constraints.
 
     Manages file rotation when size limits are exceeded to prevent individual files
@@ -702,13 +654,7 @@ def append_or_create_parquet_file(
     Returns:
         tuple: (updated_idx, (dst_chunk, dst_file)) where updated_idx is the index dict
                and (dst_chunk, dst_file) is the actual destination file the data was written to.
-
-    Raises:
-        ValueError: If aggr_root is not provided.
     """
-    if aggr_root is None:
-        raise ValueError("aggr_root must be provided.")
-
     dst_chunk, dst_file = idx["chunk"], idx["file"]
     dst_path = aggr_root / default_path.format(chunk_index=dst_chunk, file_index=dst_file)
 
@@ -752,9 +698,7 @@ def append_or_create_parquet_file(
     return idx, (dst_chunk, dst_file)
 
 
-def finalize_aggregation(
-    aggr_meta: LeRobotDatasetMetadata, all_metadata: list[LeRobotDatasetMetadata]
-) -> None:
+def finalize_aggregation(aggr_meta, all_metadata):
     """Finalizes the dataset aggregation by writing summary files and statistics.
 
     Writes the tasks file, info file with total counts and splits, and
@@ -764,16 +708,16 @@ def finalize_aggregation(
         aggr_meta: Aggregated dataset metadata.
         all_metadata: List of all source dataset metadata objects.
     """
-    logger.info("write tasks")
+    logging.info("write tasks")
     write_tasks(aggr_meta.tasks, aggr_meta.root)
 
-    logger.info("write info")
+    logging.info("write info")
     aggr_meta.info.total_tasks = len(aggr_meta.tasks)
     aggr_meta.info.total_episodes = sum(m.total_episodes for m in all_metadata)
     aggr_meta.info.total_frames = sum(m.total_frames for m in all_metadata)
     aggr_meta.info.splits = {"train": f"0:{sum(m.total_episodes for m in all_metadata)}"}
     write_info(aggr_meta.info, aggr_meta.root)
 
-    logger.info("write stats")
+    logging.info("write stats")
     aggr_meta.stats = aggregate_stats([m.stats for m in all_metadata])
     write_stats(aggr_meta.stats, aggr_meta.root)

src/lerobot/datasets/dataset_metadata.py

@@ -18,15 +18,13 @@ import logging
 from collections.abc import Callable, Iterable
 from copy import deepcopy
 from pathlib import Path
-from typing import Literal
 
 import numpy as np
 import packaging.version
 import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
-from huggingface_hub import snapshot_download, sync_bucket
-from huggingface_hub.utils import WeakFileLock
+from huggingface_hub import snapshot_download
 
 from lerobot.configs import DEPTH_METER_UNIT, VideoEncoderConfig
 from lerobot.utils.constants import DEFAULT_FEATURES, HF_LEROBOT_HOME, HF_LEROBOT_HUB_CACHE
@@ -76,7 +74,6 @@ class LeRobotDatasetMetadata:
         force_cache_sync: bool = False,
         metadata_buffer_size: int = 10,
         *,
-        repo_type: Literal["dataset", "bucket"] = "dataset",
         token: str | bool | None = None,
     ):
         """Load or download metadata for an existing LeRobot dataset.
@@ -99,53 +96,36 @@ class LeRobotDatasetMetadata:
                 even when local files exist.
             metadata_buffer_size: Number of episode metadata records to buffer
                 in memory before flushing to parquet.
-            repo_type: Repository type: "dataset" (default) or "bucket" for an
-                HF Storage Bucket streamed over hf://buckets/.
             token: Authentication token used for Hub requests. Pass a string
                 token, ``True`` to require the locally stored token, ``False``
                 to disable authentication, or ``None`` to use the Hugging Face
                 Hub default.
         """
-        if repo_type not in ("dataset", "bucket"):
-            raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {repo_type!r}")
-
         self.repo_id = repo_id
-        self.repo_type = repo_type
         self.revision = revision if revision else CODEBASE_VERSION
         self._requested_root = Path(root) if root is not None else None
-        if self._requested_root is not None:
-            self.root = self._requested_root
-        elif self.repo_type == "bucket":
-            self.root = HF_LEROBOT_HUB_CACHE / ("buckets--" + self.repo_id.replace("/", "--"))
-        else:
-            self.root = HF_LEROBOT_HOME / repo_id
+        self.root = self._requested_root if self._requested_root is not None else HF_LEROBOT_HOME / repo_id
         self._pq_writer = None
         self.latest_episode = None
         self._metadata_buffer: list[dict] = []
         self._metadata_buffer_size = metadata_buffer_size
         self._finalized = False
 
-        metadata_lock = contextlib.nullcontext()
-        if self.repo_type == "bucket":
-            self.root.parent.mkdir(parents=True, exist_ok=True)
-            metadata_lock = WeakFileLock(self.root.parent / f".{self.root.name}.lock")
+        try:
+            if force_cache_sync or (
+                self._requested_root is None and has_legacy_hub_download_metadata(self.root)
+            ):
+                raise FileNotFoundError
+            self._load_metadata()
+        except (FileNotFoundError, NotADirectoryError):
+            if is_valid_version(self.revision):
+                if token is None:
+                    self.revision = get_safe_version(self.repo_id, self.revision)
+                else:
+                    self.revision = get_safe_version(self.repo_id, self.revision, token=token)
 
-        with metadata_lock:
-            try:
-                if force_cache_sync or (
-                    self._requested_root is None and has_legacy_hub_download_metadata(self.root)
-                ):
-                    raise FileNotFoundError
-                self._load_metadata()
-            except (FileNotFoundError, NotADirectoryError):
-                if self.repo_type != "bucket" and is_valid_version(self.revision):
-                    if token is None:
-                        self.revision = get_safe_version(self.repo_id, self.revision)
-                    else:
-                        self.revision = get_safe_version(self.repo_id, self.revision, token=token)
-
-                self._pull_from_repo(allow_patterns="meta/", token=token)
-                self._load_metadata()
+            self._pull_from_repo(allow_patterns="meta/", token=token)
+            self._load_metadata()
 
     def _flush_metadata_buffer(self) -> None:
         """Write all buffered episode metadata to parquet file."""
@@ -208,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) if self.total_tasks > 0 else None
-        self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
+        self.tasks = load_tasks(self.root)
+        self.episodes = load_episodes(self.root)
         self.stats = load_stats(self.root)
 
     def ensure_readable(self) -> None:
@@ -252,16 +232,6 @@ class LeRobotDatasetMetadata:
         *,
         token: str | bool | None = None,
     ) -> None:
-        if self.repo_type == "bucket":
-            self.root.mkdir(parents=True, exist_ok=True)
-            sync_bucket(
-                f"hf://buckets/{self.repo_id}/meta",
-                str(self.root / "meta"),
-                delete=True,
-                quiet=True,
-                token=token,
-            )
-            return
         token_kwargs = {} if token is None else {"token": token}
         if self._requested_root is None:
             self.root = Path(
@@ -292,8 +262,6 @@ class LeRobotDatasetMetadata:
     @property
     def url_root(self) -> str:
         """Hugging Face Hub URL root for this dataset."""
-        if self.repo_type == "bucket":
-            return f"hf://buckets/{self.repo_id}"
         return f"hf://datasets/{self.repo_id}"
 
     @property

src/lerobot/datasets/dataset_reader.py

@@ -163,10 +163,40 @@ 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,6 +66,17 @@ 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.
 
@@ -87,11 +98,14 @@ 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=cfg.dataset.episodes,
+                episodes=episodes,
                 delta_timestamps=delta_timestamps,
                 image_transforms=image_transforms,
                 revision=cfg.dataset.revision,
@@ -104,7 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             dataset = StreamingLeRobotDataset(
                 cfg.dataset.repo_id,
                 root=cfg.dataset.root,
-                episodes=cfg.dataset.episodes,
+                episodes=episodes,
                 delta_timestamps=delta_timestamps,
                 image_transforms=image_transforms,
                 revision=cfg.dataset.revision,

src/lerobot/datasets/language_render.py

@@ -162,14 +162,28 @@ def render_sample(
     task: str | None = None,
     dataset_ctx: Any | None = None,
 ) -> RenderedMessages | None:
-    """Render the chat-style messages for a single dataset sample.
+    """Resolve one sample's bindings and render its message recipe.
 
-    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.
+    Returns ``None`` when no text or low-level action supervision applies.
     """
     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,
@@ -183,6 +197,55 @@ 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:
@@ -346,7 +409,9 @@ def _render_message_recipe(
         if turn.target:
             target_indices.append(message_idx)
 
-    if not target_indices:
+    # 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:
         return None
 
     rendered = {
@@ -403,14 +468,12 @@ def _validate_rendered(rendered: RenderedMessages) -> None:
 
     if len(streams) != len(messages):
         raise ValueError("message_streams must be aligned with messages.")
-    if not target_indices:
-        raise ValueError("Rendered samples must contain at least one target message.")
+    # 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.")
     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/datasets/pipeline_features.py

@@ -17,8 +17,8 @@ from collections.abc import Sequence
 from typing import Any
 
 from lerobot.configs import PipelineFeatureType
-from lerobot.lerobot_types import RobotAction, RobotObservation
 from lerobot.processor import DataProcessorPipeline
+from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE, OBS_STR
 from lerobot.utils.feature_utils import hw_to_dataset_features
 

src/lerobot/datasets/streaming_dataset.py

@@ -16,7 +16,6 @@
 from collections import deque
 from collections.abc import Callable, Generator, Iterable, Iterator
 from pathlib import Path
-from typing import Literal
 
 import datasets
 import numpy as np
@@ -258,16 +257,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
         return_uint8: bool = False,
         depth_output_unit: str = DEFAULT_DEPTH_UNIT,
         *,
-        repo_type: Literal["dataset", "bucket"] = "dataset",
         token: str | bool | None = None,
     ):
         """Initialize a StreamingLeRobotDataset.
 
         Args:
             repo_id (str): This is the repo id that will be used to fetch the dataset.
-            root (Path | None, optional): Local directory to use for local datasets. In bucket mode,
-                this is an optional local metadata-cache directory; parquet and video data remain remote.
-                When omitted, Hub metadata is resolved through the cache under ``$HF_LEROBOT_HOME/hub``.
+            root (Path | None, optional): Local directory to use for local datasets. When omitted, Hub
+                metadata is resolved through a revision-safe snapshot cache under
+                ``$HF_LEROBOT_HOME/hub``.
             episodes (list[int] | None, optional): If specified, this will only load episodes specified by
                 their episode_index in this list.
             image_transforms (Callable | None, optional): Transform to apply to image data.
@@ -282,8 +280,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             shuffle (bool, optional): Whether to shuffle the dataset across exhaustions. Defaults to True.
             depth_output_unit (str, optional): Physical unit depth maps are dequantized to ("m" or "mm").
                 Defaults to "mm".
-            repo_type: "dataset" (default) or "bucket" to stream from an HF Storage Bucket
-                over ``hf://buckets/``.
             token: Authentication token used while streaming this dataset from
                 the Hub. Pass a string token, ``True`` to require the locally
                 stored token, ``False`` to disable authentication, or ``None``
@@ -291,14 +287,10 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
                 on the dataset instance after initialization.
         """
         super().__init__()
-        if repo_type not in ("dataset", "bucket"):
-            raise ValueError(f"repo_type must be 'dataset' or 'bucket', got {repo_type!r}")
-
         self.repo_id = repo_id
-        self.repo_type = repo_type
-        self._requested_root = Path(root) if root is not None else None
+        self._requested_root = Path(root) if root else None
         self.root = self._requested_root if self._requested_root is not None else HF_LEROBOT_HOME / repo_id
-        self.streaming_from_local = root is not None and self.repo_type == "dataset"
+        self.streaming_from_local = root is not None
 
         self.image_transforms = image_transforms
         self.episodes = episodes
@@ -325,7 +317,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             self._requested_root,
             self.revision,
             force_cache_sync=force_cache_sync,
-            repo_type=self.repo_type,
             token=token,
         )
         self.root = self.meta.root
@@ -354,26 +345,15 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
             self.delta_timestamps = delta_timestamps
             self.delta_indices = get_delta_indices(self.delta_timestamps, self.fps)
 
-        token_kwargs = {} if token is None else {"token": token}
-        if self.repo_type == "bucket":
-            self.hf_dataset: datasets.IterableDataset = load_dataset(
-                "parquet",
-                data_files=f"hf://buckets/{self.repo_id}/data/*/*.parquet",
-                split="train",
-                streaming=self.streaming,
-                **token_kwargs,
-            )
-        else:
-            if self.streaming_from_local:
-                token_kwargs = {}
-            self.hf_dataset: datasets.IterableDataset = load_dataset(
-                self.repo_id if not self.streaming_from_local else str(self.root),
-                split="train",
-                streaming=self.streaming,
-                data_files="data/*/*.parquet",
-                revision=self.revision,
-                **token_kwargs,
-            )
+        token_kwargs = {} if token is None or self.streaming_from_local else {"token": token}
+        self.hf_dataset: datasets.IterableDataset = load_dataset(
+            self.repo_id if not self.streaming_from_local else str(self.root),
+            split="train",
+            streaming=self.streaming,
+            data_files="data/*/*.parquet",
+            revision=self.revision,
+            **token_kwargs,
+        )
 
         self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
 

src/lerobot/envs/configs.py

@@ -507,7 +507,7 @@ class MetaworldEnv(EnvConfig):
 class RoboCasaEnv(EnvConfig):
     task: str = "CloseFridge"
     fps: int = 20
-    episode_length: int | None = None
+    episode_length: int = 1000
     obs_type: str = "pixels_agent_pos"
     render_mode: str = "rgb_array"
     camera_name: str = "robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right"
@@ -560,7 +560,13 @@ class RoboCasaEnv(EnvConfig):
             kwargs["split"] = self.split
         return kwargs
 
-    def create_envs(self, n_envs: int, use_async_envs: bool = False):
+    def create_envs(
+        self,
+        n_envs: int,
+        use_async_envs: bool = False,
+        terminate_on_success: bool = True,
+        horizon: int | None = None,
+    ):
         from .robocasa import create_robocasa_envs
 
         if self.task is None:
@@ -574,6 +580,8 @@ 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

@@ -30,7 +30,7 @@ from gymnasium import spaces
 from libero.libero import benchmark, get_libero_path
 from libero.libero.envs import OffScreenRenderEnv
 
-from lerobot.lerobot_types import RobotObservation
+from lerobot.types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv, parse_camera_names
 
@@ -384,12 +384,7 @@ class LiberoEnv(gym.Env):
 
     def close(self):
         if self._env is not None:
-            try:
-                self._env.close()
-            finally:
-                # LIBERO deletes its inner env on close, so this wrapper must
-                # be recreated before the next reset.
-                self._env = None
+            self._env.close()
 
 
 def _make_env_fns(

src/lerobot/envs/metaworld.py

@@ -25,7 +25,7 @@ import metaworld.policies as policies
 import numpy as np
 from gymnasium import spaces
 
-from lerobot.lerobot_types import RobotObservation
+from lerobot.types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv
 

src/lerobot/envs/robocasa.py

@@ -25,7 +25,7 @@ import gymnasium as gym
 import numpy as np
 from gymnasium import spaces
 
-from lerobot.lerobot_types import RobotObservation
+from lerobot.types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv, parse_camera_names
 
@@ -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  # 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)
+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)
 ACTION_LOW = -1.0
 ACTION_HIGH = 1.0
 
@@ -98,30 +98,18 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
     return names, None
 
 
-def _get_task_horizon(task: str) -> int:
-    """Return the rollout horizon registered by RoboCasa for a task."""
-    from robocasa.utils.dataset_registry_utils import get_task_horizon
-
-    try:
-        return int(get_task_horizon(task))
-    except ValueError as exc:
-        raise ValueError(
-            f"No RoboCasa horizon is registered for task '{task}'. "
-            "Set `--env.episode_length=<steps>` explicitly."
-        ) from exc
-
-
 def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
     """Split a flat (12,) action vector into a RoboCasa action dict.
 
-    Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
+    Layout (openpi / robocasa.utils.env_utils.convert_action order):
+        ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
     """
     return {
-        "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],
+        "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],
     }
 
 
@@ -149,9 +137,16 @@ 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
@@ -167,7 +162,7 @@ class RoboCasaEnv(gym.Env):
 
         self.camera_name = parse_camera_names(camera_name)
 
-        self._max_episode_steps = episode_length if episode_length is not None else _get_task_horizon(task)
+        self._max_episode_steps = episode_length if episode_length is not None else 1000
 
         # Deferred — created on first reset() inside the worker subprocess
         # to avoid inheriting stale GPU/EGL contexts across fork().
@@ -223,12 +218,16 @@ 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()
@@ -243,12 +242,14 @@ 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.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.base_position", np.zeros(3)),
+                raw_obs.get("state.base_rotation", np.zeros(4)),
                 raw_obs.get("state.gripper_qpos", np.zeros(2)),
             ],
             axis=-1,
@@ -293,7 +294,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
+        terminated = done or (is_success and self.terminate_on_success)
         info.update({"task": self.task, "done": done, "is_success": is_success})
 
         observation = self._format_raw_obs(raw_obs)
@@ -326,6 +327,8 @@ 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.
 
@@ -348,6 +351,8 @@ 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)]
@@ -361,6 +366,8 @@ 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.
 
@@ -422,6 +429,8 @@ 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

@@ -28,7 +28,7 @@ import numpy as np
 import torch
 from gymnasium import spaces
 
-from lerobot.lerobot_types import RobotObservation
+from lerobot.types import RobotObservation
 from lerobot.utils.import_utils import _scipy_available
 
 from .utils import _LazyAsyncVectorEnv
@@ -384,9 +384,7 @@ 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.ndarray = np.zeros(
-            (self.observation_height, self.observation_width, 3), dtype=np.uint8
-        )
+        self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
 
         image_spaces = {
             cam: spaces.Box(

src/lerobot/envs/vlabench.py

@@ -37,7 +37,7 @@ import numpy as np
 from gymnasium import spaces
 from scipy.spatial.transform import Rotation
 
-from lerobot.lerobot_types import RobotObservation
+from lerobot.types import RobotObservation
 
 from .utils import _LazyAsyncVectorEnv
 
@@ -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.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
+            padded = 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,9 +122,6 @@ 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}
@@ -137,9 +134,6 @@ MODEL_RESOLUTION = {
     "xm430-w350": 4096,
     "xm540-w270": 4096,
     "xc430-w150": 4096,
-    "xh540-w150": 4096,
-    "xc330-t288": 4096,
-    "xc330-t181": 4096,
 }
 
 # {model: model_number}
@@ -151,9 +145,6 @@ 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}
@@ -165,9 +156,6 @@ 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 = {
@@ -178,9 +166,6 @@ 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 = {
@@ -191,9 +176,6 @@ 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,6 +104,8 @@ 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,6 +28,7 @@ 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
@@ -56,6 +57,7 @@ __all__ = [
     "PI0Config",
     "PI0FastConfig",
     "PI05Config",
+    "PI052Config",
     "SmolVLAConfig",
     "TDMPCConfig",
     "VLAJEPAConfig",

src/lerobot/policies/common/vla_utils.py

@@ -41,21 +41,20 @@ else:
 def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
     time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
 ) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
+    """Compute sine-cosine embeddings for scalar or per-action positions."""
     if dimension % 2 != 0:
         raise ValueError(f"dimension ({dimension}) must be divisible by 2")
 
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
+    if time.ndim not in (1, 2):
+        raise ValueError("The time tensor must have shape (batch_size,) or (batch_size, action_horizon).")
 
     dtype = get_safe_dtype(torch.float64, device.type)
     fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
     period = min_period * (max_period / min_period) ** fraction
 
-    # Compute the outer product
     scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+    sin_input = time[..., None] * scaling_factor
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=-1)
 
 
 def make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor:  # see openpi (exact copy)

src/lerobot/policies/eo1/processor_eo1.py

@@ -22,7 +22,6 @@ from typing import TYPE_CHECKING, Any
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.lerobot_types import TransitionKey
 from lerobot.processor import (
     ComplementaryDataProcessorStep,
     PolicyAction,
@@ -32,6 +31,7 @@ from lerobot.processor import (
     make_default_policy_processor_steps,
     make_policy_processor_pipelines,
 )
+from lerobot.types import TransitionKey
 from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.import_utils import _transformers_available, require_package
 

src/lerobot/policies/evo1/modeling_evo1.py

@@ -42,9 +42,6 @@ class Evo1Policy(PreTrainedPolicy):
     config_class = Evo1Config
     name = "evo1"
 
-    def supports_rtc(self) -> bool:
-        return True
-
     def __init__(self, config: Evo1Config, *, vlm_hub_kwargs: dict | None = None, **kwargs):
         super().__init__(config)
         config.validate_features()

src/lerobot/policies/evo1/processor_evo1.py

@@ -21,7 +21,6 @@ from typing import Any
 import torch
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -41,6 +40,7 @@ from lerobot.processor.converters import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     DONE,
@@ -302,33 +302,6 @@ def _pad_evo1_stats(
     return padded_stats
 
 
-def _refresh_evo1_normalization_steps(
-    config: Evo1Config,
-    preprocessor: PolicyProcessorPipeline,
-    postprocessor: PolicyProcessorPipeline,
-) -> None:
-    """Re-pad checkpoint-loaded (un)normalizer stats/features to EVO1's fixed widths.
-
-    Loading a checkpoint injects the raw dataset stats (unpadded to max_state_dim/max_action_dim)
-    into the (un)normalizer via the generic override path in make_pre_post_processors. Those stats
-    and their declared features must be re-padded/reshaped to EVO1's fixed widths, otherwise
-    normalization fails against the padded state/action tensors (e.g. state padded to 24 vs. 8-dim
-    LIBERO stats). Padding is a no-op when stats are already at the target width.
-    """
-    normalization_features = _evo1_normalization_features(config)
-    action_features = _evo1_action_features(config)
-    for step in preprocessor.steps:
-        if isinstance(step, NormalizerProcessorStep):
-            step.features = normalization_features
-            step.stats = _pad_evo1_stats(config, step.stats)
-            step.to(device=step.device, dtype=step.dtype)
-    for step in postprocessor.steps:
-        if isinstance(step, UnnormalizerProcessorStep):
-            step.features = action_features
-            step.stats = _pad_evo1_stats(config, step.stats)
-            step.to(device=step.device, dtype=step.dtype)
-
-
 def reconcile_evo1_processors(
     config: Evo1Config,
     preprocessor: PolicyProcessorPipeline,
@@ -336,19 +309,16 @@ def reconcile_evo1_processors(
 ) -> tuple[PolicyProcessorPipeline, PolicyProcessorPipeline]:
     """Reconcile checkpoint-loaded pipelines with the current EVO1 config.
 
-    Three things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
-    (converters are plain functions and are never serialized), eval-time CLI overrides of the
-    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`), and the
-    (un)normalizer stats/features when the generic override path injects raw, unpadded dataset
-    stats. This restores the converter, re-pads the normalization stats to EVO1's fixed widths, and
-    rebuilds the action step from the current config so those overrides take effect.
+    Two things cannot be restored from a serialized pipeline alone: the EVO1 batch converter
+    (converters are plain functions and are never serialized), and eval-time CLI overrides of the
+    action postprocessing flags (`postprocess_action_dim`, `binarize_gripper`, `gripper_*`). This
+    restores the converter and rebuilds the action step from the current config so those overrides
+    take effect.
     """
     # Pipelines reloaded from a checkpoint come back with the default batch converter, which drops
     # non-observation extras (embodiment_id, state_mask, custom task fields) needed by EVO1.
     preprocessor.to_transition = evo1_batch_to_transition
 
-    _refresh_evo1_normalization_steps(config, preprocessor, postprocessor)
-
     action_step = Evo1ActionProcessorStep(
         action_dim=_evo1_action_dim(config),
         binarize_gripper=config.binarize_gripper,

src/lerobot/policies/factory.py

@@ -28,7 +28,6 @@ if TYPE_CHECKING:
 
 from lerobot.configs import FeatureType, PreTrainedConfig
 from lerobot.envs import EnvConfig, env_to_policy_features
-from lerobot.lerobot_types import PolicyAction
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     PolicyProcessorPipeline,
@@ -38,25 +37,19 @@ from lerobot.processor import (
     transition_to_batch,
     transition_to_policy_action,
 )
+from lerobot.types import PolicyAction
 from lerobot.utils.constants import (
     ACTION,
     POLICY_POSTPROCESSOR_DEFAULT_NAME,
     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
@@ -144,6 +137,12 @@ 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
 
 
@@ -178,6 +177,10 @@ 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
 
@@ -197,12 +200,29 @@ 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=kwargs.get("preprocessor_overrides", {}),
+            overrides=preprocessor_overrides,
             to_transition=batch_to_transition,
             to_output=transition_to_batch,
             revision=pretrained_revision,
@@ -234,6 +254,11 @@ 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"),
     )
 
 
@@ -341,15 +366,12 @@ 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.
-        require_package("peft", extra="peft")
+        from peft import PeftConfig, PeftModel
 
         logging.info("Loading policy's PEFT adapter.")
 
         peft_pretrained_path = str(cfg.pretrained_path)
-        peft_config = PeftConfig.from_pretrained(
-            peft_pretrained_path,
-            revision=cfg.pretrained_revision,
-        )
+        peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
 
         kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
         if not kwargs["pretrained_name_or_path"]:
@@ -360,14 +382,9 @@ 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,
-            revision=cfg.pretrained_revision,
-            is_trainable=True,
+            policy, peft_pretrained_path, config=peft_config, is_trainable=True
         )
 
     else:
@@ -439,6 +456,7 @@ 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.
 
@@ -474,7 +492,9 @@ 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 inspect.signature(function).parameters:
+    if "dataset_meta" in 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,19 +37,13 @@ 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/groot/modeling_groot.py

@@ -68,9 +68,6 @@ class GrootPolicy(PreTrainedPolicy):
     name = "groot"
     config_class = GrootConfig
 
-    def supports_rtc(self) -> bool:
-        return True
-
     def __init__(self, config: GrootConfig, **kwargs):
         """Initialize Groot policy wrapper."""
         require_package("transformers", extra="groot")

src/lerobot/policies/groot/processor_groot.py

@@ -50,7 +50,6 @@ if TYPE_CHECKING or _datasets_available:
 else:
     LeRobotDataset = None
 
-from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
@@ -67,6 +66,7 @@ from lerobot.processor import (
     transition_to_batch,
     transition_to_policy_action,
 )
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     OBS_IMAGE,

src/lerobot/policies/molmoact2/modeling_molmoact2.py

@@ -43,22 +43,11 @@ from torch.distributions import Beta
 
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION
-from lerobot.utils.import_utils import (
-    _peft_available,
-    _scipy_available,
-    _transformers_available,
-    require_package,
-)
+from lerobot.utils.import_utils import _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__)
 
 
@@ -520,9 +509,6 @@ class MolmoAct2Policy(PreTrainedPolicy):
     config_class = MolmoAct2Config
     name = "molmoact2"
 
-    def supports_rtc(self) -> bool:
-        return self.config.inference_action_mode == "continuous"
-
     def __init__(
         self,
         config: MolmoAct2Config,
@@ -1745,11 +1731,13 @@ 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/molmoact2/processor_molmoact2.py

@@ -36,7 +36,6 @@ import torch
 from torch import Tensor
 
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
@@ -50,6 +49,7 @@ from lerobot.processor import (
     policy_action_to_transition,
     transition_to_policy_action,
 )
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
     ACTION,
     OBS_IMAGES,

src/lerobot/policies/pi0/modeling_pi0.py

@@ -749,9 +749,6 @@ class PI0Policy(PreTrainedPolicy):
     config_class = PI0Config
     name = "pi0"
 
-    def supports_rtc(self) -> bool:
-        return True
-
     def __init__(
         self,
         config: PI0Config,

src/lerobot/policies/pi05/configuration_pi05.py

@@ -58,6 +58,8 @@ class PI05Config(PreTrainedConfig):
 
     # Real-Time Chunking (RTC) configuration
     rtc_config: RTCConfig | None = None
+    # Maximum clean action-prefix length sampled during training. Zero disables trained RTC.
+    rtc_training_max_delay: int = 0
 
     image_resolution: tuple[int, int] = (
         DEFAULT_IMAGE_SIZE,
@@ -111,6 +113,11 @@ class PI05Config(PreTrainedConfig):
             raise ValueError(
                 f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
             )
+        if not 0 <= self.rtc_training_max_delay < self.chunk_size:
+            raise ValueError(
+                "rtc_training_max_delay must satisfy "
+                f"0 <= delay < chunk_size ({self.chunk_size}), got {self.rtc_training_max_delay}"
+            )
 
         if self.paligemma_variant not in ["gemma_300m", "gemma_2b"]:
             raise ValueError(f"Invalid paligemma_variant: {self.paligemma_variant}")

src/lerobot/policies/pi05/modeling_pi05.py

@@ -22,6 +22,7 @@ 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
@@ -30,6 +31,7 @@ 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,
@@ -44,20 +46,21 @@ 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 euler_integrate, sample_noise, sample_time_beta
+from ..common.flow_matching import 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
@@ -71,6 +74,110 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
+def _prepare_trained_rtc_prefix(
+    x_t: Tensor,
+    prev_chunk_left_over: Tensor | None,
+    inference_delay: int,
+    training_max_delay: int,
+) -> tuple[Tensor | None, Tensor | None]:
+    """Pad and validate a hard prefix for training-time RTC inference."""
+    if prev_chunk_left_over is None or inference_delay <= 0:
+        return None, None
+    if training_max_delay <= 0:
+        raise ValueError(
+            "RTC mode='trained' requires a checkpoint trained with policy.rtc_training_max_delay > 0."
+        )
+    if inference_delay > training_max_delay:
+        raise ValueError(
+            f"Measured RTC inference delay ({inference_delay}) exceeds the checkpoint's "
+            f"rtc_training_max_delay ({training_max_delay})."
+        )
+    if inference_delay >= x_t.shape[1]:
+        raise ValueError(
+            f"RTC inference delay ({inference_delay}) must be smaller than chunk_size ({x_t.shape[1]})."
+        )
+
+    previous = prev_chunk_left_over.to(device=x_t.device, dtype=x_t.dtype)
+    if not torch.isfinite(previous).all():
+        raise ValueError("RTC prefix contains NaN or Inf values.")
+    if previous.ndim == 2:
+        previous = previous.unsqueeze(0)
+    if previous.ndim != 3:
+        raise ValueError(f"Expected RTC prefix shape (B, T, A), got {tuple(previous.shape)}")
+    if previous.shape[0] == 1 and x_t.shape[0] > 1:
+        previous = previous.expand(x_t.shape[0], -1, -1)
+    if previous.shape[0] != x_t.shape[0]:
+        raise ValueError(
+            f"RTC prefix batch size ({previous.shape[0]}) does not match policy batch ({x_t.shape[0]})."
+        )
+    if previous.shape[1] < inference_delay:
+        raise ValueError(f"RTC prefix has {previous.shape[1]} steps, but inference_delay={inference_delay}.")
+    if previous.shape[2] > x_t.shape[2]:
+        raise ValueError(
+            f"RTC prefix action dimension ({previous.shape[2]}) exceeds model dimension ({x_t.shape[2]})."
+        )
+
+    padded_prefix = torch.zeros_like(x_t)
+    padded_prefix[:, :inference_delay, : previous.shape[2]] = previous[:, :inference_delay]
+    prefix_mask = torch.arange(x_t.shape[1], device=x_t.device) < inference_delay
+    prefix_mask = prefix_mask[None, :, None].expand(x_t.shape[0], -1, x_t.shape[2])
+    return padded_prefix, prefix_mask
+
+
+def _sample_training_rtc_prefix_mask(
+    batch_size: int,
+    action_horizon: int,
+    max_delay: int,
+    device: torch.device,
+) -> Tensor | None:
+    """Sample a clean action-prefix length independently for each training example."""
+    if max_delay <= 0:
+        return None
+    delays = torch.randint(0, max_delay + 1, (batch_size,), device=device)
+    positions = torch.arange(action_horizon, device=device)
+    return positions.unsqueeze(0) < delays.unsqueeze(1)
+
+
+def _build_flow_matching_inputs(
+    actions: Tensor,
+    noise: Tensor,
+    time: Tensor,
+    prefix_mask: Tensor | None,
+) -> tuple[Tensor, Tensor]:
+    """Keep the sampled RTC prefix clean while noising the remaining action chunk."""
+    if prefix_mask is None:
+        model_time = time
+        expanded_time = time[:, None, None]
+    else:
+        model_time = time[:, None].expand_as(prefix_mask)
+        model_time = torch.where(prefix_mask, torch.zeros_like(model_time), model_time)
+        expanded_time = model_time.unsqueeze(-1)
+    x_t = expanded_time * noise + (1 - expanded_time) * actions
+    return x_t, model_time
+
+
+def _reduce_training_rtc_loss(
+    losses: Tensor,
+    prefix_mask: Tensor | None,
+    reduction: str,
+) -> Tensor:
+    """Average flow loss over predicted postfix actions, excluding the clean RTC prefix."""
+    if reduction not in {"mean", "none"}:
+        raise ValueError(f"Unsupported loss reduction: {reduction!r}")
+    if prefix_mask is None:
+        return losses.mean() if reduction == "mean" else losses.mean(dim=(1, 2))
+
+    postfix_mask = (~prefix_mask).unsqueeze(-1).expand_as(losses)
+    if reduction == "none":
+        numerator = (losses * postfix_mask).sum(dim=(1, 2))
+        denominator = postfix_mask.sum(dim=(1, 2))
+        return numerator / denominator.clamp(min=1)
+    return (losses * postfix_mask).sum() / postfix_mask.sum().clamp(min=1)
+
+
+_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 = []
@@ -401,6 +508,12 @@ 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
@@ -444,7 +557,11 @@ 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
-        self.paligemma_with_expert.paligemma.model.vision_tower.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.gemma_expert.model.gradient_checkpointing = True
         logging.info("Enabled gradient checkpointing for PI05Pytorch model")
 
@@ -452,7 +569,11 @@ 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
-        self.paligemma_with_expert.paligemma.model.vision_tower.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.gemma_expert.model.gradient_checkpointing = False
         logging.info("Disabled gradient checkpointing for PI05Pytorch model")
 
@@ -467,6 +588,14 @@ 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)
 
@@ -488,13 +617,16 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         pad_masks = []
         att_masks = []
 
-        # Process images
-        for img, img_mask in zip(images, img_masks, strict=True):
+        if self.checkpoint_vision_embeddings:
 
-            def image_embed_func(img):
-                return self.paligemma_with_expert.embed_image(img)
+            def embed_image(img):
+                return self._apply_checkpoint(self.paligemma_with_expert.embed_image, img)
 
-            img_emb = self._apply_checkpoint(image_embed_func, 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):
             bsize, num_img_embs = img_emb.shape[:2]
 
             embs.append(img_emb)
@@ -556,19 +688,35 @@ 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))
-        att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
-        att_masks = att_masks[None, :].expand(bsize, len(att_masks))
+
+        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))
 
         return action_emb, pad_masks, att_masks, adarms_cond
 
-    def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
+    def forward(
+        self,
+        images,
+        img_masks,
+        tokens,
+        masks,
+        actions,
+        noise,
+        time,
+        prefix_mask: Tensor | None = None,
+    ) -> Tensor:
         """Do a full training forward pass and compute the loss."""
-        time_expanded = time[:, None, None]
-        x_t = time_expanded * noise + (1 - time_expanded) * actions
+        x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
         u_t = noise - actions
 
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, tokens, masks)
-        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, time)
+        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, model_time)
 
         if (
             self.paligemma_with_expert.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype
@@ -583,7 +731,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 = prepare_attention_masks_4d(att_2d_masks)
+        att_2d_masks_4d = self._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(
@@ -641,7 +789,8 @@ 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
 
-        prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
+        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)
         self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
 
         _, past_key_values = self.paligemma_with_expert.forward(
@@ -652,21 +801,78 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             use_cache=True,
         )
 
-        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"),
-        )
+        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
+        rtc_mode = "guided"
+        trained_prefix = trained_prefix_mask = None
+        if self._rtc_enabled():
+            rtc_mode = self.rtc_processor.rtc_config.mode
+            if rtc_mode == "trained":
+                training_max_delay = int(getattr(self.config, "rtc_training_max_delay", 0))
+                if training_max_delay <= 0:
+                    raise ValueError(
+                        "RTC mode='trained' requires a checkpoint trained with "
+                        "policy.rtc_training_max_delay > 0."
+                    )
+                trained_prefix, trained_prefix_mask = _prepare_trained_rtc_prefix(
+                    x_t,
+                    kwargs.get("prev_chunk_left_over"),
+                    int(kwargs.get("inference_delay") or 0),
+                    training_max_delay,
+                )
+
+        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)
+
+            denoise_timestep = time_tensor
+            if trained_prefix is not None:
+                x_t = torch.where(trained_prefix_mask, trained_prefix, x_t)
+                denoise_timestep = time_tensor[:, None].expand(bsize, x_t.shape[1]).clone()
+                denoise_timestep[trained_prefix_mask[..., 0]] = 0.0
+
+            def denoise_step_partial_call(input_x_t, current_timestep=denoise_timestep):
+                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() and rtc_mode == "guided":
+                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 trained_prefix is not None:
+                x_t = torch.where(trained_prefix_mask, trained_prefix, x_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
 
     def denoise_step(
         self,
@@ -689,7 +895,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 = prepare_attention_masks_4d(full_att_2d_masks)
+        full_att_2d_masks_4d = self._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)
@@ -713,9 +919,10 @@ class PI05Policy(PreTrainedPolicy):
 
     config_class = PI05Config
     name = "pi05"
-
-    def supports_rtc(self) -> bool:
-        return True
+    model_class = PI05Pytorch
+    eval_after_pretrained_load = False
+    show_openpi_disclaimer = True
+    use_native_pretrained_loader = False
 
     def __init__(
         self,
@@ -733,7 +940,7 @@ class PI05Policy(PreTrainedPolicy):
 
         # Initialize the core PI05 model
         self.init_rtc_processor()
-        self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
+        self.model = self.model_class(config, rtc_processor=self.rtc_processor)
 
         # Enable gradient checkpointing if requested
         if config.gradient_checkpointing:
@@ -759,16 +966,31 @@ class PI05Policy(PreTrainedPolicy):
         strict: bool = True,
         **kwargs,
     ) -> T:
-        """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"
-        )
+        """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"
+            )
         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,
@@ -782,85 +1004,41 @@ 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}.")
 
-        # 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}")
-
+        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()
         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`
@@ -940,7 +1118,10 @@ class PI05Policy(PreTrainedPolicy):
         # Create processor if config provided
         # If RTC is not enabled - we can still track the denoising data
         if self.config.rtc_config is not None:
-            self.rtc_processor = RTCProcessor(self.config.rtc_config)
+            self.rtc_processor = RTCProcessor(
+                self.config.rtc_config,
+                trained_mode_supported=int(getattr(self.config, "rtc_training_max_delay", 0)) > 0,
+            )
 
             model_value = getattr(self, "model", None)
             if model_value is not None:
@@ -1031,12 +1212,16 @@ class PI05Policy(PreTrainedPolicy):
 
         # Action queue logic for n_action_steps > 1
         if len(self._action_queue) == 0:
-            actions = self.predict_action_chunk(batch)[:, : self.config.n_action_steps]
+            action_batch = self._prepare_action_batch(batch)
+            actions = self.predict_action_chunk(action_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."""
@@ -1072,28 +1257,35 @@ class PI05Policy(PreTrainedPolicy):
 
         noise = self.model.sample_noise(actions.shape, actions.device)
         time = self.model.sample_time(actions.shape[0], actions.device)
+        prefix_mask = _sample_training_rtc_prefix_mask(
+            actions.shape[0],
+            actions.shape[1],
+            self.config.rtc_training_max_delay,
+            actions.device,
+        )
 
         # Compute loss (no separate state needed for PI05)
-        losses = self.model.forward(images, img_masks, tokens, masks, actions, noise, time)
+        losses = self.model.forward(images, img_masks, tokens, masks, actions, noise, time, prefix_mask)
 
         # Truncate losses to actual action dimensions
         original_action_dim = self.config.output_features[ACTION].shape[0]
         losses = losses[:, :, :original_action_dim]
 
-        loss_dict = {
-            "loss_per_dim": losses.mean(dim=[0, 1]).detach().cpu().numpy().tolist(),
-        }
+        if prefix_mask is None:
+            loss_per_dim = losses.mean(dim=(0, 1))
+        else:
+            postfix_mask = (~prefix_mask).unsqueeze(-1).expand_as(losses)
+            loss_per_dim = (losses * postfix_mask).sum(dim=(0, 1)) / postfix_mask.sum(dim=(0, 1)).clamp(min=1)
+        loss_dict = {"loss_per_dim": loss_per_dim.detach().cpu().numpy().tolist()}
 
         if reduction == "none":
-            # Return per-sample losses (B,) by averaging over time and action dims
-            per_sample_loss = losses.mean(dim=(1, 2))
+            per_sample_loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="none")
             loss_dict["loss"] = per_sample_loss.mean().item()
             return per_sample_loss, loss_dict
-        else:
-            # Default: return scalar mean loss
-            loss = losses.mean()
-            loss_dict["loss"] = loss.item()
-            return loss, loss_dict
+
+        loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean")
+        loss_dict["loss"] = loss.item()
+        return loss, loss_dict
 
     def _get_default_peft_targets(self) -> dict[str, any]:
         """Return default PEFT target modules for PI0.5 fine-tuning."""

src/lerobot/policies/pi05/processor_pi05.py

@@ -22,7 +22,6 @@ import numpy as np
 import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     PolicyAction,
@@ -34,6 +33,7 @@ from lerobot.processor import (
     make_default_policy_processor_steps,
     make_policy_processor_pipelines,
 )
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import OBS_STATE
 
 from .configuration_pi05 import PI05Config

src/lerobot/policies/pi052/__init__.py

@@ -0,0 +1,19 @@
+# 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

@@ -0,0 +1,172 @@
+# 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
+
+    # Training-time RTC configuration is inherited from PI05Config.
+
+    # 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

@@ -0,0 +1,522 @@
+# 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

@@ -0,0 +1,263 @@
+# 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

@@ -0,0 +1,19 @@
+# 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

@@ -0,0 +1,254 @@
+# 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

@@ -0,0 +1,164 @@
+# 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

@@ -0,0 +1,521 @@
+# 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.MEAN_STD,  # Pi0Fast uses quantiles for state
-            "ACTION": NormalizationMode.MEAN_STD,  # Pi0Fast uses quantiles for action
+            "STATE": NormalizationMode.QUANTILES,
+            "ACTION": NormalizationMode.QUANTILES,
         }
     )
 

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

@@ -24,13 +24,7 @@ import numpy as np
 import torch
 from torch import Tensor, nn
 
-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
+from lerobot.utils.import_utils import _transformers_available, require_package
 
 if TYPE_CHECKING or _transformers_available:
     from transformers import AutoProcessor, AutoTokenizer
@@ -66,6 +60,32 @@ 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."""
 
@@ -240,7 +260,6 @@ 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):
@@ -467,18 +486,12 @@ 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
 
-        # 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
+        # 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
 
         # compute cross-entropy loss
         loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
@@ -488,9 +501,7 @@ 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)
 
-        # 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)
+        fast_loss = _reduce_fast_token_loss(fast_loss_per_token, fast_action_masks.float())
 
         return {
             "ce_loss": fast_loss,
@@ -519,15 +530,7 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
         device = tokens.device
         lm_head = self.paligemma_with_expert.paligemma.lm_head
 
-        # 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]
+        # 1. Initial embedding: the prompt's existing BOS is the only BOS in the sequence.
         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
         )
@@ -539,6 +542,8 @@ 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):
@@ -556,16 +561,24 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
                 adarms_cond=[None, None],
             )
 
-            # predict next token from the very last sequence position
-            last_logits = lm_head(prefix_out[:, -1:, :])  # (B, 1, vocab_size)
-
-            if temperature > 0:
-                probs = torch.softmax(last_logits[:, -1] / temperature, dim=-1)
-                next_token = torch.multinomial(probs, num_samples=1)
+            if t == 0:
+                prediction_hidden = _gather_last_valid_language_hidden(prefix_out, masks, total_t_images)
             else:
-                next_token = torch.argmax(last_logits[:, -1], dim=-1, keepdim=True)
+                prediction_hidden = prefix_out[:, -1]
+            next_token = _sample_next_token(lm_head(prediction_hidden), temperature)
 
-            generated_action_tokens[:, t] = next_token.squeeze(-1)
+            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,
+            )
 
             # 3. Update sequence for next iteration (unless it's the last step)
             if t < max_decoding_steps - 1:
@@ -612,20 +625,14 @@ 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_in, masks_in, fast_action_tokens=None, fast_action_masks=None
+            images, img_masks, tokens, masks, fast_action_tokens=None, fast_action_masks=None
         )
 
         # Ensure correct precision (bfloat16/float32)
@@ -652,17 +659,18 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
             adarms_cond=[None, None],
         )
 
-        # 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)
+        prediction_hidden = _gather_last_valid_language_hidden(prefix_out, masks, total_t_images)
+        next_token = _sample_next_token(lm_head(prediction_hidden), temperature)
         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)
@@ -703,15 +711,19 @@ class PI0FastPytorch(nn.Module):  # see openpi `PI0Pytorch`
                 adarms_cond=[None, None],
             )
 
-            # 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)
+            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,
+            )
 
         return generated_action_tokens
 
@@ -1024,7 +1036,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[int], time_horizon: int, action_dim: int, relaxed_decoding: bool = True
+        self, token_ids: list[Tensor], time_horizon: int, action_dim: int
     ) -> np.ndarray:
         """
         Decodes action token IDs back to continuous action values using the FAST tokenizer.
@@ -1033,8 +1045,6 @@ 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.
         """
@@ -1042,40 +1052,23 @@ class PI0FastPolicy(PreTrainedPolicy):
 
         for token in token_ids:
             try:
-                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})"
+                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,
                 )
+                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(f"Error decoding tokens: {e}")
-                logging.warning(f"Tokens: {token}")
-                decoded_dct_coeff = np.zeros((time_horizon, action_dim))
+                logging.warning("Invalid FAST action sequence; returning a zero action chunk: %s", e)
+                decoded_action = np.zeros((time_horizon, action_dim))
 
-            decoded_actions.append(
-                idct(decoded_dct_coeff / self.action_tokenizer.scale, axis=0, norm="ortho")
-            )
+            decoded_actions.append(decoded_action)
 
         return np.stack(decoded_actions)
 
@@ -1105,53 +1098,28 @@ class PI0FastPolicy(PreTrainedPolicy):
         if single_sample:
             tokens = tokens.unsqueeze(0)
 
-        # 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}"
+        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,
                 )
-
-        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
-        ]
+                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))
 
         # Decode action tokens to continuous actions
         actions = self.decode_actions_with_fast(
@@ -1220,7 +1188,7 @@ class PI0FastPolicy(PreTrainedPolicy):
             )
 
         # Detokenize action tokens to continuous actions
-        action_horizon = self.config.n_action_steps
+        action_horizon = self.config.chunk_size
         action_dim = self.config.output_features[ACTION].shape[0]
 
         continuous_actions = self.detokenize_actions(

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -22,7 +22,6 @@ import numpy as np
 import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
-from lerobot.lerobot_types import EnvTransition, TransitionKey
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     ActionTokenizerProcessorStep,
@@ -35,6 +34,7 @@ from lerobot.processor import (
     make_default_policy_processor_steps,
     make_policy_processor_pipelines,
 )
+from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import OBS_STATE
 
 from .configuration_pi0_fast import PI0FastConfig
@@ -70,7 +70,7 @@ class Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(ProcessorStep):
 
         full_prompts = []
         for i, task in enumerate(tasks):
-            cleaned_text = task.strip().replace("_", " ").replace("\n", " ")
+            cleaned_text = task.strip().replace("_", " ").replace("\n", " ").lower()
             state_str = " ".join(map(str, discretized_states[i]))
             full_prompt = f"Task: {cleaned_text}, State: {state_str};\n"
             full_prompts.append(full_prompt)
@@ -92,6 +92,11 @@ 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],
@@ -136,6 +141,18 @@ 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,
@@ -149,10 +166,11 @@ def make_pi0_fast_pre_post_processors(
             padding="max_length",
         ),
         ActionTokenizerProcessorStep(
-            action_tokenizer_name=config.action_tokenizer_name,
+            action_tokenizer_name=action_tokenizer_path,
             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,6 +18,7 @@ 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
 
@@ -121,7 +122,10 @@ 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)
-        if len(x.shape) == 3:
+        # 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:
             modulation = modulation.unsqueeze(1)
         scale, shift, gate = modulation.chunk(3, dim=-1)
         normed = normed * (1 + scale.float()) + shift.float()
@@ -275,6 +279,8 @@ 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)
@@ -367,3 +373,45 @@ __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,22 +34,14 @@ 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
 
-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
+T = TypeVar("T", bound="PreTrainedPolicy")
 
 if TYPE_CHECKING:
     from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
 
-T = TypeVar("T", bound="PreTrainedPolicy")
-
 
 def _build_card_context(
     cfg: TrainPipelineConfig | None,
@@ -249,10 +241,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         """
         raise NotImplementedError
 
-    def supports_rtc(self) -> bool:
-        """Whether this policy implements Real-Time Chunking inference semantics."""
-        return False
-
     # TODO(aliberts, rcadene): split into 'forward' and 'compute_loss'?
     @abc.abstractmethod
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
@@ -350,6 +338,7 @@ 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",
         }
@@ -396,7 +385,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.
         """
-        require_package("peft", extra="peft")
+        from peft import get_peft_model
 
         # If user provided a complete config, use it directly (with overrides)
         if peft_config is not None:
@@ -467,7 +456,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         Returns:
             Preprocessed dict with renamed keys and init_type mapped to method-specific key.
         """
-        require_package("peft", extra="peft")
+        from peft import PeftType
 
         cli_overrides = cli_overrides.copy()
 
@@ -492,7 +481,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."""
-        require_package("peft", extra="peft")
+        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
 
         # Determine PEFT method type (default to LORA)
         method_type_str = cli_overrides.get("method_type") or "lora"
@@ -519,7 +508,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."""
-        require_package("peft", extra="peft")
+        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
 
         # Get method type from existing config or CLI override
         method_type_str = cli_overrides.get("method_type")

src/lerobot/policies/rtc/configuration_rtc.py

@@ -37,6 +37,10 @@ class RTCConfig:
     # Infrastructure
     enabled: bool = True
 
+    # ``guided`` is the original inference-time Jacobian guidance. ``trained``
+    # hard-inpaints a prefix and requires a compatible training-time RTC checkpoint.
+    mode: str = "guided"
+
     # Core RTC settings
     # Todo change to exp
     prefix_attention_schedule: RTCAttentionSchedule = RTCAttentionSchedule.LINEAR
@@ -49,6 +53,8 @@ class RTCConfig:
 
     def __post_init__(self):
         """Validate RTC configuration parameters."""
+        if self.mode not in {"guided", "trained"}:
+            raise ValueError(f"mode must be 'guided' or 'trained', got {self.mode!r}")
         if self.max_guidance_weight <= 0:
             raise ValueError(f"max_guidance_weight must be positive, got {self.max_guidance_weight}")
         if self.debug_maxlen <= 0: