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>

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
 

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/bring_your_own_policies.mdx

@@ -189,6 +189,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/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/language_and_recipes.mdx

@@ -141,6 +141,11 @@ 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.
+
 ## Graceful absence
 
 If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.

docs/source/pi052.mdx

@@ -0,0 +1,315 @@
+# π₀.₅ 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_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 |
+
+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 public Pi05 base 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/pi05_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   |
+
+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 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
 

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

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,

pyproject.toml

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

src/lerobot/configs/recipe.py

@@ -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_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/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/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/envs/configs.py

@@ -556,7 +556,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:
@@ -570,6 +576,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/robocasa.py

@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
 
 # Dimensions for the flat action/state vectors used by the LeRobot wrapper.
 # These correspond to the PandaOmron robot in RoboCasa365.
-OBS_STATE_DIM = 16  # 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
 
@@ -101,14 +101,15 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
 def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
     """Split a flat (12,) action vector into a RoboCasa action dict.
 
-    Layout: 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],
     }
 
 
@@ -136,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
@@ -210,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()
@@ -230,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,
@@ -280,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)
@@ -313,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.
 
@@ -335,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)]
@@ -348,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.
 
@@ -409,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/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
@@ -55,6 +56,7 @@ __all__ = [
     "PI0Config",
     "PI0FastConfig",
     "PI05Config",
+    "PI052Config",
     "SmolVLAConfig",
     "TDMPCConfig",
     "VQBeTConfig",

src/lerobot/policies/factory.py

@@ -133,6 +133,10 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
         from .pi05.modeling_pi05 import PI05Policy
 
         return PI05Policy
+    elif name == "pi052":
+        from .pi052.modeling_pi052 import PI052Policy
+
+        return PI052Policy
     elif name == "gaussian_actor":
         from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
 
@@ -193,8 +197,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
 
     Args:
         policy_type: The type of the policy. Supported types include "tdmpc",
-                     "multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor",
-                     "smolvla", "wall_x", "molmoact2", "eo1", "evo1".
+                     "multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "pi052",
+                     "gaussian_actor", "smolvla", "wall_x", "molmoact2", "eo1", "evo1".
         **kwargs: Keyword arguments to be passed to the configuration class constructor.
 
     Returns:
@@ -217,6 +221,10 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
         return PI0Config(**kwargs)
     elif policy_type == "pi05":
         return PI05Config(**kwargs)
+    elif policy_type == "pi052":
+        from .pi052.configuration_pi052 import PI052Config
+
+        return PI052Config(**kwargs)
     elif policy_type == "gaussian_actor":
         return GaussianActorConfig(**kwargs)
     elif policy_type == "smolvla":
@@ -267,6 +275,8 @@ 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 repo used for optional processor fitting; omit it to use universal tokenizers.
+    dataset_repo_id: str | None
     dataset_meta: Any | None
 
 
@@ -301,6 +311,22 @@ def make_pre_post_processors(
         NotImplementedError: If a processor factory is not implemented for the given
             policy configuration type.
     """
+    if (
+        pretrained_path
+        and getattr(policy_cfg, "type", None) in {"pi0_fast", "pi052"}
+        and getattr(policy_cfg, "auto_fit_fast_tokenizer", False)
+        and kwargs.get("dataset_repo_id") is not None
+    ):
+        from .pi052.fit_fast_tokenizer import resolve_fast_tokenizer
+
+        overrides = dict(kwargs.get("preprocessor_overrides") or {})
+        action_tokenizer_override = {
+            **overrides.get("action_tokenizer_processor", {}),
+            "action_tokenizer_name": resolve_fast_tokenizer(policy_cfg, kwargs.get("dataset_repo_id")),
+        }
+        overrides["action_tokenizer_processor"] = action_tokenizer_override
+        kwargs["preprocessor_overrides"] = overrides
+
     if pretrained_path:
         if isinstance(policy_cfg, GrootConfig):
             from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
@@ -402,6 +428,26 @@ def make_pre_post_processors(
             dataset_stats=kwargs.get("dataset_stats"),
         )
 
+    elif policy_cfg.type == "pi0_fast":
+        from .pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
+
+        processors = make_pi0_fast_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+            dataset_repo_id=kwargs.get("dataset_repo_id"),
+        )
+
+    elif policy_cfg.type == "pi052":
+        # PI052 must precede PI05 because its config subclasses PI05Config.
+        from .pi052.processor_pi052 import make_pi052_pre_post_processors
+
+        processors = make_pi052_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+            # Without a dataset repo, FAST auto-fit falls back to the universal tokenizer.
+            dataset_repo_id=kwargs.get("dataset_repo_id"),
+        )
+
     elif isinstance(policy_cfg, PI05Config):
         from .pi05.processor_pi05 import make_pi05_pre_post_processors
 

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

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import builtins
+import json
 import logging
 import math
 from collections import deque
@@ -23,6 +24,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
@@ -32,6 +34,7 @@ if TYPE_CHECKING or _transformers_available:
     from transformers.cache_utils import DynamicCache
     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,
@@ -47,6 +50,7 @@ else:
     _gated_residual = None
     layernorm_forward = None
     PaliGemmaForConditionalGenerationWithPiGemma = None
+    cached_file = None
 from lerobot.configs import PreTrainedConfig
 from lerobot.utils.constants import (
     ACTION,
@@ -66,6 +70,185 @@ 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"
+_SAFETENSORS_INDEX = "model.safetensors.index.json"
+
+
+def _resolve_weight_files(
+    pretrained_name_or_path: str | Path,
+    *,
+    force_download: bool,
+    resume_download: bool | None,
+    proxies: dict | None,
+    token: str | bool | None,
+    cache_dir: str | Path | None,
+    local_files_only: bool,
+    revision: str | None,
+) -> list[Path]:
+    model_id = str(pretrained_name_or_path)
+    local_dir = Path(model_id)
+    load_kwargs = {
+        "revision": revision,
+        "cache_dir": cache_dir,
+        "force_download": force_download,
+        "resume_download": resume_download,
+        "proxies": proxies,
+        "token": token,
+        "local_files_only": local_files_only,
+    }
+
+    if local_dir.is_dir():
+        index_path = local_dir / _SAFETENSORS_INDEX
+        single_path = local_dir / _SAFETENSORS_FILE
+    else:
+        resolved_index = cached_file(
+            model_id,
+            _SAFETENSORS_INDEX,
+            _raise_exceptions_for_missing_entries=False,
+            **load_kwargs,
+        )
+        index_path = Path(resolved_index) if resolved_index is not None else None
+        single_path = None
+        if index_path is None:
+            resolved_file = cached_file(model_id, _SAFETENSORS_FILE, **load_kwargs)
+            single_path = Path(resolved_file) if resolved_file is not None else None
+
+    if index_path is None or not index_path.is_file():
+        if single_path is None or not single_path.is_file():
+            raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
+        return [single_path]
+
+    index = json.loads(index_path.read_text())
+    shard_names = sorted(set(index.get("weight_map", {}).values()))
+    if not shard_names:
+        raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")
+    if local_dir.is_dir():
+        files = [local_dir / name for name in shard_names]
+    else:
+        files = []
+        for name in shard_names:
+            resolved_file = cached_file(model_id, name, **load_kwargs)
+            if resolved_file is None:
+                raise FileNotFoundError(f"Checkpoint shard {name!r} not found in {model_id!r}.")
+            files.append(Path(resolved_file))
+    missing = [str(path) for path in files if not path.is_file()]
+    if missing:
+        raise FileNotFoundError(f"Missing checkpoint shards: {missing}")
+    return files
+
+
+def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
+    state_dict: dict[str, Tensor] = {}
+    for path in files:
+        shard = load_file(path)
+        overlap = state_dict.keys() & shard.keys()
+        if overlap:
+            raise ValueError(f"Duplicate checkpoint keys in {path}: {sorted(overlap)[:5]}")
+        state_dict.update(shard)
+    return state_dict
+
+
 def get_safe_dtype(target_dtype, device_type):
     """Get a safe dtype for the given device type."""
     if device_type == "mps" and target_dtype == torch.float64:
@@ -82,21 +265,21 @@ def get_safe_dtype(target_dtype, device_type):
 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."""
+    """Computes 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
+    # Broadcast the frequency dimension over either (B,) or (B, H).
     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 sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
@@ -563,6 +746,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
@@ -606,7 +795,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")
 
@@ -614,7 +807,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")
 
@@ -629,10 +826,13 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
-    def _prepare_attention_masks_4d(self, att_2d_masks):
+    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, :, :]
-        return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
+        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 torch.normal(
@@ -658,13 +858,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)
@@ -734,19 +937,34 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         embs = torch.cat(embs, dim=1)
         pad_masks = torch.cat(pad_masks, dim=1)
-        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.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=embs.dtype, device=embs.device)
+            att_masks[0] = 1
+            att_masks = att_masks[None, :].expand(bsize, n)
+        else:
+            att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
+            att_masks = att_masks[None, :].expand(bsize, len(att_masks))
 
         return embs, 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
@@ -819,7 +1037,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 = self._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(
@@ -832,12 +1051,45 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         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
-            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+            if times is None:
+                time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+            else:
+                time_tensor = times[step].expand(bsize)
 
-            def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
+            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,
@@ -845,7 +1097,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
                     timestep=current_timestep,
                 )
 
-            if self._rtc_enabled():
+            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")
@@ -862,6 +1114,8 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
                 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)
@@ -913,6 +1167,9 @@ class PI05Policy(PreTrainedPolicy):
 
     config_class = PI05Config
     name = "pi05"
+    model_class = PI05Pytorch
+    eval_after_pretrained_load = False
+    show_openpi_disclaimer = True
 
     def __init__(
         self,
@@ -930,7 +1187,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:
@@ -956,16 +1213,16 @@ 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 PI05-compatible single-file or sharded safetensors checkpoints."""
+        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,
@@ -979,85 +1236,35 @@ class PI05Policy(PreTrainedPolicy):
                 **kwargs,
             )
 
-        # Initialize model without loading weights
-        # Check if dataset_stats were provided in kwargs
         model = cls(config, **kwargs)
-
-        # 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}")
-
+        files = _resolve_weight_files(
+            pretrained_name_or_path,
+            force_download=force_download,
+            resume_download=resume_download,
+            proxies=proxies,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+            revision=revision,
+        )
+        fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), 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`
@@ -1137,7 +1344,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:
@@ -1228,12 +1438,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."""
@@ -1269,28 +1483,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/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,167 @@
+# 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 configuration for recipe-driven text and action supervision."""
+
+    # Recipe / language stack ---------------------------------------------
+    recipe_path: str | None = "recipes/subtask_mem.yaml"
+    """Recipe path relative to ``src/lerobot/configs/``, or ``None`` for the plain PI0.5 prompt."""
+
+    apply_chat_template: bool = False
+    """Whether to apply a tokenizer chat template.
+
+    PaliGemma defaults to plain recipe-rendered prefixes because it is not chat-pretrained.
+    """
+
+    # Balance frequent recipe text supervision against the paper's α=10 flow weight.
+    text_loss_weight: float = 1.0
+    """LM-head cross-entropy weight; ``0`` disables text training."""
+
+    flow_loss_weight: float = 10.0
+    """Weight on action-expert flow matching relative to text supervision."""
+
+    # Backbone training ---------------------------------------------------
+    unfreeze_lm_head: bool = True
+    """Keep PaliGemma's language head trainable for hierarchical inference."""
+
+    # 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-tokenized action cross-entropy to text CE and flow matching."""
+
+    action_tokenizer_name: str = "physical-intelligence/fast"
+    """HF identifier for the FAST action tokenizer."""
+
+    max_action_tokens: int = 256
+    """Maximum number of FAST tokens per action chunk."""
+
+    fast_skip_tokens: int = 128
+    """Number of low-vocab tokens the FAST tokenizer skips to avoid
+    collisions with PaliGemma's text vocabulary."""
+
+    fast_action_loss_weight: float = 1.0
+    """Weight on FAST action-token CE relative to continuous-flow supervision."""
+
+    subtask_replan_steps: int = 0
+    """Environment steps between subtask generations during evaluation.
+
+    Non-positive values regenerate each action chunk while still refreshing the action prompt every chunk.
+    """
+
+    auto_fit_fast_tokenizer: bool = False
+    """Fit and cache a dataset-specific FAST tokenizer before training.
+
+    Disabled by default to avoid the extra dataset pass and use the universal tokenizer.
+    """
+
+    fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
+    """Where fitted FAST tokenizers are stored. ``~`` expands."""
+
+    fast_tokenizer_fit_samples: int = 1024
+    """Number of action chunks sampled when fitting FAST."""
+
+    # Knowledge insulation detaches VLM K/V from action-loss gradients (paper §III.B).
+    knowledge_insulation: bool = True
+    """Block action-loss gradients through VLM keys and values."""
+
+    # Optional training backends. Defaults preserve the eager/SDPA path.
+    use_flashrt_adarms: bool = False
+    """Use FlashRT adaptive RMSNorm kernels when available."""
+
+    use_compiled_text_ce: bool = False
+    """Compile the materialized-logits text and FAST CE path."""
+
+    use_compiled_vision: bool = False
+    """Compile the SigLIP tower for no-grad flow and inference passes."""
+
+    use_flex_attention: bool = False
+    """Use FlexAttention for amortized KI, with SDPA fallback where unsupported."""
+
+    use_manual_attention: bool = False
+    """Use materialized-logits attention for explicitly profiled KI shapes."""
+
+    manual_attention_scope: str = "all"
+    """Apply manual attention to all KI queries or only action queries."""
+
+    # 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
+    """Enable calibrated FlashRT FP8 kernels for Gemma and SigLIP MLPs.
+
+    Apply after loading with ``PI052Policy.apply_flashrt_fp8_mlp``; unavailable kernels keep BF16.
+    """
+
+    # 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.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.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,256 @@
+# 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 logging
+import os
+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_local_leader() -> bool:
+    return int(os.environ.get("LOCAL_RANK", "0")) == 0
+
+
+def _dataset_signature(
+    dataset_repo_id: str,
+    base_tokenizer_name: str,
+    n_samples: int,
+    chunk_size: int,
+) -> str:
+    """Deterministic short hash for naming the cache directory.
+
+    Keys on (dataset, base tokenizer, sample count, chunk size) so any
+    of those changing re-runs the fit. ``chunk_size`` matters because
+    the tokenizer is fit on chunks of that length.
+    """
+    h = hashlib.sha256()
+    h.update(dataset_repo_id.encode("utf-8"))
+    h.update(b"\0")
+    h.update(base_tokenizer_name.encode("utf-8"))
+    h.update(b"\0")
+    h.update(str(n_samples).encode("utf-8"))
+    h.update(b"\0")
+    h.update(str(chunk_size).encode("utf-8"))
+    return h.hexdigest()[:16]
+
+
+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,
+) -> 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)
+    sig = _dataset_signature(dataset_repo_id, base_tokenizer_name, n_samples, chunk_size)
+    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)
+
+    # Each node fits its node-local cache once; its other local ranks wait.
+    is_leader = _is_local_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,
+    )
+
+    from transformers import AutoProcessor  # noqa: PLC0415
+
+    # 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
+    from huggingface_hub import snapshot_download  # noqa: PLC0415
+
+    snap = Path(snapshot_download(repo_id=dataset_repo_id, repo_type="dataset"))
+    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}.")
+
+    # Load only episode indices and fixed-width actions across all shards.
+    tables = [_pq.read_table(f, columns=["episode_index", "action"]) 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}.")
+
+    # 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]
+
+    samples_per_episode = max(1, n_samples // max(num_episodes, 1))
+    collected = 0
+    eps_visited = 0
+    short_episodes = 0
+    ep_indices = list(ep_to_slice.keys())
+    for ep_idx in rng.permutation(ep_indices):
+        if collected >= n_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])
+            collected += 1
+            if collected >= n_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)
+    logger.info(
+        "FAST fit: collected %d chunks of shape %s from %d episodes",
+        actions.shape[0],
+        actions.shape[1:],
+        eps_visited,
+    )
+
+    # Match training-time quantile normalization so FAST sees the same bounded action space.
+    flat = actions.reshape(-1, actions.shape[-1])
+    q01 = np.quantile(flat, 0.01, axis=0)
+    q99 = np.quantile(flat, 0.99, axis=0)
+    span = np.where((q99 - q01) > 1e-6, q99 - q01, 1.0)
+    actions = np.clip((actions - q01) / span * 2.0 - 1.0, -1.0, 1.0).astype(np.float32)
+
+    base = AutoProcessor.from_pretrained(base_tokenizer_name, trust_remote_code=True)
+    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."
+        )
+
+    fitted = base.fit(actions)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    fitted.save_pretrained(str(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) -> 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
+
+    return fit_fast_tokenizer(
+        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,
+    )

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,205 @@
+# 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.
+        content = subtask
+        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
+            content = f"{subtask}, State: {discretize_state_str(state_row)};"
+
+        text_batch = _build_text_batch(
+            self.policy,
+            [{"role": "user", "content": content}],
+            add_generation_prompt=False,
+        )
+        batch = dict(observation)
+        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)
+        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:
+        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,149 @@
+# 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,
+) -> 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:
+        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),
+                max_action_tokens=config.max_action_tokens,
+                fast_skip_tokens=config.fast_skip_tokens,
+                paligemma_tokenizer_name="google/paligemma-3b-pt-224",
+            )
+        )
+
+    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,483 @@
+# 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
+
+
+@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,6 +61,9 @@ 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

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -101,6 +101,7 @@ 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,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -143,6 +144,10 @@ 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)
+
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
         AddBatchDimensionProcessorStep(),
@@ -160,7 +165,7 @@ 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,

src/lerobot/policies/pi_gemma.py

@@ -14,18 +14,27 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import torch
 from torch import nn
+from torch.nn import functional as F  # noqa: N812
 
 from lerobot.utils.import_utils import _transformers_available
 
+# Default PaliGemma SigLIP input resolution. Mirrors
+# ``pi05.configuration_pi05.DEFAULT_IMAGE_SIZE``; duplicated as a plain constant
+# to avoid importing the pi05 package here (which would create an import cycle:
+# pi_gemma -> pi05.__init__ -> modeling_pi05 -> pi_gemma).
+DEFAULT_IMAGE_SIZE = 224
+
 if TYPE_CHECKING or _transformers_available:
     from transformers.cache_utils import DynamicCache
     from transformers.masking_utils import create_causal_mask
     from transformers.modeling_layers import GradientCheckpointingLayer
     from transformers.modeling_outputs import BaseModelOutputWithPast
+    from transformers.models.auto import CONFIG_MAPPING
+    from transformers.models.gemma import modeling_gemma
     from transformers.models.gemma.modeling_gemma import (
         GemmaAttention,
         GemmaConfig,
@@ -49,6 +58,8 @@ else:
     GradientCheckpointingLayer = None
     BaseModelOutputWithPast = None
     create_causal_mask = None
+    CONFIG_MAPPING = None
+    modeling_gemma = None
 
 
 def _gated_residual(
@@ -121,7 +132,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 +289,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 +383,374 @@ __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
+
+
+# Define the complete layer computation function for gradient checkpointing
+def compute_layer_complete(
+    layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond, paligemma, gemma_expert
+):
+    models = [paligemma.model.language_model, gemma_expert.model]
+    query_states = []
+    key_states = []
+    value_states = []
+    gates = []
+    for i, hidden_states in enumerate(inputs_embeds):
+        layer = models[i].layers[layer_idx]
+        hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
+        gates.append(gate)
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+        query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        query_states.append(query_state)
+        key_states.append(key_state)
+        value_states.append(value_state)
+    # Concatenate and process attention
+    query_states = torch.cat(query_states, dim=2)
+    key_states = torch.cat(key_states, dim=2)
+    value_states = torch.cat(value_states, dim=2)
+    dummy_tensor = torch.zeros(
+        query_states.shape[0],
+        query_states.shape[2],
+        query_states.shape[-1],
+        device=query_states.device,
+        dtype=query_states.dtype,
+    )
+    cos, sin = paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
+    query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
+        query_states, key_states, cos, sin, unsqueeze_dim=1
+    )
+    batch_size = query_states.shape[0]
+    scaling = paligemma.model.language_model.layers[layer_idx].self_attn.scaling
+    att_output, _ = sdpa_attention_forward(
+        paligemma.model.language_model.layers[layer_idx].self_attn,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        scaling,
+    )
+    # Get head_dim from the current layer, not from the model
+    head_dim = paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
+    att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
+    # Process layer outputs
+    outputs_embeds = []
+    start_pos = 0
+    for i, hidden_states in enumerate(inputs_embeds):
+        layer = models[i].layers[layer_idx]
+        end_pos = start_pos + hidden_states.shape[1]
+        if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
+            att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
+        out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos])
+        # first residual
+        out_emb = _gated_residual(hidden_states, out_emb, gates[i])
+        after_first_residual = out_emb.clone()
+        out_emb, gate = layernorm_forward(layer.post_attention_layernorm, out_emb, adarms_cond[i])
+        # Convert to bfloat16 if the next layer (mlp) uses bfloat16
+        if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
+            out_emb = out_emb.to(dtype=torch.bfloat16)
+        out_emb = layer.mlp(out_emb)
+        # second residual
+        out_emb = _gated_residual(after_first_residual, out_emb, gate)
+        outputs_embeds.append(out_emb)
+        start_pos = end_pos
+    return outputs_embeds
+
+
+class GemmaVariantConfig:  # see openpi `gemma.py: Config`
+    """Configuration for Gemma model variants."""
+
+    def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
+        self.width = width
+        self.depth = depth
+        self.mlp_dim = mlp_dim
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+
+
+def get_gemma_config(variant: str) -> GemmaVariantConfig:  # see openpi `gemma.py: get_config`
+    """Returns config for specified gemma variant."""
+    if variant == "gemma_300m":
+        return GemmaVariantConfig(
+            width=1024,
+            depth=18,
+            mlp_dim=4096,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+        )
+    elif variant == "gemma_2b":
+        return GemmaVariantConfig(
+            width=2048,
+            depth=18,
+            mlp_dim=16_384,
+            num_heads=8,
+            num_kv_heads=1,
+            head_dim=256,
+        )
+    else:
+        raise ValueError(f"Unknown variant: {variant}")
+
+
+class PaliGemmaWithExpertModel(
+    nn.Module
+):  # see openpi `gemma_pytorch.py: PaliGemmaWithExpertModel` this class is almost a exact copy of PaliGemmaWithExpertModel in openpi
+    """PaliGemma model with action expert for PI05."""
+
+    def __init__(
+        self,
+        vlm_config,
+        action_expert_config,
+        use_adarms=None,
+        precision: Literal["bfloat16", "float32"] = "bfloat16",
+        image_size: int = DEFAULT_IMAGE_SIZE,
+        freeze_vision_encoder: bool = False,
+        train_expert_only: bool = False,
+    ):
+        if use_adarms is None:
+            use_adarms = [False, False]
+        super().__init__()
+        self.freeze_vision_encoder = freeze_vision_encoder
+        self.train_expert_only = train_expert_only
+
+        vlm_config_hf = CONFIG_MAPPING["paligemma"]()
+        vlm_config_hf._vocab_size = 257152  # noqa: SLF001
+        vlm_config_hf.image_token_index = 257152
+        vlm_config_hf.text_config.hidden_size = vlm_config.width
+        vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
+        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
+        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
+        vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
+        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
+        vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
+        vlm_config_hf.text_config.dtype = "float32"
+        vlm_config_hf.text_config.vocab_size = 257152
+        vlm_config_hf.text_config.use_adarms = use_adarms[0]
+        vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None
+        vlm_config_hf.vision_config.image_size = image_size
+        vlm_config_hf.vision_config.intermediate_size = 4304
+        vlm_config_hf.vision_config.projection_dim = 2048
+        vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
+        vlm_config_hf.vision_config.dtype = "float32"
+
+        action_expert_config_hf = CONFIG_MAPPING["gemma"](
+            head_dim=action_expert_config.head_dim,
+            hidden_size=action_expert_config.width,
+            intermediate_size=action_expert_config.mlp_dim,
+            num_attention_heads=action_expert_config.num_heads,
+            num_hidden_layers=action_expert_config.depth,
+            num_key_value_heads=action_expert_config.num_kv_heads,
+            vocab_size=257152,
+            hidden_activation="gelu_pytorch_tanh",
+            dtype="float32",
+            use_adarms=use_adarms[1],
+            adarms_cond_dim=action_expert_config.width if use_adarms[1] else None,
+        )
+
+        self.paligemma = PaliGemmaForConditionalGenerationWithPiGemma(config=vlm_config_hf)
+        self.gemma_expert = PiGemmaForCausalLM(config=action_expert_config_hf)
+        self.gemma_expert.model.embed_tokens = None
+
+        self.to_bfloat16_for_selected_params(precision)
+        self._set_requires_grad()
+
+    def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"):
+        if precision == "bfloat16":
+            self.to(dtype=torch.bfloat16)
+        elif precision == "float32":
+            self.to(dtype=torch.float32)
+            return
+        else:
+            raise ValueError(f"Invalid precision: {precision}")
+
+        # Keep full vision path in float32 so we never toggle (toggle causes optimizer
+        # "same dtype" error). Saves memory vs full float32; more memory than only 3 params.
+        params_to_keep_float32 = [
+            "vision_tower",
+            "multi_modal_projector",
+            "lm_head",
+            "input_layernorm",
+            "post_attention_layernorm",
+            "model.norm",
+        ]
+
+        for name, param in self.named_parameters():
+            if any(selector in name for selector in params_to_keep_float32):
+                param.data = param.data.to(dtype=torch.float32)
+
+    def _set_requires_grad(self):
+        if self.freeze_vision_encoder:
+            self.paligemma.model.vision_tower.eval()
+            for param in self.paligemma.model.vision_tower.parameters():
+                param.requires_grad = False
+        if self.train_expert_only:
+            self.paligemma.eval()
+            for param in self.paligemma.parameters():
+                param.requires_grad = False
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        if self.freeze_vision_encoder:
+            self.paligemma.model.vision_tower.eval()
+        if self.train_expert_only:
+            self.paligemma.eval()
+
+    def embed_image(self, image: torch.Tensor):
+        # Vision tower and multi_modal_projector are kept in float32 (params_to_keep_float32).
+        out_dtype = image.dtype
+        if image.dtype != torch.float32:
+            image = image.to(torch.float32)
+        image_outputs = self.paligemma.model.get_image_features(image)
+        # OpenPI / big_vision convention: image (soft) tokens are NOT scaled by the
+        # Gemma embedder normalizer (sqrt(hidden_size)) — only text tokens are. lerobot/pi05_base
+        # was trained in this regime, so scaling image features here over-scales them ~45x and
+        # breaks the pretrained vision-language alignment. Keep image features un-normalized.
+        features = image_outputs.pooler_output
+        if features.dtype != out_dtype:
+            features = features.to(out_dtype)
+        return features
+
+    def embed_language_tokens(self, tokens: torch.Tensor):
+        return self.paligemma.model.language_model.embed_tokens(tokens)
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: list[torch.FloatTensor] | None = None,
+        inputs_embeds: list[torch.FloatTensor] | None = None,
+        use_cache: bool | None = None,
+        adarms_cond: list[torch.Tensor] | None = None,
+    ):
+        if adarms_cond is None:
+            adarms_cond = [None, None]
+        if inputs_embeds[1] is None:
+            prefix_output = self.paligemma.model.language_model.forward(
+                inputs_embeds=inputs_embeds[0],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                adarms_cond=adarms_cond[0] if adarms_cond is not None else None,
+            )
+            prefix_past_key_values = prefix_output.past_key_values
+            prefix_output = prefix_output.last_hidden_state
+            suffix_output = None
+        elif inputs_embeds[0] is None:
+            suffix_output = self.gemma_expert.model.forward(
+                inputs_embeds=inputs_embeds[1],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                adarms_cond=adarms_cond[1] if adarms_cond is not None else None,
+            )
+            suffix_output = suffix_output.last_hidden_state
+            prefix_output = None
+            prefix_past_key_values = None
+        else:
+            models = [self.paligemma.model.language_model, self.gemma_expert.model]
+            num_layers = self.paligemma.config.text_config.num_hidden_layers
+
+            # Check if gradient checkpointing is enabled for any of the models
+            use_gradient_checkpointing = (
+                hasattr(self.gemma_expert.model, "gradient_checkpointing")
+                and self.gemma_expert.model.gradient_checkpointing
+                and self.training
+            ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
+
+            # Process all layers with gradient checkpointing if enabled
+            for layer_idx in range(num_layers):
+                if use_gradient_checkpointing:
+                    inputs_embeds = torch.utils.checkpoint.checkpoint(
+                        compute_layer_complete,
+                        layer_idx,
+                        inputs_embeds,
+                        attention_mask,
+                        position_ids,
+                        adarms_cond,
+                        use_reentrant=False,
+                        preserve_rng_state=False,
+                        paligemma=self.paligemma,
+                        gemma_expert=self.gemma_expert,
+                    )
+                else:
+                    inputs_embeds = compute_layer_complete(
+                        layer_idx,
+                        inputs_embeds,
+                        attention_mask,
+                        position_ids,
+                        adarms_cond,
+                        paligemma=self.paligemma,
+                        gemma_expert=self.gemma_expert,
+                    )
+
+            # final norm
+            def compute_final_norms(inputs_embeds, adarms_cond):
+                outputs_embeds = []
+                for i, hidden_states in enumerate(inputs_embeds):
+                    out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
+                    outputs_embeds.append(out_emb)
+                return outputs_embeds
+
+            # Apply gradient checkpointing to final norm if enabled
+            if use_gradient_checkpointing:
+                outputs_embeds = torch.utils.checkpoint.checkpoint(
+                    compute_final_norms,
+                    inputs_embeds,
+                    adarms_cond,
+                    use_reentrant=False,
+                    preserve_rng_state=False,
+                )
+            else:
+                outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond)
+
+            prefix_output = outputs_embeds[0]
+            suffix_output = outputs_embeds[1]
+            prefix_past_key_values = None
+
+        return [prefix_output, suffix_output], prefix_past_key_values

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:

src/lerobot/policies/rtc/modeling_rtc.py

@@ -42,7 +42,12 @@ class RTCProcessor:
     prefix attention, and adaptive chunk processing.
     """
 
-    def __init__(self, rtc_config: RTCConfig):
+    def __init__(self, rtc_config: RTCConfig, *, trained_mode_supported: bool = False):
+        if rtc_config.enabled and rtc_config.mode == "trained" and not trained_mode_supported:
+            raise ValueError(
+                "RTC mode='trained' requires a PI05-compatible checkpoint trained with "
+                "rtc_training_max_delay > 0."
+            )
         self.rtc_config = rtc_config
 
         self.tracker = None

src/lerobot/processor/batch_processor.py

@@ -175,9 +175,6 @@ class AddBatchDimensionComplementaryDataStep(ComplementaryDataProcessorStep):
             if isinstance(task_index_value, Tensor) and task_index_value.dim() == 0:
                 complementary_data["task_index"] = task_index_value.unsqueeze(0)
 
-        complementary_data.pop("language_persistent", None)
-        complementary_data.pop("language_events", None)
-
         if "messages" in complementary_data:
             messages = complementary_data["messages"]
             if isinstance(messages, list) and (not messages or isinstance(messages[0], dict)):

src/lerobot/processor/render_messages_processor.py

@@ -16,7 +16,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import asdict, dataclass
 from typing import Any
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -32,17 +32,18 @@ from .pipeline import ProcessorStep, ProcessorStepRegistry
 @dataclass
 @ProcessorStepRegistry.register(name="render_messages_processor")
 class RenderMessagesStep(ProcessorStep):
-    """Processor step that turns raw language columns into rendered chat messages.
-
-    Reads ``language_persistent`` and ``language_events`` from the transition's
-    complementary data, renders them through ``recipe`` at the sample timestamp,
-    and replaces the raw columns with the resulting ``messages`` /
-    ``message_streams`` / ``target_message_indices`` keys.
-    """
+    """Render language columns into recipe-defined messages and supervision metadata."""
 
     recipe: TrainingRecipe
     dataset_ctx: Any | None = None
 
+    def __post_init__(self) -> None:
+        if isinstance(self.recipe, dict):
+            self.recipe = TrainingRecipe.from_dict(self.recipe)
+
+    def get_config(self) -> dict[str, Any]:
+        return {"recipe": asdict(self.recipe)}
+
     def __call__(self, transition: EnvTransition) -> EnvTransition | None:
         """Render messages for a single transition; return ``None`` to drop it."""
         complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
@@ -50,7 +51,17 @@ class RenderMessagesStep(ProcessorStep):
         events = complementary_data.get(LANGUAGE_EVENTS) or []
 
         if not persistent and not events:
-            return transition
+            rendered = _fallback_low_level_render(complementary_data.get("task"))
+            if rendered is None:
+                return transition
+            new_transition = transition.copy()
+            new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
+            new_complementary_data.update(rendered)
+            new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
+            return new_transition
+
+        if _is_batched_language(persistent) or _is_batched_language(events):
+            return self._call_batch(transition, complementary_data, persistent, events)
 
         timestamp = complementary_data.get("timestamp")
         if timestamp is None:
@@ -67,18 +78,147 @@ class RenderMessagesStep(ProcessorStep):
             dataset_ctx=self.dataset_ctx,
         )
         if rendered is None:
-            return None
+            rendered = _fallback_low_level_render(complementary_data.get("task"))
+            if rendered is None:
+                return None
 
         new_transition = transition.copy()
-        new_complementary_data = dict(complementary_data)
+        new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
         new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
         new_complementary_data.pop(LANGUAGE_EVENTS, None)
         new_complementary_data.update(rendered)
         new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
         return new_transition
 
+    def _call_batch(
+        self,
+        transition: EnvTransition,
+        complementary_data: dict[str, Any],
+        persistent_batch: list,
+        events_batch: list,
+    ) -> EnvTransition | None:
+        timestamp = complementary_data.get("timestamp")
+        if timestamp is None:
+            raise KeyError("RenderMessagesStep requires sample timestamp in complementary data.")
+
+        batch_size = max(len(persistent_batch), len(events_batch))
+        messages: list[list[dict[str, Any]]] = []
+        message_streams: list[list[str | None]] = []
+        target_message_indices: list[list[int]] = []
+        keep_indices: list[int] = []
+
+        for i in range(batch_size):
+            rendered = render_sample(
+                recipe=self.recipe,
+                persistent=persistent_batch[i] if i < len(persistent_batch) else [],
+                events=events_batch[i] if i < len(events_batch) else [],
+                t=_batch_value(timestamp, i),
+                sample_idx=int(_batch_value(complementary_data.get("index", 0), i)),
+                task=_batch_value(complementary_data.get("task"), i),
+                dataset_ctx=self.dataset_ctx,
+            )
+            if rendered is None:
+                rendered = _fallback_low_level_render(_batch_value(complementary_data.get("task"), i))
+                if rendered is None:
+                    continue
+            keep_indices.append(i)
+            messages.append(rendered["messages"])
+            message_streams.append(rendered["message_streams"])
+            target_message_indices.append(rendered["target_message_indices"])
+
+        if not messages:
+            return None
+
+        new_transition = (
+            _select_batch_indices(transition, keep_indices)
+            if len(keep_indices) != batch_size
+            else transition.copy()
+        )
+        new_complementary_data = dict(new_transition.get(TransitionKey.COMPLEMENTARY_DATA) or {})
+        new_complementary_data.pop(LANGUAGE_PERSISTENT, None)
+        new_complementary_data.pop(LANGUAGE_EVENTS, None)
+        new_complementary_data["messages"] = messages
+        new_complementary_data["message_streams"] = message_streams
+        new_complementary_data["target_message_indices"] = target_message_indices
+        new_transition[TransitionKey.COMPLEMENTARY_DATA] = new_complementary_data
+        return new_transition
+
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
         """Pass features through unchanged; rendering only touches complementary data."""
         return features
+
+
+def _scalar(value: Any) -> float | int:
+    """Unwrap a tensor/array/single-element list into a Python scalar."""
+    if hasattr(value, "item"):
+        return value.item()
+    if isinstance(value, list):
+        if len(value) != 1:
+            raise ValueError(f"Expected a scalar, got list of length {len(value)}: {value!r}")
+        return _scalar(value[0])
+    return value
+
+
+def _is_batched_language(value: Any) -> bool:
+    return isinstance(value, list) and bool(value) and isinstance(value[0], list)
+
+
+def _batch_value(value: Any, index: int) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, list):
+        return value[index]
+    if hasattr(value, "ndim") and value.ndim > 0:
+        return _scalar(value[index])
+    return _scalar(value)
+
+
+def _select_batch_indices(transition: EnvTransition, indices: list[int]) -> EnvTransition:
+    selected = transition.copy()
+    for key in (TransitionKey.OBSERVATION, TransitionKey.COMPLEMENTARY_DATA):
+        data = selected.get(key)
+        if isinstance(data, dict):
+            selected[key] = {k: _select_value(v, indices) for k, v in data.items()}
+    action = selected.get(TransitionKey.ACTION)
+    if action is not None:
+        selected[TransitionKey.ACTION] = _select_value(action, indices)
+    return selected
+
+
+def _select_value(value: Any, indices: list[int]) -> Any:
+    if isinstance(value, list) and len(value) >= len(indices):
+        return [value[i] for i in indices]
+    if hasattr(value, "index_select") and hasattr(value, "new_tensor") and getattr(value, "ndim", 0) > 0:
+        return value.index_select(0, value.new_tensor(indices).long())
+    return value
+
+
+def _fallback_low_level_render(task: Any) -> dict[str, Any] | None:
+    """Keep action-only samples trainable when no recipe branch matches."""
+    if hasattr(task, "item"):
+        task = task.item()
+    if isinstance(task, list):
+        messages = []
+        message_streams = []
+        target_message_indices = []
+        for t in task:
+            rendered = _fallback_low_level_render(t)
+            if rendered is None:
+                return None
+            messages.append(rendered["messages"])
+            message_streams.append(rendered["message_streams"])
+            target_message_indices.append(rendered["target_message_indices"])
+        return {
+            "messages": messages,
+            "message_streams": message_streams,
+            "target_message_indices": target_message_indices,
+        }
+    if not isinstance(task, str) or not task:
+        return None
+    return {
+        "messages": [{"role": "user", "content": task}],
+        "message_streams": ["low_level"],
+        "target_message_indices": [],
+    }

src/lerobot/processor/tokenizer_processor.py

@@ -32,6 +32,7 @@ import torch
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.types import EnvTransition, RobotObservation, TransitionKey
 from lerobot.utils.constants import (
+    ACTION_CODE_TOKEN_MASK,
     ACTION_TOKEN_MASK,
     ACTION_TOKENS,
     OBS_LANGUAGE_ATTENTION_MASK,
@@ -412,14 +413,15 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             # During inference, no action is available, skip tokenization
             return new_transition
 
-        # Tokenize and get both tokens and mask
-        tokens, mask = self._tokenize_action(action)
+        # Tokenize and get masks for the full formatted sequence and the discrete action codes.
+        tokens, mask, code_mask = self._tokenize_action(action)
 
         # Store mask in complementary data
         complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
         if complementary_data is None:
             complementary_data = {}
         complementary_data[ACTION_TOKEN_MASK] = mask
+        complementary_data[ACTION_CODE_TOKEN_MASK] = code_mask
         complementary_data[ACTION_TOKENS] = tokens
         new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
         return new_transition
@@ -430,7 +432,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
         """
         return self._paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens - tokens
 
-    def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """
         Tokenizes the action tensor and creates a mask.
 
@@ -459,6 +461,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
         # The fast tokenizer expects action data and returns token IDs
         tokens_list = []
         masks_list = []
+        code_masks_list = []
 
         for i in range(batch_size):
             # Tokenize single action (move to CPU first as tokenizer uses scipy which requires numpy)
@@ -476,19 +479,26 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             if tokens.dim() > 1:
                 tokens = tokens.flatten()
 
+            action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens)
             bos_id = self._paligemma_tokenizer.bos_token_id
-            # add bos
+            prompt_tokens = torch.tensor(
+                self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
+                device=action.device,
+            )
+            end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device)
+
+            code_start = 1 + len(prompt_tokens)
+            code_end = code_start + len(action_code_tokens)
             tokens = torch.cat(
                 [
                     torch.tensor([bos_id], device=action.device),
-                    torch.tensor(
-                        self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
-                        device=action.device,
-                    ),
-                    self._act_tokens_to_paligemma_tokens(tokens),
-                    torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device),
+                    prompt_tokens,
+                    action_code_tokens,
+                    end_tokens,
                 ]
             )
+            code_mask = torch.zeros(len(tokens), dtype=torch.bool, device=action.device)
+            code_mask[code_start:code_end] = True
 
             # Truncate or pad to max_action_tokens
             if len(tokens) > self.max_action_tokens:
@@ -497,44 +507,49 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
                     "Consider increasing the `max_action_tokens` in your model config if this happens frequently."
                 )
                 tokens = tokens[: self.max_action_tokens]
+                code_mask = code_mask[: self.max_action_tokens]
                 mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device)
             else:
+                pad_len = self.max_action_tokens - len(tokens)
                 mask = torch.cat(
                     [
                         torch.ones(len(tokens), dtype=torch.bool, device=action.device),
-                        torch.zeros(
-                            self.max_action_tokens - len(tokens), dtype=torch.bool, device=action.device
-                        ),
+                        torch.zeros(pad_len, dtype=torch.bool, device=action.device),
                     ]
                 )
+                code_mask = torch.nn.functional.pad(code_mask, (0, pad_len), value=False)
                 # Pad tokens with zeros
-                tokens = torch.nn.functional.pad(tokens, (0, self.max_action_tokens - len(tokens)), value=0)
+                tokens = torch.nn.functional.pad(tokens, (0, pad_len), value=0)
 
             tokens_list.append(tokens)
             masks_list.append(mask)
+            code_masks_list.append(code_mask)
 
         # Stack into batched tensors
         tokens_batch = torch.stack(tokens_list, dim=0)  # (B, max_action_tokens)
         masks_batch = torch.stack(masks_list, dim=0)  # (B, max_action_tokens)
+        code_masks_batch = torch.stack(code_masks_list, dim=0)  # (B, max_action_tokens)
 
         # Remove batch dimension if input was single sample
         if single_sample:
             tokens_batch = tokens_batch.squeeze(0)
             masks_batch = masks_batch.squeeze(0)
+            code_masks_batch = code_masks_batch.squeeze(0)
 
         # Move to the same device as the input
         if device is not None:
             tokens_batch = tokens_batch.to(device)
             masks_batch = masks_batch.to(device)
+            code_masks_batch = code_masks_batch.to(device)
 
-        return tokens_batch, masks_batch
+        return tokens_batch, masks_batch, code_masks_batch
 
     def action(self, action: torch.Tensor) -> torch.Tensor:
         """
         This method is not used since we override __call__.
         Required by ActionProcessorStep ABC.
         """
-        tokens, _ = self._tokenize_action(action)
+        tokens, _, _ = self._tokenize_action(action)
         return tokens
 
     def get_config(self) -> dict[str, Any]:
@@ -550,6 +565,8 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
         config = {
             "trust_remote_code": self.trust_remote_code,
             "max_action_tokens": self.max_action_tokens,
+            "fast_skip_tokens": self.fast_skip_tokens,
+            "paligemma_tokenizer_name": self.paligemma_tokenizer_name,
         }
 
         # Only save tokenizer_name if it was used to create the tokenizer

src/lerobot/robots/utils.py

@@ -21,6 +21,8 @@ from lerobot.utils.import_utils import make_device_from_device_class
 from .config import RobotConfig
 from .robot import Robot
 
+logger = logging.getLogger(__name__)
+
 
 def make_robot_from_config(config: RobotConfig) -> Robot:
     # TODO(Steven): Consider just using the make_device_from_device_class for all types
@@ -118,7 +120,7 @@ def ensure_safe_goal_position(
             }
 
     if warnings_dict:
-        logging.warning(
+        logger.warning(
             "Relative goal position magnitude had to be clamped to be safe.\n"
             f"{pformat(warnings_dict, indent=4)}"
         )

src/lerobot/rollout/context.py

@@ -46,6 +46,7 @@ from lerobot.processor import (
 from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
 from lerobot.robots import make_robot_from_config
 from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
+from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
 
 from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
@@ -60,6 +61,35 @@ from .robot_wrapper import ThreadSafeRobot
 logger = logging.getLogger(__name__)
 
 
+def _validate_trained_rtc_rollout_config(policy_config, inference_config: RTCInferenceConfig) -> None:
+    """Fail fast when rollout cannot retain every trained RTC prefix."""
+    rtc = inference_config.rtc
+    if not rtc.enabled or rtc.mode != "trained":
+        return
+    if policy_config.type not in {"pi05", "pi052"}:
+        raise ValueError(
+            "--inference.rtc.mode=trained currently requires a PI05-compatible checkpoint; "
+            f"got policy type {policy_config.type!r}."
+        )
+
+    training_max_delay = int(getattr(policy_config, "rtc_training_max_delay", 0))
+    if training_max_delay <= 0:
+        raise ValueError(
+            "--inference.rtc.mode=trained requires a checkpoint trained with "
+            "--policy.rtc_training_max_delay > 0."
+        )
+    if rtc.execution_horizon < training_max_delay:
+        raise ValueError(
+            f"--inference.rtc.execution_horizon ({rtc.execution_horizon}) must be at least the "
+            f"checkpoint's rtc_training_max_delay ({training_max_delay})."
+        )
+    if inference_config.queue_threshold < training_max_delay:
+        raise ValueError(
+            f"--inference.queue_threshold ({inference_config.queue_threshold}) must be at least the "
+            f"checkpoint's rtc_training_max_delay ({training_max_delay})."
+        )
+
+
 def _resolve_action_key_order(
     policy_action_names: list[str] | None, dataset_action_names: list[str]
 ) -> list[str]:
@@ -80,6 +110,26 @@ def _resolve_action_key_order(
     return policy_action_names
 
 
+def _align_relative_state_feature_order(
+    hw_features: dict[str, dict], policy_action_names: list[str] | None
+) -> dict[str, dict]:
+    """Align policy-facing state with named relative-action dimensions."""
+    if not policy_action_names or OBS_STATE not in hw_features:
+        return hw_features
+
+    state_feature = hw_features[OBS_STATE]
+    state_names = state_feature.get("names")
+    if not state_names or len(state_names) != len(policy_action_names):
+        return hw_features
+    if set(state_names) != set(policy_action_names) or state_names == policy_action_names:
+        return hw_features
+
+    aligned = dict(hw_features)
+    aligned[OBS_STATE] = {**state_feature, "names": list(policy_action_names)}
+    logger.info("Aligned relative-action state order with checkpoint action names")
+    return aligned
+
+
 # ---------------------------------------------------------------------------
 # Sub-contexts
 # ---------------------------------------------------------------------------
@@ -178,6 +228,9 @@ def build_rollout_context(
     policy_config = cfg.policy
     policy_class = get_policy_class(policy_config.type)
 
+    if is_rtc:
+        _validate_trained_rtc_rollout_config(policy_config, cfg.inference)
+
     if hasattr(policy_config, "compile_model"):
         policy_config.compile_model = cfg.use_torch_compile
 
@@ -399,10 +452,22 @@ def build_rollout_context(
         },
     )
 
-    if isinstance(cfg.inference, SyncInferenceConfig) and any(
-        isinstance(step, RelativeActionsProcessorStep) and step.enabled
-        for step in getattr(preprocessor, "steps", ())
-    ):
+    relative_action_step = next(
+        (
+            step
+            for step in getattr(preprocessor, "steps", ())
+            if isinstance(step, RelativeActionsProcessorStep) and step.enabled
+        ),
+        None,
+    )
+    if relative_action_step is not None:
+        relative_action_names = relative_action_step.action_names or policy_action_names
+        hw_features = _align_relative_state_feature_order(
+            hw_features,
+            list(relative_action_names) if relative_action_names else None,
+        )
+
+    if isinstance(cfg.inference, SyncInferenceConfig) and relative_action_step is not None:
         raise NotImplementedError(
             "SyncInferenceEngine does not support policies with relative actions for now."
             "Use --inference.type=rtc or remove relative action processor steps from the policy pipeline."

src/lerobot/rollout/inference/rtc.py

@@ -57,6 +57,18 @@ _RTC_MAX_CONSECUTIVE_ERRORS: int = 10
 _RTC_JOIN_TIMEOUT_S: float = 3.0
 
 
+class _FatalRTCInferenceError(RuntimeError):
+    """Base class for RTC errors that cannot become valid after a retry."""
+
+
+class _TrainedRTCDelayExceededError(_FatalRTCInferenceError):
+    """Raised when measured latency exceeds a trained RTC checkpoint's support."""
+
+
+class _TrainedRTCPrefixUnavailableError(_FatalRTCInferenceError):
+    """Raised when the queue cannot provide the prefix used for conditioning."""
+
+
 # ---------------------------------------------------------------------------
 # RTC helpers
 # ---------------------------------------------------------------------------
@@ -76,6 +88,50 @@ def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int
     return padded
 
 
+def _trained_rtc_chunk_can_merge(
+    *,
+    conditioned_delay: int,
+    measured_delay: int,
+    training_max_delay: int,
+    has_previous_actions: bool,
+) -> bool:
+    """Check that a trained RTC chunk covers the overlap observed during inference."""
+    if not has_previous_actions:
+        return True
+    if measured_delay > training_max_delay:
+        raise _TrainedRTCDelayExceededError(
+            f"Measured RTC inference delay ({measured_delay}) exceeds the checkpoint's "
+            f"rtc_training_max_delay ({training_max_delay})."
+        )
+    return measured_delay <= conditioned_delay
+
+
+def _estimate_rtc_delay(
+    *,
+    latency: float,
+    time_per_step: float,
+    mode: str,
+    training_max_delay: int,
+    has_previous_actions: bool,
+) -> int:
+    """Estimate overlap, using the trained capacity to bootstrap the first transition."""
+    if latency:
+        return math.ceil(latency / time_per_step)
+    if mode == "trained" and has_previous_actions:
+        return training_max_delay
+    return 0
+
+
+def _validate_trained_rtc_prefix_available(*, conditioned_delay: int, available_steps: int) -> None:
+    """Reject hard-prefix inference when the real queue is shorter than its delay."""
+    if conditioned_delay > available_steps:
+        raise _TrainedRTCPrefixUnavailableError(
+            f"Trained RTC needs {conditioned_delay} committed prefix actions, but the queue has "
+            f"only {available_steps}. Increase --inference.queue_threshold and "
+            "--inference.rtc.execution_horizon."
+        )
+
+
 # ---------------------------------------------------------------------------
 # RTCInferenceEngine
 # ---------------------------------------------------------------------------
@@ -272,9 +328,23 @@ class RTCInferenceEngine(InferenceEngine):
                         current_time = time.perf_counter()
                         idx_before = queue.get_action_index()
                         prev_actions = queue.get_left_over()
+                        has_previous_actions = prev_actions is not None and prev_actions.numel() > 0
 
+                        training_max_delay = int(getattr(self._policy.config, "rtc_training_max_delay", 0))
                         latency = latency_tracker.max()
-                        delay = math.ceil(latency / time_per_chunk) if latency else 0
+                        delay = _estimate_rtc_delay(
+                            latency=latency,
+                            time_per_step=time_per_chunk,
+                            mode=self._rtc_config.mode,
+                            training_max_delay=training_max_delay,
+                            has_previous_actions=has_previous_actions,
+                        )
+                        if self._rtc_config.mode == "trained" and delay > 0:
+                            available_steps = 0 if prev_actions is None else prev_actions.shape[0]
+                            _validate_trained_rtc_prefix_available(
+                                conditioned_delay=delay,
+                                available_steps=available_steps,
+                            )
 
                         obs_batch = build_dataset_frame(self._hw_features, obs, prefix="observation")
                         obs_batch = prepare_observation_for_inference(
@@ -316,11 +386,32 @@ class RTCInferenceEngine(InferenceEngine):
                         inference_count += 1
                         consecutive_errors = 0
                         is_warmup = self._use_torch_compile and inference_count <= warmup_required
-                        if is_warmup:
+                        is_initial_trained_chunk = (
+                            self._rtc_config.mode == "trained" and not has_previous_actions
+                        )
+                        if is_warmup or is_initial_trained_chunk:
                             latency_tracker.reset()
                         else:
                             latency_tracker.add(new_latency)
 
+                        if (
+                            not is_warmup
+                            and self._rtc_config.mode == "trained"
+                            and not _trained_rtc_chunk_can_merge(
+                                conditioned_delay=delay,
+                                measured_delay=new_delay,
+                                training_max_delay=training_max_delay,
+                                has_previous_actions=has_previous_actions,
+                            )
+                        ):
+                            logger.warning(
+                                "Discarding trained RTC chunk: measured delay %d exceeded "
+                                "conditioned delay %d; retrying with updated latency",
+                                new_delay,
+                                delay,
+                            )
+                            continue
+
                         queue.merge(original, processed, new_delay, idx_before)
 
                         if (
@@ -333,6 +424,8 @@ class RTCInferenceEngine(InferenceEngine):
 
                         logger.debug("RTC inference latency=%.2fs, queue=%d", new_latency, queue.qsize())
 
+                    except _FatalRTCInferenceError:
+                        raise
                     except Exception as e:
                         consecutive_errors += 1
                         logger.error(

src/lerobot/runtime/__init__.py

@@ -0,0 +1,38 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Policy-agnostic runtime for language-conditioned policies.
+
+Adapters registered in :mod:`lerobot.runtime.registry` are served by ``lerobot-rollout --language``.
+"""
+
+from .adapter import BaseLanguageAdapter, GenerationConfig, LanguageDiagnostics
+from .language_runtime import (
+    LanguageConditionedPolicyAdapter,
+    LanguageConditionedRuntime,
+    RuntimeState,
+    Tick,
+    TickClock,
+)
+
+__all__ = [
+    "BaseLanguageAdapter",
+    "GenerationConfig",
+    "LanguageConditionedPolicyAdapter",
+    "LanguageConditionedRuntime",
+    "LanguageDiagnostics",
+    "RuntimeState",
+    "Tick",
+    "TickClock",
+]

src/lerobot/runtime/adapter.py

@@ -0,0 +1,165 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Policy adapters for the language runtime.
+
+The base adapter owns generation control and diagnostics while subclasses provide policy-specific actions and text.
+"""
+
+from __future__ import annotations
+
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+from .language_runtime import RuntimeState
+
+_SAY_RE = re.compile(r"<\s*say\s*>(.*?)<\s*/\s*say\s*>", re.IGNORECASE | re.DOTALL)
+
+
+@dataclass
+class GenerationConfig:
+    """Text-generation settings fixed for the adapter's lifetime."""
+
+    min_new_tokens: int = 0
+    temperature: float = 0.0
+    top_p: float = 1.0
+    chunks_per_regen: int = 1  # regenerate the language context every N action chunks
+    enable_memory: bool = True  # generate a running memory note on subtask change
+    enable_subtask: bool = True  # generate the low-level subtask (off => use the given text directly)
+
+
+@dataclass
+class LanguageDiagnostics:
+    """Runtime-panel generation counters keyed by text kind."""
+
+    last_raw: dict[str, str] = field(default_factory=dict)
+    empty: dict[str, int] = field(default_factory=dict)
+    repeat: int = 0
+
+    def _bump(self, table: dict[str, int], kind: str) -> int:
+        table[kind] = table.get(kind, 0) + 1
+        return table[kind]
+
+
+class BaseLanguageAdapter(ABC):
+    """Batteries-included adapter: generic high-level control, policy primitives abstract."""
+
+    def __init__(self, policy: Any, gen: GenerationConfig | None = None) -> None:
+        self.policy = policy
+        self.gen = gen or GenerationConfig()
+        self.diag = LanguageDiagnostics()
+        self._chunks_until_regen = 0
+
+    @abstractmethod
+    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
+        """Produce an action chunk from the observation + current language context."""
+
+    @abstractmethod
+    def generate_text(
+        self,
+        kind: str,
+        observation: dict[str, Any] | None,
+        state: RuntimeState,
+        user_text: str | None = None,
+    ) -> str:
+        """Generate one text stream (``kind``) and return the decoded string."""
+
+    def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
+        """Throttled regeneration of the language context (subtask / memory / ...)."""
+        if self._chunks_until_regen > 0:
+            self._chunks_until_regen -= 1
+            return
+        self._chunks_until_regen = max(1, self.gen.chunks_per_regen) - 1
+        self._regenerate_context(observation, state)
+
+    def handle_interjection(
+        self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
+    ) -> None:
+        """React to a mid-run user message by regenerating the plan."""
+        out = self.generate_text("interjection", observation, state, user_text=user_text)
+        plan = self.plan_from_text(out)
+        if plan:
+            state.set_context("plan", plan, label="plan")
+
+    def plan_from_text(self, text: str) -> str:
+        """Strip ``<say>`` speech markers from a generated plan."""
+        plan, _speech = split_plan_and_say(text)
+        return plan
+
+    def _regenerate_context(self, observation: dict[str, Any] | None, state: RuntimeState) -> None:
+        """Default hierarchy: regenerate the subtask, then memory when it changes.
+
+        Override for a policy with a different language hierarchy.
+        """
+        if not self.gen.enable_subtask:
+            # Preserve operator-provided subtasks in direct mode.
+            return
+        subtask = self._generate_filtered("subtask", observation, state)
+        if subtask is None:
+            return
+        previous = state.language_context.get("subtask")
+        if not state.set_context("subtask", subtask, label="subtask"):
+            self.diag.repeat += 1
+            return
+        self.diag.repeat = 0
+        if previous:
+            state.extra["prior_subtask"] = previous
+        if not self.gen.enable_memory:
+            return
+        memory = self._generate_filtered("memory", observation, state)
+        if memory is not None:
+            state.set_context("memory", memory, label="memory")
+
+    def _generate_filtered(
+        self, kind: str, observation: dict[str, Any] | None, state: RuntimeState
+    ) -> str | None:
+        """Generate one ``kind``, record diagnostics, and drop empty output."""
+        text = self.generate_text(kind, observation, state)
+        self.diag.last_raw[kind] = text or ""
+        if not text:
+            count = self.diag._bump(self.diag.empty, kind)
+            if count == 1 or count % 5 == 0:
+                state.log(f"  [info] {kind} gen returned empty (x{count})")
+            return None
+        return text
+
+
+class DirectTaskPolicyAdapter(BaseLanguageAdapter):
+    """Adapter for flat policies whose preprocessors condition actions on the operator's task."""
+
+    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any:
+        return self.policy.predict_action_chunk(observation)
+
+    def generate_text(
+        self,
+        kind: str,
+        observation: dict[str, Any] | None,
+        state: RuntimeState,
+        user_text: str | None = None,
+    ) -> str:
+        return ""
+
+
+def split_plan_and_say(text: str) -> tuple[str, str]:
+    """Split ``plan <say>speech</say>`` into ``(plan, speech)``."""
+    if not text:
+        return "", ""
+    match = _SAY_RE.search(text)
+    if not match:
+        return text.strip(), ""
+    speech = match.group(1).strip().strip('"').strip("'")
+    plan = (text[: match.start()] + text[match.end() :]).strip()
+    return plan, speech

src/lerobot/runtime/language_runtime.py

@@ -0,0 +1,349 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Small reusable runtime for language-conditioned robot policies."""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from collections import deque
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RuntimeState:
+    """Explicit state shared by the runtime and policy adapter."""
+
+    task: str = ""
+    language_context: dict[str, str] = field(default_factory=dict)
+    action_queue: deque[Any] = field(default_factory=deque)
+    events: set[str] = field(default_factory=set)
+    log_lines: list[str] = field(default_factory=list)
+    mode: str = "action"
+    stop: bool = False
+    tick: Tick | None = None
+    actions_dispatched: int = 0
+    action_deadline: float | None = None
+    extra: dict[str, Any] = field(default_factory=dict)
+    revision: int = 0
+    lock: Any = field(default_factory=threading.RLock, repr=False)
+
+    def emit(self, event_name: str) -> None:
+        self.events.add(event_name)
+
+    def take_event(self, event_name: str) -> bool:
+        if event_name not in self.events:
+            return False
+        self.events.remove(event_name)
+        return True
+
+    def log(self, line: str) -> None:
+        self.log_lines.append(line)
+
+    def set_context(self, key: str, value: str | None, *, label: str | None = None) -> bool:
+        with self.lock:
+            previous = self.language_context.get(key)
+            if previous == value:
+                return False
+            if value is None:
+                self.language_context.pop(key, None)
+            else:
+                self.language_context[key] = value
+            self.revision += 1
+            if label is not None and value:
+                self.log(f"  {label}: {value}")
+            return True
+
+    def get(self, key: str, default: Any = None) -> Any:
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def setdefault(self, key: str, default: Any = None) -> Any:
+        current = self.get(key, None)
+        if current is not None:
+            return current
+        self[key] = default
+        return default
+
+    def __getitem__(self, key: str) -> Any:
+        if hasattr(self, key):
+            return getattr(self, key)
+        if key in self.extra:
+            return self.extra[key]
+        raise KeyError(key)
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        with self.lock:
+            if hasattr(self, key):
+                if key == "mode" and self.mode != value:
+                    self.revision += 1
+                setattr(self, key, value)
+            else:
+                self.extra[key] = value
+
+
+class LanguageConditionedPolicyAdapter(Protocol):
+    """Runtime policy contract, implemented directly or through ``BaseLanguageAdapter``."""
+
+    def select_action(self, observation: dict[str, Any], state: RuntimeState) -> Any: ...
+
+    def update_language_state(self, observation: dict[str, Any] | None, state: RuntimeState) -> None: ...
+
+    def handle_interjection(
+        self, user_text: str, observation: dict[str, Any] | None, state: RuntimeState
+    ) -> None: ...
+
+
+@dataclass
+class Tick:
+    index: int
+    monotonic_seconds: float
+
+
+@dataclass
+class TickClock:
+    max_rate_hz: float = 50.0
+    _index: int = field(default=0, init=False)
+    _last_seconds: float | None = field(default=None, init=False)
+
+    def advance(self) -> Tick:
+        period = 1.0 / max(self.max_rate_hz, 0.1)
+        now = time.monotonic()
+        if self._last_seconds is not None:
+            sleep_for = (self._last_seconds + period) - now
+            if sleep_for > 0:
+                time.sleep(sleep_for)
+                now = time.monotonic()
+        self._last_seconds = now
+        self._index += 1
+        return Tick(index=self._index, monotonic_seconds=now)
+
+
+@dataclass
+class _RateGate:
+    hz: float
+    _last_seconds: float | None = None
+
+    def due(self, tick: Tick, *, force: bool = False) -> bool:
+        if force:
+            self._last_seconds = tick.monotonic_seconds
+            return True
+        period = 1.0 / max(self.hz, 1e-6)
+        if self._last_seconds is None or tick.monotonic_seconds - self._last_seconds >= period:
+            self._last_seconds = tick.monotonic_seconds
+            return True
+        return False
+
+    def rearm(self) -> None:
+        self._last_seconds = None
+
+
+@dataclass
+class LanguageConditionedRuntime:
+    """Generic tick loop for language-conditioned robot policies."""
+
+    policy_adapter: LanguageConditionedPolicyAdapter
+    observation_provider: Callable[[], dict[str, Any] | None] | None = None
+    action_executor: Callable[[Any], None] | None = None
+    event_collector: Callable[[RuntimeState], None] | None = None
+    chunk_hz: float = 4.0
+    ctrl_hz: float = 50.0
+    high_level_hz: float = 1.0
+    max_rate_hz: float = 50.0
+
+    state: RuntimeState = field(default_factory=RuntimeState)
+    _chunk_gate: _RateGate = field(init=False)
+    _ctrl_gate: _RateGate = field(init=False)
+    _language_gate: _RateGate = field(init=False)
+    _stop: bool = field(default=False, init=False)
+    _last_dispatch_seconds: float | None = field(default=None, init=False)
+
+    def __post_init__(self) -> None:
+        self._chunk_gate = _RateGate(self.chunk_hz)
+        self._ctrl_gate = _RateGate(self.ctrl_hz)
+        self._language_gate = _RateGate(self.high_level_hz)
+
+    @property
+    def policy(self) -> Any:
+        return getattr(self.policy_adapter, "policy", self.policy_adapter)
+
+    def set_task(self, task: str) -> None:
+        with self.state.lock:
+            if self.state.task != task:
+                self.state.revision += 1
+            self.state.task = task
+            self.state.log(f"Task: {task}")
+
+    def stop(self) -> None:
+        self._stop = True
+        self.state.stop = True
+
+    def run(self, *, max_ticks: int | None = None) -> None:
+        clock = TickClock(max_rate_hz=self.max_rate_hz)
+        while not self._stop:
+            tick = clock.advance()
+            self._run_tick(tick)
+            self._flush_logs()
+            if self.state.stop:
+                self._stop = True
+            if max_ticks is not None and tick.index >= max_ticks:
+                break
+        self._on_shutdown()
+
+    def step_once(self) -> list[str]:
+        previous = self.state.tick.index if self.state.tick is not None else 0
+        tick = Tick(index=previous + 1, monotonic_seconds=time.monotonic())
+        self._run_tick(tick, force_rates=True)
+        return list(self.state.log_lines)
+
+    def _run_tick(self, tick: Tick, *, force_rates: bool = False) -> None:
+        self.state.tick = tick
+        self.state.log_lines = []
+        if self.event_collector is not None:
+            self.event_collector(self.state)
+        self._handle_action_deadline()
+        if self.state.stop:
+            return
+        self.maybe_update_language_state(force=force_rates)
+        self.maybe_handle_user_events()
+        self.maybe_enqueue_action_chunk(force=force_rates)
+        self.dispatch_action(force=force_rates)
+        self.state.events.clear()
+
+    def _current_observation(self) -> dict[str, Any] | None:
+        if self.observation_provider is None:
+            return None
+        try:
+            return self.observation_provider()
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("observation_provider failed: %s", exc)
+            return None
+
+    def maybe_update_language_state(self, *, force: bool = False) -> None:
+        if self.state.mode != "action" or not self.state.task:
+            return
+        if self.state.action_queue:
+            self._language_gate.rearm()
+            return
+        if self.state.tick is None or not self._language_gate.due(self.state.tick, force=force):
+            return
+        observation = self._current_observation()
+        try:
+            self.policy_adapter.update_language_state(observation, self.state)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("language update failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
+            self.state.log(f"  [warn] language update failed: {type(exc).__name__}: {exc}")
+
+    def maybe_handle_user_events(self) -> None:
+        if self.state.take_event("user_interjection"):
+            self._handle_user_interjection()
+
+    def _handle_user_interjection(self) -> None:
+        text = str(self.state.extra.get("recent_interjection") or "")
+        if not text:
+            return
+        observation = self._current_observation()
+        self.policy_adapter.handle_interjection(text, observation, self.state)
+        self.state.extra["recent_interjection"] = None
+
+    def maybe_enqueue_action_chunk(self, *, force: bool = False) -> None:
+        with self.state.lock:
+            if self.state.mode != "action" or not self.state.task:
+                return
+            if self.state.action_queue:
+                return
+            if self.state.tick is None or not self._chunk_gate.due(self.state.tick, force=force):
+                return
+            revision = self.state.revision
+        observation = self._current_observation()
+        if observation is None:
+            return
+        try:
+            chunk = self.policy_adapter.select_action(observation, self.state)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("select_action failed: %s", exc, exc_info=logger.isEnabledFor(logging.DEBUG))
+            self.state.log(f"  [warn] select_action failed: {type(exc).__name__}: {exc}")
+            return
+        with self.state.lock:
+            if (
+                self.state.revision != revision
+                or self.state.mode != "action"
+                or self.state.stop
+                or self._stop
+            ):
+                logger.info("Discarded an action chunk invalidated during inference.")
+                return
+            self._enqueue_chunk(chunk)
+
+    def _enqueue_chunk(self, chunk: Any) -> None:
+        if chunk is None:
+            return
+        chunk_iter = chunk[0] if getattr(chunk, "ndim", None) == 3 else chunk
+        if getattr(chunk_iter, "ndim", None) == 1:
+            chunk_iter = chunk_iter.unsqueeze(0)
+        for step in chunk_iter:
+            self.state.action_queue.append(step.unsqueeze(0) if hasattr(step, "unsqueeze") else step)
+        try:
+            self.state.extra["last_chunk_size"] = int(chunk_iter.shape[0])
+        except Exception:  # noqa: BLE001
+            self.state.extra["last_chunk_size"] = len(self.state.action_queue)
+
+    def dispatch_action(self, *, force: bool = False) -> None:
+        if self.state.mode != "action":
+            self._last_dispatch_seconds = None
+            return
+        if self.state.tick is None or not self._ctrl_gate.due(self.state.tick, force=force):
+            return
+        queue = self.state.action_queue
+        if not queue:
+            self._last_dispatch_seconds = None
+            return
+        now = time.monotonic()
+        if self._last_dispatch_seconds is None or self.ctrl_hz <= 0:
+            n_to_pop = 1
+        else:
+            n_to_pop = max(1, min(len(queue), int(round((now - self._last_dispatch_seconds) * self.ctrl_hz))))
+        self._last_dispatch_seconds = now
+        latest = None
+        for _ in range(n_to_pop):
+            if not queue:
+                break
+            latest = queue.popleft()
+            self.state.actions_dispatched += 1
+        if latest is not None and self.action_executor is not None:
+            self.action_executor(latest)
+
+    def _handle_action_deadline(self) -> None:
+        deadline = self.state.action_deadline
+        if self.state.mode == "action" and deadline is not None and time.monotonic() >= deadline:
+            self.state.mode = "paused"
+            self.state.action_deadline = None
+            self.state.action_queue.clear()
+            self.state.log("timed action elapsed — paused")
+
+    def _flush_logs(self) -> None:
+        for line in self.state.log_lines:
+            print(f"[runtime] {line}", flush=True)
+
+    def _on_shutdown(self) -> None:
+        self.state.action_queue.clear()
+        print("[runtime] stopped", flush=True)

src/lerobot/runtime/registry.py

@@ -0,0 +1,39 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Lazy mapping from policy types to language-runtime adapters."""
+
+from __future__ import annotations
+
+import importlib
+from collections.abc import Callable
+from typing import Any
+
+_ADAPTERS: dict[str, str] = {
+    "pi052": "lerobot.policies.pi052.inference.pi052_adapter:PI052PolicyAdapter",
+    "pi05": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
+    "molmoact2": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
+}
+
+
+def get_language_adapter_factory(policy_type: str) -> Callable[..., Any]:
+    """Return the adapter class registered for ``policy_type``."""
+    spec = _ADAPTERS.get(policy_type)
+    if spec is None:
+        raise ValueError(
+            f"No language-runtime adapter registered for policy type {policy_type!r}. "
+            f"Registered: {sorted(_ADAPTERS)}. Add an entry to lerobot.runtime.registry."
+        )
+    module_path, class_name = spec.split(":")
+    return getattr(importlib.import_module(module_path), class_name)

src/lerobot/runtime/sim_robocasa.py

@@ -0,0 +1,406 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""RoboCasa backend for interactive language-conditioned rollouts.
+
+It reuses the eval observation/action pipeline while prompts control a persistent selected scene.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+from lerobot.utils.io_utils import StreamingVideoWriter
+from lerobot.utils.video_annotation import annotate_frame
+
+logger = logging.getLogger(__name__)
+
+
+def _short_cam_name(cam: str) -> str:
+    """Human-friendly view label for a RoboCasa camera name."""
+    c = cam.replace("robot0_", "")
+    return {
+        "agentview_left": "left",
+        "agentview_right": "right",
+        "eye_in_hand": "wrist",
+    }.get(c, c)
+
+
+def _label_panel(img: np.ndarray, label: str) -> np.ndarray:
+    """Draw a small camera-view label in the bottom-left corner of a panel."""
+    try:
+        import cv2  # noqa: PLC0415
+    except ImportError:
+        return img
+    y = img.shape[0] - 6
+    cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (0, 0, 0), 3, cv2.LINE_AA)
+    cv2.putText(img, label, (5, y), cv2.FONT_HERSHEY_SIMPLEX, 0.45, (255, 255, 0), 1, cv2.LINE_AA)
+    return img
+
+
+# Two workers avoid broken single-worker EGL rendering; only env 0 is displayed.
+_SIM_N_ENVS = 2
+
+
+def create_sim_env(
+    *,
+    task: str,
+    split: str | None,
+    obj_registries: list[str],
+    seed: int | None,
+    render_size: int = 384,
+) -> tuple[Any, dict]:
+    """Create and reset the vectorized RoboCasa environment before CUDA initializes.
+
+    Two workers keep EGL stable, while only env 0 is driven and displayed.
+    """
+    from lerobot.envs.configs import RoboCasaEnv as RoboCasaEnvConfig  # noqa: PLC0415
+
+    # The policy resizes inputs, so render_size only affects display quality and cost.
+    env_cfg = RoboCasaEnvConfig(
+        task=task,
+        split=split,
+        obj_registries=list(obj_registries),
+        observation_height=render_size,
+        observation_width=render_size,
+    )
+    # Keep one kitchen alive across sequential prompts.
+    envs = env_cfg.create_envs(
+        n_envs=_SIM_N_ENVS,
+        use_async_envs=True,
+        terminate_on_success=False,
+        horizon=100_000,
+    )
+    env = envs[next(iter(envs))][0]
+    logger.info("[sim] resetting RoboCasa scene task=%r split=%r (n_envs=%d)", task, split, _SIM_N_ENVS)
+    seeds = None if seed is None else [seed + i for i in range(_SIM_N_ENVS)]
+    obs, _ = env.reset(seed=seeds)
+    return env, obs
+
+
+def start_mjpeg_server(port: int, get_frame: Callable[[], np.ndarray | None]) -> Any:
+    """Start an MJPEG server that shows a placeholder until ``get_frame`` returns frames."""
+    import io  # noqa: PLC0415
+    import threading  # noqa: PLC0415
+    import time  # noqa: PLC0415
+    from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer  # noqa: PLC0415
+
+    from PIL import Image  # noqa: PLC0415
+
+    _placeholder = Image.new("RGB", (256, 256), (17, 17, 17))
+
+    class _Handler(BaseHTTPRequestHandler):
+        def log_message(self, *args):  # silence per-request logging
+            pass
+
+        def do_GET(self):  # noqa: N802
+            if self.path in ("/", "/index.html"):
+                self.send_response(200)
+                self.send_header("Content-Type", "text/html")
+                self.end_headers()
+                self.wfile.write(
+                    b"<html><body style='margin:0;background:#111;text-align:center'>"
+                    b"<img src='/stream' style='max-width:100vw;max-height:100vh;"
+                    b"image-rendering:pixelated'></body></html>"
+                )
+                return
+            if self.path != "/stream":
+                self.send_response(404)
+                self.end_headers()
+                return
+            self.send_response(200)
+            self.send_header("Content-Type", "multipart/x-mixed-replace; boundary=frame")
+            self.end_headers()
+            try:
+                while True:
+                    frame = get_frame()
+                    buf = io.BytesIO()
+                    img = Image.fromarray(frame) if frame is not None else _placeholder
+                    img.save(buf, format="JPEG", quality=80)
+                    data = buf.getvalue()
+                    self.wfile.write(
+                        b"--frame\r\nContent-Type: image/jpeg\r\nContent-Length: "
+                        + str(len(data)).encode()
+                        + b"\r\n\r\n"
+                        + data
+                        + b"\r\n"
+                    )
+                    time.sleep(0.05)
+            except (BrokenPipeError, ConnectionResetError):
+                pass
+
+    try:
+        # Bind all interfaces intentionally so the viewer remains reachable
+        # through the documented SSH port-forwarding workflow.
+        server = ThreadingHTTPServer(("0.0.0.0", port), _Handler)  # nosec B104
+    except OSError as exc:
+        logger.warning("[sim] could not start live stream on port %d: %s", port, exc)
+        print(f"[runtime] WARNING: live stream port {port} unavailable ({exc})", flush=True)
+        return None
+    threading.Thread(target=server.serve_forever, daemon=True, name="sim-mjpeg").start()
+    print(
+        f"[runtime] live view: http://localhost:{port}  "
+        f"(over SSH: ssh -L {port}:localhost:{port} <host>) — loading until scene is ready",
+        flush=True,
+    )
+    return server
+
+
+class RoboCasaSimBackend:
+    """Expose a RoboCasa environment through the runtime observation/action contract.
+
+    The environment must be created before the policy initializes CUDA.
+    """
+
+    def __init__(
+        self,
+        *,
+        env: Any,
+        last_obs: dict,
+        task: str,
+        seed: int | None,
+        device: str,
+        preprocessor: Any,
+        postprocessor: Any,
+        record: bool = True,
+        output_dir: str | None = None,
+        view_cams: list[str] | None = None,
+    ) -> None:
+        self.env = env
+        self._last_obs = last_obs
+        self._scene_task = task
+        self._view_cams = view_cams or [
+            "robot0_agentview_left",
+            "robot0_eye_in_hand",
+            "robot0_agentview_right",
+        ]
+        self.device = torch.device(device) if isinstance(device, str) else device
+        self.preprocessor = preprocessor
+        self.postprocessor = postprocessor
+        self.seed = seed
+        self.record = record
+        self.output_dir = Path(output_dir) if output_dir else Path("outputs/runtime_sim")
+
+        self._video_writer: StreamingVideoWriter | None = None
+        self._video_path: Path | None = None
+        self._live_counter = 0
+        self._latest_frame: np.ndarray | None = None
+        self._stream_server: Any = None
+        self._reset_count = 0
+        # Bind these after runtime construction for live annotations.
+        self._task_getter: Callable[[], str | None] | None = None
+        self._subtask_getter: Callable[[], str | None] | None = None
+        self._memory_getter: Callable[[], str | None] | None = None
+        logger.info("[sim] scene ready — task_description=%r", self._scene_description())
+
+    def bind_runtime(self, runtime: Any) -> None:
+        """Wire live task/subtask/memory getters from the runtime state."""
+        self._task_getter = lambda: runtime.state.get("task")
+        self._subtask_getter = lambda: runtime.state.language_context.get("subtask")
+        self._memory_getter = lambda: (runtime.state.get("language_context") or {}).get("memory")
+
+    def _scene_description(self) -> str:
+        try:
+            return str(self.env.get_attr("task_description")[0]) or self._scene_task
+        except Exception:  # noqa: BLE001
+            return self._scene_task
+
+    def _current_task(self) -> str:
+        task = self._task_getter() if self._task_getter else None
+        return task or self._scene_description() or self._scene_task
+
+    def reset_scene(self) -> None:
+        """Re-roll the kitchen: reset the env to a fresh scene (new layout/style).
+
+        Uses a new seed each call so ``/reset`` explores different kitchens.
+        """
+        self._reset_count += 1
+        n = self.env.num_envs
+        if self.seed is None:
+            seeds = None
+        else:
+            base = self.seed + self._reset_count * 1000
+            seeds = [base + i for i in range(n)]
+        obs, _ = self.env.reset(seed=seeds)
+        self._last_obs = obs
+        logger.info("[sim] scene reset (#%d)", self._reset_count)
+
+    def _env0_obs(self) -> dict:
+        """Slice env 0 out of the batched vec-env observation (batch of 1)."""
+        raw = self._last_obs or {}
+        pixels = raw.get("pixels")
+        out: dict[str, Any] = {}
+        if isinstance(pixels, dict):
+            out["pixels"] = {k: np.asarray(v)[0:1] for k, v in pixels.items()}
+        agent_pos = raw.get("agent_pos")
+        if agent_pos is not None:
+            out["agent_pos"] = np.asarray(agent_pos)[0:1]
+        return out
+
+    def observation_provider(self) -> dict | None:
+        from lerobot.envs.utils import preprocess_observation  # noqa: PLC0415
+
+        try:
+            obs = preprocess_observation(self._env0_obs())
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("[sim] preprocess_observation failed: %s", exc)
+            return None
+        # The adapter later replaces this recipe input with its generated subtask.
+        obs["task"] = [self._current_task()]
+        if self.preprocessor is not None:
+            try:
+                obs = self.preprocessor(obs)
+            except Exception as exc:  # noqa: BLE001
+                logger.warning("[sim] preprocessor failed: %s", exc)
+                return None
+        return {
+            k: (v.to(self.device) if isinstance(v, torch.Tensor) else v)
+            for k, v in obs.items()
+            if isinstance(k, str) and k.startswith("observation.")
+        }
+
+    def action_executor(self, action: Any) -> None:
+        try:
+            if self.postprocessor is not None:
+                action = self.postprocessor(action)
+            if isinstance(action, torch.Tensor):
+                if action.ndim > 1 and action.shape[0] == 1:
+                    action = action.squeeze(0)
+                action = action.detach().to("cpu").numpy()
+            # Tile env 0's action because the extra workers exist only for EGL stability.
+            action_row = np.asarray(action, dtype=np.float32).reshape(-1)
+            action_np = np.tile(action_row, (self.env.num_envs, 1))
+            obs, _reward, terminated, truncated, _info = self.env.step(action_np)
+            self._last_obs = obs
+            self._capture_frame()
+            # AsyncVectorEnv resets terminated sub-environments automatically.
+            if bool(np.any(terminated)) or bool(np.any(truncated)):
+                logger.info("[sim] episode ended — scene auto-reset")
+        except Exception as exc:  # noqa: BLE001
+            logger.error("[sim] env.step failed: %s", exc, exc_info=True)
+
+    def _multiview_frame(self) -> np.ndarray | None:
+        """Label and compose env 0's existing observation views without extra rendering."""
+        pixels = (self._last_obs or {}).get("pixels")
+        if not isinstance(pixels, dict) or not pixels:
+            return None
+        panels: list[np.ndarray] = []
+        for cam in self._view_cams:
+            v = pixels.get(cam)
+            if v is None:
+                continue
+            img = np.asarray(v)
+            if img.ndim == 4:  # (n_envs, H, W, C) -> env 0
+                img = img[0]
+            if img.ndim != 3 or img.shape[-1] != 3:
+                continue
+            panels.append(_label_panel(np.ascontiguousarray(img.astype(np.uint8)), _short_cam_name(cam)))
+        if not panels:
+            return None
+        h = min(p.shape[0] for p in panels)
+        panels = [p[:h] for p in panels]
+        return np.concatenate(panels, axis=1)
+
+    def _capture_frame(self) -> None:
+        frame = self._multiview_frame()
+        if frame is None:  # fallback to single env.render()
+            try:
+                rendered = self.env.call("render")[0]
+                if isinstance(rendered, np.ndarray) and rendered.ndim == 3:
+                    frame = rendered
+            except Exception as exc:  # noqa: BLE001
+                logger.debug("[sim] render failed: %s", exc)
+        if frame is None:
+            return
+        subtask = self._subtask_getter() if self._subtask_getter else None
+        memory = self._memory_getter() if self._memory_getter else None
+        annotated = annotate_frame(
+            frame,
+            (("Task", self._current_task()), ("Subtask", subtask), ("Memory", memory)),
+        )
+        self._latest_frame = annotated  # served by the live MJPEG stream
+        self._write_live_frame(annotated)
+        if self.record:
+            self._write_recording_frame(annotated)
+
+    def _write_recording_frame(self, frame: np.ndarray) -> None:
+        try:
+            if self._video_writer is None:
+                self.output_dir.mkdir(parents=True, exist_ok=True)
+                stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+                self._video_path = self.output_dir / f"sim_{stamp}.mp4"
+                fps = int((getattr(self.env, "metadata", None) or {}).get("render_fps", 20))
+                self._video_writer = StreamingVideoWriter(self._video_path, fps)
+            self._video_writer.add_frame(frame)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("[sim] video encoding failed: %s", exc)
+            self.record = False
+
+    def _write_live_frame(self, frame: np.ndarray) -> None:
+        """Write a rolling latest.png every few frames for live viewing over SSH.
+
+        Open ``{output_dir}/latest.png`` in an editor/viewer and refresh to watch
+        the rollout in near-real-time without a GUI window. Written atomically
+        (temp + replace) so a reader never sees a half-written file.
+        """
+        self._live_counter += 1
+        if self._live_counter % 3 != 0:
+            return
+        try:
+            import os  # noqa: PLC0415
+
+            from PIL import Image  # noqa: PLC0415
+
+            self.output_dir.mkdir(parents=True, exist_ok=True)
+            tmp = self.output_dir / ".latest.tmp.png"
+            Image.fromarray(frame).save(tmp)
+            os.replace(tmp, self.output_dir / "latest.png")
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("[sim] live frame write failed: %s", exc)
+
+    def _flush_video(self) -> None:
+        if self._video_writer is None:
+            return
+        writer = self._video_writer
+        self._video_writer = None
+        try:
+            writer.close()
+            logger.info("[sim] wrote video (%d frames) to %s", writer.frames_written, self._video_path)
+            print(f"[runtime] sim video saved to {self._video_path}", flush=True)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning("[sim] video close failed: %s", exc)
+
+    def attach_stream_server(self, server: Any) -> None:
+        """Attach an already-running MJPEG server so disconnect() can stop it."""
+        self._stream_server = server
+
+    def disconnect(self) -> None:
+        """Match the robot backend's cleanup contract."""
+        if self._stream_server is not None:
+            try:
+                self._stream_server.shutdown()
+            except Exception as exc:  # noqa: BLE001
+                logger.debug("[sim] stream server shutdown raised %s", exc)
+        self._flush_video()
+        try:
+            self.env.close()
+        except Exception as exc:  # noqa: BLE001
+            logger.debug("[sim] env.close raised %s", exc)

src/lerobot/scripts/lerobot_eval.py

@@ -94,6 +94,19 @@ from lerobot.utils.utils import (
     init_logging,
     inside_slurm,
 )
+from lerobot.utils.video_annotation import annotate_frame
+
+
+def _annotate_eval_frames(frames: np.ndarray, task: str | None, subtask: str | None) -> np.ndarray:
+    """Overlay the high-level task and predicted subtask onto rendered frames.
+
+    ``frames`` is ``(n_envs, H, W, C)`` uint8. Best-effort: if OpenCV isn't
+    available the frames are returned unchanged so eval never fails over a
+    visualization concern.
+    """
+    if frames.ndim != 4 or frames.shape[-1] != 3:
+        return frames
+    return np.stack([annotate_frame(frame, (("Task", task), ("Subtask", subtask))) for frame in frames])
 
 
 def _env_features_to_dataset_features(env_features: dict) -> dict:
@@ -474,11 +487,36 @@ def eval_policy(
             return
         n_to_render_now = min(max_episodes_rendered - n_episodes_rendered, env.num_envs)
         if isinstance(env, gym.vector.SyncVectorEnv):
-            ep_frames.append(np.stack([env.envs[i].render() for i in range(n_to_render_now)]))  # noqa: B023
+            frames = np.stack([env.envs[i].render() for i in range(n_to_render_now)])  # noqa: B023
         elif hasattr(env, "call"):
             # Here we must render all frames and discard any we don't need.
             # Covers AsyncVectorEnv and _LazyAsyncVectorEnv (which wraps one).
-            ep_frames.append(np.stack(env.call("render")[:n_to_render_now]))
+            frames = np.stack(env.call("render")[:n_to_render_now])
+        else:
+            return
+
+        # Overlay the high-level task and (for hierarchical policies like
+        # pi052) the predicted low-level subtask onto each frame. Both are
+        # best-effort: missing values just skip that line.
+        try:
+            tasks = list(env.call("task_description"))
+        except (AttributeError, NotImplementedError):
+            try:
+                tasks = list(env.call("task"))
+            except (AttributeError, NotImplementedError):
+                tasks = None
+        subtasks = getattr(policy, "last_subtasks", None)
+        annotated = []
+        for i in range(frames.shape[0]):
+            subtask_i = subtasks[i] if subtasks is not None and i < len(subtasks) else None
+            annotated.append(
+                _annotate_eval_frames(
+                    frames[i : i + 1],
+                    tasks[i] if tasks is not None and i < len(tasks) else None,
+                    subtask_i,
+                )[0]
+            )
+        ep_frames.append(np.stack(annotated))
 
     if max_episodes_rendered > 0:
         video_paths: list[str] = []

src/lerobot/scripts/lerobot_rollout.py

@@ -151,6 +151,7 @@ Usage examples
 """
 
 import logging
+import sys
 
 from lerobot.cameras.opencv import OpenCVCameraConfig  # noqa: F401
 from lerobot.cameras.realsense import RealSenseCameraConfig  # noqa: F401
@@ -241,10 +242,69 @@ def rollout(cfg: RolloutConfig):
     logger.info("Rollout finished")
 
 
-def main():
-    """CLI entry point for ``lerobot-rollout``."""
+_LANGUAGE_RUNTIME_FLAGS = {
+    "--language",
+    "--no_robot",
+    "--sim",
+    "--direct_subtask",
+    "--sim.direct_subtask",
+    "--disable_memory",
+    "--fp8",
+}
+_LANGUAGE_RUNTIME_PREFIXES = (
+    "--sim.",
+    "--chunk_hz",
+    "--ctrl_hz",
+    "--high_level_hz",
+    "--subtask_chunks_per_gen",
+    "--text_min_new_tokens",
+    "--text_temperature",
+    "--text_top_p",
+)
+
+
+def _uses_language_runtime(argv: list[str]) -> bool:
+    """Return whether *argv* selects the interactive language runtime.
+
+    ``--language`` is the explicit selector for real-robot runs whose other
+    options overlap with the standard rollout CLI. Language-only options also
+    select it automatically, which keeps the former language-runtime examples
+    working after replacing their command name with ``lerobot-rollout``.
+    """
+    return any(
+        arg.split("=", 1)[0] in _LANGUAGE_RUNTIME_FLAGS or arg.startswith(_LANGUAGE_RUNTIME_PREFIXES)
+        for arg in argv
+    )
+
+
+def main(argv: list[str] | None = None):
+    """CLI entry point for ``lerobot-rollout``.
+
+    Standard policy deployment continues through :class:`RolloutConfig`.
+    Interactive language-conditioned and RoboCasa runs share this entry point
+    and are selected with ``--language`` or any language-runtime-only option.
+    """
     register_third_party_plugins()
-    rollout()
+    cli_args = list(sys.argv[1:] if argv is None else argv)
+    if _uses_language_runtime(cli_args):
+        from lerobot.runtime.cli import run as run_language_runtime
+
+        # ``--language`` is a dispatcher flag, not part of the runtime's own
+        # argparse surface. All other arguments pass through unchanged.
+        runtime_args = [arg for arg in cli_args if arg != "--language"]
+        return run_language_runtime(runtime_args, prog="lerobot-rollout")
+
+    if argv is None:
+        return rollout()
+
+    # draccus reads sys.argv. Supporting an explicit argv keeps this entry
+    # point easy to smoke-test and mirrors the language-runtime branch above.
+    previous_argv = sys.argv
+    try:
+        sys.argv = [previous_argv[0], *cli_args]
+        return rollout()
+    finally:
+        sys.argv = previous_argv
 
 
 if __name__ == "__main__":

src/lerobot/scripts/lerobot_train.py

@@ -20,9 +20,11 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
+import os
 import sys
 import time
 from contextlib import nullcontext
+from datetime import timedelta
 from pprint import pformat
 from typing import TYPE_CHECKING, Any
 
@@ -81,6 +83,7 @@ def update_policy(
     lr_scheduler=None,
     lock=None,
     sample_weighter=None,
+    log_metrics: bool = True,
 ) -> tuple[MetricsTracker, dict | None]:
     """
     Performs a single training step to update the policy's weights.
@@ -98,6 +101,7 @@ def update_policy(
         lr_scheduler: An optional learning rate scheduler.
         lock: An optional lock for thread-safe optimizer updates.
         sample_weighter: Optional SampleWeighter instance for per-sample loss weighting.
+        log_metrics: Whether to synchronize and record GPU metrics this step.
 
     Returns:
         A tuple containing:
@@ -165,12 +169,20 @@ def update_policy(
     if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"):
         accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
 
-    train_metrics.loss = loss.item()
-    train_metrics.grad_norm = grad_norm.item()
     train_metrics.lr = optimizer.param_groups[0]["lr"]
-    train_metrics.update_s = time.perf_counter() - start_time
     if torch.cuda.is_available():
         train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
+    train_metrics.accumulate_tensor("loss", loss)
+    train_metrics.accumulate_tensor("grad_norm", grad_norm)
+    train_metrics.update_s = time.perf_counter() - start_time
+    # Synchronize accumulated GPU metrics only when logging.
+    if log_metrics:
+        train_metrics.materialize_tensors()
+        # Materialize detached loss components during the same logging synchronization.
+        if output_dict:
+            output_dict = {
+                k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items()
+            }
     return train_metrics, output_dict
 
 
@@ -198,7 +210,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     require_package("accelerate", extra="training")
     from accelerate import Accelerator
-    from accelerate.utils import DistributedDataParallelKwargs, DistributedType
+    from accelerate.utils import DistributedDataParallelKwargs, DistributedType, InitProcessGroupKwargs
 
     cfg.validate()
 
@@ -207,7 +219,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes
     # We set find_unused_parameters=True to handle models with conditional computation
     if accelerator is None:
-        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+        # Static graphs restore DDP overlap when conditional parameter usage is stable.
+        # Environment flags retain the existing defaults.
+        ddp_find_unused = os.environ.get("LEROBOT_DDP_FIND_UNUSED", "1") == "1"
+        ddp_static_graph = os.environ.get("LEROBOT_DDP_STATIC_GRAPH", "0") == "1"
+        ddp_kwargs = DistributedDataParallelKwargs(
+            find_unused_parameters=ddp_find_unused and not ddp_static_graph,
+            static_graph=ddp_static_graph,
+        )
+        # Allow rank 0 enough time to index large datasets before other ranks leave the barrier.
+        ipg_kwargs = InitProcessGroupKwargs(timeout=timedelta(hours=2))
         # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
         # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
         force_cpu = cfg.trainable_config.device == "cpu"
@@ -217,7 +238,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         accelerator = Accelerator(
             step_scheduler_with_optimizer=False,
             mixed_precision=mixed_precision,
-            kwargs_handlers=[ddp_kwargs],
+            kwargs_handlers=[ddp_kwargs, ipg_kwargs],
             cpu=force_cpu,
         )
 
@@ -313,6 +334,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     active_cfg = cfg.trainable_config
     processor_pretrained_path = active_cfg.pretrained_path
+    # A weight checkpoint may contain PI05 or differently configured PI052 processors.
+    if cfg.policy.type == "pi052" and processor_pretrained_path is not None and not cfg.resume:
+        logging.warning(
+            "pi052 is loading pretrained weights from %s, but building processors from the current "
+            "pi052 config so recipe text labels and FAST action labels are generated.",
+            processor_pretrained_path,
+        )
+        processor_pretrained_path = None
 
     processor_kwargs = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
@@ -321,6 +350,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
+    if cfg.policy.type in {"pi0_fast", "pi052"}:
+        processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
+
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
         preprocessor_overrides = {
             "device_processor": {"device": device.type},
@@ -417,13 +449,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         # same permutation. accelerate then shards it disjointly across ranks via BatchSamplerShard
         # without needing a `generator` attribute to synchronize an RNG, and resume is sample-exact.
         shuffle = False
+        from_indices = dataset.meta.episodes["dataset_from_index"]
+        to_indices = dataset.meta.episodes["dataset_to_index"]
+        seed = cfg.seed if cfg.seed is not None else 0
+
         sampler = EpisodeAwareSampler(
-            dataset.meta.episodes["dataset_from_index"],
-            dataset.meta.episodes["dataset_to_index"],
+            from_indices,
+            to_indices,
             episode_indices_to_use=dataset.episodes,
             drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
             shuffle=True,
-            seed=cfg.seed if cfg.seed is not None else 0,
+            seed=seed,
             absolute_to_relative_idx=dataset.absolute_to_relative_idx,
         )
         if cfg.resume and step > 0:
@@ -461,6 +497,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # declares language columns; otherwise stay on PyTorch's default
     # collate so non-language training runs are unaffected.
     collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
+    # Allow spawn/forkserver workers where forking large rank processes exhausts memory.
+    mp_context = os.environ.get("LEROBOT_DATALOADER_MP_CONTEXT") or None
     dataloader = torch.utils.data.DataLoader(
         dataset,
         num_workers=cfg.num_workers,
@@ -472,6 +510,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         collate_fn=collate_fn,
         prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
         persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
+        multiprocessing_context=mp_context if cfg.num_workers > 0 else None,
     )
 
     # Build eval dataloader if a held-out split exists
@@ -572,6 +611,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
+        # Synchronize GPU metrics only for updates that will be logged.
+        log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0
+
         train_tracker, output_dict = update_policy(
             train_tracker,
             policy,
@@ -581,6 +623,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             accelerator=accelerator,
             lr_scheduler=lr_scheduler,
             sample_weighter=sample_weighter,
+            log_metrics=log_metrics,
         )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
@@ -680,10 +723,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             if is_main_process:
                 step_id = get_step_identifier(step, cfg.steps)
                 logging.info(f"Eval policy at step {step}")
+                eval_target_policy = accelerator.unwrap_model(policy)
                 with torch.no_grad(), accelerator.autocast():
                     eval_info = eval_policy_all(
                         envs=eval_env,  # dict[suite][task_id] -> vec_env
-                        policy=accelerator.unwrap_model(policy),
+                        policy=eval_target_policy,
                         env_preprocessor=env_preprocessor,
                         env_postprocessor=env_postprocessor,
                         preprocessor=preprocessor,

src/lerobot/utils/collate.py

@@ -22,7 +22,7 @@ from torch.utils.data._utils.collate import default_collate
 
 from lerobot.datasets.language import LANGUAGE_COLUMNS
 
-_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"}
+_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices", *LANGUAGE_COLUMNS}
 
 
 def lerobot_collate_fn(batch: list[dict[str, Any] | None]) -> dict[str, Any] | None:

src/lerobot/utils/constants.py

@@ -34,6 +34,7 @@ ACTION = "action"
 ACTION_PREFIX = ACTION + "."
 ACTION_TOKENS = ACTION + ".tokens"
 ACTION_TOKEN_MASK = ACTION + ".token_mask"
+ACTION_CODE_TOKEN_MASK = ACTION + ".code_token_mask"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"

src/lerobot/utils/io_utils.py

@@ -23,6 +23,46 @@ logger = logging.getLogger(__name__)
 JsonLike = str | int | float | bool | None | list["JsonLike"] | dict[str, "JsonLike"] | tuple["JsonLike", ...]
 
 
+class StreamingVideoWriter:
+    """Incrementally encode RGB frames to an MP4 without retaining them in memory."""
+
+    def __init__(self, video_path: str | Path, fps: int) -> None:
+        from .import_utils import require_package
+
+        require_package("av", extra="av-dep")
+        import av
+
+        self._av = av
+        self._container = av.open(str(video_path), mode="w")
+        self._stream = self._container.add_stream("libx264", rate=fps)
+        self._shape: tuple[int, int] | None = None
+        self.frames_written = 0
+
+    def add_frame(self, frame_array) -> None:
+        orig_height, orig_width = frame_array.shape[:2]
+        height = orig_height - orig_height % 2
+        width = orig_width - orig_width % 2
+        if self._shape is None:
+            self._shape = (height, width)
+            self._stream.width = width
+            self._stream.height = height
+            self._stream.pix_fmt = "yuv420p"
+        elif self._shape != (height, width):
+            raise ValueError(f"Video frame shape changed from {self._shape} to {(height, width)}")
+        frame = self._av.VideoFrame.from_ndarray(frame_array[:height, :width], format="rgb24")
+        for packet in self._stream.encode(frame):
+            self._container.mux(packet)
+        self.frames_written += 1
+
+    def close(self) -> None:
+        if self._container is None:
+            return
+        for packet in self._stream.encode():
+            self._container.mux(packet)
+        self._container.close()
+        self._container = None
+
+
 def load_json(fpath: Path) -> Any:
     """Load data from a JSON file.
 
@@ -58,36 +98,12 @@ def write_video(video_path: str | Path, stacked_frames: list, fps: int) -> None:
         stacked_frames: List of HWC uint8 numpy arrays (RGB).
         fps: Frames per second for the output video.
     """
-    from .import_utils import require_package
-
-    require_package("av", extra="av-dep")
-    import av
-
-    with av.open(str(video_path), mode="w") as container:
-        orig_height, orig_width = stacked_frames[0].shape[:2]
-        # yuv420p requires even dimensions; crop by one pixel if needed
-        height = orig_height if orig_height % 2 == 0 else orig_height - 1
-        width = orig_width if orig_width % 2 == 0 else orig_width - 1
-        if height != orig_height or width != orig_width:
-            logger.warning(
-                "Frame dimensions %dx%d are not even; cropping to %dx%d for yuv420p compatibility.",
-                orig_width,
-                orig_height,
-                width,
-                height,
-            )
-        stream = container.add_stream("libx264", rate=fps)
-        stream.width = width
-        stream.height = height
-        stream.pix_fmt = "yuv420p"
+    writer = StreamingVideoWriter(video_path, fps)
+    try:
         for frame_array in stacked_frames:
-            if height != orig_height or width != orig_width:
-                frame_array = frame_array[:height, :width]
-            frame = av.VideoFrame.from_ndarray(frame_array, format="rgb24")
-            for packet in stream.encode(frame):
-                container.mux(packet)
-        for packet in stream.encode():
-            container.mux(packet)
+            writer.add_frame(frame_array)
+    finally:
+        writer.close()
 
 
 def deserialize_json_into_object[T: JsonLike](fpath: Path, obj: T) -> T:

src/lerobot/utils/logging_utils.py

@@ -104,6 +104,8 @@ class MetricsTracker:
         "episodes",
         "epochs",
         "accelerator",
+        "_tensor_sums",
+        "_tensor_counts",
     ]
 
     def __init__(
@@ -129,6 +131,8 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
         self.accelerator = accelerator
+        self._tensor_sums: dict[str, torch.Tensor] = {}
+        self._tensor_counts: dict[str, int] = {}
 
     def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
         if name in self.__dict__:
@@ -156,6 +160,22 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
 
+    def accumulate_tensor(self, name: str, value: torch.Tensor) -> None:
+        """Accumulate a detached metric on-device until the next logging step."""
+        if name not in self.metrics:
+            raise KeyError(f"Unknown metric {name!r}.")
+        value = value.detach()
+        self._tensor_sums[name] = self._tensor_sums.get(name, torch.zeros_like(value)) + value
+        self._tensor_counts[name] = self._tensor_counts.get(name, 0) + 1
+
+    def materialize_tensors(self) -> None:
+        """Transfer pending tensor averages to their meters with one sync per metric."""
+        for name, total in self._tensor_sums.items():
+            count = self._tensor_counts[name]
+            self.metrics[name].update((total / count).item(), n=count)
+        self._tensor_sums.clear()
+        self._tensor_counts.clear()
+
     def reduce_across_ranks(self) -> None:
         """
         Synchronises the running averages of every metric whose ``reduction`` is not ``"none"``
@@ -176,10 +196,21 @@ class MetricsTracker:
         if not buckets:
             return
 
+        # NB: don't use ``accelerator.reduce(..., reduction="max")`` — accelerate only implements
+        # "sum"/"mean" (it always all-reduces with SUM and divides for "mean"), so "max" silently
+        # returns the SUM across ranks, inflating every "max" metric by ``num_processes`` (e.g. a
+        # 3.5s step reported as 28s on 8 GPUs). Gather per-rank values and reduce them explicitly.
         device = self.accelerator.device
+        num_processes = self.accelerator.num_processes
         for reduction, names in buckets.items():
-            tensor = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
-            reduced = self.accelerator.reduce(tensor, reduction=reduction)
+            local = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
+            gathered = self.accelerator.gather(local).view(num_processes, len(names))
+            if reduction == "max":
+                reduced = gathered.amax(dim=0)
+            elif reduction == "sum":
+                reduced = gathered.sum(dim=0)
+            else:  # "mean"
+                reduced = gathered.mean(dim=0)
             for name, value in zip(names, reduced.tolist(), strict=True):
                 meter = self.metrics[name]
                 # Preserve avg == sum / count so a later .update() on this meter accumulates
@@ -216,3 +247,5 @@ class MetricsTracker:
         """Resets average meters."""
         for m in self.metrics.values():
             m.reset()
+        self._tensor_sums.clear()
+        self._tensor_counts.clear()

src/lerobot/utils/rerun_visualization.py

@@ -38,7 +38,10 @@ def _is_scalar(x):
 
 
 def init_rerun(
-    session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
+    session_name: str = "lerobot_control_loop",
+    ip: str | None = None,
+    port: int | None = None,
+    web_port: int | None = None,
 ) -> None:
     """
     Initializes the Rerun SDK for visualizing the control loop.
@@ -47,6 +50,7 @@ def init_rerun(
         session_name: Name of the Rerun session.
         ip: Optional IP for connecting to a Rerun server.
         port: Optional port for connecting to a Rerun server.
+        web_port: Serve a headless web viewer on this port, using ``port`` for gRPC.
     """
 
     require_package("rerun-sdk", extra="viz", import_name="rerun")
@@ -60,6 +64,10 @@ def init_rerun(
     memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%")
     if ip and port:
         rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy")
+    elif web_port is not None:
+        grpc_port = port or 9876
+        url = rr.serve_grpc(grpc_port=grpc_port)
+        rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url)
     else:
         rr.spawn(memory_limit=memory_limit)
 

src/lerobot/utils/video_annotation.py

@@ -0,0 +1,71 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Best-effort text overlays shared by evaluation and interactive rollouts."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+import numpy as np
+
+
+def annotate_frame(frame: np.ndarray, fields: Iterable[tuple[str, str | None]]) -> np.ndarray:
+    """Return an RGB frame annotated with the non-empty labeled ``fields``."""
+    if frame.ndim != 3 or frame.shape[-1] != 3:
+        return frame
+    try:
+        import cv2  # noqa: PLC0415
+    except ImportError:
+        return frame
+
+    text_rows = [f"{label}: {value}" for label, value in fields if value]
+    if not text_rows:
+        return frame
+
+    image = np.ascontiguousarray(frame).copy()
+    font, scale, thickness, margin = cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1, 6
+    max_width = image.shape[1] - 2 * margin
+    lines: list[str] = []
+    for text in text_rows:
+        current = ""
+        for word in text.split():
+            candidate = f"{current} {word}".strip()
+            width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
+            if width > max_width and current:
+                lines.append(current)
+                current = word
+            else:
+                current = candidate
+        if current:
+            lines.append(current)
+
+    line_height = 20
+    header_height = min(image.shape[0], len(lines) * line_height + 6)
+    backdrop = image.copy()
+    cv2.rectangle(backdrop, (0, 0), (image.shape[1], header_height), (0, 0, 0), -1)
+    cv2.addWeighted(backdrop, 0.55, image, 0.45, 0, dst=image)
+
+    for index, line in enumerate(lines):
+        cv2.putText(
+            image,
+            line,
+            (margin, 18 + index * line_height),
+            font,
+            scale,
+            (255, 255, 255),
+            thickness,
+            cv2.LINE_AA,
+        )
+    return image

tests/configs/test_recipe.py

@@ -29,6 +29,13 @@ def test_message_recipe_validates_unknown_binding():
         )
 
 
+def test_canonical_recipe_loads():
+    """The canonical PI052 blend YAML loads + validates."""
+    recipe = TrainingRecipe.from_yaml(Path("src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml"))
+    assert recipe.blend is not None
+    assert sum(c.weight for c in recipe.blend.values()) == pytest.approx(1.0)
+
+
 def test_message_turn_requires_a_stream():
     """Every turn must declare a stream — None is rejected at construction.
 

tests/datasets/test_language_render.py

@@ -343,6 +343,84 @@ def test_resolve_task_explicit_override_beats_rephrasings():
     assert rendered["messages"][0]["content"] == "explicit override wins"
 
 
+def test_flow_only_low_level_recipe_renders_without_target():
+    """Regression: a flow-only ``low_level`` recipe has no ``target`` turn —
+    its supervision is the action-expert flow loss, not text-CE. It must
+    still render (not ``None``), otherwise every blend draw of it is dropped
+    and the action expert never receives a flow loss."""
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(
+                role="user",
+                content="${subtask}",
+                stream="low_level",
+                if_present="subtask",
+            ),
+        ],
+        bindings={"subtask": "active_at(t, style=subtask)"},
+    )
+
+    rendered = render_sample(
+        recipe=recipe,
+        persistent=PERSISTENT,
+        events=[],
+        t=0.5,
+        sample_idx=0,
+        task="clean kitchen",
+    )
+
+    assert rendered is not None
+    assert rendered["messages"] == [{"role": "user", "content": "subtask 0"}]
+    assert rendered["message_streams"] == ["low_level"]
+    assert rendered["target_message_indices"] == []
+
+
+def test_vqa_frame_is_consumed_over_the_weighted_blend():
+    """A frame carrying a VQA annotation renders the ``ask_vqa*`` sub-recipe
+    even when its blend weight is tiny — VQA annotations are sparse and must
+    never be wasted on a subtask/action draw."""
+    recipe = TrainingRecipe(
+        blend={
+            "high_level_subtask": TrainingRecipe(
+                weight=0.99,
+                messages=[
+                    MessageTurn(role="user", content="${task}", stream="high_level"),
+                    MessageTurn(role="assistant", content="a subtask", stream="high_level", target=True),
+                ],
+            ),
+            "ask_vqa_top": TrainingRecipe(
+                weight=0.01,
+                bindings={
+                    "vqa_query": "emitted_at(t, style=vqa, role=user, camera=observation.images.top)",
+                    "vqa": "emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)",
+                },
+                messages=[
+                    MessageTurn(
+                        role="user", content="${vqa_query}", stream="high_level", if_present="vqa_query"
+                    ),
+                    MessageTurn(
+                        role="assistant",
+                        content="${vqa}",
+                        stream="high_level",
+                        target=True,
+                        if_present="vqa",
+                    ),
+                ],
+            ),
+        }
+    )
+    # A frame WITH a vqa event renders VQA on every sample_idx, despite the
+    # ask_vqa weight being only 0.01.
+    for sample_idx in range(20):
+        rendered = render_sample(
+            recipe=recipe, persistent=PERSISTENT, events=EVENTS_AT_1, t=1.0, sample_idx=sample_idx, task="x"
+        )
+        assert rendered["messages"][-1]["content"] == '{"count": 2}', sample_idx
+    # A frame WITHOUT a vqa event falls back to the normal weighted blend.
+    rendered = render_sample(recipe=recipe, persistent=PERSISTENT, events=[], t=1.0, sample_idx=0, task="x")
+    assert rendered["messages"][-1]["content"] == "a subtask"
+
+
 def test_emitted_at_persistent_tolerates_small_timestamp_drift():
     """Persistent ``emitted_at`` should match within EMITTED_AT_TOLERANCE_S
     so callers that derive ``t`` arithmetically (``frame_idx / fps``) still

tests/datasets/test_sampler.py

@@ -25,7 +25,7 @@ from datasets import Dataset  # noqa: E402
 from lerobot.datasets.io_utils import (
     hf_transform_to_torch,
 )
-from lerobot.datasets.sampler import EpisodeAwareSampler
+from lerobot.datasets.sampler import EpisodeAwareSampler, compute_sampler_state
 
 
 def calculate_episode_data_index(hf_dataset: Dataset) -> dict[str, torch.Tensor]:
@@ -154,8 +154,6 @@ def test_partial_episode_drop_warns(caplog):
 
 # --- seeded (seed, epoch) shuffling, resume, and state ---
 
-from lerobot.datasets.sampler import compute_sampler_state  # noqa: E402
-
 EPISODE_BOUNDS = ([0, 2, 3], [2, 3, 6])  # episodes of 2, 1 and 3 frames
 
 

tests/policies/pi052/test_pi052_attention_masking.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Attention-masking tests for the PI052 (π0.5 v2) text head.
+
+Regression coverage for the text-CE collapse bug: PaliGemma's
+``embed_prefix`` flags every language token ``att=0``, which
+``make_att_2d_masks`` turns into one fully *bidirectional* block. Under
+that mask the text cross-entropy degenerates into a copy task — a
+supervised target token attends to the tokens it is trained to predict —
+and the LM head never learns causal generation, so ``select_message``
+collapses at inference.
+
+``_mark_target_span_causal`` sets ``att=1`` on the supervised target
+language positions so each target token attends causally among the
+targets while staying bidirectional to images + the user prompt. These
+tests pin that behaviour for the PaliGemma prefix layout.
+"""
+
+import pytest
+import torch
+
+# modeling_pi052 / modeling_pi05 import transformers transitively.
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks  # noqa: E402
+from lerobot.policies.pi052.modeling_pi052 import (  # noqa: E402
+    _mark_target_span_causal,
+    _shifted_lin_ce,
+)
+
+
+def _shifted_ce(logits, labels):
+    """Adapter: ``_shifted_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ).
+
+    An identity ``lm_head_weight`` makes the computed logits equal ``logits``.
+    Liger's Triton kernel is GPU-only, so inputs run on CUDA; the loss is
+    returned on CPU so grad still flows back to the CPU ``logits`` leaf.
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("Liger fused CE requires CUDA")
+    vocab_size = logits.shape[-1]
+    eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
+    return _shifted_lin_ce(logits.cuda(), eye, labels.cuda()).cpu()
+
+
+# Synthetic prefix: two image tokens, three prompt tokens, and four supervised target tokens.
+# Text labels mask the prompt with -100 and cover the target through the prefix end.
+N_IMAGE = 2
+N_PROMPT = 3
+N_TARGET = 4
+LANG_START = N_IMAGE
+LANG_END = N_IMAGE + N_PROMPT + N_TARGET  # = prefix length
+PREFIX_LEN = LANG_END
+
+
+def _embed_prefix_att_masks() -> torch.Tensor:
+    """Mimic PaliGemma ``embed_prefix``: images + lang all att=0."""
+    return torch.zeros(1, PREFIX_LEN, dtype=torch.bool)
+
+
+def _text_labels() -> torch.Tensor:
+    """-100 over the prompt span, real ids over the target span."""
+    labels = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
+    labels[0, N_PROMPT:] = torch.arange(10, 10 + N_TARGET)
+    return labels
+
+
+def _attends(prefix_att_masks: torch.Tensor) -> torch.Tensor:
+    """2D boolean attendance matrix; ``[i, j]`` True ⇒ i attends to j."""
+    pad = torch.ones(1, PREFIX_LEN, dtype=torch.bool)
+    return make_att_2d_masks(pad, prefix_att_masks)[0]
+
+
+def test_mark_sets_att_on_targets_only():
+    """Only the supervised target language positions flip to att=1."""
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
+    expected = [False] * PREFIX_LEN
+    for i in range(LANG_START + N_PROMPT, LANG_END):  # target span
+        expected[i] = True
+    assert marked[0].tolist() == expected
+
+
+def test_target_tokens_attend_causally_among_themselves():
+    """A target token must NOT attend to later targets, but must attend
+    to earlier ones — genuine causal next-token prediction."""
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
+    attends = _attends(marked)
+    tgt = range(LANG_START + N_PROMPT, LANG_END)
+    for i in tgt:
+        for j in tgt:
+            if j > i:
+                assert not attends[i, j], f"target {i} must not see future target {j}"
+            else:
+                assert attends[i, j], f"target {i} must see earlier/self target {j}"
+
+
+def test_target_tokens_attend_prompt_and_images_bidirectionally():
+    """Targets keep full visibility of images + the user prompt."""
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), _text_labels(), LANG_START, LANG_END)
+    attends = _attends(marked)
+    context = list(range(0, LANG_START + N_PROMPT))  # images + prompt
+    for i in range(LANG_START + N_PROMPT, LANG_END):
+        for j in context:
+            assert attends[i, j], f"target {i} must attend context {j}"
+
+
+def test_non_target_subtask_stays_bidirectional():
+    """A flow-only / non-target language span (all -100 labels) leaves the
+    mask untouched — the action expert reads it bidirectionally."""
+    all_ignored = torch.full((1, N_PROMPT + N_TARGET), -100, dtype=torch.long)
+    marked = _mark_target_span_causal(_embed_prefix_att_masks(), all_ignored, LANG_START, LANG_END)
+    assert torch.equal(marked, _embed_prefix_att_masks())
+
+
+def test_unmarked_mask_is_bidirectional_the_bug():
+    """Documents the bug the fix prevents: without ``_mark_target_span_causal``
+    a target token attends *bidirectionally* to later targets — the
+    text-CE can copy the answer it is trained to predict."""
+    attends = _attends(_embed_prefix_att_masks())
+    first_tgt = LANG_START + N_PROMPT
+    last_tgt = LANG_END - 1
+    assert attends[first_tgt, last_tgt], (
+        "raw embed_prefix mask is bidirectional over language — the first "
+        "target token can see the last, which is the collapse bug"
+    )
+
+
+def test_shifted_ce_returns_zero_when_no_text_positions_are_supervised():
+    pytest.importorskip("liger_kernel")
+    logits = torch.randn(2, 4, 8, requires_grad=True)
+    labels = torch.full((2, 4), -100, dtype=torch.long)
+
+    loss = _shifted_ce(logits, labels)
+
+    assert loss.item() == 0
+    loss.backward()
+    assert logits.grad is not None

tests/policies/pi052/test_pi052_bucketed_ce.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi052.modeling_pi052 import _lin_ce_flat, _shifted_lin_ce
+
+
+def test_shifted_ce_none_retains_distinct_per_sample_losses():
+    hidden = torch.tensor(
+        [
+            [[8.0, 0.0], [0.0, 8.0], [0.0, 0.0]],
+            [[0.0, 8.0], [8.0, 0.0], [0.0, 0.0]],
+        ]
+    )
+    labels = torch.tensor([[0, 0, 1], [0, 0, 1]])
+    losses = _shifted_lin_ce(hidden, torch.eye(2), labels, reduction="none")
+
+    assert losses.shape == (2,)
+    assert losses[0] < losses[1]
+
+
+def test_checkpoint_resolution_forwards_explicit_hub_options(monkeypatch, tmp_path):
+    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
+
+    checkpoint = tmp_path / "model.safetensors"
+    checkpoint.touch()
+    calls = []
+
+    def fake_cached_file(model_id, filename, **kwargs):
+        calls.append((model_id, filename, kwargs))
+        return None if filename.endswith("index.json") else str(checkpoint)
+
+    monkeypatch.setattr(modeling_pi05, "cached_file", fake_cached_file)
+    files = modeling_pi05._resolve_weight_files(
+        "org/model",
+        force_download=True,
+        resume_download=True,
+        proxies={"https": "proxy"},
+        token="secret",
+        cache_dir=tmp_path / "cache",
+        local_files_only=True,
+        revision="commit",
+    )
+
+    assert files == [checkpoint]
+    for _model_id, _filename, kwargs in calls:
+        assert kwargs["revision"] == "commit"
+        assert kwargs["cache_dir"] == tmp_path / "cache"
+        assert kwargs["force_download"] is True
+        assert kwargs["resume_download"] is True
+        assert kwargs["proxies"] == {"https": "proxy"}
+        assert kwargs["token"] == "secret"
+        assert kwargs["local_files_only"] is True
+
+
+def test_checkpoint_resolution_rejects_local_directory_without_weights(tmp_path):
+    import lerobot.policies.pi05.modeling_pi05 as modeling_pi05
+
+    with pytest.raises(FileNotFoundError, match="model.safetensors"):
+        modeling_pi05._resolve_weight_files(
+            tmp_path,
+            force_download=False,
+            resume_download=None,
+            proxies=None,
+            token=None,
+            cache_dir=None,
+            local_files_only=False,
+            revision=None,
+        )
+
+
+@pytest.mark.parametrize("z_loss_weight", [0.0, 1e-4])
+@pytest.mark.parametrize("rows,valid_rows", [(24, 9), (48, 25)])
+def test_bucketed_ce_matches_dense_loss_and_gradients(z_loss_weight, rows, valid_rows):
+    generator = torch.Generator().manual_seed(23)
+    hidden_size, vocab_size = 7, 19
+    hidden_ref = torch.randn(rows, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True)
+    weight_ref = torch.randn(
+        vocab_size, hidden_size, generator=generator, dtype=torch.float64, requires_grad=True
+    )
+    labels = torch.full((rows,), -100, dtype=torch.long)
+    valid_indices = torch.randperm(rows, generator=generator)[:valid_rows]
+    labels[valid_indices] = torch.randint(0, vocab_size, (valid_rows,), generator=generator)
+    hidden_bucketed = hidden_ref.detach().clone().requires_grad_(True)
+    weight_bucketed = weight_ref.detach().clone().requires_grad_(True)
+
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    loss_ref = _lin_ce_flat(hidden_ref, weight_ref, labels, z_loss_weight=z_loss_weight)
+    old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
+    modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
+    try:
+        loss_bucketed = _lin_ce_flat(
+            hidden_bucketed,
+            weight_bucketed,
+            labels,
+            z_loss_weight=z_loss_weight,
+        )
+    finally:
+        modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
+
+    loss_ref.backward()
+    loss_bucketed.backward()
+
+    torch.testing.assert_close(loss_bucketed, loss_ref, rtol=1e-6, atol=1e-6)
+    torch.testing.assert_close(hidden_bucketed.grad, hidden_ref.grad, rtol=1e-12, atol=1e-12)
+    torch.testing.assert_close(weight_bucketed.grad, weight_ref.grad, rtol=1e-12, atol=1e-12)
+
+
+def test_bucketed_ce_all_ignored_preserves_zero_gradients():
+    hidden = torch.randn(24, 7, dtype=torch.float64, requires_grad=True)
+    weight = torch.randn(19, 7, dtype=torch.float64, requires_grad=True)
+    labels = torch.full((24,), -100, dtype=torch.long)
+
+    import lerobot.policies.pi052.modeling_pi052 as modeling_pi052
+
+    old_limit = modeling_pi052._LOGITS_CE_MAX_POSITIONS
+    modeling_pi052._LOGITS_CE_MAX_POSITIONS = 16
+    try:
+        loss = _lin_ce_flat(hidden, weight, labels)
+    finally:
+        modeling_pi052._LOGITS_CE_MAX_POSITIONS = old_limit
+    loss.backward()
+
+    assert loss.item() == 0.0
+    assert hidden.grad is not None
+    assert weight.grad is not None
+    assert torch.count_nonzero(hidden.grad) == 0
+    assert torch.count_nonzero(weight.grad) == 0

tests/policies/pi052/test_pi052_fast_action_loss.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Regression tests for PI052 FAST action-code supervision."""
+
+import pytest
+import torch
+from torch.nn import functional as F  # noqa: N812
+
+pytest.importorskip("transformers")
+pytest.importorskip("liger_kernel")
+
+from lerobot.policies.pi052.modeling_pi052 import _fast_lin_ce  # noqa: E402
+
+
+def _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t):
+    """Adapter: ``_fast_lin_ce`` is Liger-fused (hidden @ lm_head_weightᵀ).
+
+    Feeding an identity ``lm_head_weight`` makes the computed logits equal the
+    provided ``logits``, so these regression tests exercise the masking/gating
+    logic exactly as before the fused-CE refactor. Liger's Triton kernel is
+    GPU-only, so inputs are moved to CUDA and the loss is returned on CPU
+    (keeping grad flowing back to the CPU ``logits`` leaf).
+    """
+    if not torch.cuda.is_available():
+        pytest.skip("Liger fused CE requires CUDA")
+    vocab_size = logits.shape[-1]
+    eye = torch.eye(vocab_size, dtype=logits.dtype, device="cuda")
+    predict = predict_actions_t.cuda() if predict_actions_t is not None else None
+    loss = _fast_lin_ce(logits.cuda(), eye, action_tokens.cuda(), action_code_mask.cuda(), predict)
+    return loss.cpu()
+
+
+def test_fast_ce_supervises_only_discrete_action_codes():
+    """Wrapper tokens can be wrong without affecting the FAST action-code loss."""
+    vocab_size = 8
+    action_tokens = torch.tensor([[1, 2, 3, 4, 5, 0]])
+    action_code_mask = torch.tensor([[False, False, True, True, False, False]])
+
+    logits = torch.zeros(1, action_tokens.shape[1], vocab_size)
+    # Deliberately bad wrapper-token predictions. These should be ignored.
+    logits[0, 0, 7] = 10.0  # target would be token 2
+    logits[0, 3, 7] = 10.0  # target would be delimiter token 5
+    # Correct action-code predictions: hidden t predicts target t + 1.
+    logits[0, 1, 3] = 10.0
+    logits[0, 2, 4] = 10.0
+
+    loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t=None)
+    expected = F.cross_entropy(
+        torch.stack([logits[0, 1], logits[0, 2]]),
+        torch.tensor([3, 4]),
+        reduction="mean",
+    )
+
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
+    assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
+
+
+def test_fast_ce_masks_non_action_samples():
+    """Recipe samples with predict_actions=False do not contribute FAST loss."""
+    vocab_size = 8
+    action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
+    action_code_mask = torch.tensor([[False, False, True, True], [False, False, True, True]])
+    predict_actions = torch.tensor([True, False])
+
+    logits = torch.zeros(2, action_tokens.shape[1], vocab_size)
+    logits[0, 1, 3] = 10.0
+    logits[0, 2, 4] = 10.0
+    # Bad predictions in the masked sample should not matter.
+    logits[1, 1, 7] = 10.0
+    logits[1, 2, 7] = 10.0
+
+    loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions)
+    expected = F.cross_entropy(
+        torch.stack([logits[0, 1], logits[0, 2]]),
+        torch.tensor([3, 4]),
+        reduction="mean",
+    )
+
+    # Allow the fused GPU kernel's ~1e-7 difference on small losses.
+    assert torch.allclose(loss, expected, atol=1e-5, rtol=1e-3)
+
+
+def test_fast_ce_returns_zero_when_no_action_code_positions_are_valid():
+    logits = torch.randn(2, 4, 8, requires_grad=True)
+    action_tokens = torch.tensor([[1, 2, 3, 4], [1, 2, 5, 6]])
+    action_code_mask = torch.zeros_like(action_tokens, dtype=torch.bool)
+
+    loss = _fast_ce(logits, action_tokens, action_code_mask, predict_actions_t=None)
+
+    assert loss.item() == 0
+    loss.backward()
+    assert logits.grad is not None

tests/policies/pi052/test_pi052_flex_attention.py

@@ -0,0 +1,65 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+import lerobot.policies.pi052.modeling_pi052 as modeling_pi052  # noqa: E402
+from lerobot.policies.pi052.configuration_pi052 import PI052Config  # noqa: E402
+
+
+def test_flex_backend_skips_non_cuda_without_initializing(monkeypatch):
+    monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
+    monkeypatch.setattr(torch, "compile", lambda *args, **kwargs: pytest.fail("torch.compile was called"))
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_properties",
+        lambda *args, **kwargs: pytest.fail("CUDA properties were queried"),
+    )
+
+    assert modeling_pi052._get_flex_fns(torch.device("cpu")) is None
+    assert modeling_pi052._get_flex_kernel_options(torch.device("cpu")) is None
+    assert modeling_pi052._flex_fns is None
+
+
+def test_flex_initialization_failure_falls_back(monkeypatch, caplog):
+    monkeypatch.setattr(modeling_pi052, "_flex_fns", None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+
+    def fail_compile(*args, **kwargs):
+        raise RuntimeError("compile failed")
+
+    monkeypatch.setattr(torch, "compile", fail_compile)
+
+    with caplog.at_level(logging.WARNING, logger=modeling_pi052.__name__):
+        assert modeling_pi052._get_flex_fns(torch.device("cuda", 0)) is None
+
+    assert modeling_pi052._flex_fns is False
+    assert "FlexAttention unavailable" in caplog.text
+
+
+def test_flex_rejects_single_repeat_configuration():
+    with pytest.raises(ValueError, match="use_flex_attention requires flow_num_repeats > 1"):
+        PI052Config(use_flex_attention=True, flow_num_repeats=1)
+
+
+def test_flex_accepts_amortized_repeat_configuration():
+    config = PI052Config(use_flex_attention=True, flow_num_repeats=5)
+    assert config.use_flex_attention

tests/policies/pi052/test_pi052_import.py

@@ -0,0 +1,27 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import subprocess
+import sys
+
+
+def test_pi052_config_import_does_not_load_model_or_dataset_processor():
+    code = """
+import sys
+from lerobot.policies import PI052Config
+assert PI052Config.__name__ == "PI052Config"
+assert "lerobot.policies.pi052.modeling_pi052" not in sys.modules
+assert "lerobot.policies.pi052.processor_pi052" not in sys.modules
+"""
+    subprocess.run([sys.executable, "-c", code], check=True)

tests/policies/pi052/test_pi052_runtime_adapter.py

@@ -0,0 +1,85 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+from lerobot.policies.pi052.inference.pi052_adapter import PI052PolicyAdapter
+from lerobot.runtime import RuntimeState
+from lerobot.runtime.adapter import split_plan_and_say
+
+
+def test_pi052_adapter_builds_recipe_prompts_from_runtime_state():
+    adapter = PI052PolicyAdapter(policy=object())
+    state = RuntimeState(
+        task="clean the kitchen",
+        language_context={"memory": "cup moved", "plan": "pick then place"},
+        extra={"prior_subtask": "pick the cup"},
+    )
+
+    assert adapter.build_messages("subtask", state) == [{"role": "user", "content": "clean the kitchen"}]
+    assert adapter.build_messages("memory", state) == [
+        {"role": "user", "content": "clean the kitchen"},
+        {"role": "assistant", "content": "Previous memory: cup moved"},
+        {"role": "user", "content": "Completed subtask: pick the cup"},
+    ]
+    assert adapter.build_messages("interjection", state, user_text="wait") == [
+        {"role": "user", "content": "clean the kitchen"},
+        {"role": "assistant", "content": "Previous plan:\npick then place"},
+        {"role": "user", "content": "wait"},
+    ]
+
+
+def test_pi052_adapter_strips_say_markers_from_plan_text():
+    adapter = PI052PolicyAdapter(policy=object())
+    text = "Move to the sink. <say>heading to the sink</say>"
+
+    assert split_plan_and_say(text) == ("Move to the sink.", "heading to the sink")
+    assert adapter.plan_from_text(text) == "Move to the sink."
+
+
+def test_rollout_language_cli_smoke_does_not_load_model(monkeypatch):
+    """lerobot-rollout dispatches language flags to the adapter-based runtime."""
+    from lerobot.runtime import cli
+    from lerobot.scripts import lerobot_rollout
+
+    fake_policy = SimpleNamespace(config=SimpleNamespace(device="cpu", type="pi052"))
+
+    monkeypatch.setattr(
+        cli,
+        "_load_policy_and_preprocessor",
+        lambda policy_path, **kwargs: (fake_policy, None, None),
+    )
+    monkeypatch.setattr(cli, "_run_repl", lambda runtime, **kwargs: 0)
+
+    assert lerobot_rollout.main(["--policy.path=fake", "--no_robot", "--task=clean", "--max_ticks=0"]) == 0
+
+
+def test_rollout_language_dispatch_preserves_standard_molmoact2_path(monkeypatch):
+    """MolmoAct2 only opts into open prompting when a language flag is present."""
+    from lerobot.scripts import lerobot_rollout
+
+    standard = [
+        "--policy.path=lerobot/MolmoAct2-SO100_101-LeRobot",
+        "--robot.type=so101_follower",
+        "--task=pick up the cube",
+    ]
+    assert not lerobot_rollout._uses_language_runtime(standard)
+    assert lerobot_rollout._uses_language_runtime([*standard, "--direct_subtask"])
+    assert lerobot_rollout._uses_language_runtime(["--policy.path=lerobot/pi052_robocasa", "--sim"])
+
+    standard_calls = []
+    monkeypatch.setattr(lerobot_rollout, "register_third_party_plugins", lambda: None)
+    monkeypatch.setattr(lerobot_rollout, "rollout", lambda: standard_calls.append(True))
+    lerobot_rollout.main(standard)
+    assert standard_calls == [True]

tests/policies/pi052/test_pi052_sdpa_attention.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Numerical-parity tests for the SDPA attention port.
+
+``pi05`` / ``pi052`` replaced the per-layer call from
+``modeling_gemma.eager_attention_forward`` with
+``sdpa_attention_forward`` (PyTorch SDPA + GQA repeat). The forward
+output must be bit-equivalent (within bf16 tolerance) on the masks
+this model actually uses — block-bidirectional with an arbitrary
+additive bias — otherwise we silently change training behaviour.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from transformers.models.gemma import modeling_gemma  # noqa: E402
+
+from lerobot.policies.pi052.modeling_pi052 import make_att_2d_masks  # noqa: E402
+from lerobot.policies.pi_gemma import sdpa_attention_forward  # noqa: E402
+from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE  # noqa: E402
+
+
+def _mock_self_attn(num_kv_groups: int, training: bool = False):
+    """Bare module surface that both forwards read."""
+    return SimpleNamespace(
+        num_key_value_groups=num_kv_groups,
+        training=training,
+    )
+
+
+def _build_inputs(
+    bsize: int,
+    num_heads: int,
+    num_kv_heads: int,
+    seq_len: int,
+    head_dim: int,
+    dtype: torch.dtype,
+    seed: int = 0,
+):
+    g = torch.Generator(device="cpu").manual_seed(seed)
+    q = torch.randn(bsize, num_heads, seq_len, head_dim, dtype=dtype, generator=g)
+    k = torch.randn(bsize, num_kv_heads, seq_len, head_dim, dtype=dtype, generator=g)
+    v = torch.randn(bsize, num_kv_heads, seq_len, head_dim, dtype=dtype, generator=g)
+    return q, k, v
+
+
+def _block_bidirectional_mask(
+    bsize: int, seq_len: int, block_sizes: list[int], dtype: torch.dtype
+) -> torch.Tensor:
+    """Mimic ``_prepare_attention_masks_4d`` on a block layout that
+    matches ``[images, language, suffix]`` from ``embed_prefix`` +
+    ``embed_suffix``: every block bidirectional internally, later
+    blocks visible to earlier ones via the cumulative-block rule.
+    """
+    assert sum(block_sizes) == seq_len
+    att_marks = []
+    for i, n in enumerate(block_sizes):
+        att_marks += [1 if i > 0 else 0] + [0] * (n - 1)
+    pad = torch.ones(bsize, seq_len, dtype=torch.bool)
+    att = torch.tensor(att_marks, dtype=torch.bool)[None].expand(bsize, seq_len)
+    att_2d = make_att_2d_masks(pad, att)
+    bias = torch.where(
+        att_2d[:, None, :, :],
+        torch.zeros((), dtype=dtype),
+        torch.tensor(OPENPI_ATTENTION_MASK_VALUE, dtype=dtype),
+    )
+    return bias
+
+
+@pytest.mark.parametrize(
+    "num_heads,num_kv_heads,head_dim",
+    [
+        (8, 1, 256),  # gemma_2b / paligemma config
+        (8, 8, 64),  # MHA control (no GQA repeat)
+    ],
+)
+def test_sdpa_parity_with_eager_block_bidirectional(num_heads, num_kv_heads, head_dim):
+    """SDPA forward output matches the eager softmax(QK^T)@V on the
+    block-bidirectional mask layout pi05 actually uses."""
+    bsize, seq_len = 2, 13
+    block_sizes = [4, 5, 4]  # images, language, suffix-style blocks
+    dtype = torch.float32  # cpu math kernel — keep fp32 for tight tol
+    scaling = head_dim**-0.5
+
+    q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, dtype)
+    mask = _block_bidirectional_mask(bsize, seq_len, block_sizes, dtype)
+
+    module = _mock_self_attn(num_heads // num_kv_heads)
+
+    out_eager, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
+    out_sdpa, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
+    assert out_eager.shape == out_sdpa.shape
+    torch.testing.assert_close(out_sdpa, out_eager, atol=1e-5, rtol=1e-4)
+
+
+def test_sdpa_parity_bf16():
+    """bf16 path — looser tolerance, must still match eager."""
+    bsize, num_heads, num_kv_heads, seq_len, head_dim = 2, 8, 1, 17, 256
+    scaling = head_dim**-0.5
+    q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, torch.bfloat16)
+    mask = _block_bidirectional_mask(bsize, seq_len, [5, 6, 6], torch.bfloat16)
+    module = _mock_self_attn(num_heads // num_kv_heads)
+
+    out_eager, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
+    out_sdpa, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
+    torch.testing.assert_close(out_sdpa, out_eager, atol=2e-2, rtol=2e-2)
+
+
+def test_sdpa_parity_backward():
+    """Gradients flow through SDPA and match the eager path within
+    bf16 tolerance — critical for any training-side parity claim."""
+    bsize, num_heads, num_kv_heads, seq_len, head_dim = 1, 4, 2, 9, 32
+    scaling = head_dim**-0.5
+    q, k, v = _build_inputs(bsize, num_heads, num_kv_heads, seq_len, head_dim, torch.float32)
+    q.requires_grad_(True)
+    k.requires_grad_(True)
+    v.requires_grad_(True)
+    mask = _block_bidirectional_mask(bsize, seq_len, [3, 3, 3], torch.float32)
+    module = _mock_self_attn(num_heads // num_kv_heads)
+
+    out_e, _ = modeling_gemma.eager_attention_forward(module, q, k, v, mask, scaling)
+    g_q_e, g_k_e, g_v_e = torch.autograd.grad(out_e.sum(), [q, k, v])
+
+    out_s, _ = sdpa_attention_forward(module, q, k, v, mask, scaling)
+    g_q_s, g_k_s, g_v_s = torch.autograd.grad(out_s.sum(), [q, k, v])
+
+    torch.testing.assert_close(g_q_s, g_q_e, atol=1e-5, rtol=1e-4)
+    torch.testing.assert_close(g_k_s, g_k_e, atol=1e-5, rtol=1e-4)
+    torch.testing.assert_close(g_v_s, g_v_e, atol=1e-5, rtol=1e-4)

tests/policies/pi052/test_pi052_text_processor.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for PI052's text tokenizer.
+
+Covers ``say`` tool-call flattening (PaliGemma's flat prompt has no
+structured tool calls, so a ``say`` call must be serialized into a
+``<say>...</say>`` text marker) and EOS-termination supervision (the
+supervised target span must end with an EOS token so the LM head learns
+to stop instead of rambling to ``max_length`` at inference).
+"""
+
+import torch
+
+from lerobot.configs.recipe import MessageTurn, TrainingRecipe
+from lerobot.policies.pi052.text_processor_pi052 import (
+    PI052TextTokenizerStep,
+    _flatten_say_tool_calls,
+    _format_messages,
+)
+from lerobot.processor import PolicyProcessorPipeline
+from lerobot.processor.render_messages_processor import RenderMessagesStep
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import (
+    OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_TOKENS,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
+
+
+def _say_call(text):
+    return {"type": "function", "function": {"name": "say", "arguments": {"text": text}}}
+
+
+def test_flatten_appends_say_marker_and_drops_tool_calls():
+    msg = {"role": "assistant", "content": "Heading to the cube.", "tool_calls": [_say_call("On it!")]}
+    out = _flatten_say_tool_calls(msg)
+    assert "tool_calls" not in out
+    assert out["content"] == "Heading to the cube.\n<say>On it!</say>"
+
+
+def test_flatten_marker_only_when_content_empty_or_none():
+    out = _flatten_say_tool_calls({"role": "assistant", "tool_calls": [_say_call("hi")]})
+    assert out["content"] == "<say>hi</say>"
+
+
+def test_flatten_accepts_json_string_arguments():
+    call = {"type": "function", "function": {"name": "say", "arguments": '{"text": "hello there"}'}}
+    out = _flatten_say_tool_calls({"role": "assistant", "content": "p", "tool_calls": [call]})
+    assert out["content"] == "p\n<say>hello there</say>"
+
+
+def test_flatten_leaves_messages_without_tool_calls_untouched():
+    msg = {"role": "assistant", "content": "just a plan"}
+    assert _flatten_say_tool_calls(msg) == msg
+
+
+def test_flatten_drops_non_say_tool_calls_but_keeps_content():
+    weather = {"type": "function", "function": {"name": "check_weather", "arguments": {}}}
+    out = _flatten_say_tool_calls({"role": "assistant", "content": "plan only", "tool_calls": [weather]})
+    assert out["content"] == "plan only"
+    assert "tool_calls" not in out
+
+
+def test_format_messages_appends_eos_to_target_turns_only():
+    msgs = [
+        {"role": "user", "content": "pick cube"},
+        {"role": "assistant", "content": "move to cube"},
+    ]
+    prompt, spans = _format_messages(msgs, target_indices=[1], eos_token="<eos>")
+    # EOS is appended to the supervised target (assistant) turn only.
+    assert prompt == "User: pick cube\nAssistant: move to cube<eos>\n"
+    # The user span is unchanged; the target span covers content + EOS.
+    assert prompt[spans[0][0] : spans[0][1]] == "pick cube"
+    assert prompt[spans[1][0] : spans[1][1]] == "move to cube<eos>"
+
+
+def test_format_messages_without_eos_args_is_unchanged():
+    """Inference callers omit target_indices / eos_token — no EOS baked in."""
+    prompt, spans = _format_messages([{"role": "user", "content": "hi"}])
+    assert prompt == "User: hi\n"
+    assert prompt[spans[0][0] : spans[0][1]] == "hi"
+
+
+def test_pi052_steps_roundtrip_through_standard_pipeline_loader(tmp_path):
+    recipe = TrainingRecipe(messages=[MessageTurn(role="user", content="${task}", stream="low_level")])
+    pipeline = PolicyProcessorPipeline(
+        steps=[
+            RenderMessagesStep(recipe),
+            PI052TextTokenizerStep(
+                tokenizer_name="custom-tokenizer",
+                max_length=77,
+                plan_dropout_prob=0.2,
+                dropout_seed=3,
+            ),
+        ],
+        name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+    )
+    pipeline.save_pretrained(tmp_path)
+
+    loaded = PolicyProcessorPipeline.from_pretrained(
+        tmp_path, config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
+    )
+
+    assert loaded.steps[0].recipe == recipe
+    assert loaded.steps[1].tokenizer_name == "custom-tokenizer"
+    assert loaded.steps[1].max_length == 77
+    assert loaded.steps[1].plan_dropout_prob == 0.2
+    assert loaded.steps[1].dropout_seed == 3
+
+
+def _eos_char_id() -> int:
+    """Token id _CharTokenizer assigns to its 1-char EOS."""
+    return ord("\x1f") % 251 + 1
+
+
+def test_pi052_text_tokenizer_supervises_eos_at_target_end():
+    """The appended EOS is the last supervised label on a target turn —
+    that's the signal that teaches the LM head to stop. The trailing
+    newline right after it stays unsupervised (-100)."""
+    step = PI052TextTokenizerStep(max_length=64)
+    step._tokenizer = _CharTokenizer()
+    transition = {
+        TransitionKey.OBSERVATION: {},
+        TransitionKey.COMPLEMENTARY_DATA: {
+            "messages": [
+                {"role": "user", "content": "pick cube"},
+                {"role": "assistant", "content": "move to cube"},
+            ],
+            "target_message_indices": [1],
+            "message_streams": ["high_level", "high_level"],
+            "index": torch.tensor(10),
+        },
+    }
+    out = step(transition)
+    ids = out[TransitionKey.OBSERVATION][OBS_LANGUAGE_TOKENS][0]
+    labels = out[TransitionKey.COMPLEMENTARY_DATA]["text_labels"][0]
+
+    supervised = (labels != -100).nonzero().flatten().tolist()
+    assert supervised, "target turn produced no supervised labels"
+    last = supervised[-1]
+    # The last supervised token is the appended EOS.
+    assert int(ids[last]) == _eos_char_id()
+    assert int(labels[last]) == _eos_char_id()
+    # The token right after the EOS (the trailing newline) is NOT supervised.
+    assert int(labels[last + 1]) == -100
+
+
+class _CharTokenizer:
+    pad_token_id = 0
+    eos_token = "\x1f"  # unit separator — a 1-char "EOS" for testing
+
+    def __call__(
+        self,
+        text,
+        max_length,
+        padding,
+        truncation,
+        return_tensors,
+        return_offsets_mapping,
+        padding_side,
+    ):
+        ids = [ord(c) % 251 + 1 for c in text[:max_length]]
+        offsets = [(i, i + 1) for i in range(len(ids))]
+        attention = [1] * len(ids)
+        if padding == "max_length" and len(ids) < max_length:
+            pad = max_length - len(ids)
+            ids += [self.pad_token_id] * pad
+            offsets += [(0, 0)] * pad
+            attention += [0] * pad
+        return {
+            "input_ids": torch.tensor([ids], dtype=torch.long),
+            "attention_mask": torch.tensor([attention], dtype=torch.long),
+            "offset_mapping": torch.tensor([offsets], dtype=torch.long),
+        }
+
+    def decode(self, token_ids, skip_special_tokens=False):
+        return "".join(chr(max(int(i) - 1, 0)) for i in token_ids if int(i) != self.pad_token_id)
+
+
+def test_pi052_text_tokenizer_handles_batched_rendered_messages():
+    step = PI052TextTokenizerStep(max_length=64)
+    step._tokenizer = _CharTokenizer()
+
+    transition = {
+        TransitionKey.OBSERVATION: {},
+        TransitionKey.COMPLEMENTARY_DATA: {
+            "messages": [
+                [
+                    {"role": "user", "content": "pick cube"},
+                    {"role": "assistant", "content": "move to cube"},
+                ],
+                [{"role": "user", "content": "open drawer"}],
+            ],
+            "target_message_indices": [[1], []],
+            "message_streams": [["high_level", "high_level"], ["low_level"]],
+            "index": torch.tensor([10, 11]),
+        },
+    }
+
+    out = step(transition)
+    obs = out[TransitionKey.OBSERVATION]
+    comp = out[TransitionKey.COMPLEMENTARY_DATA]
+
+    assert obs[OBS_LANGUAGE_TOKENS].shape == (2, 64)
+    assert obs[OBS_LANGUAGE_ATTENTION_MASK].shape == (2, 64)
+    assert comp["text_labels"].shape == (2, 64)
+    assert comp["predict_actions"].tolist() == [False, True]
+    assert (comp["text_labels"][0] != -100).any()
+    assert not (comp["text_labels"][1] != -100).any()

tests/policies/pi052/test_pi052_training_time_rtc.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+
+"""Unit coverage for Pi052 training-time RTC conditioning."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05.modeling_pi05 import (  # noqa: E402
+    _build_flow_matching_inputs,
+    _prepare_trained_rtc_prefix,
+    create_sinusoidal_pos_embedding,
+)
+from lerobot.policies.pi052.configuration_pi052 import PI052Config  # noqa: E402
+from lerobot.policies.pi052.modeling_pi052 import (  # noqa: E402
+    PI05Pytorch as PI052Pytorch,
+    _flow_loss_per_sample,
+    _reduce_flow_loss,
+)
+
+
+def test_training_rtc_uses_clean_prefix_and_per_token_time():
+    actions = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]])
+    noise = torch.tensor([[[10.0], [20.0], [30.0], [40.0]]])
+    time = torch.tensor([0.25])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
+
+    assert model_time.tolist() == [[0.0, 0.0, 0.25, 0.25]]
+    assert torch.equal(x_t[:, :2], actions[:, :2])
+    assert torch.equal(x_t[:, 2:], 0.25 * noise[:, 2:] + 0.75 * actions[:, 2:])
+
+
+def test_training_rtc_loss_averages_over_postfix_only():
+    flow_loss = torch.tensor([[[100.0], [100.0], [2.0], [4.0]]])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    per_sample = _flow_loss_per_sample(flow_loss, prefix_mask)
+
+    assert per_sample.tolist() == [3.0]
+
+
+def test_training_rtc_mean_loss_uses_global_postfix_normalization():
+    flow_loss = torch.tensor(
+        [
+            [[1.0], [1.0], [1.0], [1.0]],
+            [[9.0], [9.0], [9.0], [9.0]],
+        ]
+    )
+    prefix_mask = torch.tensor(
+        [
+            [False, False, False, False],
+            [True, True, True, False],
+        ]
+    )
+
+    loss = _reduce_flow_loss(flow_loss, prefix_mask, predict_actions_t=None, reduction="mean")
+
+    assert loss.item() == pytest.approx(13 / 5)
+
+
+def test_per_token_time_embedding_preserves_action_axis():
+    time = torch.tensor([[0.0, 0.5, 1.0]])
+
+    embedding = create_sinusoidal_pos_embedding(time, 8, 4e-3, 4.0, time.device)
+
+    assert embedding.shape == (1, 3, 8)
+    assert not torch.equal(embedding[:, 0], embedding[:, 1])
+
+
+def test_action_expert_embeds_per_token_flow_times():
+    model = PI052Pytorch.__new__(PI052Pytorch)
+    nn.Module.__init__(model)
+    model.config = SimpleNamespace(chunk_size=3, min_period=4e-3, max_period=4.0)
+    model.gradient_checkpointing_enabled = False
+    model.action_in_proj = nn.Linear(2, 8)
+    model.time_mlp_in = nn.Linear(8, 8)
+    model.time_mlp_out = nn.Linear(8, 8)
+
+    suffix, _, _, adarms_cond = model.embed_suffix(
+        torch.randn(2, 3, 2),
+        torch.tensor([[0.0, 0.5, 0.5], [0.0, 0.0, 0.5]]),
+    )
+
+    assert suffix.shape == (2, 3, 8)
+    assert adarms_cond.shape == (2, 3, 8)
+
+
+def test_trained_rtc_prefix_is_padded_and_masked():
+    x_t = torch.randn(1, 5, 4)
+    previous = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+
+    prefix, mask = _prepare_trained_rtc_prefix(x_t, previous, inference_delay=2, training_max_delay=3)
+
+    assert prefix.shape == x_t.shape
+    assert mask.shape == x_t.shape
+    assert torch.equal(prefix[0, :2, :2], previous[:2])
+    assert torch.count_nonzero(prefix[0, :2, 2:]) == 0
+    assert mask[0, :2].all()
+    assert not mask[0, 2:].any()
+
+
+def test_trained_rtc_rejects_delay_outside_training_distribution():
+    with pytest.raises(ValueError, match="exceeds the checkpoint"):
+        _prepare_trained_rtc_prefix(
+            torch.randn(1, 5, 4),
+            torch.randn(4, 2),
+            inference_delay=4,
+            training_max_delay=3,
+        )
+
+
+@pytest.mark.parametrize("non_finite", [float("nan"), float("inf")])
+def test_trained_rtc_rejects_non_finite_prefix(non_finite):
+    previous = torch.randn(4, 2)
+    previous[0, 0] = non_finite
+
+    with pytest.raises(ValueError, match="NaN or Inf"):
+        _prepare_trained_rtc_prefix(
+            torch.randn(1, 5, 4),
+            previous,
+            inference_delay=2,
+            training_max_delay=3,
+        )
+
+
+@pytest.mark.parametrize("max_delay", [-1, 5])
+def test_pi052_config_rejects_invalid_training_rtc_delay(max_delay):
+    with pytest.raises(ValueError, match="rtc_training_max_delay"):
+        PI052Config(chunk_size=5, n_action_steps=5, rtc_training_max_delay=max_delay)

tests/policies/pi052/test_pi052_vision_checkpointing.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import MethodType, SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi052.modeling_pi052 import PI05Pytorch
+
+
+class _MockVisionTower:
+    def __init__(self):
+        self.enable_kwargs = None
+        self.disable_calls = 0
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.enable_kwargs = kwargs
+
+    def gradient_checkpointing_disable(self):
+        self.disable_calls += 1
+
+
+def _checkpoint_model():
+    tower = _MockVisionTower()
+    language_model = SimpleNamespace(gradient_checkpointing=False)
+    expert_model = SimpleNamespace(gradient_checkpointing=False)
+    model = PI05Pytorch.__new__(PI05Pytorch)
+    nn.Module.__init__(model)
+    model.gradient_checkpointing_enabled = False
+    model.paligemma_with_expert = SimpleNamespace(
+        paligemma=SimpleNamespace(model=SimpleNamespace(language_model=language_model, vision_tower=tower)),
+        gemma_expert=SimpleNamespace(model=expert_model),
+    )
+    return model, tower, language_model, expert_model
+
+
+def test_gradient_checkpointing_uses_vision_tower_layer_api():
+    model, tower, language_model, expert_model = _checkpoint_model()
+
+    PI05Pytorch.gradient_checkpointing_enable(model)
+
+    assert model.gradient_checkpointing_enabled
+    assert language_model.gradient_checkpointing
+    assert expert_model.gradient_checkpointing
+    assert tower.enable_kwargs == {"gradient_checkpointing_kwargs": {"use_reentrant": False}}
+
+    PI05Pytorch.gradient_checkpointing_disable(model)
+
+    assert not model.gradient_checkpointing_enabled
+    assert not language_model.gradient_checkpointing
+    assert not expert_model.gradient_checkpointing
+    assert tower.disable_calls == 1
+
+
+def test_siglip_layers_recompute_individually():
+    from transformers.models.siglip.configuration_siglip import SiglipVisionConfig
+    from transformers.models.siglip.modeling_siglip import SiglipVisionModel
+
+    config = SiglipVisionConfig(
+        hidden_size=16,
+        intermediate_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        num_channels=3,
+        image_size=16,
+        patch_size=8,
+    )
+    tower = SiglipVisionModel(config).train()
+    tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    calls = [0] * config.num_hidden_layers
+
+    for index, layer in enumerate(tower.vision_model.encoder.layers):
+        original_forward = layer.forward
+
+        def counted_forward(self, *args, _index=index, _forward=original_forward, **kwargs):
+            calls[_index] += 1
+            return _forward(*args, **kwargs)
+
+        layer.forward = MethodType(counted_forward, layer)
+
+    pixels = torch.randn(2, config.num_channels, config.image_size, config.image_size)
+    tower(pixels).last_hidden_state.sum().backward()
+
+    assert calls == [2] * config.num_hidden_layers
+
+
+def test_embed_prefix_does_not_wrap_the_whole_vision_tower_checkpoint():
+    model = PI05Pytorch.__new__(PI05Pytorch)
+    nn.Module.__init__(model)
+    model.config = SimpleNamespace()
+    model.gradient_checkpointing_enabled = True
+    model.train()
+
+    image_calls = []
+
+    def embed_image(image):
+        image_calls.append(image.shape)
+        return image[:, :1, 0, :2]
+
+    def embed_language_tokens(tokens):
+        return tokens.to(torch.float32).unsqueeze(-1).expand(*tokens.shape, 2)
+
+    model.paligemma_with_expert = SimpleNamespace(
+        embed_image=embed_image,
+        embed_language_tokens=embed_language_tokens,
+    )
+    outer_checkpoint_calls = []
+
+    def apply_checkpoint(func, value):
+        outer_checkpoint_calls.append(value.shape)
+        return func(value)
+
+    model._apply_checkpoint = apply_checkpoint
+
+    images = [torch.randn(2, 3, 4, 4), torch.randn(2, 3, 4, 4)]
+    image_masks = [torch.ones(2, dtype=torch.bool) for _ in images]
+    tokens = torch.ones(2, 3, dtype=torch.long)
+    token_masks = torch.ones_like(tokens, dtype=torch.bool)
+
+    embeddings, _, _ = model.embed_prefix(images, image_masks, tokens, token_masks)
+
+    assert image_calls == [image.shape for image in images]
+    assert outer_checkpoint_calls == [tokens.shape]
+    assert embeddings.shape == (2, 5, 2)

tests/policies/pi0_fast/test_pi0_fast_tokenizer_fit.py

@@ -0,0 +1,98 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+
+import pytest
+
+from lerobot.policies import factory
+from lerobot.policies.pi0_fast.configuration_pi0_fast import PI0FastConfig
+from lerobot.policies.pi052 import fit_fast_tokenizer as fit_module
+
+
+def test_pi0_fast_resolves_dataset_specific_tokenizer(monkeypatch, tmp_path):
+    config = PI0FastConfig(
+        auto_fit_fast_tokenizer=True,
+        action_tokenizer_name="base-tokenizer",
+        fast_tokenizer_cache_dir=str(tmp_path),
+        fast_tokenizer_fit_samples=17,
+        chunk_size=12,
+        n_action_steps=12,
+    )
+    received = {}
+
+    def fake_fit(**kwargs):
+        received.update(kwargs)
+        return "/cache/fitted-tokenizer"
+
+    monkeypatch.setattr(fit_module, "fit_fast_tokenizer", fake_fit)
+
+    assert fit_module.resolve_fast_tokenizer(config, "user/dataset") == "/cache/fitted-tokenizer"
+    assert received == {
+        "dataset_repo_id": "user/dataset",
+        "cache_dir": tmp_path,
+        "base_tokenizer_name": "base-tokenizer",
+        "n_samples": 17,
+        "chunk_size": 12,
+    }
+
+
+def test_fast_fit_failure_is_not_silently_replaced(monkeypatch, tmp_path):
+    config = PI0FastConfig(auto_fit_fast_tokenizer=True, fast_tokenizer_cache_dir=str(tmp_path))
+    monkeypatch.setattr(
+        fit_module,
+        "fit_fast_tokenizer",
+        lambda **kwargs: (_ for _ in ()).throw(RuntimeError("fit failed")),
+    )
+
+    with pytest.raises(RuntimeError, match="fit failed"):
+        fit_module.resolve_fast_tokenizer(config, "user/dataset")
+
+
+def test_each_node_uses_its_local_rank_zero_as_fit_leader(monkeypatch):
+    monkeypatch.setenv("RANK", "8")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    assert fit_module._is_local_leader()
+
+    monkeypatch.setenv("LOCAL_RANK", "1")
+    assert not fit_module._is_local_leader()
+
+
+def test_pretrained_pi0_fast_overrides_only_fitted_tokenizer(monkeypatch):
+    config = PI0FastConfig(auto_fit_fast_tokenizer=True)
+    calls = []
+
+    monkeypatch.setattr(
+        fit_module,
+        "resolve_fast_tokenizer",
+        lambda config, dataset_repo_id: "/cache/fitted-tokenizer",
+    )
+
+    def fake_from_pretrained(cls, *args, **kwargs):
+        calls.append(kwargs)
+        return SimpleNamespace(steps=[])
+
+    monkeypatch.setattr(factory.PolicyProcessorPipeline, "from_pretrained", classmethod(fake_from_pretrained))
+
+    factory.make_pre_post_processors(
+        config,
+        pretrained_path="checkpoint",
+        dataset_repo_id="user/dataset",
+    )
+
+    assert calls[0]["overrides"] == {
+        "action_tokenizer_processor": {"action_tokenizer_name": "/cache/fitted-tokenizer"}
+    }

tests/policies/pi0_pi05/test_pi05.py

@@ -16,8 +16,12 @@
 
 """Test script to verify PI0.5 (pi05) support in PI0 policy"""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
+from safetensors.torch import save_file
+from torch import nn
 
 pytest.importorskip("transformers")
 
@@ -31,6 +35,26 @@ from lerobot.utils.random_utils import set_seed
 from tests.utils import require_cuda, require_hf_token  # noqa: E402
 
 
+class _CheckpointPolicy(PI05Policy):
+    def __init__(self, config, **kwargs):
+        nn.Module.__init__(self)
+        self.config = config
+        self.loaded_state_dict = None
+
+    def load_state_dict(self, state_dict, strict=True, assign=False):
+        self.loaded_state_dict = state_dict
+        return [], []
+
+
+def test_from_pretrained_loads_existing_single_file_checkpoint(tmp_path):
+    save_file({"weight": torch.tensor([1.0])}, tmp_path / "model.safetensors")
+
+    policy = _CheckpointPolicy.from_pretrained(tmp_path, config=SimpleNamespace())
+
+    assert policy.loaded_state_dict is not None
+    torch.testing.assert_close(policy.loaded_state_dict["model.weight"], torch.tensor([1.0]))
+
+
 @require_cuda
 @require_hf_token
 def test_policy_instantiation():

tests/policies/pi0_pi05/test_pi05_training_time_rtc.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.modeling_pi05 import (  # noqa: E402
+    _build_flow_matching_inputs,
+    _reduce_training_rtc_loss,
+)
+
+
+def test_pi05_training_rtc_uses_clean_prefix_and_per_token_time():
+    actions = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]])
+    noise = torch.tensor([[[10.0], [20.0], [30.0], [40.0]]])
+    time = torch.tensor([0.25])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
+
+    assert model_time.tolist() == [[0.0, 0.0, 0.25, 0.25]]
+    assert torch.equal(x_t[:, :2], actions[:, :2])
+    assert torch.equal(x_t[:, 2:], 0.25 * noise[:, 2:] + 0.75 * actions[:, 2:])
+
+
+def test_pi05_training_rtc_loss_excludes_clean_prefix():
+    losses = torch.tensor([[[100.0], [100.0], [2.0], [4.0]]])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean")
+
+    assert loss.item() == pytest.approx(3.0)
+
+
+@pytest.mark.parametrize("max_delay", [-1, 5])
+def test_pi05_config_rejects_invalid_training_rtc_delay(max_delay):
+    with pytest.raises(ValueError, match="rtc_training_max_delay"):
+        PI05Config(chunk_size=5, n_action_steps=5, rtc_training_max_delay=max_delay)

tests/policies/rtc/test_configuration_rtc.py

@@ -16,6 +16,8 @@
 
 """Tests for RTC configuration module."""
 
+import pytest
+
 from lerobot.configs.types import RTCAttentionSchedule
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
 
@@ -27,6 +29,7 @@ def test_rtc_config_default_initialization():
     config = RTCConfig()
 
     assert config.enabled is True
+    assert config.mode == "guided"
     assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
     assert config.max_guidance_weight == 10.0
     assert config.execution_horizon == 10
@@ -34,10 +37,16 @@ def test_rtc_config_default_initialization():
     assert config.debug_maxlen == 100
 
 
+def test_rtc_config_rejects_unknown_mode():
+    with pytest.raises(ValueError, match="mode must be"):
+        RTCConfig(mode="unknown")
+
+
 def test_rtc_config_custom_initialization():
     """Test RTCConfig initializes with custom values."""
     config = RTCConfig(
         enabled=True,
+        mode="trained",
         prefix_attention_schedule=RTCAttentionSchedule.EXP,
         max_guidance_weight=5.0,
         execution_horizon=20,
@@ -46,6 +55,7 @@ def test_rtc_config_custom_initialization():
     )
 
     assert config.enabled is True
+    assert config.mode == "trained"
     assert config.prefix_attention_schedule == RTCAttentionSchedule.EXP
     assert config.max_guidance_weight == 5.0
     assert config.execution_horizon == 20

tests/policies/rtc/test_modeling_rtc.py

@@ -93,6 +93,19 @@ def test_rtc_processor_initialization_without_debug(rtc_config_debug_disabled):
     assert processor.tracker is None
 
 
+def test_rtc_processor_rejects_trained_mode_when_policy_does_not_support_it():
+    config = RTCConfig(mode="trained")
+
+    with pytest.raises(ValueError, match="requires a PI05-compatible checkpoint"):
+        RTCProcessor(config)
+
+    processor = RTCProcessor(config, trained_mode_supported=True)
+    assert processor.rtc_config.mode == "trained"
+
+    disabled = RTCProcessor(RTCConfig(enabled=False, mode="trained"))
+    assert disabled.rtc_config.enabled is False
+
+
 # ====================== Tracker Proxy Methods Tests ======================
 
 

tests/processor/test_render_messages_processor.py

@@ -12,7 +12,9 @@ from lerobot.processor.render_messages_processor import RenderMessagesStep  # no
 from lerobot.types import TransitionKey  # noqa: E402
 
 
-def test_render_messages_step_noops_without_language_columns():
+def test_render_messages_step_renders_task_fallback_without_language_columns():
+    """No language columns + a task string → low-level task fallback render,
+    matching what the policy sees at eval time on unannotated observations."""
     recipe = TrainingRecipe(
         messages=[
             MessageTurn(role="user", content="${task}", stream="high_level"),
@@ -21,6 +23,24 @@ def test_render_messages_step_noops_without_language_columns():
     )
     transition = create_transition(complementary_data={"task": "do it"})
 
+    out = RenderMessagesStep(recipe)(transition)
+    data = out[TransitionKey.COMPLEMENTARY_DATA]
+
+    assert data["messages"] == [{"role": "user", "content": "do it"}]
+    assert data["message_streams"] == ["low_level"]
+    assert data["target_message_indices"] == []
+    assert data["task"] == "do it"
+
+
+def test_render_messages_step_noops_without_language_columns_or_task():
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(role="user", content="${task}", stream="high_level"),
+            MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
+        ]
+    )
+    transition = create_transition(complementary_data={})
+
     assert RenderMessagesStep(recipe)(transition) == transition
 
 
@@ -58,3 +78,70 @@ def test_render_messages_step_renders_and_drops_raw_language():
     assert data["messages"][-1]["content"] == "reach carefully"
     assert data["message_streams"] == ["high_level", "low_level"]
     assert data["target_message_indices"] == [1]
+
+
+def test_render_messages_step_falls_back_to_low_level_task_when_recipe_misses():
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(
+                role="assistant",
+                content="${subtask}",
+                stream="high_level",
+                target=True,
+                if_present="subtask",
+            ),
+        ]
+    )
+    transition = create_transition(
+        complementary_data={
+            "task": "pick the cube",
+            "timestamp": torch.tensor(0.0),
+            "index": torch.tensor(7),
+            "language_persistent": [],
+            "language_events": [{"style": "unmatched", "timestamp": 0.0}],
+        }
+    )
+
+    out = RenderMessagesStep(recipe)(transition)
+    data = out[TransitionKey.COMPLEMENTARY_DATA]
+
+    assert data["messages"] == [{"role": "user", "content": "pick the cube"}]
+    assert data["message_streams"] == ["low_level"]
+    assert data["target_message_indices"] == []
+
+
+def test_render_messages_step_falls_back_per_sample_in_batched_language():
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(
+                role="assistant",
+                content="${subtask}",
+                stream="high_level",
+                target=True,
+                if_present="subtask",
+            ),
+        ]
+    )
+    transition = create_transition(
+        action=torch.arange(4).reshape(2, 2),
+        complementary_data={
+            "task": ["pick the cube", "open the drawer"],
+            "timestamp": torch.tensor([0.0, 1.0]),
+            "index": torch.tensor([7, 8]),
+            "language_persistent": [[], []],
+            "language_events": [
+                [{"style": "unmatched", "timestamp": 0.0}],
+                [{"style": "unmatched", "timestamp": 1.0}],
+            ],
+        },
+    )
+
+    out = RenderMessagesStep(recipe)(transition)
+    data = out[TransitionKey.COMPLEMENTARY_DATA]
+
+    assert data["messages"] == [
+        [{"role": "user", "content": "pick the cube"}],
+        [{"role": "user", "content": "open the drawer"}],
+    ]
+    assert data["message_streams"] == [["low_level"], ["low_level"]]
+    assert data["target_message_indices"] == [[], []]

tests/processor/test_tokenizer_processor.py

@@ -25,7 +25,7 @@ import pytest
 import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
-from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep
+from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, TokenizerProcessorStep
 from lerobot.processor.converters import create_transition, identity_transition
 from lerobot.types import TransitionKey
 from lerobot.utils.constants import (
@@ -88,6 +88,24 @@ class MockTokenizer:
         return result
 
 
+def test_action_tokenizer_config_preserves_token_mapping():
+    processor = object.__new__(ActionTokenizerProcessorStep)
+    processor.trust_remote_code = True
+    processor.max_action_tokens = 384
+    processor.fast_skip_tokens = 64
+    processor.paligemma_tokenizer_name = "custom/paligemma"
+    processor.action_tokenizer_name = "custom/fast"
+    processor.action_tokenizer_input_object = None
+
+    assert processor.get_config() == {
+        "trust_remote_code": True,
+        "max_action_tokens": 384,
+        "fast_skip_tokens": 64,
+        "paligemma_tokenizer_name": "custom/paligemma",
+        "action_tokenizer_name": "custom/fast",
+    }
+
+
 @pytest.fixture
 def mock_tokenizer():
     """Provide a mock tokenizer for testing."""

tests/runtime/test_adapter.py

@@ -0,0 +1,105 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from lerobot.runtime import RuntimeState
+from lerobot.runtime.adapter import (
+    BaseLanguageAdapter,
+    DirectTaskPolicyAdapter,
+    GenerationConfig,
+)
+
+
+class ScriptedAdapter(BaseLanguageAdapter):
+    """Base adapter whose text generation returns queued strings per kind."""
+
+    def __init__(self, scripts, gen=None):
+        super().__init__(policy=object(), gen=gen)
+        self.scripts = {k: list(v) for k, v in scripts.items()}
+        self.calls = []
+
+    def select_action(self, observation, state):
+        return None
+
+    def generate_text(self, kind, observation, state, user_text=None):
+        self.calls.append(kind)
+        queue = self.scripts.get(kind, [])
+        return queue.pop(0) if queue else ""
+
+
+def test_cascade_sets_subtask_then_memory():
+    adapter = ScriptedAdapter({"subtask": ["pick the red cup"], "memory": ["the cup is grasped"]})
+    state = RuntimeState(task="clean")
+
+    adapter.update_language_state(None, state)
+
+    assert state.language_context["subtask"] == "pick the red cup"
+    assert state.language_context["memory"] == "the cup is grasped"
+    assert adapter.calls == ["subtask", "memory"]
+
+
+def test_nonempty_generation_is_used_verbatim():
+    adapter = ScriptedAdapter({"subtask": [":::: ::"], "memory": ["memory"]})
+    state = RuntimeState(task="clean")
+
+    adapter.update_language_state(None, state)
+
+    assert state.language_context["subtask"] == ":::: ::"
+    assert state.language_context["memory"] == "memory"
+    assert adapter.calls == ["subtask", "memory"]
+
+
+def test_throttle_regenerates_every_n_chunks():
+    adapter = ScriptedAdapter(
+        {
+            "subtask": ["pick the first cup", "pick the second cup"],
+            "memory": ["memory one two three", "memory four five six"],
+        },
+        gen=GenerationConfig(chunks_per_regen=2),
+    )
+    state = RuntimeState(task="clean")
+
+    adapter.update_language_state(None, state)  # generates
+    assert state.language_context["subtask"] == "pick the first cup"
+    adapter.update_language_state(None, state)  # throttled — no generation
+    assert state.language_context["subtask"] == "pick the first cup"
+    adapter.update_language_state(None, state)  # generates again
+    assert state.language_context["subtask"] == "pick the second cup"
+
+
+def test_handle_interjection_sets_plan_and_strips_say():
+    adapter = ScriptedAdapter({"interjection": ["turn to the left now <say>heading left</say>"]})
+    state = RuntimeState(task="clean")
+
+    adapter.handle_interjection("turn", None, state)
+
+    assert state.language_context["plan"] == "turn to the left now"
+
+
+def test_direct_task_adapter_delegates_action_chunk():
+    class Policy:
+        def predict_action_chunk(self, observation):
+            return ("chunk", observation)
+
+    observation = {"task": "pick up the cube"}
+    adapter = DirectTaskPolicyAdapter(Policy())
+
+    assert adapter.select_action(observation, RuntimeState()) == ("chunk", observation)
+    assert adapter.generate_text("subtask", observation, RuntimeState()) == ""
+
+
+def test_flat_policy_registry_reuses_direct_task_adapter():
+    from lerobot.runtime.registry import get_language_adapter_factory
+
+    assert get_language_adapter_factory("pi05") is DirectTaskPolicyAdapter
+    assert get_language_adapter_factory("molmoact2") is DirectTaskPolicyAdapter

tests/runtime/test_cli.py

@@ -0,0 +1,75 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from lerobot.runtime.cli import _build_rollout_runtime_io, _parse_args
+
+
+def test_parse_args_preserves_rollout_robot_overrides():
+    args = _parse_args(
+        [
+            "--policy.path=checkpoint",
+            "--robot.type=so101_follower",
+            "--robot.calibration_dir=/tmp/calibration",
+        ]
+    )
+
+    assert args.robot_type == "so101_follower"
+    assert "--robot.calibration_dir=/tmp/calibration" in args.raw_argv
+
+
+def test_parse_args_rejects_removed_dataset_replay_flags():
+    with pytest.raises(SystemExit):
+        _parse_args(["--policy.path=checkpoint", "--dataset.repo_id=dataset"])
+
+
+def test_rollout_runtime_io_uses_context_processors():
+    robot = MagicMock()
+    robot.robot_type = "mock_robot"
+    robot.cameras = {}
+    robot.get_observation.return_value = {"joint.pos": 1.5}
+    ctx = SimpleNamespace(
+        hardware=SimpleNamespace(robot_wrapper=robot),
+        runtime=SimpleNamespace(cfg=SimpleNamespace(device="cpu")),
+        processors=SimpleNamespace(
+            robot_observation_processor=lambda observation: observation,
+            robot_action_processor=lambda pair: pair[0],
+        ),
+        policy=SimpleNamespace(
+            preprocessor=lambda observation: observation,
+            postprocessor=lambda action: action,
+        ),
+        data=SimpleNamespace(
+            dataset_features={
+                "observation.state": {
+                    "dtype": "float32",
+                    "shape": (1,),
+                    "names": ["joint.pos"],
+                },
+                "action": {"dtype": "float32", "shape": (1,), "names": ["joint.pos"]},
+            }
+        ),
+    )
+    provider, executor = _build_rollout_runtime_io(ctx, rerun_log=False, get_task=lambda: "move")
+
+    observation = provider()
+    executor(torch.tensor([[2.0]]))
+
+    assert observation["observation.state"].shape == (1, 1)
+    robot.send_action.assert_called_once_with({"joint.pos": 2.0})

tests/runtime/test_language_runtime.py

@@ -0,0 +1,100 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import threading
+import time
+
+from lerobot.runtime import LanguageConditionedRuntime, Tick
+
+
+class FakeAdapter:
+    def __init__(self):
+        self.updated = False
+        self.interjections = []
+
+    def select_action(self, observation, state):
+        assert observation == {"observation.state": 1}
+        assert state.task == "clean"
+        return ["a0", "a1"]
+
+    def update_language_state(self, observation, state):
+        self.updated = True
+        state.set_context("subtask", "pick cup", label="subtask")
+
+    def handle_interjection(self, user_text, observation, state):
+        self.interjections.append(user_text)
+        state.set_context("plan", "new plan", label="plan")
+
+
+def test_runtime_tick_updates_language_enqueues_and_dispatches_action():
+    adapter = FakeAdapter()
+    executed = []
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=adapter,
+        observation_provider=lambda: {"observation.state": 1},
+        action_executor=executed.append,
+    )
+    runtime.set_task("clean")
+
+    logs = runtime.step_once()
+
+    assert adapter.updated
+    assert runtime.state.language_context["subtask"] == "pick cup"
+    assert executed == ["a0"]
+    assert list(runtime.state.action_queue) == ["a1"]
+    assert "  subtask: pick cup" in logs
+
+
+def test_runtime_handles_user_interjection():
+    adapter = FakeAdapter()
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=adapter,
+        observation_provider=lambda: {"observation.state": 1},
+    )
+    runtime.set_task("clean")
+    runtime.state.extra["recent_interjection"] = "please say ok"
+    runtime.state.emit("user_interjection")
+
+    runtime.step_once()
+
+    assert "please say ok" in adapter.interjections
+    assert runtime.state.language_context["plan"] == "new plan"
+
+
+def test_prompt_change_discards_in_flight_action_chunk():
+    started = threading.Event()
+    release = threading.Event()
+
+    class BlockingAdapter(FakeAdapter):
+        def select_action(self, observation, state):
+            started.set()
+            assert release.wait(timeout=2)
+            return ["stale"]
+
+    runtime = LanguageConditionedRuntime(
+        policy_adapter=BlockingAdapter(),
+        observation_provider=lambda: {"observation.state": 1},
+    )
+    runtime.set_task("old task")
+    runtime.state.tick = Tick(index=1, monotonic_seconds=time.monotonic())
+    inference = threading.Thread(target=runtime.maybe_enqueue_action_chunk, kwargs={"force": True})
+    inference.start()
+    assert started.wait(timeout=2)
+
+    runtime.set_task("new task")
+    release.set()
+    inference.join(timeout=2)
+
+    assert not inference.is_alive()
+    assert list(runtime.state.action_queue) == []

tests/runtime/test_sim_robocasa.py

@@ -0,0 +1,82 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+
+from lerobot.runtime.sim_robocasa import RoboCasaSimBackend
+from lerobot.utils.video_annotation import annotate_frame
+
+
+def test_overlay_draws_each_label_once(monkeypatch):
+    put_text_calls = []
+    rectangle_calls = []
+
+    def put_text(image, text, origin, font, scale, color, thickness, line_type):
+        put_text_calls.append((text, color, thickness))
+        return image
+
+    def rectangle(image, start, end, color, thickness):
+        rectangle_calls.append((start, end, color, thickness))
+        return image
+
+    def add_weighted(src1, alpha, src2, beta, gamma, *, dst):
+        dst[:] = src1 * alpha + src2 * beta + gamma
+        return dst
+
+    fake_cv2 = SimpleNamespace(
+        FONT_HERSHEY_SIMPLEX=0,
+        LINE_AA=16,
+        getTextSize=lambda text, font, scale, thickness: ((len(text) * 7, 10), 0),
+        putText=put_text,
+        rectangle=rectangle,
+        addWeighted=add_weighted,
+    )
+    monkeypatch.setitem(sys.modules, "cv2", fake_cv2)
+
+    frame = np.full((120, 480, 3), 200, dtype=np.uint8)
+    annotated = annotate_frame(
+        frame,
+        (("Task", "close the fridge"), ("Subtask", "reach for the handle"), ("Memory", None)),
+    )
+
+    assert [call[0] for call in put_text_calls] == [
+        "Task: close the fridge",
+        "Subtask: reach for the handle",
+    ]
+    assert all(color == (255, 255, 255) and thickness == 1 for _, color, thickness in put_text_calls)
+    assert len(rectangle_calls) == 1
+    assert not np.shares_memory(annotated, frame)
+
+
+def test_capture_updates_live_frame_when_recording_is_disabled(monkeypatch):
+    backend = object.__new__(RoboCasaSimBackend)
+    frame = np.full((8, 8, 3), 42, dtype=np.uint8)
+    written = []
+    backend.record = False
+    backend.runtime_state = None
+    backend._multiview_frame = lambda: frame
+    backend._current_task = lambda: "task"
+    backend._subtask_getter = None
+    backend._memory_getter = None
+    backend._latest_frame = None
+    backend._write_live_frame = written.append
+    monkeypatch.setattr("lerobot.runtime.sim_robocasa.annotate_frame", lambda image, labels: image)
+
+    backend._capture_frame()
+
+    assert backend._latest_frame is frame
+    assert written == [frame]

tests/test_rollout.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import dataclasses
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -98,6 +99,131 @@ def test_inference_config_types():
     assert rtc.rtc is not None
 
 
+def test_trained_rtc_retries_chunk_when_measured_delay_exceeds_conditioning():
+    from lerobot.rollout.inference.rtc import _trained_rtc_chunk_can_merge
+
+    assert not _trained_rtc_chunk_can_merge(
+        conditioned_delay=2,
+        measured_delay=3,
+        training_max_delay=4,
+        has_previous_actions=True,
+    )
+    assert _trained_rtc_chunk_can_merge(
+        conditioned_delay=2,
+        measured_delay=5,
+        training_max_delay=4,
+        has_previous_actions=False,
+    )
+
+
+def test_trained_rtc_bootstraps_first_overlap_with_checkpoint_capacity():
+    from lerobot.rollout.inference.rtc import _estimate_rtc_delay
+
+    assert (
+        _estimate_rtc_delay(
+            latency=0,
+            time_per_step=1 / 30,
+            mode="trained",
+            training_max_delay=10,
+            has_previous_actions=False,
+        )
+        == 0
+    )
+    assert (
+        _estimate_rtc_delay(
+            latency=0,
+            time_per_step=1 / 30,
+            mode="trained",
+            training_max_delay=10,
+            has_previous_actions=True,
+        )
+        == 10
+    )
+
+
+def test_trained_rtc_rejects_measured_delay_above_checkpoint_support():
+    from lerobot.rollout.inference.rtc import (
+        _trained_rtc_chunk_can_merge,
+        _TrainedRTCDelayExceededError,
+    )
+
+    with pytest.raises(_TrainedRTCDelayExceededError, match="rtc_training_max_delay"):
+        _trained_rtc_chunk_can_merge(
+            conditioned_delay=3,
+            measured_delay=5,
+            training_max_delay=4,
+            has_previous_actions=True,
+        )
+
+
+def test_trained_rtc_rejects_prefix_shorter_than_conditioned_delay():
+    from lerobot.rollout.inference.rtc import (
+        _TrainedRTCPrefixUnavailableError,
+        _validate_trained_rtc_prefix_available,
+    )
+
+    with pytest.raises(_TrainedRTCPrefixUnavailableError, match="only 2"):
+        _validate_trained_rtc_prefix_available(conditioned_delay=4, available_steps=2)
+
+
+@pytest.mark.parametrize(
+    ("execution_horizon", "queue_threshold", "match"),
+    [
+        (3, 4, "execution_horizon"),
+        (4, 3, "queue_threshold"),
+    ],
+)
+def test_trained_rtc_rollout_requires_capacity_for_max_delay(execution_horizon, queue_threshold, match):
+    from lerobot.policies.rtc.configuration_rtc import RTCConfig
+    from lerobot.rollout.context import _validate_trained_rtc_rollout_config
+    from lerobot.rollout.inference import RTCInferenceConfig
+
+    policy_config = SimpleNamespace(type="pi052", rtc_training_max_delay=4)
+    inference_config = RTCInferenceConfig(
+        rtc=RTCConfig(mode="trained", execution_horizon=execution_horizon),
+        queue_threshold=queue_threshold,
+    )
+
+    with pytest.raises(ValueError, match=match):
+        _validate_trained_rtc_rollout_config(policy_config, inference_config)
+
+
+def test_relative_state_order_follows_checkpoint_action_names():
+    from lerobot.rollout.context import _align_relative_state_feature_order
+    from lerobot.utils.constants import OBS_STATE
+    from lerobot.utils.feature_utils import build_dataset_frame
+
+    hw_features = {
+        OBS_STATE: {
+            "dtype": "float32",
+            "shape": (4,),
+            "names": ["left_joint.pos", "left_gripper.pos", "right_joint.pos", "right_gripper.pos"],
+        }
+    }
+    checkpoint_order = [
+        "right_joint.pos",
+        "right_gripper.pos",
+        "left_joint.pos",
+        "left_gripper.pos",
+    ]
+
+    aligned = _align_relative_state_feature_order(hw_features, checkpoint_order)
+    frame = build_dataset_frame(
+        aligned,
+        {
+            "left_joint.pos": 1.0,
+            "left_gripper.pos": 2.0,
+            "right_joint.pos": 3.0,
+            "right_gripper.pos": 4.0,
+        },
+        prefix="observation",
+    )
+
+    assert aligned[OBS_STATE]["names"] == checkpoint_order
+    assert frame[OBS_STATE].tolist() == [3.0, 4.0, 1.0, 2.0]
+    assert hw_features[OBS_STATE]["names"][0] == "left_joint.pos"
+
+
 def test_sentry_config_defaults():
     from lerobot.rollout import SentryStrategyConfig
 

tests/utils/test_logging_utils.py

@@ -37,6 +37,12 @@ class MockAccelerator:
             return self._reduce_fn(tensor, reduction)
         return tensor
 
+    def gather(self, tensor):
+        if self._reduce_fn is None:
+            return tensor.repeat(self.num_processes)
+        reduced = self._reduce_fn(tensor, "max")
+        return torch.cat([tensor.repeat(self.num_processes - 1), reduced])
+
 
 def test_average_meter_initialization():
     meter = AverageMeter("loss", ":.2f")
@@ -168,6 +174,18 @@ def test_metrics_tracker_reset_averages(mock_metrics):
     assert tracker.accuracy.avg == 0.0
 
 
+def test_metrics_tracker_materializes_full_tensor_window(mock_metrics):
+    tracker = MetricsTracker(batch_size=2, num_frames=10, num_episodes=2, metrics=mock_metrics)
+    tracker.accumulate_tensor("loss", torch.tensor(1.0))
+    tracker.accumulate_tensor("loss", torch.tensor(3.0))
+
+    assert tracker.loss.count == 0
+    tracker.materialize_tensors()
+
+    assert tracker.loss.avg == pytest.approx(2.0)
+    assert tracker.loss.count == 2
+
+
 def test_average_meter_invalid_reduction():
     with pytest.raises(ValueError):
         AverageMeter("loss", reduction="median")

uv.lock

@@ -1129,7 +1129,7 @@ name = "decord"
 version = "0.6.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "numpy", marker = "(platform_machine != 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" },
+    { name = "numpy", marker = "(platform_machine != 'arm64' and platform_machine != 's390x' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'linux')" },
 ]
 wheels = [
     { url = "https://files.pythonhosted.org/packages/11/79/936af42edf90a7bd4e41a6cac89c913d4b47fa48a26b042d5129a9242ee3/decord-0.6.0-py3-none-manylinux2010_x86_64.whl", hash = "sha256:51997f20be8958e23b7c4061ba45d0efcd86bffd5fe81c695d0befee0d442976", size = 13602299, upload-time = "2021-06-14T21:30:55.486Z" },
@@ -2910,6 +2910,7 @@ all = [
     { name = "ruff" },
     { name = "scikit-image" },
     { name = "scipy" },
+    { name = "sentencepiece" },
     { name = "teleop" },
     { name = "timm" },
     { name = "torchcodec", marker = "(platform_machine == 'arm64' and sys_platform == 'darwin') or (platform_machine == 'AMD64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'arm64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux') or sys_platform == 'win32'" },
@@ -3158,6 +3159,7 @@ phone = [
 ]
 pi = [
     { name = "scipy" },
+    { name = "sentencepiece" },
     { name = "transformers" },
 ]
 placo-dep = [
@@ -3218,6 +3220,9 @@ sarm = [
 scipy-dep = [
     { name = "scipy" },
 ]
+sentencepiece-dep = [
+    { name = "sentencepiece" },
+]
 smolvla = [
     { name = "accelerate" },
     { name = "num2words" },
@@ -3418,6 +3423,7 @@ requires-dist = [
     { name = "lerobot", extras = ["scipy-dep"], marker = "extra == 'phone'" },
     { name = "lerobot", extras = ["scipy-dep"], marker = "extra == 'pi'" },
     { name = "lerobot", extras = ["scipy-dep"], marker = "extra == 'wallx'" },
+    { name = "lerobot", extras = ["sentencepiece-dep"], marker = "extra == 'pi'" },
     { name = "lerobot", extras = ["smolvla"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["test"], marker = "extra == 'all'" },
     { name = "lerobot", extras = ["timm-dep"], marker = "extra == 'groot'" },
@@ -3493,6 +3499,7 @@ requires-dist = [
     { name = "scikit-image", marker = "extra == 'video-benchmark'", specifier = ">=0.23.2,<0.26.0" },
     { name = "scipy", marker = "extra == 'all'", specifier = ">=1.14.0,<2.0.0" },
     { name = "scipy", marker = "extra == 'scipy-dep'", specifier = ">=1.14.0,<2.0.0" },
+    { name = "sentencepiece", marker = "extra == 'sentencepiece-dep'", specifier = ">=0.2.0,<0.3.0" },
     { name = "setuptools", specifier = ">=71.0.0,<81.0.0" },
     { name = "teleop", marker = "extra == 'phone'", specifier = ">=0.1.0,<0.2.0" },
     { name = "termcolor", specifier = ">=2.4.0,<4.0.0" },
@@ -3509,7 +3516,7 @@ requires-dist = [
     { name = "transformers", marker = "extra == 'transformers-dep'", specifier = ">=5.4.0,<5.6.0" },
     { name = "wandb", marker = "extra == 'training'", specifier = ">=0.24.0,<0.28.0" },
 ]
-provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "timm-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "evo1", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
+provides-extras = ["dataset", "training", "hardware", "viz", "core-scripts", "evaluation", "dataset-viz", "av-dep", "pygame-dep", "placo-dep", "transformers-dep", "sentencepiece-dep", "grpcio-dep", "accelerate-dep", "can-dep", "peft-dep", "scipy-dep", "diffusers-dep", "qwen-vl-utils-dep", "matplotlib-dep", "pyserial-dep", "deepdiff-dep", "pynput-dep", "pyzmq-dep", "motorbridge-dep", "motorbridge-smart-servo-dep", "timm-dep", "feetech", "dynamixel", "damiao", "robstride", "openarms", "gamepad", "hopejr", "lekiwi", "unitree-g1", "reachy2", "rebot", "kinematics", "intelrealsense", "phone", "diffusion", "wallx", "pi", "molmoact2", "smolvla", "multi-task-dit", "groot", "sarm", "robometer", "topreward", "xvla", "eo1", "fastwam", "evo1", "hilserl", "vla-jepa", "lingbot-va", "async", "peft", "annotations", "dev", "notebook", "test", "video-benchmark", "aloha", "pusht", "libero", "metaworld", "all"]
 
 [[package]]
 name = "librt"
@@ -6231,6 +6238,54 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1c/78/504fdd027da3b84ff1aecd9f6957e65f35134534ccc6da8628eb71e76d3f/send2trash-2.1.0-py3-none-any.whl", hash = "sha256:0da2f112e6d6bb22de6aa6daa7e144831a4febf2a87261451c4ad849fe9a873c", size = 17610, upload-time = "2026-01-14T06:27:35.218Z" },
 ]
 
+[[package]]
+name = "sentencepiece"
+version = "0.2.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" },
+    { url = "https://files.pythonhosted.org/packages/88/7e/ff23008899a58678e98c6ff592bf4d368eee5a71af96d0df6b38a039dd4f/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:56dd39a3c4d6493db3cdca7e8cc68c6b633f0d4195495cbadfcf5af8a22d05a6", size = 1325651, upload-time = "2025-08-12T06:59:41.536Z" },
+    { url = "https://files.pythonhosted.org/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" },
+    { url = "https://files.pythonhosted.org/packages/89/fa/d3d5ebcba3cb9e6d3775a096251860c41a6bc53a1b9461151df83fe93255/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99f955df238021bf11f0fc37cdb54fd5e5b5f7fd30ecc3d93fb48b6815437167", size = 1316273, upload-time = "2025-08-12T06:59:44.476Z" },
+    { url = "https://files.pythonhosted.org/packages/04/88/14f2f4a2b922d8b39be45bf63d79e6cd3a9b2f248b2fcb98a69b12af12f5/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cdfecef430d985f1c2bcbfff3defd1d95dae876fbd0173376012d2d7d24044b", size = 1387881, upload-time = "2025-08-12T06:59:46.09Z" },
+    { url = "https://files.pythonhosted.org/packages/fd/b8/903e5ccb77b4ef140605d5d71b4f9e0ad95d456d6184688073ed11712809/sentencepiece-0.2.1-cp312-cp312-win32.whl", hash = "sha256:a483fd29a34c3e34c39ac5556b0a90942bec253d260235729e50976f5dba1068", size = 999540, upload-time = "2025-08-12T06:59:48.023Z" },
+    { url = "https://files.pythonhosted.org/packages/2d/81/92df5673c067148c2545b1bfe49adfd775bcc3a169a047f5a0e6575ddaca/sentencepiece-0.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:4cdc7c36234fda305e85c32949c5211faaf8dd886096c7cea289ddc12a2d02de", size = 1054671, upload-time = "2025-08-12T06:59:49.895Z" },
+    { url = "https://files.pythonhosted.org/packages/fe/02/c5e3bc518655d714622bec87d83db9cdba1cd0619a4a04e2109751c4f47f/sentencepiece-0.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:daeb5e9e9fcad012324807856113708614d534f596d5008638eb9b40112cd9e4", size = 1033923, upload-time = "2025-08-12T06:59:51.952Z" },
+    { url = "https://files.pythonhosted.org/packages/ba/4a/85fbe1706d4d04a7e826b53f327c4b80f849cf1c7b7c5e31a20a97d8f28b/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:dcd8161eee7b41aae57ded06272905dbd680a0a04b91edd0f64790c796b2f706", size = 1943150, upload-time = "2025-08-12T06:59:53.588Z" },
+    { url = "https://files.pythonhosted.org/packages/c2/83/4cfb393e287509fc2155480b9d184706ef8d9fa8cbf5505d02a5792bf220/sentencepiece-0.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c6c8f42949f419ff8c7e9960dbadcfbc982d7b5efc2f6748210d3dd53a7de062", size = 1325651, upload-time = "2025-08-12T06:59:55.073Z" },
+    { url = "https://files.pythonhosted.org/packages/8d/de/5a007fb53b1ab0aafc69d11a5a3dd72a289d5a3e78dcf2c3a3d9b14ffe93/sentencepiece-0.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:097f3394e99456e9e4efba1737c3749d7e23563dd1588ce71a3d007f25475fff", size = 1253641, upload-time = "2025-08-12T06:59:56.562Z" },
+    { url = "https://files.pythonhosted.org/packages/2c/d2/f552be5928105588f4f4d66ee37dd4c61460d8097e62d0e2e0eec41bc61d/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d7b670879c370d350557edabadbad1f6561a9e6968126e6debca4029e5547820", size = 1316271, upload-time = "2025-08-12T06:59:58.109Z" },
+    { url = "https://files.pythonhosted.org/packages/96/df/0cfe748ace5485be740fed9476dee7877f109da32ed0d280312c94ec259f/sentencepiece-0.2.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c7f0fd2f2693309e6628aeeb2e2faf6edd221134dfccac3308ca0de01f8dab47", size = 1387882, upload-time = "2025-08-12T07:00:00.701Z" },
+    { url = "https://files.pythonhosted.org/packages/ac/dd/f7774d42a881ced8e1739f393ab1e82ece39fc9abd4779e28050c2e975b5/sentencepiece-0.2.1-cp313-cp313-win32.whl", hash = "sha256:92b3816aa2339355fda2c8c4e021a5de92180b00aaccaf5e2808972e77a4b22f", size = 999541, upload-time = "2025-08-12T07:00:02.709Z" },
+    { url = "https://files.pythonhosted.org/packages/dd/e9/932b9eae6fd7019548321eee1ab8d5e3b3d1294df9d9a0c9ac517c7b636d/sentencepiece-0.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:10ed3dab2044c47f7a2e7b4969b0c430420cdd45735d78c8f853191fa0e3148b", size = 1054669, upload-time = "2025-08-12T07:00:04.915Z" },
+    { url = "https://files.pythonhosted.org/packages/c9/3a/76488a00ea7d6931689cda28726a1447d66bf1a4837943489314593d5596/sentencepiece-0.2.1-cp313-cp313-win_arm64.whl", hash = "sha256:ac650534e2251083c5f75dde4ff28896ce7c8904133dc8fef42780f4d5588fcd", size = 1033922, upload-time = "2025-08-12T07:00:06.496Z" },
+    { url = "https://files.pythonhosted.org/packages/4a/b6/08fe2ce819e02ccb0296f4843e3f195764ce9829cbda61b7513f29b95718/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_universal2.whl", hash = "sha256:8dd4b477a7b069648d19363aad0cab9bad2f4e83b2d179be668efa672500dc94", size = 1946052, upload-time = "2025-08-12T07:00:08.136Z" },
+    { url = "https://files.pythonhosted.org/packages/ab/d9/1ea0e740591ff4c6fc2b6eb1d7510d02f3fb885093f19b2f3abd1363b402/sentencepiece-0.2.1-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:0c0f672da370cc490e4c59d89e12289778310a0e71d176c541e4834759e1ae07", size = 1327408, upload-time = "2025-08-12T07:00:09.572Z" },
+    { url = "https://files.pythonhosted.org/packages/99/7e/1fb26e8a21613f6200e1ab88824d5d203714162cf2883248b517deb500b7/sentencepiece-0.2.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:ad8493bea8432dae8d6830365352350f3b4144415a1d09c4c8cb8d30cf3b6c3c", size = 1254857, upload-time = "2025-08-12T07:00:11.021Z" },
+    { url = "https://files.pythonhosted.org/packages/bc/85/c72fd1f3c7a6010544d6ae07f8ddb38b5e2a7e33bd4318f87266c0bbafbf/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b81a24733726e3678d2db63619acc5a8dccd074f7aa7a54ecd5ca33ca6d2d596", size = 1315722, upload-time = "2025-08-12T07:00:12.989Z" },
+    { url = "https://files.pythonhosted.org/packages/4a/e8/661e5bd82a8aa641fd6c1020bd0e890ef73230a2b7215ddf9c8cd8e941c2/sentencepiece-0.2.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0a81799d0a68d618e89063fb423c3001a034c893069135ffe51fee439ae474d6", size = 1387452, upload-time = "2025-08-12T07:00:15.088Z" },
+    { url = "https://files.pythonhosted.org/packages/99/5e/ae66c361023a470afcbc1fbb8da722c72ea678a2fcd9a18f1a12598c7501/sentencepiece-0.2.1-cp313-cp313t-win32.whl", hash = "sha256:89a3ea015517c42c0341d0d962f3e6aaf2cf10d71b1932d475c44ba48d00aa2b", size = 1002501, upload-time = "2025-08-12T07:00:16.966Z" },
+    { url = "https://files.pythonhosted.org/packages/c1/03/d332828c4ff764e16c1b56c2c8f9a33488bbe796b53fb6b9c4205ddbf167/sentencepiece-0.2.1-cp313-cp313t-win_amd64.whl", hash = "sha256:33f068c9382dc2e7c228eedfd8163b52baa86bb92f50d0488bf2b7da7032e484", size = 1057555, upload-time = "2025-08-12T07:00:18.573Z" },
+    { url = "https://files.pythonhosted.org/packages/88/14/5aee0bf0864df9bd82bd59e7711362908e4935e3f9cdc1f57246b5d5c9b9/sentencepiece-0.2.1-cp313-cp313t-win_arm64.whl", hash = "sha256:b3616ad246f360e52c85781e47682d31abfb6554c779e42b65333d4b5f44ecc0", size = 1036042, upload-time = "2025-08-12T07:00:20.209Z" },
+    { url = "https://files.pythonhosted.org/packages/24/9c/89eb8b2052f720a612478baf11c8227dcf1dc28cd4ea4c0c19506b5af2a2/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:5d0350b686c320068702116276cfb26c066dc7e65cfef173980b11bb4d606719", size = 1943147, upload-time = "2025-08-12T07:00:21.809Z" },
+    { url = "https://files.pythonhosted.org/packages/82/0b/a1432bc87f97c2ace36386ca23e8bd3b91fb40581b5e6148d24b24186419/sentencepiece-0.2.1-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:c7f54a31cde6fa5cb030370566f68152a742f433f8d2be458463d06c208aef33", size = 1325624, upload-time = "2025-08-12T07:00:23.289Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/99/bbe054ebb5a5039457c590e0a4156ed073fb0fe9ce4f7523404dd5b37463/sentencepiece-0.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c83b85ab2d6576607f31df77ff86f28182be4a8de6d175d2c33ca609925f5da1", size = 1253670, upload-time = "2025-08-12T07:00:24.69Z" },
+    { url = "https://files.pythonhosted.org/packages/19/ad/d5c7075f701bd97971d7c2ac2904f227566f51ef0838dfbdfdccb58cd212/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1855f57db07b51fb51ed6c9c452f570624d2b169b36f0f79ef71a6e6c618cd8b", size = 1316247, upload-time = "2025-08-12T07:00:26.435Z" },
+    { url = "https://files.pythonhosted.org/packages/fb/03/35fbe5f3d9a7435eebd0b473e09584bd3cc354ce118b960445b060d33781/sentencepiece-0.2.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:01e6912125cb45d3792f530a4d38f8e21bf884d6b4d4ade1b2de5cf7a8d2a52b", size = 1387894, upload-time = "2025-08-12T07:00:28.339Z" },
+    { url = "https://files.pythonhosted.org/packages/dc/aa/956ef729aafb6c8f9c443104c9636489093bb5c61d6b90fc27aa1a865574/sentencepiece-0.2.1-cp314-cp314-win32.whl", hash = "sha256:c415c9de1447e0a74ae3fdb2e52f967cb544113a3a5ce3a194df185cbc1f962f", size = 1096698, upload-time = "2025-08-12T07:00:29.764Z" },
+    { url = "https://files.pythonhosted.org/packages/b8/cb/fe400d8836952cc535c81a0ce47dc6875160e5fedb71d2d9ff0e9894c2a6/sentencepiece-0.2.1-cp314-cp314-win_amd64.whl", hash = "sha256:881b2e44b14fc19feade3cbed314be37de639fc415375cefaa5bc81a4be137fd", size = 1155115, upload-time = "2025-08-12T07:00:32.865Z" },
+    { url = "https://files.pythonhosted.org/packages/32/89/047921cf70f36c7b6b6390876b2399b3633ab73b8d0cb857e5a964238941/sentencepiece-0.2.1-cp314-cp314-win_arm64.whl", hash = "sha256:2005242a16d2dc3ac5fe18aa7667549134d37854823df4c4db244752453b78a8", size = 1133890, upload-time = "2025-08-12T07:00:34.763Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/11/5b414b9fae6255b5fb1e22e2ed3dc3a72d3a694e5703910e640ac78346bb/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_universal2.whl", hash = "sha256:a19adcec27c524cb7069a1c741060add95f942d1cbf7ad0d104dffa0a7d28a2b", size = 1946081, upload-time = "2025-08-12T07:00:36.97Z" },
+    { url = "https://files.pythonhosted.org/packages/77/eb/7a5682bb25824db8545f8e5662e7f3e32d72a508fdce086029d89695106b/sentencepiece-0.2.1-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:e37e4b4c4a11662b5db521def4e44d4d30ae69a1743241412a93ae40fdcab4bb", size = 1327406, upload-time = "2025-08-12T07:00:38.669Z" },
+    { url = "https://files.pythonhosted.org/packages/03/b0/811dae8fb9f2784e138785d481469788f2e0d0c109c5737372454415f55f/sentencepiece-0.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:477c81505db072b3ab627e7eab972ea1025331bd3a92bacbf798df2b75ea86ec", size = 1254846, upload-time = "2025-08-12T07:00:40.611Z" },
+    { url = "https://files.pythonhosted.org/packages/ef/23/195b2e7ec85ebb6a547969f60b723c7aca5a75800ece6cc3f41da872d14e/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:010f025a544ef770bb395091d57cb94deb9652d8972e0d09f71d85d5a0816c8c", size = 1315721, upload-time = "2025-08-12T07:00:42.914Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/aa/553dbe4178b5f23eb28e59393dddd64186178b56b81d9b8d5c3ff1c28395/sentencepiece-0.2.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:733e59ff1794d26db706cd41fc2d7ca5f6c64a820709cb801dc0ea31780d64ab", size = 1387458, upload-time = "2025-08-12T07:00:44.56Z" },
+    { url = "https://files.pythonhosted.org/packages/66/7c/08ff0012507297a4dd74a5420fdc0eb9e3e80f4e88cab1538d7f28db303d/sentencepiece-0.2.1-cp314-cp314t-win32.whl", hash = "sha256:d3233770f78e637dc8b1fda2cd7c3b99ec77e7505041934188a4e7fe751de3b0", size = 1099765, upload-time = "2025-08-12T07:00:46.058Z" },
+    { url = "https://files.pythonhosted.org/packages/91/d5/2a69e1ce15881beb9ddfc7e3f998322f5cedcd5e4d244cb74dade9441663/sentencepiece-0.2.1-cp314-cp314t-win_amd64.whl", hash = "sha256:5e4366c97b68218fd30ea72d70c525e6e78a6c0a88650f57ac4c43c63b234a9d", size = 1157807, upload-time = "2025-08-12T07:00:47.673Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/16/54f611fcfc2d1c46cbe3ec4169780b2cfa7cf63708ef2b71611136db7513/sentencepiece-0.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:105e36e75cbac1292642045458e8da677b2342dcd33df503e640f0b457cb6751", size = 1136264, upload-time = "2025-08-12T07:00:49.485Z" },
+]
+
 [[package]]
 name = "sentry-sdk"
 version = "2.64.0"