refactor(pi052): load native base checkpoint

Keep the parity tests compatible with the repository's current Ruff hooks and remove the requested migration warning from the docs.

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

README.md

@@ -102,10 +102,10 @@ lerobot-train \
 ```
 
 | 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) |
+| **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)                                                                                                                                                                                                                                                                                                           |
 

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

@@ -191,6 +191,162 @@ def make_my_policy_pre_post_processors(
 
 ---
 
+## Adding high- and low-level language control
+
+The policy API above is sufficient for training and standard evaluation. To use a language-conditioned policy with interactive `lerobot-rollout`, also register a runtime adapter. The adapter keeps policy-specific prompting and tokenization out of the generic control loop.
+
+The runtime supports two policy shapes:
+
+| Policy shape     | Behavior                                                                | Adapter                                        |
+| ---------------- | ----------------------------------------------------------------------- | ---------------------------------------------- |
+| Low-level / flat | The operator's task or subtask directly conditions action prediction.   | Reuse `DirectTaskPolicyAdapter`.               |
+| High + low level | The policy generates subtasks or memory, then conditions actions on it. | Subclass `BaseLanguageAdapter`, as PI052 does. |
+
+During a rollout, `RuntimeState` stores the high-level task and the active language context:
+
+```text
+task ──> adapter.generate_text("subtask", ...) ──> state.language_context["subtask"]
+                                                       │
+observation ──> processors ──> adapter.select_action() ─┴─> action chunk ──> robot
+```
+
+The generic runtime handles generation frequency, pause/resume, prompt replacement, action queues, and dispatch. The adapter only translates between that runtime contract and your policy.
+
+### Low-level policies
+
+If your policy already consumes the live task through its normal preprocessor and implements `predict_action_chunk`, register the shared direct adapter. PI0.5 and MolmoAct2 use this path:
+
+```python
+# src/lerobot/runtime/registry.py
+_ADAPTERS = {
+    # ...
+    "my_policy": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
+}
+```
+
+Run it with direct-subtask mode so the operator supplies the instruction used by the action policy:
+
+```bash
+lerobot-rollout \
+    --language \
+    --policy.path=user/my_policy_checkpoint \
+    --robot.type=so101_follower \
+    --robot.port=/dev/ttyACM0 \
+    --direct_subtask
+```
+
+The rollout context builds the observation batch with the current instruction before `DirectTaskPolicyAdapter` calls `policy.predict_action_chunk(observation)`. No text-generation method is required.
+
+### Hierarchical policies
+
+For a policy that generates language and actions, subclass [`BaseLanguageAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/runtime/adapter.py) and implement two methods:
+
+- `generate_text(kind, observation, state, user_text=None) -> str` generates a `subtask`, `memory`, or interjection response.
+- `select_action(observation, state)` builds the low-level prompt from the active context and returns an action chunk.
+
+This abbreviated adapter follows [`PI052PolicyAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py):
+
+```python
+# inference/my_policy_adapter.py
+from typing import Any
+
+from lerobot.runtime import RuntimeState
+from lerobot.runtime.adapter import BaseLanguageAdapter
+from lerobot.utils.constants import (
+    OBS_LANGUAGE_ATTENTION_MASK,
+    OBS_LANGUAGE_TOKENS,
+)
+
+
+class MyPolicyAdapter(BaseLanguageAdapter):
+    def select_action(self, observation: dict[str, Any], state: RuntimeState):
+        instruction = state.language_context.get("subtask") or state.task or ""
+        tokens, attention_mask = tokenize_instruction(instruction)
+
+        batch = dict(observation)
+        batch[OBS_LANGUAGE_TOKENS] = tokens
+        batch[OBS_LANGUAGE_ATTENTION_MASK] = attention_mask
+        return self.policy.predict_action_chunk(batch)
+
+    def generate_text(
+        self,
+        kind: str,
+        observation: dict[str, Any] | None,
+        state: RuntimeState,
+        user_text: str | None = None,
+    ) -> str:
+        messages = self.build_messages(kind, state, user_text)
+        batch, tokenizer = tokenize_messages(messages, observation)
+        return self.policy.select_message(
+            batch,
+            tokenizer=tokenizer,
+            min_new_tokens=self.gen.min_new_tokens,
+            temperature=self.gen.temperature,
+            top_p=self.gen.top_p,
+        )
+
+    def build_messages(
+        self, kind: str, state: RuntimeState, user_text: str | None
+    ) -> list[dict[str, str]]:
+        if kind == "subtask":
+            return [{"role": "user", "content": state.task or ""}]
+        if kind == "memory":
+            return [
+                {"role": "user", "content": state.task or ""},
+                {
+                    "role": "user",
+                    "content": f"Completed subtask: {state.extra.get('prior_subtask', '')}",
+                },
+            ]
+        if kind == "interjection":
+            return [
+                {"role": "user", "content": state.task or ""},
+                {"role": "user", "content": user_text or ""},
+            ]
+        raise ValueError(f"Unsupported text kind: {kind}")
+```
+
+`tokenize_instruction` and `tokenize_messages` are policy-specific helpers. They must reproduce the prompt format used during training; PI052, for example, adds the discretized robot state to its low-level subtask prompt and uses the same PaliGemma formatting for `select_message`.
+
+`BaseLanguageAdapter` provides the default hierarchy: regenerate a subtask at action-chunk boundaries, update memory when the subtask changes, and handle user interjections. Override `_regenerate_context` only if your policy uses a different hierarchy.
+
+Register the adapter with a lazy import so importing LeRobot does not load the model or its optional dependencies:
+
+```python
+# src/lerobot/runtime/registry.py
+_ADAPTERS = {
+    # ...
+    "my_policy": "lerobot.policies.my_policy.inference.my_policy_adapter:MyPolicyAdapter",
+}
+```
+
+The key must match the policy's registered type. Once registered, the same checkpoint works through the shared entry point:
+
+```bash
+lerobot-rollout \
+    --language \
+    --policy.path=user/my_hierarchical_checkpoint \
+    --robot.type=so101_follower \
+    --robot.port=/dev/ttyACM0 \
+    --task="put the cup in the sink"
+```
+
+For RoboCasa-compatible policies, replace the robot arguments with `--sim --sim.task=<task>`. Without `--direct_subtask`, the adapter generates the low-level subtask; with it, the operator bypasses high-level generation and supplies each subtask.
+
+### Keep training and deployment aligned
+
+The adapter is intentionally small, but its prompts are part of the model contract:
+
+- Use the same tokenizer, role formatting, special tokens, image ordering, and state encoding as training.
+- Condition `select_action` on `state.language_context["subtask"]`, falling back to `state.task` for direct or not-yet-generated prompts.
+- Return a full action chunk from `select_action`; the runtime handles control-rate dispatch.
+- Keep optional model dependencies inside lazy imports.
+- Test adapter selection, generated-message routing, action-batch construction, and direct-subtask behavior with a lightweight fake policy.
+
+PI052 is the complete in-tree reference: its [processor](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/processor_pi052.py) renders the training recipe, its [policy](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/modeling_pi052.py) exposes text and action generation, and its [adapter](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py) reconstructs those same prompts at deployment.
+
+---
+
 ## Path A: Out-of-tree plugin
 
 The fastest way to ship a policy: package it as a standalone Python distribution and install it alongside LeRobot. No PR required, you own the release cycle, and you can publish to PyPI under your own namespace.

docs/source/cameras.mdx

@@ -136,10 +136,6 @@ config = RealSenseCameraConfig(
     height=480,
     color_mode=ColorMode.RGB,
     use_depth=True,
-    # Optional fixed color controls. Omit them to leave the current sensor settings unchanged.
-    exposure=120,
-    gain=64,
-    white_balance=4600,
     rotation=Cv2Rotation.NO_ROTATION
 )
 
@@ -158,15 +154,6 @@ finally:
 ```
 <!-- prettier-ignore-end -->
 
-Manual color controls disable the corresponding automatic exposure or white-balance mode. Their
-supported ranges vary by camera model; an invalid value raises an error at connection time that
-includes the range reported by the sensor. Requesting an unsupported control also raises an error.
-Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
-require `use_rgb=True`.
-
-On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
-exposure or gain also affects the depth stream.
-
 </hfoption>
 </hfoptions>
 

docs/source/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.

docs/source/language_and_recipes.mdx

@@ -141,6 +141,17 @@ sample["target_message_indices"]
 
 The renderer does not apply a tokenizer chat template. Policy processors decide how to serialize the messages for their backbone, which keeps the same dataset usable across SmolVLA, Pi0.5, and any future VLM that expects OpenAI-style chat messages.
 
+## Blends
+
+Blend recipes select one weighted sub-recipe deterministically from the sample index.
+`recipes/subtask_mem.yaml` trains the compact core blend — high-level subtask prediction, low-level execution, and memory. `recipes/subtask_mem_vqa_speech.yaml` is the fuller variant that also adds VQA and spoken interjection responses.
+
+A message recipe with a supervised assistant turn on the `low_level` stream trains
+the π0.5 paper's joint sequence instead of a blend: the target span gets text CE
+while also conditioning the action losses in the same forward.
+`recipes/subtask_joint.yaml` is the provided example; pair it with
+`--policy.joint_subtask_conditioning=true` at inference.
+
 ## Graceful absence
 
 If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.

docs/source/pi052.mdx

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

docs/source/pi0fast.mdx

@@ -110,15 +110,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.n_action_steps`               | Number of decoded action steps to execute          | `50`                            |
 | `--policy.max_action_tokens`            | Maximum number of FAST tokens per action chunk     | `256`                           |
 | `--policy.action_tokenizer_name`        | FAST tokenizer to use                              | `lerobot/fast-action-tokenizer` |
+| `--policy.auto_fit_fast_tokenizer=true` | Fit and cache a tokenizer for the training dataset | `false`                         |
 | `--policy.compile_model=true`           | Enable torch.compile for faster training           | `false`                         |
 
+Set `--policy.auto_fit_fast_tokenizer=true` to sample action chunks from the
+training dataset and cache a fitted tokenizer under
+`~/.cache/lerobot/fast_tokenizers`. This also works when fine-tuning with
+`--policy.path`; leave it disabled to retain the checkpoint's tokenizer.
+
 ## Inference
 
 ### KV-Caching for Fast Inference
@@ -151,6 +157,9 @@ actions = policy.predict_action_chunk(batch)
 
 The model takes images, text instructions, and robot state as input, and outputs discrete FAST tokens that are decoded back to continuous actions.
 
+PI0-FAST always decodes a complete `chunk_size` action chunk. `n_action_steps` controls only
+how many actions from that chunk are executed before the policy predicts again.
+
 ## Configuration Options
 
 | Parameter            | Description                                     | Default    |

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]"]
@@ -494,19 +495,6 @@ ignore_errors = true
 module = "lerobot.envs.*"
 ignore_errors = false
 
-[[tool.mypy.overrides]]
-module = "lerobot.annotations.*"
-ignore_errors = false
-disallow_untyped_defs = true
-disallow_incomplete_defs = true
-check_untyped_defs = true
-
-[[tool.mypy.overrides]]
-module = "lerobot.transforms.*"
-ignore_errors = false
-disallow_untyped_defs = true
-disallow_incomplete_defs = true
-check_untyped_defs = true
 
 # [[tool.mypy.overrides]]
 # module = "lerobot.utils.*"

src/lerobot/cameras/opencv/camera_opencv.py

@@ -120,22 +120,14 @@ class OpenCVCamera(Camera):
         self.rotation: int | None = get_cv2_rotation(config.rotation)
         self.backend: int = config.backend
 
-        self.capture_width: int | None = None
-        self.capture_height: int | None = None
-        self._reset_connection_settings()
-
-    def __str__(self) -> str:
-        return f"{self.__class__.__name__}({self.index_or_path})"
-
-    def _reset_connection_settings(self) -> None:
-        """Restore settings that may have been auto-detected during a failed connection."""
-        self.fps = self.config.fps
-        self.width = self.config.width
-        self.height = self.config.height
+        if self.height and self.width:
             self.capture_width, self.capture_height = self.width, self.height
             if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
                 self.capture_width, self.capture_height = self.height, self.width
 
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}({self.index_or_path})"
+
     @property
     def is_connected(self) -> bool:
         """Checks if the camera is currently connected and opened."""
@@ -172,7 +164,6 @@ class OpenCVCamera(Camera):
                 f"Failed to open {self}.Run `lerobot-find-cameras opencv` to find available cameras."
             )
 
-        try:
         self._configure_capture_settings()
         self._start_read_thread()
 
@@ -184,13 +175,6 @@ class OpenCVCamera(Camera):
             with self.frame_lock:
                 if self.latest_frame is None:
                     raise ConnectionError(f"{self} failed to capture frames during warmup.")
-        except BaseException:
-            try:
-                self._cleanup_resources()
-            except Exception:
-                logger.exception(f"Failed to fully clean up {self} after connect() failed.")
-            self._reset_connection_settings()
-            raise
 
         logger.info(f"{self} connected.")
 
@@ -328,7 +312,6 @@ class OpenCVCamera(Camera):
 
         for target in targets_to_scan:
             camera = cv2.VideoCapture(target)
-            try:
             if camera.isOpened():
                 default_width = int(camera.get(cv2.CAP_PROP_FRAME_WIDTH))
                 default_height = int(camera.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -338,9 +321,7 @@ class OpenCVCamera(Camera):
                 # Get FOURCC code and convert to string
                 default_fourcc_code = camera.get(cv2.CAP_PROP_FOURCC)
                 default_fourcc_code_int = int(default_fourcc_code)
-                    default_fourcc = "".join(
-                        [chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)]
-                    )
+                default_fourcc = "".join([chr((default_fourcc_code_int >> 8 * i) & 0xFF) for i in range(4)])
 
                 camera_info = {
                     "name": f"OpenCV Camera @ {target}",
@@ -357,7 +338,6 @@ class OpenCVCamera(Camera):
                 }
 
                 found_cameras_info.append(camera_info)
-            finally:
                 camera.release()
 
         return found_cameras_info
@@ -516,26 +496,6 @@ class OpenCVCamera(Camera):
             self.latest_timestamp = None
             self.new_frame_event.clear()
 
-    def _cleanup_resources(self) -> None:
-        """Stop background reads and release the capture, including after partial setup."""
-        read_thread = self.thread
-        videocapture = self.videocapture
-
-        try:
-            self._stop_read_thread()
-        finally:
-            self.videocapture = None
-            try:
-                if videocapture is not None:
-                    videocapture.release()
-            finally:
-                # Releasing the device may unblock a hardware read that outlived
-                # the first bounded join in _stop_read_thread().
-                if read_thread is not None and read_thread.is_alive():
-                    read_thread.join(timeout=2.0)
-                    if read_thread.is_alive():  # pragma: no cover
-                        logger.warning(f"{self} read thread remained alive after releasing the capture.")
-
     @check_if_not_connected
     def async_read(self, timeout_ms: float = 200) -> NDArray[Any]:
         """
@@ -626,6 +586,16 @@ class OpenCVCamera(Camera):
         if not self.is_connected and self.thread is None:
             raise DeviceNotConnectedError(f"{self} not connected.")
 
-        self._cleanup_resources()
+        if self.thread is not None:
+            self._stop_read_thread()
+
+        if self.videocapture is not None:
+            self.videocapture.release()
+            self.videocapture = None
+
+        with self.frame_lock:
+            self.latest_frame = None
+            self.latest_timestamp = None
+            self.new_frame_event.clear()
 
         logger.info(f"{self} disconnected.")

src/lerobot/cameras/realsense/camera_realsense.py

@@ -121,9 +121,6 @@ class RealSenseCamera(Camera):
 
         self.config = config
 
-        self.width: int | None = config.width
-        self.height: int | None = config.height
-
         if config.serial_number_or_name.isdigit():
             self.serial_number = config.serial_number_or_name
         else:
@@ -134,9 +131,6 @@ class RealSenseCamera(Camera):
         self.use_rgb = config.use_rgb
         self.use_depth = config.use_depth
         self.warmup_s = config.warmup_s
-        self.exposure: int | None = config.exposure
-        self.gain: int | None = config.gain
-        self.white_balance: int | None = config.white_balance
 
         self.rs_pipeline: rs.pipeline | None = None
         self.rs_profile: rs.pipeline_profile | None = None
@@ -151,23 +145,14 @@ class RealSenseCamera(Camera):
 
         self.rotation: int | None = get_cv2_rotation(config.rotation)
 
-        self.capture_width: int | None = None
-        self.capture_height: int | None = None
-        self._reset_connection_settings()
-
-    def __str__(self) -> str:
-        return f"{self.__class__.__name__}({self.serial_number})"
-
-    def _reset_connection_settings(self) -> None:
-        """Restore settings that may have been auto-detected during a failed connection."""
-        self.fps = self.config.fps
-        self.width = self.config.width
-        self.height = self.config.height
-        self.warmup_s = self.config.warmup_s
+        if self.height and self.width:
             self.capture_width, self.capture_height = self.width, self.height
             if self.rotation in [cv2.ROTATE_90_CLOCKWISE, cv2.ROTATE_90_COUNTERCLOCKWISE]:
                 self.capture_width, self.capture_height = self.height, self.width
 
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}({self.serial_number})"
+
     @property
     def is_connected(self) -> bool:
         """Checks if the camera pipeline is started and streams are active."""
@@ -187,8 +172,7 @@ class RealSenseCamera(Camera):
 
         Raises:
             DeviceAlreadyConnectedError: If the camera is already connected.
-            ValueError: If the configuration is invalid, a requested sensor option is unsupported,
-                or a requested sensor value is invalid.
+            ValueError: If the configuration is invalid (e.g., missing serial/name, name not unique).
             ConnectionError: If the camera is found but fails to start the pipeline or no RealSense devices are detected at all.
             RuntimeError: If the pipeline starts but fails to apply requested settings.
         """
@@ -206,9 +190,7 @@ class RealSenseCamera(Camera):
                 f"Failed to open {self}.Run `lerobot-find-cameras realsense` to find available cameras."
             ) from e
 
-        try:
         self._configure_capture_settings()
-            self._configure_sensor_options()
         self._start_read_thread()
 
         # NOTE(Steven/Caroline): Enforcing at least one second of warmup as RS cameras need a bit of time before the first read. If we don't wait, the first read from the warmup will raise.
@@ -224,13 +206,6 @@ class RealSenseCamera(Camera):
                 self.use_depth and self.latest_depth_frame is None
             ):
                 raise ConnectionError(f"{self} failed to capture frames during warmup.")
-        except BaseException:
-            try:
-                self._cleanup_resources()
-            except Exception:
-                logger.exception(f"Failed to fully clean up {self} after connect() failed.")
-            self._reset_connection_settings()
-            raise
 
         logger.info(f"{self} connected.")
 
@@ -364,111 +339,6 @@ class RealSenseCamera(Camera):
         self.new_frame_event.clear()
         return self._async_read(timeout_ms=10000, read_depth=read_depth)
 
-    def _get_color_sensor(self) -> "rs.sensor":
-        """Returns the sensor that controls the color stream.
-
-        Most RealSense cameras expose "RGB Camera" for color. The D405 has no
-        separate RGB module — its color stream comes from "Stereo Module".
-        We try RGB Camera first, then fall back to Stereo Module.
-        """
-        if self.rs_profile is None:
-            raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
-
-        device = self.rs_profile.get_device()
-        sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
-
-        for name in ("RGB Camera", "Stereo Module"):
-            if name in sensors:
-                return sensors[name]
-
-        available = list(sensors.keys())
-        raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
-
-    def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
-        """Sets a sensor option, re-raising range errors with actionable diagnostics."""
-        try:
-            sensor.set_option(option, value)
-        except Exception as e:
-            range_info = ""
-            try:
-                option_range = sensor.get_option_range(option)
-                range_info = (
-                    f" (supported range: min={option_range.min}, max={option_range.max}, "
-                    f"step={option_range.step}, default={option_range.default})"
-                )
-            except Exception:
-                range_info = " (option range unavailable)"
-            raise ValueError(
-                f"{self}: failed to set {label} to {value}{range_info}. Original error: {e}"
-            ) from e
-
-    def _configure_sensor_options(self) -> None:
-        """Applies manual sensor options (exposure, gain, white balance) to the color sensor.
-
-        When exposure or gain is set, auto-exposure is disabled first. When white_balance
-        is set, auto white balance is disabled first. An omitted option is left unchanged,
-        and configuration is skipped entirely if all options are omitted.
-
-        Raises:
-            ValueError: If the sensor does not support a requested option or a requested
-                value is invalid. Invalid-value errors include the option name, requested
-                value, and supported range when available.
-        """
-        if self.exposure is None and self.gain is None and self.white_balance is None:
-            return
-
-        color_sensor = self._get_color_sensor()
-
-        requested_options = (
-            (rs.option.exposure, self.exposure, "exposure"),
-            (rs.option.gain, self.gain, "gain"),
-            (rs.option.white_balance, self.white_balance, "white balance"),
-        )
-        unsupported_options = [
-            label
-            for option, value, label in requested_options
-            if value is not None and not color_sensor.supports(option)
-        ]
-        if unsupported_options:
-            raise ValueError(
-                f"{self}: color sensor does not support requested manual options: {unsupported_options}."
-            )
-
-        manual_exposure_requested = self.exposure is not None or self.gain is not None
-        if manual_exposure_requested:
-            if color_sensor.supports(rs.option.enable_auto_exposure):
-                self._set_sensor_option(color_sensor, rs.option.enable_auto_exposure, 0, "auto-exposure")
-                logger.info(f"{self} auto-exposure disabled.")
-            else:
-                logger.warning(
-                    f"{self} sensor does not support disabling auto-exposure; "
-                    "applying manual exposure/gain directly."
-                )
-
-        if self.exposure is not None:
-            self._set_sensor_option(color_sensor, rs.option.exposure, self.exposure, "exposure")
-            logger.info(f"{self} exposure set to {self.exposure}.")
-
-        if self.gain is not None:
-            self._set_sensor_option(color_sensor, rs.option.gain, self.gain, "gain")
-            logger.info(f"{self} gain set to {self.gain}.")
-
-        if self.white_balance is not None:
-            if color_sensor.supports(rs.option.enable_auto_white_balance):
-                self._set_sensor_option(
-                    color_sensor, rs.option.enable_auto_white_balance, 0, "auto white balance"
-                )
-                logger.info(f"{self} auto white balance disabled.")
-            else:
-                logger.warning(
-                    f"{self} sensor does not support disabling auto white balance; "
-                    "applying manual white balance directly."
-                )
-            self._set_sensor_option(
-                color_sensor, rs.option.white_balance, self.white_balance, "white balance"
-            )
-            logger.info(f"{self} white balance set to {self.white_balance}.")
-
     @check_if_not_connected
     def read_depth(self, timeout_ms: int = 200) -> NDArray[Any]:
         """
@@ -671,27 +541,6 @@ class RealSenseCamera(Camera):
             self.latest_timestamp = None
             self.new_frame_event.clear()
 
-    def _cleanup_resources(self) -> None:
-        """Stop background reads and stop the pipeline, including after partial setup."""
-        read_thread = self.thread
-        rs_pipeline = self.rs_pipeline
-
-        try:
-            self._stop_read_thread()
-        finally:
-            self.rs_pipeline = None
-            self.rs_profile = None
-            try:
-                if rs_pipeline is not None:
-                    rs_pipeline.stop()
-            finally:
-                # Stopping the pipeline may unblock a hardware read that outlived
-                # the first bounded join in _stop_read_thread().
-                if read_thread is not None and read_thread.is_alive():
-                    read_thread.join(timeout=2.0)
-                    if read_thread.is_alive():  # pragma: no cover
-                        logger.warning(f"{self} read thread remained alive after stopping the pipeline.")
-
     def _async_read(self, timeout_ms: float, read_depth: bool = False) -> NDArray[Any]:
         """Shared helper for :meth:`async_read`/:meth:`async_read_depth`: return the latest buffered frame."""
         if self.thread is None or not self.thread.is_alive():
@@ -835,5 +684,18 @@ class RealSenseCamera(Camera):
                 f"Attempted to disconnect {self}, but it appears already disconnected."
             )
 
-        self._cleanup_resources()
+        if self.thread is not None:
+            self._stop_read_thread()
+
+        if self.rs_pipeline is not None:
+            self.rs_pipeline.stop()
+            self.rs_pipeline = None
+            self.rs_profile = None
+
+        with self.frame_lock:
+            self.latest_color_frame = None
+            self.latest_depth_frame = None
+            self.latest_timestamp = None
+            self.new_frame_event.clear()
+
         logger.info(f"{self} disconnected.")

src/lerobot/cameras/realsense/configuration_realsense.py

@@ -46,17 +46,6 @@ class RealSenseCameraConfig(CameraConfig):
         use_depth: Whether to enable depth stream. Defaults to False.
         rotation: Image rotation setting (0°, 90°, 180°, or 270°). Defaults to no rotation.
         warmup_s: Time reading frames before returning from connect (in seconds)
-        exposure: Manual exposure value for the color sensor. When set, auto-exposure is
-            disabled and this fixed value is used. Valid ranges are camera-model specific
-            and reported if the value is rejected. Defaults to None (leave unchanged).
-        gain: Manual gain value for the color sensor. When set, auto-exposure is disabled
-            and this fixed gain is used, which also freezes exposure at its current value
-            when no exposure is configured. Valid ranges are camera-model specific and
-            reported if the value is rejected. Defaults to None (leave unchanged).
-        white_balance: Manual white balance value for the color sensor. When set, auto
-            white balance is disabled and this fixed value is used. Valid ranges are
-            camera-model specific and reported if the value is rejected. Defaults to None
-            (leave unchanged).
 
     Note:
         - Either name or serial_number must be specified.
@@ -72,9 +61,6 @@ class RealSenseCameraConfig(CameraConfig):
     use_depth: bool = False
     rotation: Cv2Rotation = Cv2Rotation.NO_ROTATION
     warmup_s: int = 1
-    exposure: int | None = None
-    gain: int | None = None
-    white_balance: int | None = None
 
     def __post_init__(self) -> None:
         self.color_mode = ColorMode(self.color_mode)
@@ -83,18 +69,6 @@ class RealSenseCameraConfig(CameraConfig):
         if not self.use_rgb and not self.use_depth:
             raise ValueError("At least one of `use_rgb` or `use_depth` must be enabled.")
 
-        manual_color_options = {
-            "exposure": self.exposure,
-            "gain": self.gain,
-            "white_balance": self.white_balance,
-        }
-        configured_color_options = [name for name, value in manual_color_options.items() if value is not None]
-        if configured_color_options and not self.use_rgb:
-            raise ValueError(
-                "Manual color sensor options require `use_rgb=True`. "
-                f"Configured options: {configured_color_options}."
-            )
-
         values = (self.fps, self.width, self.height)
         if any(v is not None for v in values) and any(v is None for v in values):
             raise ValueError(

src/lerobot/configs/dataset.py

@@ -71,19 +71,13 @@ class DatasetRecordConfig:
     # Number of threads per encoder instance. None = auto (codec default).
     # Lower values reduce CPU usage, maps to 'lp' (via svtav1-params) for libsvtav1 and 'threads' for h264/hevc..
     encoder_threads: int | None = None
-    # Skip appending the date-time tag to repo_id, keeping the user-provided name as-is
-    # (e.g. self-managed versioned names intended for a later `lerobot-edit-dataset merge`).
-    no_stamp: bool = False
 
     def stamp_repo_id(self) -> None:
         """Append a date-time tag to ``repo_id`` so each recording session gets a unique name.
 
         Must be called explicitly at dataset *creation* time — not on resume,
         where the existing ``repo_id`` (already stamped) must be preserved.
-        No-op when ``no_stamp`` is set, preserving a user-managed ``repo_id``.
         """
-        if self.no_stamp:
-            return
         if self.repo_id:
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             self.repo_id = f"{self.repo_id}_{timestamp}"

src/lerobot/configs/default.py

@@ -33,6 +33,8 @@ class DatasetConfig:
     # looked up under $HF_LEROBOT_HOME/repo_id and Hub downloads use a revision-safe cache under $HF_LEROBOT_HOME/hub.
     root: str | None = None
     episodes: list[int] | None = None
+    # Episode indices to drop (e.g. corrupt or heterogeneous ones). Applied on top of `episodes`.
+    exclude_episodes: list[int] | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None
     use_imagenet_stats: bool = True
@@ -62,6 +64,10 @@ class DatasetConfig:
             if len(self.episodes) != len(set(self.episodes)):
                 duplicates = sorted({ep for ep in self.episodes if self.episodes.count(ep) > 1})
                 raise ValueError(f"Episode indices contain duplicates: {duplicates}")
+        if self.exclude_episodes is not None and any(ep < 0 for ep in self.exclude_episodes):
+            raise ValueError(
+                f"exclude_episodes must be non-negative, got: {[ep for ep in self.exclude_episodes if ep < 0]}"
+            )
 
 
 @dataclass

src/lerobot/configs/recipe.py

@@ -78,7 +78,7 @@ class MessageTurn:
             raise ValueError(f"Unsupported message stream: {self.stream!r}")
         if self.content is None and self.tool_calls_from is None:
             raise ValueError("MessageTurn.content is required unless tool_calls_from is set.")
-        if self.content is not None and not isinstance(self.content, (str, list)):
+        if self.content is not None and not isinstance(self.content, str | list):
             raise TypeError("MessageTurn.content must be a string, a list of HF-style blocks, or None.")
         if isinstance(self.content, list):
             for block in self.content:
@@ -147,7 +147,7 @@ class TrainingRecipe:
         return cls.from_dict(data)
 
     def _validate_message_recipe(self) -> None:
-        """Ensure every templated binding is known and at least one turn is a target."""
+        """Validate bindings and require text or low-level action supervision."""
         assert self.messages is not None
         known_bindings = set(DEFAULT_BINDINGS) | set(self.bindings or {}) | {"task"}
 
@@ -156,8 +156,14 @@ class TrainingRecipe:
             if missing:
                 raise ValueError(f"MessageTurn references unknown binding(s): {sorted(missing)}")
 
-        if not any(turn.target for turn in self.messages):
-            raise ValueError("Message recipes must contain at least one target turn.")
+        has_target = any(turn.target for turn in self.messages)
+        has_low_level = any(turn.stream == "low_level" for turn in self.messages)
+        if not (has_target or has_low_level):
+            raise ValueError(
+                "Message recipes must contain at least one supervised turn — "
+                "either ``target: true`` (text CE) or ``stream: low_level`` "
+                "(flow/action loss)."
+            )
 
     def _validate_blend_recipe(self) -> None:
         """Ensure each blend component is a non-empty, weighted message recipe."""

src/lerobot/configs/recipes/subtask.yaml

@@ -0,0 +1,16 @@
+# Predicts subtasks from tasks and trains subtask-conditioned action flow without memory or plans.
+# Requires `subtask` annotations; samples with missing `if_present` bindings do not render.
+
+blend:
+
+  high_level_subtask:
+    weight: 0.30
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
+
+  low_level_execution:
+    weight: 0.70
+    messages:
+      # The low-level stream trains action flow on the generated or annotated subtask.
+      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}

src/lerobot/configs/recipes/subtask_joint.yaml

@@ -0,0 +1,13 @@
+# Paper-style joint sequence (pi0.5 §IV-B): one sample supervises the subtask
+# text with CE and, because the assistant turn is part of the prefix, conditions
+# the FAST and flow action losses on the same annotated subtask in one forward.
+# The supervised span is attended causally; the action losses see task + subtask.
+#
+# Pair with `--policy.joint_subtask_conditioning=true` at inference so the flow
+# prefix reproduces this layout (task turn with state + causal generated subtask).
+# Samples without a `subtask` annotation fall back to a plain task-prompt
+# low-level sample via `if_present`.
+
+messages:
+  - {role: user, content: "${task}", stream: low_level}
+  - {role: assistant, content: "${subtask}", stream: low_level, target: true, if_present: subtask}

src/lerobot/configs/recipes/subtask_mem.yaml

@@ -0,0 +1,30 @@
+# Trains subtask prediction, subtask-conditioned action flow, and memory updates without plans.
+# Requires `subtask` and `memory`; missing `if_present` bindings skip the affected sub-recipe.
+
+blend:
+
+  high_level_subtask:
+    weight: 0.25
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
+
+  low_level_execution:
+    weight: 0.60
+    messages:
+      # The low-level stream trains action flow on the generated or annotated subtask.
+      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
+
+  memory_update:
+    # `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
+    # Inference controls update timing through `subtask_change` events.
+    weight: 0.15
+    bindings:
+      prior_memory: "nth_prev(style=memory, offset=1)"
+      current_memory: "active_at(t, style=memory)"
+      completed_subtask: "nth_prev(style=subtask, offset=1)"
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
+      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
+      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}

src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml

@@ -0,0 +1,70 @@
+# Adds memory, spoken interjection responses, and camera-grounded VQA to subtask/action training.
+# Missing optional annotations skip only their sub-recipe; `say` tool calls tokenize as `<say>...</say>`.
+
+blend:
+
+  high_level_subtask:
+    weight: 0.25
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: assistant, content: "${subtask}", stream: high_level, target: true, if_present: subtask}
+
+  low_level_execution:
+    weight: 0.40
+    messages:
+      # The low-level stream trains action flow on the generated or annotated subtask.
+      - {role: user, content: "${subtask}", stream: low_level, if_present: subtask}
+
+  memory_update:
+    # `active_at` densifies sparse boundaries while preserving the prior-memory/subtask mapping.
+    # Inference controls update timing through `subtask_change` events.
+    weight: 0.10
+    bindings:
+      prior_memory: "nth_prev(style=memory, offset=1)"
+      current_memory: "active_at(t, style=memory)"
+      completed_subtask: "nth_prev(style=subtask, offset=1)"
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: assistant, content: "Previous memory: ${prior_memory}", stream: high_level, if_present: prior_memory}
+      - {role: user, content: "Completed subtask: ${completed_subtask}", stream: high_level, if_present: completed_subtask}
+      - {role: assistant, content: "${current_memory}", stream: high_level, target: true, if_present: current_memory}
+
+  user_interjection_response:
+    weight: 0.10
+    bindings:
+      interjection: "emitted_at(t, style=interjection)"
+      speech: "emitted_at(t, role=assistant, tool_name=say)"
+    messages:
+      - {role: user, content: "${task}", stream: high_level}
+      - {role: user, content: "${interjection}", stream: high_level, if_present: interjection}
+      # The assistant target is a `say` tool call flattened to a `<say>...</say>` marker.
+      - {role: assistant, stream: high_level, target: true, if_present: speech, tool_calls_from: speech}
+
+  # Each camera uses a separate VQA sub-recipe for view-specific binding.
+  ask_vqa_top:
+    weight: 0.075
+    bindings:
+      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.front)"
+      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.front)"
+    messages:
+      - role: user
+        stream: high_level
+        if_present: vqa_query
+        content:
+          - {type: image, feature: observation.images.front}
+          - {type: text, text: "${vqa_query}"}
+      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}
+
+  ask_vqa_wrist:
+    weight: 0.075
+    bindings:
+      vqa_query: "emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
+      vqa: "emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
+    messages:
+      - role: user
+        stream: high_level
+        if_present: vqa_query
+        content:
+          - {type: image, feature: observation.images.wrist}
+          - {type: text, text: "${vqa_query}"}
+      - {role: assistant, content: "${vqa}", stream: high_level, target: true, if_present: vqa}

src/lerobot/datasets/dataset_metadata.py

@@ -188,8 +188,8 @@ class LeRobotDatasetMetadata:
     def _load_metadata(self):
         self.info = load_info(self.root)
         check_version_compatibility(self.repo_id, self._version, CODEBASE_VERSION)
-        self.tasks = load_tasks(self.root) if self.total_tasks > 0 else None
-        self.episodes = load_episodes(self.root) if self.total_episodes > 0 else None
+        self.tasks = load_tasks(self.root)
+        self.episodes = load_episodes(self.root)
         self.stats = load_stats(self.root)
 
     def ensure_readable(self) -> None:

src/lerobot/datasets/dataset_reader.py

@@ -163,10 +163,40 @@ class DatasetReader:
     def _load_hf_dataset(self) -> datasets.Dataset:
         """hf_dataset contains all the observations, states, actions, rewards, etc."""
         features = get_hf_features_from_features(self._meta.features)
+        # Annotated datasets may have language columns absent from metadata.
+        # Extend the schema before the strict Parquet cast.
+        features = self._extend_features_with_language_columns(features)
         hf_dataset = load_nested_dataset(self.root / "data", features=features, episodes=self.episodes)
         hf_dataset.set_transform(hf_transform_to_torch)
         return hf_dataset
 
+    def _extend_features_with_language_columns(self, features: datasets.Features) -> datasets.Features:
+        """Register language columns found in Parquet but missing from metadata."""
+        # Leave empty datasets to fail through the normal loading path.
+        try:
+            sample = next((self.root / "data").glob("*/*.parquet"))
+        except StopIteration:
+            return features
+
+        from pyarrow import parquet as _pq  # noqa: PLC0415
+
+        schema_names = set(_pq.read_schema(sample).names)
+        from .language import (  # noqa: PLC0415
+            LANGUAGE_EVENTS,
+            LANGUAGE_PERSISTENT,
+            language_events_column_feature,
+            language_persistent_column_feature,
+        )
+
+        extra: dict[str, object] = {}
+        if LANGUAGE_PERSISTENT in schema_names and LANGUAGE_PERSISTENT not in features:
+            extra[LANGUAGE_PERSISTENT] = language_persistent_column_feature()
+        if LANGUAGE_EVENTS in schema_names and LANGUAGE_EVENTS not in features:
+            extra[LANGUAGE_EVENTS] = language_events_column_feature()
+        if not extra:
+            return features
+        return datasets.Features({**features, **extra})
+
     def _check_cached_episodes_sufficient(self) -> bool:
         """Check if the cached dataset contains all requested episodes and their video files."""
         if self.hf_dataset is None or len(self.hf_dataset) == 0:

src/lerobot/datasets/factory.py

@@ -66,6 +66,17 @@ def resolve_delta_timestamps(
     return delta_timestamps
 
 
+def _resolve_episodes(
+    episodes: list[int] | None, exclude_episodes: list[int] | None, total_episodes: int
+) -> list[int] | None:
+    """Apply an episode exclusion list on top of an optional allowlist."""
+    if not exclude_episodes:
+        return episodes
+    base = episodes if episodes is not None else list(range(total_episodes))
+    excluded = set(exclude_episodes)
+    return [episode for episode in base if episode not in excluded]
+
+
 def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDataset:
     """Handles the logic of setting up delta timestamps and image transforms before creating a dataset.
 
@@ -87,11 +98,14 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
         )
         delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
+        episodes = _resolve_episodes(
+            cfg.dataset.episodes, cfg.dataset.exclude_episodes, ds_meta.total_episodes
+        )
         if not cfg.dataset.streaming:
             dataset = LeRobotDataset(
                 cfg.dataset.repo_id,
                 root=cfg.dataset.root,
-                episodes=cfg.dataset.episodes,
+                episodes=episodes,
                 delta_timestamps=delta_timestamps,
                 image_transforms=image_transforms,
                 revision=cfg.dataset.revision,
@@ -104,7 +118,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
             dataset = StreamingLeRobotDataset(
                 cfg.dataset.repo_id,
                 root=cfg.dataset.root,
-                episodes=cfg.dataset.episodes,
+                episodes=episodes,
                 delta_timestamps=delta_timestamps,
                 image_transforms=image_transforms,
                 revision=cfg.dataset.revision,

src/lerobot/datasets/language_render.py

@@ -162,14 +162,28 @@ def render_sample(
     task: str | None = None,
     dataset_ctx: Any | None = None,
 ) -> RenderedMessages | None:
-    """Render the chat-style messages for a single dataset sample.
+    """Resolve one sample's bindings and render its message recipe.
 
-    Resolves the recipe's bindings against ``persistent`` and ``events`` rows
-    at frame timestamp ``t``, then expands the recipe's message templates.
-    Returns ``None`` if the resolved sample contains no target message.
+    Returns ``None`` when no text or low-level action supervision applies.
     """
     persistent_rows = _normalize_rows(persistent or [])
     event_rows = _normalize_rows(events or [])
+
+    # Route sparse VQA frames to a matching view-specific component before weighted selection.
+    # This avoids dropping annotated frames or selecting VQA without annotations.
+    if recipe.blend is not None:
+        vqa_rendered = _render_vqa_if_present(
+            recipe,
+            persistent=persistent_rows,
+            events=event_rows,
+            t=t,
+            sample_idx=sample_idx,
+            task=task,
+            dataset_ctx=dataset_ctx,
+        )
+        if vqa_rendered is not None:
+            return vqa_rendered
+
     selected_recipe = _select_recipe(recipe, sample_idx)
     bindings = _resolve_bindings(
         selected_recipe,
@@ -183,6 +197,55 @@ def render_sample(
     return _render_message_recipe(selected_recipe, bindings)
 
 
+def _render_vqa_if_present(
+    recipe: TrainingRecipe,
+    *,
+    persistent: Sequence[LanguageRow],
+    events: Sequence[LanguageRow],
+    t: float,
+    sample_idx: int,
+    task: str | None,
+    dataset_ctx: Any | None,
+) -> RenderedMessages | None:
+    """Render a matching VQA component, or return ``None`` for normal selection.
+
+    Multiple matching views are selected deterministically by relative weight.
+    """
+    assert recipe.blend is not None
+    renderable: list[tuple[float, RenderedMessages]] = []
+    for name, component in recipe.blend.items():
+        if not name.startswith("ask_vqa"):
+            continue
+        bindings = _resolve_bindings(
+            component,
+            persistent=persistent,
+            events=events,
+            t=t,
+            sample_idx=sample_idx,
+            task=task,
+            dataset_ctx=dataset_ctx,
+        )
+        rendered = _render_message_recipe(component, bindings)
+        if rendered is not None:
+            renderable.append((float(component.weight or 0.0), rendered))
+
+    if not renderable:
+        return None
+    if len(renderable) == 1:
+        return renderable[0][1]
+
+    # Choose among matching cameras by relative weight, or uniformly when all weights are zero.
+    total = sum(w for w, _ in renderable) or float(len(renderable))
+    digest = hashlib.blake2b(f"vqa:{sample_idx}".encode(), digest_size=8).digest()
+    draw = int.from_bytes(digest, "big") / 2**64 * total
+    cumulative = 0.0
+    for w, rendered in renderable:
+        cumulative += w or (total / len(renderable))
+        if draw < cumulative:
+            return rendered
+    return renderable[-1][1]
+
+
 def _select_recipe(recipe: TrainingRecipe, sample_idx: int) -> TrainingRecipe:
     """Pick a deterministic blend component for ``sample_idx`` (or return ``recipe``)."""
     if recipe.blend is None:
@@ -346,7 +409,9 @@ def _render_message_recipe(
         if turn.target:
             target_indices.append(message_idx)
 
-    if not target_indices:
+    # Keep samples with either text targets or low-level action supervision.
+    has_low_level = any(stream == "low_level" for stream in streams)
+    if not target_indices and not has_low_level:
         return None
 
     rendered = {
@@ -403,14 +468,12 @@ def _validate_rendered(rendered: RenderedMessages) -> None:
 
     if len(streams) != len(messages):
         raise ValueError("message_streams must be aligned with messages.")
-    if not target_indices:
-        raise ValueError("Rendered samples must contain at least one target message.")
+    # Require text or low-level action supervision.
+    if not target_indices and not any(s == "low_level" for s in streams):
+        raise ValueError("Rendered samples must contain a target message or a low_level-stream message.")
     for idx in target_indices:
         if idx < 0 or idx >= len(messages):
             raise ValueError(f"Target message index {idx} is out of bounds.")
-    # ``stream`` is enforced non-None at MessageTurn construction time
-    # (see ``MessageTurn.__post_init__``), so a missing stream here would
-    # mean the dataclass invariant was bypassed; no need to re-check.
 
 
 def _nth_relative(

src/lerobot/envs/configs.py

@@ -560,7 +560,13 @@ class RoboCasaEnv(EnvConfig):
             kwargs["split"] = self.split
         return kwargs
 
-    def create_envs(self, n_envs: int, use_async_envs: bool = False):
+    def create_envs(
+        self,
+        n_envs: int,
+        use_async_envs: bool = False,
+        terminate_on_success: bool = True,
+        horizon: int | None = None,
+    ):
         from .robocasa import create_robocasa_envs
 
         if self.task is None:
@@ -574,6 +580,8 @@ class RoboCasaEnv(EnvConfig):
             env_cls=env_cls,
             episode_length=self.episode_length,
             obj_registries=tuple(self.obj_registries),
+            terminate_on_success=terminate_on_success,
+            horizon=horizon,
         )
 
 

src/lerobot/envs/libero.py

@@ -384,12 +384,7 @@ class LiberoEnv(gym.Env):
 
     def close(self):
         if self._env is not None:
-            try:
             self._env.close()
-            finally:
-                # LIBERO deletes its inner env on close, so this wrapper must
-                # be recreated before the next reset.
-                self._env = None
 
 
 def _make_env_fns(

src/lerobot/envs/robocasa.py

@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
 
 # Dimensions for the flat action/state vectors used by the LeRobot wrapper.
 # These correspond to the PandaOmron robot in RoboCasa365.
-OBS_STATE_DIM = 16  # 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/envs/robotwin.py

@@ -384,9 +384,7 @@ class RoboTwinEnv(gym.Env):
 
         self._env: Any | None = None  # deferred — created on first reset() inside worker
         self._step_count: int = 0
-        self._black_frame: np.ndarray = np.zeros(
-            (self.observation_height, self.observation_width, 3), dtype=np.uint8
-        )
+        self._black_frame = np.zeros((self.observation_height, self.observation_width, 3), dtype=np.uint8)
 
         image_spaces = {
             cam: spaces.Box(

src/lerobot/envs/vlabench.py

@@ -373,7 +373,7 @@ class VLABenchEnv(gym.Env):
 
         if action.shape[0] != 7:
             # Unknown layout — fall back to zero-pad so the sim doesn't crash.
-            padded: np.ndarray = np.zeros(ctrl_dim, dtype=np.float64)
+            padded = np.zeros(ctrl_dim, dtype=np.float64)
             padded[: min(action.shape[0], ctrl_dim)] = action[:ctrl_dim]
             return padded
 

src/lerobot/motors/dynamixel/tables.py

@@ -122,9 +122,6 @@ MODEL_ENCODING_TABLE = {
     "xm430-w350": X_SERIES_ENCODINGS_TABLE,
     "xm540-w270": X_SERIES_ENCODINGS_TABLE,
     "xc430-w150": X_SERIES_ENCODINGS_TABLE,
-    "xh540-w150": X_SERIES_ENCODINGS_TABLE,
-    "xc330-t288": X_SERIES_ENCODINGS_TABLE,
-    "xc330-t181": X_SERIES_ENCODINGS_TABLE,
 }
 
 # {model: model_resolution}
@@ -137,9 +134,6 @@ MODEL_RESOLUTION = {
     "xm430-w350": 4096,
     "xm540-w270": 4096,
     "xc430-w150": 4096,
-    "xh540-w150": 4096,
-    "xc330-t288": 4096,
-    "xc330-t181": 4096,
 }
 
 # {model: model_number}
@@ -151,9 +145,6 @@ MODEL_NUMBER_TABLE = {
     "xm430-w350": 1020,
     "xm540-w270": 1120,
     "xc430-w150": 1070,
-    "xh540-w150": 1110,
-    "xc330-t288": 1220,
-    "xc330-t181": 1210,
 }
 
 # {model: available_operating_modes}
@@ -165,9 +156,6 @@ MODEL_OPERATING_MODES = {
     "xm430-w350": [0, 1, 3, 4, 5, 16],
     "xm540-w270": [0, 1, 3, 4, 5, 16],
     "xc430-w150": [1, 3, 4, 16],
-    "xh540-w150": [0, 1, 3, 4, 5, 16],
-    "xc330-t288": [0, 1, 3, 4, 5, 16],
-    "xc330-t181": [0, 1, 3, 4, 5, 16],
 }
 
 MODEL_CONTROL_TABLE = {
@@ -178,9 +166,6 @@ MODEL_CONTROL_TABLE = {
     "xm430-w350": X_SERIES_CONTROL_TABLE,
     "xm540-w270": X_SERIES_CONTROL_TABLE,
     "xc430-w150": X_SERIES_CONTROL_TABLE,
-    "xh540-w150": X_SERIES_CONTROL_TABLE,
-    "xc330-t288": X_SERIES_CONTROL_TABLE,
-    "xc330-t181": X_SERIES_CONTROL_TABLE,
 }
 
 MODEL_BAUDRATE_TABLE = {
@@ -191,9 +176,6 @@ MODEL_BAUDRATE_TABLE = {
     "xm430-w350": X_SERIES_BAUDRATE_TABLE,
     "xm540-w270": X_SERIES_BAUDRATE_TABLE,
     "xc430-w150": X_SERIES_BAUDRATE_TABLE,
-    "xh540-w150": X_SERIES_BAUDRATE_TABLE,
-    "xc330-t288": X_SERIES_BAUDRATE_TABLE,
-    "xc330-t181": X_SERIES_BAUDRATE_TABLE,
 }
 
 AVAILABLE_BAUDRATES = [

src/lerobot/optim/optimizers.py

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

src/lerobot/policies/__init__.py

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

src/lerobot/policies/factory.py

@@ -44,19 +44,12 @@ from lerobot.utils.constants import (
     POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
 from lerobot.utils.feature_utils import dataset_to_policy_features
-from lerobot.utils.import_utils import _peft_available, require_package
 
 from .evo1.configuration_evo1 import Evo1Config
 from .groot.configuration_groot import GrootConfig
 from .pretrained import PreTrainedPolicy
 from .utils import validate_visual_features_consistency
 
-if TYPE_CHECKING or _peft_available:
-    from peft import PeftConfig, PeftModel
-else:
-    PeftConfig = None
-    PeftModel = None
-
 
 def _reconnect_relative_absolute_steps(
     preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -144,6 +137,12 @@ class ProcessorConfigKwargs(TypedDict, total=False):
     preprocessor_overrides: dict[str, Any] | None
     postprocessor_overrides: dict[str, Any] | None
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None
+    # Dataset source used by policies that optionally fit processor artifacts.
+    dataset_repo_id: str | None
+    dataset_root: str | None
+    dataset_revision: str | None
+    dataset_episodes: list[int] | None
+    dataset_exclude_episodes: list[int] | None
     dataset_meta: Any | None
 
 
@@ -178,6 +177,10 @@ def make_pre_post_processors(
         ValueError: If no processor factory exists for the given policy configuration type.
     """
     if pretrained_path:
+        # Register the PI052-only stateful tokenizer step before deserializing its pipeline.
+        if policy_cfg.type == "pi052":
+            from .pi052 import processor_pi052 as _processor_pi052  # noqa: F401
+
         if isinstance(policy_cfg, GrootConfig):
             from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
 
@@ -197,12 +200,29 @@ def make_pre_post_processors(
                 ),
             )
 
+        preprocessor_overrides = dict(kwargs.get("preprocessor_overrides") or {})
+        if policy_cfg.type == "pi0_fast" and getattr(policy_cfg, "auto_fit_fast_tokenizer", False):
+            from .pi052.fit_fast_tokenizer import resolve_fast_tokenizer
+
+            fitted_tokenizer = resolve_fast_tokenizer(
+                policy_cfg,
+                kwargs.get("dataset_repo_id"),
+                kwargs.get("dataset_root"),
+                kwargs.get("dataset_stats"),
+                kwargs.get("dataset_revision"),
+                kwargs.get("dataset_episodes"),
+                kwargs.get("dataset_exclude_episodes"),
+            )
+            tokenizer_overrides = dict(preprocessor_overrides.get("action_tokenizer_processor") or {})
+            tokenizer_overrides["action_tokenizer_name"] = fitted_tokenizer
+            preprocessor_overrides["action_tokenizer_processor"] = tokenizer_overrides
+
         preprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_model_name_or_path=pretrained_path,
             config_filename=kwargs.get(
                 "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
             ),
-            overrides=kwargs.get("preprocessor_overrides", {}),
+            overrides=preprocessor_overrides,
             to_transition=batch_to_transition,
             to_output=transition_to_batch,
             revision=pretrained_revision,
@@ -234,6 +254,11 @@ def make_pre_post_processors(
         config=policy_cfg,
         dataset_stats=kwargs.get("dataset_stats"),
         dataset_meta=kwargs.get("dataset_meta"),
+        dataset_repo_id=kwargs.get("dataset_repo_id"),
+        dataset_root=kwargs.get("dataset_root"),
+        dataset_revision=kwargs.get("dataset_revision"),
+        episodes=kwargs.get("dataset_episodes"),
+        exclude_episodes=kwargs.get("dataset_exclude_episodes"),
     )
 
 
@@ -341,15 +366,12 @@ def make_policy(
         # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
         # of the adapter and the adapter's config contains the path to the base policy. So we need the
         # adapter config first, then load the correct policy and then apply PEFT.
-        require_package("peft", extra="peft")
+        from peft import PeftConfig, PeftModel
 
         logging.info("Loading policy's PEFT adapter.")
 
         peft_pretrained_path = str(cfg.pretrained_path)
-        peft_config = PeftConfig.from_pretrained(
-            peft_pretrained_path,
-            revision=cfg.pretrained_revision,
-        )
+        peft_config = PeftConfig.from_pretrained(peft_pretrained_path)
 
         kwargs["pretrained_name_or_path"] = peft_config.base_model_name_or_path
         if not kwargs["pretrained_name_or_path"]:
@@ -360,14 +382,9 @@ def make_policy(
                 "the adapter was trained."
             )
 
-        kwargs["revision"] = peft_config.revision
         policy = policy_cls.from_pretrained(**kwargs)
         policy = PeftModel.from_pretrained(
-            policy,
-            peft_pretrained_path,
-            config=peft_config,
-            revision=cfg.pretrained_revision,
-            is_trainable=True,
+            policy, peft_pretrained_path, config=peft_config, is_trainable=True
         )
 
     else:
@@ -439,6 +456,7 @@ def _make_processors_from_policy_config(
     config: PreTrainedConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
     dataset_meta: Any | None = None,
+    **optional_kwargs: Any,
 ) -> tuple[Any, Any]:
     """Create pre- and post-processors from a policy configuration using dynamic imports.
 
@@ -474,7 +492,9 @@ def _make_processors_from_policy_config(
     function = getattr(module, function_name, None)
     if function is None:
         raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
+    parameters = inspect.signature(function).parameters
     call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
-    if "dataset_meta" in inspect.signature(function).parameters:
+    if "dataset_meta" in parameters:
         call_kwargs["dataset_meta"] = dataset_meta
+    call_kwargs.update({name: value for name, value in optional_kwargs.items() if name in parameters})
     return function(config, **call_kwargs)

src/lerobot/policies/gaussian_actor/configuration_gaussian_actor.py

@@ -37,19 +37,13 @@ def is_image_feature(key: str) -> bool:
 @dataclass
 class ConcurrencyConfig:
     """Configuration for the concurrency of the actor and learner.
-
     Possible values are:
     - "threads": Use threads for the actor and learner.
     - "processes": Use processes for the actor and learner.
-
-    ``multiprocessing_context`` selects the process-wide start method when
-    processes are used. Set it to ``None`` to preserve Python's default or a
-    method already selected by the embedding application.
     """
 
     actor: str = "threads"
     learner: str = "threads"
-    multiprocessing_context: str | None = "spawn"
 
 
 @dataclass

src/lerobot/policies/molmoact2/modeling_molmoact2.py

@@ -43,22 +43,11 @@ from torch.distributions import Beta
 
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION
-from lerobot.utils.import_utils import (
-    _peft_available,
-    _scipy_available,
-    _transformers_available,
-    require_package,
-)
+from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package
 
 from ..rtc.modeling_rtc import RTCProcessor
 from .configuration_molmoact2 import MolmoAct2Config
 
-if TYPE_CHECKING or _peft_available:
-    from peft import LoraConfig, get_peft_model
-else:
-    LoraConfig = None
-    get_peft_model = None
-
 logger = logging.getLogger(__name__)
 
 
@@ -1742,11 +1731,13 @@ class MolmoAct2Policy(PreTrainedPolicy):
 
     def _build_inner_lora_config(self):
         require_package("peft", extra="molmoact2")
+        from peft import LoraConfig
 
         return LoraConfig(**self._get_inner_peft_targets())
 
     def _apply_lora_adapters(self) -> None:
         require_package("peft", extra="molmoact2")
+        from peft import get_peft_model
 
         peft_config = self._build_inner_lora_config()
         self._validate_peft_config(peft_config)

src/lerobot/policies/pi05/modeling_pi05.py

@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
 
 import torch
 import torch.nn.functional as F  # noqa: N812
+from safetensors.torch import load_file
 from torch import Tensor, nn
 
 from lerobot.utils.import_utils import _transformers_available, require_package
@@ -30,6 +31,7 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 if TYPE_CHECKING or _transformers_available:
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
+    from transformers.utils import cached_file
 
     from ..pi_gemma import (
         PaliGemmaForConditionalGenerationWithPiGemma,
@@ -44,20 +46,21 @@ else:
     _gated_residual = None
     layernorm_forward = None
     PaliGemmaForConditionalGenerationWithPiGemma = None
+    cached_file = None
 from lerobot.configs import PreTrainedConfig
 from lerobot.utils.constants import (
     ACTION,
     OBS_LANGUAGE_ATTENTION_MASK,
     OBS_LANGUAGE_TOKENS,
+    OPENPI_ATTENTION_MASK_VALUE,
 )
 
-from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
+from ..common.flow_matching import sample_noise, sample_time_beta
 from ..common.vla_utils import (
     clone_past_key_values,
     create_sinusoidal_pos_embedding,
     make_att_2d_masks,
     pad_vector,
-    prepare_attention_masks_4d,
     resize_with_pad_torch,
 )
 from ..pretrained import PreTrainedPolicy, T
@@ -71,6 +74,9 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
+_SAFETENSORS_FILE = "model.safetensors"
+
+
 # Define the complete layer computation function for gradient checkpointing
 def compute_layer_complete(inputs_embeds, attention_mask, position_ids, adarms_cond, layers, rotary_emb):
     query_states = []
@@ -401,6 +407,12 @@ class PaliGemmaWithExpertModel(
 class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
     """Core PI05 PyTorch model."""
 
+    use_hf_vision_checkpointing_api = False
+    checkpoint_vision_embeddings = True
+    use_typed_attention_masks = False
+    use_on_device_suffix_mask = False
+    precompute_denoise_times = False
+
     def __init__(self, config: PI05Config, rtc_processor: RTCProcessor | None = None):
         super().__init__()
         self.config = config
@@ -444,7 +456,11 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         """Enable gradient checkpointing for memory optimization."""
         self.gradient_checkpointing_enabled = True
         self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
-        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = True
+        vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
+        if self.use_hf_vision_checkpointing_api:
+            vision_tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+        else:
+            vision_tower.gradient_checkpointing = True
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
         logging.info("Enabled gradient checkpointing for PI05Pytorch model")
 
@@ -452,7 +468,11 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         """Disable gradient checkpointing."""
         self.gradient_checkpointing_enabled = False
         self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = False
-        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = False
+        vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
+        if self.use_hf_vision_checkpointing_api:
+            vision_tower.gradient_checkpointing_disable()
+        else:
+            vision_tower.gradient_checkpointing = False
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
         logging.info("Disabled gradient checkpointing for PI05Pytorch model")
 
@@ -467,6 +487,14 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
+    def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
+        """Helper method to prepare 4D attention masks for transformer."""
+        att_2d_masks_4d = att_2d_masks[:, None, :, :]
+        result = torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
+        if dtype is not None:
+            result = result.to(dtype=dtype)
+        return result
+
     def sample_noise(self, shape, device):
         return sample_noise(shape, device)
 
@@ -488,13 +516,16 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         pad_masks = []
         att_masks = []
 
-        # Process images
-        for img, img_mask in zip(images, img_masks, strict=True):
+        if self.checkpoint_vision_embeddings:
 
-            def image_embed_func(img):
-                return self.paligemma_with_expert.embed_image(img)
+            def embed_image(img):
+                return self._apply_checkpoint(self.paligemma_with_expert.embed_image, img)
 
-            img_emb = self._apply_checkpoint(image_embed_func, img)
+            img_embs = [embed_image(img) for img in images]
+        else:
+            img_embs = [self.paligemma_with_expert.embed_image(img) for img in images]
+
+        for img_emb, img_mask in zip(img_embs, img_masks, strict=True):
             bsize, num_img_embs = img_emb.shape[:2]
 
             embs.append(img_emb)
@@ -556,6 +587,13 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         # Set attention masks so that image, language and state inputs do not attend to action tokens
         att_masks += [1] + ([0] * (self.config.chunk_size - 1))
+
+        if self.use_on_device_suffix_mask:
+            n = len(att_masks)
+            att_masks = torch.zeros(n, dtype=action_emb.dtype, device=action_emb.device)
+            att_masks[0] = 1
+            att_masks = att_masks[None, :].expand(bsize, n)
+        else:
             att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
             att_masks = att_masks[None, :].expand(bsize, len(att_masks))
 
@@ -583,7 +621,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         att_2d_masks = make_att_2d_masks(pad_masks, att_masks)
         position_ids = torch.cumsum(pad_masks, dim=1) - 1
 
-        att_2d_masks_4d = prepare_attention_masks_4d(att_2d_masks)
+        att_2d_masks_4d = self._prepare_attention_masks_4d(att_2d_masks)
 
         def forward_func(prefix_embs, suffix_embs, att_2d_masks_4d, position_ids, adarms_cond):
             (_, suffix_out), _ = self.paligemma_with_expert.forward(
@@ -641,7 +679,8 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
         prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
 
-        prefix_att_2d_masks_4d = prepare_attention_masks_4d(prefix_att_2d_masks)
+        mask_dtype = prefix_embs.dtype if self.use_typed_attention_masks else None
+        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks, dtype=mask_dtype)
         self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
 
         _, past_key_values = self.paligemma_with_expert.forward(
@@ -652,22 +691,53 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             use_cache=True,
         )
 
-        return euler_integrate(
-            lambda input_x_t, current_timestep: self.denoise_step(
+        dt = -1.0 / num_steps
+
+        times = None
+        if self.precompute_denoise_times:
+            times = torch.tensor(
+                [1.0 + step * dt for step in range(num_steps)], dtype=torch.float32, device=device
+            )
+
+        x_t = noise
+        for step in range(num_steps):
+            time = 1.0 + step * dt
+            if times is None:
+                time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+            else:
+                time_tensor = times[step].expand(bsize)
+
+            def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
+                return self.denoise_step(
                     prefix_pad_masks=prefix_pad_masks,
                     past_key_values=past_key_values,
                     x_t=input_x_t,
                     timestep=current_timestep,
-            ),
-            noise,
-            num_steps,
-            rtc_processor=self.rtc_processor,
-            rtc_enabled=self._rtc_enabled(),
-            inference_delay=kwargs.get("inference_delay"),
-            prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
-            execution_horizon=kwargs.get("execution_horizon"),
                 )
 
+            if self._rtc_enabled():
+                inference_delay = kwargs.get("inference_delay")
+                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
+                execution_horizon = kwargs.get("execution_horizon")
+
+                v_t = self.rtc_processor.denoise_step(
+                    x_t=x_t,
+                    prev_chunk_left_over=prev_chunk_left_over,
+                    inference_delay=inference_delay,
+                    time=time,
+                    original_denoise_step_partial=denoise_step_partial_call,
+                    execution_horizon=execution_horizon,
+                )
+            else:
+                v_t = denoise_step_partial_call(x_t)
+
+            x_t = x_t + dt * v_t
+
+            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
+                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
+
+        return x_t
+
     def denoise_step(
         self,
         prefix_pad_masks,
@@ -689,7 +759,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]
         position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1
 
-        full_att_2d_masks_4d = prepare_attention_masks_4d(full_att_2d_masks)
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(full_att_2d_masks)
         self.paligemma_with_expert.gemma_expert.model.config._attn_implementation = "eager"  # noqa: SLF001
 
         past_key_values = clone_past_key_values(past_key_values)
@@ -713,6 +783,10 @@ class PI05Policy(PreTrainedPolicy):
 
     config_class = PI05Config
     name = "pi05"
+    model_class = PI05Pytorch
+    eval_after_pretrained_load = False
+    show_openpi_disclaimer = True
+    use_native_pretrained_loader = False
 
     def __init__(
         self,
@@ -730,7 +804,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:
@@ -756,7 +830,23 @@ class PI05Policy(PreTrainedPolicy):
         strict: bool = True,
         **kwargs,
     ) -> T:
-        """Override the from_pretrained method to handle key remapping and display important disclaimer."""
+        """Load a native LeRobot checkpoint or convert the PI05 base checkpoint."""
+        if cls.use_native_pretrained_loader:
+            return super().from_pretrained(
+                pretrained_name_or_path,
+                config=config,
+                force_download=force_download,
+                resume_download=resume_download,
+                proxies=proxies,
+                token=token,
+                cache_dir=cache_dir,
+                local_files_only=local_files_only,
+                revision=revision,
+                strict=strict,
+                **kwargs,
+            )
+
+        if cls.show_openpi_disclaimer:
             print(
                 "The PI05 model is a direct port of the OpenPI implementation. \n"
                 "This implementation follows the original OpenPI structure for compatibility. \n"
@@ -765,7 +855,6 @@ class PI05Policy(PreTrainedPolicy):
         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,
@@ -779,85 +868,41 @@ 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
-
+        model_id = str(pretrained_name_or_path)
         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),
+            model_id,
+            _SAFETENSORS_FILE,
+            _raise_exceptions_for_missing_entries=False,
+            force_download=force_download,
+            resume_download=resume_download,
+            proxies=proxies,
+            token=token,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+            revision=revision,
         )
-                from safetensors.torch import load_file
+        if resolved_file is None:
+            raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
 
-                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
+        fixed_state_dict = model._fix_pytorch_state_dict_keys(load_file(resolved_file), model.config)
+        remapped_state_dict = {
+            key if key.startswith("model.") else f"model.{key}": value
+            for key, value in fixed_state_dict.items()
+        }
+        remapped_state_dict = model._prepare_pretrained_state_dict(remapped_state_dict)
         missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
-
         if missing_keys:
-                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")
-
+            logging.warning("Missing %s checkpoint keys: %s", cls.name, missing_keys)
         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}")
-
+            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`
@@ -1028,12 +1073,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."""

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

src/lerobot/policies/pi052/fit_fast_tokenizer.py

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

src/lerobot/policies/pi052/flashrt_fp8.py

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

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

@@ -0,0 +1,19 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PI052 adapter for the policy-agnostic language runtime."""
+
+from .pi052_adapter import PI052PolicyAdapter
+
+__all__ = ["PI052PolicyAdapter"]

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

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

src/lerobot/policies/pi052/processor_pi052.py

@@ -0,0 +1,164 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PI052 processor factory with optional recipe rendering and text tokenization.
+
+Without a recipe it delegates to the standard PI0.5 pipeline.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from lerobot.configs.recipe import TrainingRecipe
+from lerobot.processor import (
+    AbsoluteActionsProcessorStep,
+    ActionTokenizerProcessorStep,
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
+    PolicyAction,
+    PolicyProcessorPipeline,
+    RelativeActionsProcessorStep,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
+)
+
+# Import directly to keep optional language dependencies out of ``lerobot.processor``.
+from lerobot.processor.render_messages_processor import RenderMessagesStep
+from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+
+from ..pi05.processor_pi05 import make_pi05_pre_post_processors
+from .configuration_pi052 import PI052Config
+from .text_processor_pi052 import PI052TextTokenizerStep
+
+
+def make_pi052_pre_post_processors(
+    config: PI052Config,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    dataset_repo_id: str | None = None,
+    dataset_root: str | None = None,
+    dataset_revision: str | None = None,
+    episodes: list[int] | None = None,
+    exclude_episodes: list[int] | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """Build PI0.5-v2's pre/post-processor pipelines.
+
+    Falls through to π0.5's stock pipeline when ``recipe_path`` is unset.
+    """
+    if not config.recipe_path:
+        if getattr(config, "enable_fast_action_loss", False):
+            raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
+        return make_pi05_pre_post_processors(config, dataset_stats=dataset_stats)
+
+    recipe = _load_recipe(config.recipe_path)
+
+    relative_step = RelativeActionsProcessorStep(
+        enabled=config.use_relative_actions,
+        exclude_joints=getattr(config, "relative_exclude_joints", []),
+        action_names=getattr(config, "action_feature_names", None),
+    )
+
+    input_steps = [
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        relative_step,
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+        RenderMessagesStep(recipe=recipe),
+        PI052TextTokenizerStep(
+            tokenizer_name="google/paligemma-3b-pt-224",
+            max_length=config.tokenizer_max_length,
+            plan_dropout_prob=getattr(config, "plan_dropout_prob", 0.0),
+            memory_dropout_prob=getattr(config, "memory_dropout_prob", 0.0),
+            subtask_dropout_prob=getattr(config, "subtask_dropout_prob", 0.0),
+        ),
+    ]
+
+    # Add FAST action-token supervision only when explicitly enabled.
+    if getattr(config, "enable_fast_action_loss", False):
+        from .fit_fast_tokenizer import resolve_fast_tokenizer  # noqa: PLC0415
+
+        input_steps.append(
+            ActionTokenizerProcessorStep(
+                action_tokenizer_name=resolve_fast_tokenizer(
+                    config,
+                    dataset_repo_id,
+                    dataset_root,
+                    dataset_stats,
+                    dataset_revision,
+                    episodes,
+                    exclude_episodes,
+                ),
+                max_action_tokens=config.max_action_tokens,
+                fast_skip_tokens=config.fast_skip_tokens,
+                paligemma_tokenizer_name="google/paligemma-3b-pt-224",
+                allow_truncation=False,
+            )
+        )
+
+    input_steps.append(DeviceProcessorStep(device=config.device))
+
+    output_steps = [
+        UnnormalizerProcessorStep(
+            features=config.output_features,
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+        AbsoluteActionsProcessorStep(
+            enabled=config.use_relative_actions,
+            relative_step=relative_step,
+        ),
+        DeviceProcessorStep(device="cpu"),
+    ]
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )
+
+
+def _load_recipe(path_str: str) -> TrainingRecipe:
+    """Resolve ``path_str`` to a ``TrainingRecipe``.
+
+    Accepts an absolute path or a path relative to
+    ``src/lerobot/configs/``.
+    """
+    p = Path(path_str)
+    if not p.is_absolute() and not p.exists():
+        from lerobot.configs import recipe as _recipe_module  # noqa: PLC0415
+
+        configs_dir = Path(_recipe_module.__file__).resolve().parent
+        candidate = configs_dir / path_str
+        if candidate.exists():
+            p = candidate
+    return TrainingRecipe.from_yaml(p)

src/lerobot/policies/pi052/text_processor_pi052.py

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

src/lerobot/policies/pi0_fast/configuration_pi0_fast.py

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

src/lerobot/policies/pi0_fast/modeling_pi0_fast.py

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

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

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

src/lerobot/policies/pi_gemma.py

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

src/lerobot/policies/pretrained.py

@@ -34,22 +34,14 @@ from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
 from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
-from lerobot.utils.import_utils import _peft_available, require_package
 
 from .utils import log_model_loading_keys
 
-if TYPE_CHECKING or _peft_available:
-    from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
-else:
-    PEFT_TYPE_TO_CONFIG_MAPPING = None
-    PeftType = None
-    get_peft_model = None
+T = TypeVar("T", bound="PreTrainedPolicy")
 
 if TYPE_CHECKING:
     from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
 
-T = TypeVar("T", bound="PreTrainedPolicy")
-
 
 def _build_card_context(
     cfg: TrainPipelineConfig | None,
@@ -346,6 +338,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             "smolvla": "lerobot/smolvla_base",
             "pi0": "lerobot/pi0_base",
             "pi05": "lerobot/pi05_base",
+            "pi052": "lerobot/pi052_base",
             "pi0_fast": "lerobot/pi0fast-base",
             "xvla": "lerobot/xvla-base",
         }
@@ -392,7 +385,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
             peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
                 These are merged with policy defaults to build the final config.
         """
-        require_package("peft", extra="peft")
+        from peft import get_peft_model
 
         # If user provided a complete config, use it directly (with overrides)
         if peft_config is not None:
@@ -463,7 +456,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         Returns:
             Preprocessed dict with renamed keys and init_type mapped to method-specific key.
         """
-        require_package("peft", extra="peft")
+        from peft import PeftType
 
         cli_overrides = cli_overrides.copy()
 
@@ -488,7 +481,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
 
     def _build_peft_config(self, cli_overrides: dict):
         """Build a PEFT config from policy defaults and CLI overrides."""
-        require_package("peft", extra="peft")
+        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
 
         # Determine PEFT method type (default to LORA)
         method_type_str = cli_overrides.get("method_type") or "lora"
@@ -515,7 +508,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
 
     def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
         """Apply CLI overrides to an existing PEFT config."""
-        require_package("peft", extra="peft")
+        from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType
 
         # Get method type from existing config or CLI override
         method_type_str = cli_overrides.get("method_type")

src/lerobot/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/delta_action_processor.py

@@ -132,20 +132,10 @@ class MapDeltaActionToRobotActionStep(RobotActionProcessorStep):
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        for axis in ["x", "y", "z"]:
+        for axis in ["x", "y", "z", "gripper"]:
             features[PipelineFeatureType.ACTION].pop(f"delta_{axis}", None)
-        features[PipelineFeatureType.ACTION].pop("gripper", None)
 
-        for feat in [
-            "enabled",
-            "target_x",
-            "target_y",
-            "target_z",
-            "target_wx",
-            "target_wy",
-            "target_wz",
-            "gripper_vel",
-        ]:
+        for feat in ["enabled", "target_x", "target_y", "target_z", "target_wx", "target_wy", "target_wz"]:
             features[PipelineFeatureType.ACTION][f"{feat}"] = PolicyFeature(
                 type=FeatureType.ACTION, shape=(1,)
             )

src/lerobot/processor/pipeline.py

@@ -41,7 +41,7 @@ from pathlib import Path
 from typing import Any, TypedDict, TypeVar, cast
 
 import torch
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, snapshot_download
 from safetensors.torch import load_file, save_file
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -205,6 +205,10 @@ class ProcessorStep(ABC):
         """
         return None
 
+    def save_artifacts(self, save_directory: Path) -> dict[str, str]:
+        """Save non-tensor assets and map constructor arguments to relative paths."""
+        return {}
+
     def reset(self) -> None:
         """Resets the internal state of the processor step, if any."""
         return None
@@ -549,6 +553,22 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         pipeline_config = self.get_config()
         pipeline_state_dict = self.state_dict()
 
+        for processor_step, step_entry in zip(self.steps, pipeline_config["steps"], strict=True):
+            artifacts = processor_step.save_artifacts(save_directory)
+            if artifacts:
+                for config_key, relative_path in artifacts.items():
+                    artifact_path = Path(relative_path)
+                    if artifact_path.is_absolute() or ".." in artifact_path.parts:
+                        raise ValueError(
+                            f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
+                        )
+                    if not (save_directory / artifact_path).exists():
+                        raise FileNotFoundError(
+                            f"Processor step did not save declared artifact '{relative_path}'"
+                        )
+                    step_entry["config"][config_key] = artifact_path.as_posix()
+                step_entry["artifacts"] = artifacts
+
         for state_key, step_state_dict in pipeline_state_dict.items():
             state_filename = f"{state_key}.safetensors"
             save_file(step_state_dict, save_directory / state_filename)
@@ -733,7 +753,13 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
 
         # 3. Build steps with overrides
         steps, validated_overrides = cls._build_steps_with_overrides(
-            loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs, is_local_source
+            loaded_config,
+            overrides or {},
+            model_id,
+            base_path,
+            config_filename,
+            hub_download_kwargs,
+            is_local_source,
         )
 
         # 4. Validate that all overrides were used
@@ -922,6 +948,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         overrides: dict[str, Any],
         model_id: str,
         base_path: Path | None,
+        config_filename: str,
         hub_download_kwargs: dict[str, Any],
         is_local_source: bool = False,
     ) -> tuple[list[ProcessorStep], set[str]]:
@@ -976,15 +1003,68 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
             ImportError: If a step class cannot be imported or found in registry
             ValueError: If a step cannot be instantiated with its configuration
         """
+        loaded_config = deepcopy(loaded_config)
+        cls._resolve_artifact_paths(
+            loaded_config,
+            model_id,
+            base_path,
+            config_filename,
+            hub_download_kwargs,
+        )
         steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
 
         for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
             cls._load_step_state(
-                step_instance, step_entry, model_id, base_path, hub_download_kwargs, is_local_source
+                step_instance,
+                step_entry,
+                model_id,
+                base_path,
+                config_filename,
+                hub_download_kwargs,
+                is_local_source,
             )
 
         return steps, remaining_override_keys
 
+    @classmethod
+    def _resolve_artifact_paths(
+        cls,
+        loaded_config: dict[str, Any],
+        model_id: str,
+        base_path: Path | None,
+        config_filename: str,
+        hub_download_kwargs: dict[str, Any],
+    ) -> None:
+        """Resolve declared relative processor artifacts before step construction."""
+        is_local = Path(model_id).is_dir() or Path(model_id).is_file()
+
+        for step_entry in loaded_config["steps"]:
+            artifacts = step_entry.get("artifacts", {})
+            for config_key, relative_path in artifacts.items():
+                artifact_path = Path(relative_path)
+                if artifact_path.is_absolute() or ".." in artifact_path.parts:
+                    raise ValueError(
+                        f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
+                    )
+
+                resolved_path = base_path / artifact_path if base_path is not None else artifact_path
+                if not resolved_path.exists() and not is_local:
+                    repository_path = Path(config_filename).parent / artifact_path
+                    snapshot_download(
+                        repo_id=model_id,
+                        repo_type="model",
+                        allow_patterns=f"{repository_path.as_posix()}/**",
+                        **hub_download_kwargs,
+                    )
+
+                if not resolved_path.exists():
+                    step_name = step_entry.get("registry_name", step_entry.get("class", "unknown"))
+                    raise FileNotFoundError(
+                        f"Missing processor artifact '{relative_path}' for step '{step_name}' "
+                        f"next to '{config_filename}'. Checkpoint artifacts are incomplete."
+                    )
+                step_entry["config"][config_key] = str(resolved_path)
+
     @classmethod
     def _build_steps_from_config(
         cls,
@@ -1144,6 +1224,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         step_entry: dict[str, Any],
         model_id: str,
         base_path: Path | None,
+        config_filename: str,
         hub_download_kwargs: dict[str, Any],
         is_local_source: bool = False,
     ) -> None:
@@ -1209,7 +1290,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
             # Download from Hub
             state_path = hf_hub_download(
                 repo_id=model_id,
-                filename=state_filename,
+                filename=(Path(config_filename).parent / state_filename).as_posix(),
                 repo_type="model",
                 **hub_download_kwargs,
             )

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:
+            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:
@@ -66,19 +77,148 @@ class RenderMessagesStep(ProcessorStep):
             task=complementary_data.get("task"),
             dataset_ctx=self.dataset_ctx,
         )
+        if rendered is None:
+            rendered = _fallback_low_level_render(complementary_data.get("task"))
             if rendered is None:
                 return None
 
         new_transition = transition.copy()
-        new_complementary_data = dict(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

@@ -25,6 +25,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -32,6 +33,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,
@@ -136,7 +138,7 @@ class TokenizerProcessorStep(ObservationProcessorStep):
         # Standardize to a list of strings for the tokenizer
         if isinstance(task, str):
             return [task]
-        elif isinstance(task, (list, tuple)) and all(isinstance(t, str) for t in task):
+        elif isinstance(task, list | tuple) and all(isinstance(t, str) for t in task):
             return list(task)
 
         return None
@@ -349,6 +351,8 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
     max_action_tokens: int = 256
     fast_skip_tokens: int = 128
     paligemma_tokenizer_name: str = "google/paligemma-3b-pt-224"
+    allow_truncation: bool = True
+    prepend_bos: bool = True
     # Internal tokenizer instance (not part of the config)
     action_tokenizer: Any = field(default=None, init=False, repr=False)
     _paligemma_tokenizer: Any = field(default=None, init=False, repr=False)
@@ -412,14 +416,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 +435,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 +464,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,65 +482,79 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             if tokens.dim() > 1:
                 tokens = tokens.flatten()
 
-            bos_id = self._paligemma_tokenizer.bos_token_id
-            # add bos
-            tokens = torch.cat(
-                [
-                    torch.tensor([bos_id], device=action.device),
-                    torch.tensor(
+            action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens)
+            prompt_tokens = torch.tensor(
                 self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
                 device=action.device,
-                    ),
-                    self._act_tokens_to_paligemma_tokens(tokens),
-                    torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device),
-                ]
             )
+            end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device)
+
+            token_parts = []
+            if self.prepend_bos:
+                token_parts.append(
+                    torch.tensor([self._paligemma_tokenizer.bos_token_id], device=action.device)
+                )
+            code_start = sum(len(part) for part in token_parts) + len(prompt_tokens)
+            code_end = code_start + len(action_code_tokens)
+            tokens = torch.cat([*token_parts, prompt_tokens, action_code_tokens, end_tokens])
+            code_mask = torch.zeros(len(tokens), dtype=torch.bool, device=action.device)
+            code_mask[code_start:code_end] = True
 
             # Truncate or pad to max_action_tokens
             if len(tokens) > self.max_action_tokens:
+                if not self.allow_truncation:
+                    raise ValueError(
+                        f"FAST action sequence has {len(tokens)} tokens, exceeding "
+                        f"max_action_tokens={self.max_action_tokens}."
+                    )
                 logging.warning(
                     f"Token length ({len(tokens)}) exceeds max length ({self.max_action_tokens}), truncating. "
                     "Consider increasing the `max_action_tokens` in your model config if this happens frequently."
                 )
                 tokens = tokens[: self.max_action_tokens]
+                code_mask = code_mask[: self.max_action_tokens]
                 mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device)
             else:
+                pad_len = self.max_action_tokens - len(tokens)
                 mask = torch.cat(
                     [
                         torch.ones(len(tokens), dtype=torch.bool, device=action.device),
-                        torch.zeros(
-                            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 +570,10 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
         config = {
             "trust_remote_code": self.trust_remote_code,
             "max_action_tokens": self.max_action_tokens,
+            "fast_skip_tokens": self.fast_skip_tokens,
+            "paligemma_tokenizer_name": self.paligemma_tokenizer_name,
+            "allow_truncation": self.allow_truncation,
+            "prepend_bos": self.prepend_bos,
         }
 
         # Only save tokenizer_name if it was used to create the tokenizer
@@ -558,6 +582,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
 
         return config
 
+    def save_artifacts(self, save_directory: Path) -> dict[str, str]:
+        artifact_path = Path("action_tokenizer")
+        save_pretrained = getattr(self.action_tokenizer, "save_pretrained", None)
+        if save_pretrained is None:
+            raise TypeError("Action tokenizer must implement save_pretrained() to save a portable pipeline.")
+        save_pretrained(save_directory / artifact_path)
+        return {"action_tokenizer_name": artifact_path.as_posix()}
+
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:

src/lerobot/rl/actor.py

@@ -91,7 +91,7 @@ from lerobot.robots import so_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so_leader  # noqa: F401
 from lerobot.teleoperators.utils import TeleopEvents
 from lerobot.utils.device_utils import get_safe_torch_device
-from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
+from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.robot_utils import precise_sleep
 from lerobot.utils.transition import (
@@ -124,7 +124,9 @@ def actor_cli(cfg: TrainRLServerPipelineConfig):
     cfg.validate()
     display_pid = False
     if not use_threads(cfg):
-        ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
+        import torch.multiprocessing as mp
+
+        mp.set_start_method("spawn")
         display_pid = True
 
     # Create logs directory to ensure it exists

src/lerobot/rl/learner.py

@@ -102,7 +102,7 @@ from lerobot.utils.constants import (
 )
 from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.io_utils import load_json, write_json
-from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
+from lerobot.utils.process import ProcessSignalHandler
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
     format_big_number,
@@ -123,7 +123,9 @@ def train_cli(cfg: TrainRLServerPipelineConfig):
     # Fail fast with a friendly error if the optional ``hilserl`` extra is missing.
     require_package("grpcio", extra="hilserl", import_name="grpc")
     if not use_threads(cfg):
-        ensure_multiprocessing_start_method(cfg.policy.concurrency.multiprocessing_context)
+        import torch.multiprocessing as mp
+
+        mp.set_start_method("spawn")
 
     # Use the job_name from the config
     train(

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

@@ -24,7 +24,6 @@ from __future__ import annotations
 import logging
 from dataclasses import dataclass, field
 from threading import Event
-from typing import TYPE_CHECKING
 
 import torch
 
@@ -48,7 +47,6 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
 from lerobot.robots import make_robot_from_config
 from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
 from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
-from lerobot.utils.import_utils import _peft_available, require_package
 
 from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
 from .inference import (
@@ -59,12 +57,6 @@ from .inference import (
 )
 from .robot_wrapper import ThreadSafeRobot
 
-if TYPE_CHECKING or _peft_available:
-    from peft import PeftConfig, PeftModel
-else:
-    PeftConfig = None
-    PeftModel = None
-
 logger = logging.getLogger(__name__)
 
 
@@ -179,7 +171,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy
             revision=pretrained_revision,
         )
 
-    require_package("peft", extra="peft")
+    from peft import PeftConfig, PeftModel
 
     peft_path = policy_config.pretrained_path
     peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
@@ -302,22 +294,12 @@ def build_rollout_context(
     # ``observation_features`` values are either a tuple (camera shape) or the
     # ``float`` type itself used as a sentinel for scalar motor features —
     # see ``dict[str, type | tuple]`` annotation on ``Robot.observation_features``.
-    # Keep cameras (tuple) plus both joint-position (.pos) and base-velocity (.vel)
-    # scalar state features. LeKiwi's observation.state is 9-dim (6 arm .pos +
-    # x/y/theta.vel) and the policy was trained/normalized on all 9; the old .pos-only
-    # filter fed a 6-dim state into a 9-dim normalizer → RuntimeError (size 6 vs 9).
-    # Pure-arm robots have no .vel state keys, so this is a no-op for them.
     observation_features_hw = {
         k: v
         for k, v in all_obs_features.items()
-        if isinstance(v, tuple) or (v is float and k.endswith((".pos", ".vel")))
+        if isinstance(v, tuple) or (v is float and k.endswith(".pos"))
     }
-    # Keep both joint-position (.pos) and base-velocity (.vel) action features so
-    # mobile manipulators command the base too (e.g. LeKiwi: 6 arm .pos +
-    # x/y/theta.vel = 9-dim action). Pure-arm robots have no .vel keys, so this is
-    # a no-op for them. Without the .vel keys the base velocities are silently
-    # dropped from dataset_features[ACTION]/ordered_action_keys and the base never moves.
-    action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith((".pos", ".vel"))}
+    action_features_hw = {k: v for k, v in robot.action_features.items() if k.endswith(".pos")}
 
     # The action side is always needed: sync inference reads action names from
     # ``dataset_features[ACTION]`` to map policy tensors back to robot actions.

src/lerobot/runtime/__init__.py

@@ -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/convert_dataset_v21_to_v30.py

@@ -94,8 +94,6 @@ from lerobot.datasets.video_utils import concatenate_video_files, get_video_dura
 from lerobot.utils.constants import HF_LEROBOT_HOME
 from lerobot.utils.utils import flatten_dict, init_logging
 
-logger = logging.getLogger(__name__)
-
 V21 = "v2.1"
 V30 = "v3.0"
 
@@ -478,11 +476,11 @@ def convert_dataset(
     # First check if the dataset already has a v3.0 version
     if root is None and not force_conversion:
         try:
-            logger.info("Trying to download v3.0 version of the dataset from the hub...")
+            print("Trying to download v3.0 version of the dataset from the hub...")
             snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id)
             return
         except Exception:
-            logger.info("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
+            print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.")
 
     # Set root based on whether local dataset path is provided
     use_local_dataset = False
@@ -490,7 +488,7 @@ def convert_dataset(
     if root.exists():
         validate_local_dataset_version(root)
         use_local_dataset = True
-        logger.info(f"Using local dataset at {root}")
+        print(f"Using local dataset at {root}")
 
     old_root = root.parent / f"{root.name}_old"
     new_root = root.parent / f"{root.name}_v30"
@@ -525,7 +523,7 @@ def convert_dataset(
         try:
             hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset")
         except (HTTPError, RevisionNotFoundError) as e:
-            logger.warning(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
+            print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})")
             pass
         hub_api.delete_files(
             delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"],

src/lerobot/scripts/lerobot_annotate.py

@@ -154,14 +154,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
     repo_id = cfg.new_repo_id or cfg.repo_id
     commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
     api = HfApi()
-    logger.info(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...")
+    print(f"[lerobot-annotate] creating/locating dataset repo {repo_id}...", flush=True)
     api.create_repo(
         repo_id=repo_id,
         repo_type="dataset",
         private=cfg.push_private,
         exist_ok=True,
     )
-    logger.info(f"[lerobot-annotate] uploading {root} -> {repo_id}...")
+    print(f"[lerobot-annotate] uploading {root} -> {repo_id}...", flush=True)
     commit_info = api.upload_folder(
         folder_path=str(root),
         repo_id=repo_id,
@@ -172,7 +172,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         # at the source dataset; a fresh card is generated below instead.
         ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
     )
-    logger.info(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}")
+    print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
 
     dataset_info = load_info(root)
     card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
@@ -200,13 +200,14 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         with suppress(RevisionNotFoundError):
             api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
         api.create_tag(**tag_kwargs)
-        logger.info(f"[lerobot-annotate] tagged {repo_id} as {version_tag}")
+        print(f"[lerobot-annotate] tagged {repo_id} as {version_tag}", flush=True)
     except Exception as exc:  # noqa: BLE001
-        logger.warning(
+        print(
             f"[lerobot-annotate] WARNING: could not create tag {version_tag!r} on {repo_id}: {exc}. "
             "Dataset is uploaded but ``LeRobotDataset`` won't be able to load it until it's tagged. "
             "Run: from huggingface_hub import HfApi; "
-            f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)"
+            f"HfApi().create_tag({repo_id!r}, tag={version_tag!r}, repo_type='dataset', exist_ok=True)",
+            flush=True,
         )
 
 

src/lerobot/scripts/lerobot_dataset_viz.py

@@ -89,8 +89,6 @@ from lerobot.datasets import LeRobotDataset
 from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
 from lerobot.utils.utils import init_logging
 
-logger = logging.getLogger(__name__)
-
 DEFAULT_FOXGLOVE_PORT = 8765
 DEFAULT_RERUN_PORT = 9090
 
@@ -301,7 +299,7 @@ def visualize_dataset(
             while True:
                 time.sleep(1)
         except KeyboardInterrupt:
-            logger.info("Ctrl-C received. Exiting.")
+            print("Ctrl-C received. Exiting.")
 
 
 def main():

src/lerobot/scripts/lerobot_eval.py

@@ -62,7 +62,7 @@ from dataclasses import asdict
 from functools import partial
 from pathlib import Path
 from pprint import pformat
-from typing import TYPE_CHECKING, Any, TypedDict
+from typing import Any, TypedDict
 
 import einops
 import gymnasium as gym
@@ -87,21 +87,26 @@ from lerobot.processor import PolicyProcessorPipeline
 from lerobot.types import PolicyAction
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
 from lerobot.utils.device_utils import get_safe_torch_device
-from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
+from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.io_utils import write_video
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
     init_logging,
     inside_slurm,
 )
-
-if TYPE_CHECKING or _peft_available:
-    from peft import PeftModel
-else:
-    PeftModel = None
+from lerobot.utils.video_annotation import annotate_frame
 
 
-logger = logging.getLogger(__name__)
+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:
@@ -452,11 +457,13 @@ def eval_policy(
         exc = ValueError(
             f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
         )
-        if not _peft_available:
-            raise exc
-        require_package("peft", extra="peft")
+        try:
+            from peft import PeftModel
+
             if not isinstance(policy, PeftModel):
                 raise exc
+        except ImportError:
+            raise exc from None
 
     start = time.time()
     # Preserve the mode for direct callers. eval_policy_all scopes the mode
@@ -483,11 +490,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] = []
@@ -564,7 +596,7 @@ def eval_policy(
         if seeds:
             all_seeds.extend(seeds)
         else:
-            all_seeds.extend([None] * env.num_envs)
+            all_seeds.append(None)
 
         # FIXME: episode_data is either None or it doesn't exist
         if return_episode_data:
@@ -802,13 +834,13 @@ def eval_main(cfg: EvalPipelineConfig):
             recording_repo_id=cfg.eval.recording_repo_id,
             recording_private=cfg.eval.recording_private,
         )
-        logger.info("Overall Aggregated Metrics:")
-        logger.info(info["overall"])
+        print("Overall Aggregated Metrics:")
+        print(info["overall"])
 
         # Print per-suite stats
         for task_group, task_group_info in info.items():
-            logger.info(f"\nAggregated Metrics for {task_group}:")
-            logger.info(task_group_info)
+            print(f"\nAggregated Metrics for {task_group}:")
+            print(task_group_info)
     # Close all vec envs
     close_envs(envs)
 

src/lerobot/scripts/lerobot_find_cameras.py

@@ -40,7 +40,6 @@ from PIL import Image
 from lerobot.cameras import ColorMode
 from lerobot.cameras.opencv import OpenCVCamera, OpenCVCameraConfig
 from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
-from lerobot.utils.utils import init_logging
 
 logger = logging.getLogger(__name__)
 
@@ -286,8 +285,6 @@ def save_images_from_all_cameras(
 
 
 def main():
-    init_logging()
-
     parser = argparse.ArgumentParser(
         description="Unified camera utility script for listing cameras and capturing images."
     )

src/lerobot/scripts/lerobot_rollout.py

@@ -151,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
@@ -165,7 +166,6 @@ from lerobot.robots import (  # noqa: F401
     earthrover_mini_plus,
     hope_jr,
     koch_follower,
-    lekiwi,
     omx_follower,
     openarm_follower,
     reachy2,
@@ -242,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,10 +20,11 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
+import os
 import sys
 import time
-from collections.abc import Iterator
-from contextlib import contextmanager, nullcontext
+from contextlib import nullcontext
+from datetime import timedelta
 from pprint import pformat
 from typing import TYPE_CHECKING, Any
 
@@ -58,7 +59,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
 from lerobot.rewards import make_reward_pre_post_processors
 from lerobot.utils.collate import lerobot_collate_fn
-from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
+from lerobot.utils.import_utils import register_third_party_plugins
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
 from lerobot.utils.random_utils import set_seed
 from lerobot.utils.utils import (
@@ -69,28 +70,9 @@ from lerobot.utils.utils import (
     inside_slurm,
 )
 
-if TYPE_CHECKING or _peft_available:
-    from peft import PeftModel
-else:
-    PeftModel = None
-
 from .lerobot_eval import eval_policy_all
 
 
-@contextmanager
-def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, Any]]]:
-    """Create evaluation environments for one run and always dispose of them."""
-    envs = make_env(
-        cfg.env,
-        n_envs=cfg.eval.batch_size,
-        use_async_envs=cfg.eval.use_async_envs,
-    )
-    try:
-        yield envs
-    finally:
-        close_envs(envs)
-
-
 def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
     """Return worker-only DataLoader options, disabling them for single-process loading."""
     workers_enabled = cfg.num_workers > 0
@@ -111,6 +93,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.
@@ -128,6 +111,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:
@@ -195,12 +179,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()
+            }
     # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
     if output_dict:
         train_metrics.update_metrics(output_dict)
@@ -227,9 +219,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if cfg.job.is_remote:
         return submit_to_hf(cfg)
 
+    from lerobot.utils.import_utils import require_package
+
     require_package("accelerate", extra="training")
     from accelerate import Accelerator
-    from accelerate.utils import DistributedDataParallelKwargs, DistributedType
+    from accelerate.utils import DistributedDataParallelKwargs, DistributedType, InitProcessGroupKwargs
 
     cfg.validate()
 
@@ -238,7 +232,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"
@@ -248,7 +251,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,
         )
 
@@ -295,6 +298,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if not is_main_process:
         dataset, eval_dataset = make_train_eval_datasets(cfg)
 
+    # Create environment used for evaluating checkpoints during training on simulation data.
+    # On real-world data, no need to create an environment as evaluations are done outside train.py,
+    # using the eval.py instead, with gym_dora environment and dora-rs.
+    eval_env = None
+    if cfg.env_eval_freq > 0 and cfg.env is not None and is_main_process:
+        logging.info("Creating env")
+        eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
+
     if cfg.is_reward_model_training:
         if is_main_process:
             logging.info("Creating reward model")
@@ -322,7 +333,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if cfg.peft is not None:
         if cfg.is_reward_model_training:
             raise ValueError("PEFT is only supported for policy training. ")
-        require_package("peft", extra="peft")
+        from peft import PeftModel
 
         if isinstance(policy, PeftModel):
             logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
@@ -336,6 +347,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:
@@ -344,6 +363,13 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
+    if cfg.policy.type in {"pi0_fast", "pi052"}:
+        processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
+        processor_kwargs["dataset_revision"] = cfg.dataset.revision
+        processor_kwargs["dataset_episodes"] = cfg.dataset.episodes
+        processor_kwargs["dataset_exclude_episodes"] = cfg.dataset.exclude_episodes
+        processor_kwargs["dataset_root"] = cfg.dataset.root
+
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
         preprocessor_overrides = {
             "device_processor": {"device": device.type},
@@ -440,13 +466,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:
@@ -593,7 +623,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        train_tracker, _ = update_policy(
+        # 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,
             batch,
@@ -602,6 +635,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
@@ -702,10 +736,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}")
-                with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast():
+                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,
@@ -750,6 +785,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if is_main_process:
         progbar.close()
 
+    if eval_env:
+        close_envs(eval_env)
+
     is_fsdp = accelerator.distributed_type == DistributedType.FSDP
     model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None
     if is_main_process:

src/lerobot/scripts/lerobot_train_tokenizer.py

@@ -45,7 +45,6 @@ lerobot-train-tokenizer \
 """
 
 import json
-import logging
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -64,9 +63,6 @@ else:
 from lerobot.configs import NormalizationMode, parser
 from lerobot.datasets import LeRobotDataset
 from lerobot.utils.constants import ACTION, OBS_STATE
-from lerobot.utils.utils import init_logging
-
-logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -278,8 +274,11 @@ def process_episode(args):
 
         return action_chunks
 
-    except Exception:
-        logger.exception("Error processing episode %s", ep_idx)
+    except Exception as e:
+        print(f"Error processing episode {ep_idx}: {e}")
+        import traceback
+
+        traceback.print_exc()
         return None
 
 
@@ -301,10 +300,10 @@ def train_fast_tokenizer(
     Returns:
         Trained FAST tokenizer
     """
-    logger.info(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
-    logger.info(f"Action chunk shape: {action_chunks.shape}")
-    logger.info(f"Vocab size: {vocab_size}")
-    logger.info(f"DCT scale: {scale}")
+    print(f"Training FAST tokenizer on {len(action_chunks)} action chunks...")
+    print(f"Action chunk shape: {action_chunks.shape}")
+    print(f"Vocab size: {vocab_size}")
+    print(f"DCT scale: {scale}")
 
     # download the tokenizer source code (not pretrained weights)
     # we'll train a new tokenizer on our own data
@@ -315,7 +314,7 @@ def train_fast_tokenizer(
 
     # train the new tokenizer on our action data using .fit()
     # this trains the BPE tokenizer on DCT coefficients
-    logger.info("Training new tokenizer (this may take a few minutes)...")
+    print("Training new tokenizer (this may take a few minutes)...")
     tokenizer = base_tokenizer.fit(
         action_data_list,
         scale=scale,
@@ -323,21 +322,21 @@ def train_fast_tokenizer(
         time_horizon=action_chunks.shape[1],  # action_horizon
         action_dim=action_chunks.shape[2],  # encoded dimensions
     )
-    logger.info("✓ Tokenizer training complete!")
+    print("✓ Tokenizer training complete!")
 
     # validate it works
     sample_chunk = action_chunks[0]
     encoded = tokenizer(sample_chunk[None])[0]
     if isinstance(encoded, list):
         encoded = np.array(encoded)
-    logger.info(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
+    print(f"Sample encoding: {len(encoded)} tokens for chunk shape {sample_chunk.shape}")
 
     return tokenizer
 
 
 def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
     """Compute compression statistics."""
-    logger.info("\nComputing compression statistics...")
+    print("\nComputing compression statistics...")
 
     # sample for stats (use max 1000 chunks for speed)
     sample_size = min(1000, len(action_chunks))
@@ -367,12 +366,12 @@ def compute_compression_stats(tokenizer, action_chunks: np.ndarray):
         "max_token_length": float(np.max(token_lengths)),
     }
 
-    logger.info("Compression Statistics:")
-    logger.info(f"  Average compression ratio: {stats['compression_ratio']:.2f}x")
-    logger.info(f"  Mean token length: {stats['mean_token_length']:.1f}")
-    logger.info(f"  P99 token length: {stats['p99_token_length']:.0f}")
-    logger.info(f"  Min token length: {stats['min_token_length']:.0f}")
-    logger.info(f"  Max token length: {stats['max_token_length']:.0f}")
+    print("Compression Statistics:")
+    print(f"  Average compression ratio: {stats['compression_ratio']:.2f}x")
+    print(f"  Mean token length: {stats['mean_token_length']:.1f}")
+    print(f"  P99 token length: {stats['p99_token_length']:.0f}")
+    print(f"  Min token length: {stats['min_token_length']:.0f}")
+    print(f"  Max token length: {stats['max_token_length']:.0f}")
 
     return stats
 
@@ -386,9 +385,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
         cfg: TokenizerTrainingConfig dataclass with all configuration parameters
     """
     # load dataset
-    logger.info(f"Loading dataset: {cfg.repo_id}")
+    print(f"Loading dataset: {cfg.repo_id}")
     dataset = LeRobotDataset(repo_id=cfg.repo_id, root=cfg.root)
-    logger.info(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
+    print(f"Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
 
     # parse normalization mode
     try:
@@ -398,7 +397,7 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
             f"Invalid normalization_mode: {cfg.normalization_mode}. "
             f"Must be one of: {', '.join([m.value for m in NormalizationMode])}"
         ) from err
-    logger.info(f"Normalization mode: {norm_mode.value}")
+    print(f"Normalization mode: {norm_mode.value}")
 
     # parse encoded dimensions
     encoded_dim_ranges = []
@@ -407,38 +406,38 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
         encoded_dim_ranges.append((start, end))
 
     total_encoded_dims = sum(end - start for start, end in encoded_dim_ranges)
-    logger.info(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
+    print(f"Encoding {total_encoded_dims} dimensions: {cfg.encoded_dims}")
 
     # parse relative dimensions
     relative_dim_list = None
     if cfg.relative_dims is not None and cfg.relative_dims.strip():
         relative_dim_list = [int(d.strip()) for d in cfg.relative_dims.split(",")]
-        logger.info(f"Relative dimensions: {relative_dim_list}")
+        print(f"Relative dimensions: {relative_dim_list}")
     else:
-        logger.info("No relative dimensions specified")
+        print("No relative dimensions specified")
 
-    logger.info(f"Use relative transform: {cfg.use_relative_transform}")
+    print(f"Use relative transform: {cfg.use_relative_transform}")
     if cfg.use_relative_transform and (relative_dim_list is None or len(relative_dim_list) == 0):
-        logger.warning(
+        print(
             "Warning: use_relative_transform=True but no relative_dims specified. "
             "No relative transform will be applied."
         )
 
-    logger.info(f"Action horizon: {cfg.action_horizon}")
-    logger.info(f"State key: {cfg.state_key}")
+    print(f"Action horizon: {cfg.action_horizon}")
+    print(f"State key: {cfg.state_key}")
 
     # determine episodes to process
     num_episodes = dataset.num_episodes
     if cfg.max_episodes is not None:
         num_episodes = min(cfg.max_episodes, num_episodes)
 
-    logger.info(f"Processing {num_episodes} episodes...")
+    print(f"Processing {num_episodes} episodes...")
 
     # process episodes sequentially (to avoid pickling issues with dataset)
     all_chunks = []
     for ep_idx in range(num_episodes):
         if ep_idx % 10 == 0:
-            logger.info(f"  Processing episode {ep_idx}/{num_episodes}...")
+            print(f"  Processing episode {ep_idx}/{num_episodes}...")
 
         chunks = process_episode(
             (
@@ -456,19 +455,19 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
 
     # concatenate all chunks
     all_chunks = np.concatenate(all_chunks, axis=0)
-    logger.info(f"Collected {len(all_chunks)} action chunks")
+    print(f"Collected {len(all_chunks)} action chunks")
 
     # extract only encoded dimensions FIRST (before normalization)
     encoded_chunks = []
     for start, end in encoded_dim_ranges:
         encoded_chunks.append(all_chunks[:, :, start:end])
     encoded_chunks = np.concatenate(encoded_chunks, axis=-1)  # [N, H, D_encoded]
-    logger.info(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
+    print(f"Extracted {encoded_chunks.shape[-1]} encoded dimensions")
 
     # apply normalization to encoded dimensions
-    logger.info("\nBefore normalization - overall stats:")
-    logger.info(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
-    logger.info(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
+    print("\nBefore normalization - overall stats:")
+    print(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
+    print(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
 
     # get normalization stats from dataset
     norm_stats = dataset.meta.stats
@@ -490,9 +489,9 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
                     encoded_stats[stat_name] = stat_array[encoded_dim_indices]
 
         if encoded_stats:
-            logger.info(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
+            print(f"\nNormalization stats for encoded dimensions (mode: {norm_mode.value}):")
             for stat_name, stat_values in encoded_stats.items():
-                logger.info(
+                print(
                     f"  {stat_name}: shape={stat_values.shape}, "
                     f"range=[{np.min(stat_values):.4f}, {np.max(stat_values):.4f}]"
                 )
@@ -500,27 +499,27 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
             # apply normalization based on mode
             try:
                 encoded_chunks = apply_normalization(encoded_chunks, encoded_stats, norm_mode, eps=1e-8)
-                logger.info(f"\nApplied {norm_mode.value} normalization")
+                print(f"\nApplied {norm_mode.value} normalization")
             except ValueError as e:
-                logger.warning(f"Warning: {e}. Using raw actions without normalization.")
+                print(f"Warning: {e}. Using raw actions without normalization.")
 
-            logger.info("\nAfter normalization - overall stats:")
-            logger.info(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
-            logger.info(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
+            print("\nAfter normalization - overall stats:")
+            print(f"  Min: {np.min(encoded_chunks):.4f}, Max: {np.max(encoded_chunks):.4f}")
+            print(f"  Mean: {np.mean(encoded_chunks):.4f}, Std: {np.std(encoded_chunks):.4f}")
 
-            logger.info("\nPer-dimension stats (after normalization):")
+            print("\nPer-dimension stats (after normalization):")
             for d in range(encoded_chunks.shape[-1]):
                 dim_data = encoded_chunks[:, :, d]
-                logger.info(
+                print(
                     f"  Dim {d}: min={np.min(dim_data):7.4f}, max={np.max(dim_data):7.4f}, "
                     f"mean={np.mean(dim_data):7.4f}, std={np.std(dim_data):7.4f}"
                 )
         else:
-            logger.warning("Warning: Could not extract stats for encoded dimensions, using raw actions")
+            print("Warning: Could not extract stats for encoded dimensions, using raw actions")
     else:
-        logger.warning("Warning: No normalization stats found in dataset, using raw actions")
+        print("Warning: No normalization stats found in dataset, using raw actions")
 
-    logger.info(f"Encoded chunks shape: {encoded_chunks.shape}")
+    print(f"Encoded chunks shape: {encoded_chunks.shape}")
 
     # train FAST tokenizer
     tokenizer = train_fast_tokenizer(
@@ -562,8 +561,8 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
     with open(output_path / "metadata.json", "w") as f:
         json.dump(metadata, f, indent=2)
 
-    logger.info(f"\nSaved FAST tokenizer to {output_path}")
-    logger.info(f"Metadata: {json.dumps(metadata, indent=2)}")
+    print(f"\nSaved FAST tokenizer to {output_path}")
+    print(f"Metadata: {json.dumps(metadata, indent=2)}")
 
     # push to Hugging Face Hub if requested
     if cfg.push_to_hub:
@@ -571,10 +570,10 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
         hub_repo_id = cfg.hub_repo_id
         if hub_repo_id is None:
             hub_repo_id = output_path.name
-            logger.info(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
+            print(f"\nNo hub_repo_id provided, using: {hub_repo_id}")
 
-        logger.info(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
-        logger.info(f"   Private: {cfg.hub_private}")
+        print(f"\nPushing tokenizer to Hugging Face Hub: {hub_repo_id}")
+        print(f"   Private: {cfg.hub_private}")
 
         try:
             # use the tokenizer's push_to_hub method
@@ -594,15 +593,14 @@ def train_tokenizer(cfg: TokenizerTrainingConfig):
                 commit_message="Upload tokenizer metadata",
             )
 
-            logger.info(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
+            print(f"Successfully pushed tokenizer to: https://huggingface.co/{hub_repo_id}")
         except Exception as e:
-            logger.error(f"Error pushing to hub: {e}")
-            logger.error("   Make sure you're logged in with `huggingface-cli login`")
+            print(f"Error pushing to hub: {e}")
+            print("   Make sure you're logged in with `huggingface-cli login`")
 
 
 def main():
     """CLI entry point that parses arguments and runs the tokenizer training."""
-    init_logging()
     train_tokenizer()
 
 

src/lerobot/transforms/transforms.py

@@ -41,7 +41,7 @@ class RandomSubsetApply(Transform):
 
     def __init__(
         self,
-        transforms: Sequence[Callable[..., Any]],
+        transforms: Sequence[Callable],
         p: list[float] | None = None,
         n_subset: int | None = None,
         random_order: bool = False,
@@ -50,7 +50,7 @@ class RandomSubsetApply(Transform):
         if not isinstance(transforms, Sequence):
             raise TypeError("Argument transforms should be a sequence of callables")
         if p is None:
-            p = [1.0] * len(transforms)
+            p = [1] * len(transforms)
         elif len(p) != len(transforms):
             raise ValueError(
                 f"Length of p doesn't match the number of transforms: {len(p)} != {len(transforms)}"
@@ -69,7 +69,7 @@ class RandomSubsetApply(Transform):
         self.n_subset = n_subset
         self.random_order = random_order
 
-        self.selected_transforms: list[Callable[..., Any]] = []
+        self.selected_transforms = None
 
     def forward(self, *inputs: Any) -> Any:
         needs_unpacking = len(inputs) > 1
@@ -119,7 +119,7 @@ class SharpnessJitter(Transform):
         super().__init__()
         self.sharpness = self._check_input(sharpness)
 
-    def _check_input(self, sharpness: float | Sequence[float]) -> tuple[float, float]:
+    def _check_input(self, sharpness):
         if isinstance(sharpness, (int | float)):
             if sharpness < 0:
                 raise ValueError("If sharpness is a single number, it must be non negative.")
@@ -215,7 +215,7 @@ class ImageTransformsConfig:
     )
 
 
-def make_transform_from_config(cfg: ImageTransformConfig) -> Transform:
+def make_transform_from_config(cfg: ImageTransformConfig):
     if cfg.type == "SharpnessJitter":
         return SharpnessJitter(**cfg.kwargs)
 
@@ -236,8 +236,8 @@ class ImageTransforms(Transform):
         super().__init__()
         self._cfg = cfg
 
-        self.weights: list[float] = []
-        self.transforms: dict[str, Transform] = {}
+        self.weights = []
+        self.transforms = {}
         for tf_name, tf_cfg in cfg.tfs.items():
             if tf_cfg.weight <= 0.0:
                 continue

src/lerobot/utils/collate.py

@@ -22,7 +22,7 @@ from torch.utils.data._utils.collate import default_collate
 
 from lerobot.datasets.language import LANGUAGE_COLUMNS
 
-_PYTHON_LIST_KEYS = {"messages", "message_streams", "target_message_indices"}
+_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

@@ -26,6 +26,7 @@ OBS_IMAGES = OBS_IMAGE + "s"
 OBS_LANGUAGE = OBS_STR + ".language"
 OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
 OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
+OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks"
 OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
 OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
 OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
@@ -34,6 +35,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

@@ -105,6 +105,8 @@ class MetricsTracker:
         "epochs",
         "accelerator",
         "_caller_metrics",
+        "_tensor_sums",
+        "_tensor_counts",
     ]
 
     def __init__(
@@ -133,6 +135,8 @@ class MetricsTracker:
         # Meter names the caller registered up front. update_metrics() leaves these untouched, so a
         # policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter.
         self._caller_metrics: set[str] = set(self.metrics)
+        self._tensor_sums: dict[str, torch.Tensor] = {}
+        self._tensor_counts: dict[str, int] = {}
 
     def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
         if name in self.__dict__:
@@ -160,6 +164,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 update_metrics(self, values: dict[str, Any]) -> None:
         """Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
 
@@ -167,7 +187,7 @@ class MetricsTracker:
         Caller-registered metrics (those passed to the constructor) are never overridden.
         """
         for name, value in values.items():
-            if isinstance(value, bool) or not isinstance(value, (int, float)):
+            if isinstance(value, bool) or not isinstance(value, int | float):
                 continue
             if name in self._caller_metrics:
                 continue
@@ -235,3 +255,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/process.py

@@ -16,39 +16,11 @@
 # limitations under the License.
 
 import logging
-import multiprocessing
 import os
 import signal
 import sys
 
 
-def ensure_multiprocessing_start_method(start_method: str | None) -> None:
-    """Set a multiprocessing start method once, or verify the existing method matches.
-
-    Passing ``None`` leaves Python's process-wide default untouched. This is useful
-    when LeRobot is embedded in an application that owns multiprocessing setup.
-    """
-    if start_method is None:
-        return
-
-    available_methods = multiprocessing.get_all_start_methods()
-    if start_method not in available_methods:
-        raise ValueError(
-            f"Multiprocessing start method must be one of {available_methods} on this platform, "
-            f"got {start_method!r}."
-        )
-
-    current_method = multiprocessing.get_start_method(allow_none=True)
-    if current_method is None:
-        multiprocessing.set_start_method(start_method)
-    elif current_method != start_method:
-        raise RuntimeError(
-            f"Multiprocessing start method is already {current_method!r}; cannot change it to "
-            f"{start_method!r}. Set the configured multiprocessing context to null to keep the "
-            "application's existing method, or launch LeRobot in a fresh process."
-        )
-
-
 class ProcessSignalHandler:
     """Utility class to attach graceful shutdown signal handlers.
 

src/lerobot/utils/rerun_visualization.py

@@ -38,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/utils.py

@@ -133,13 +133,10 @@ def say(text: str, blocking: bool = False):
     else:
         raise RuntimeError("Unsupported operating system for text-to-speech.")
 
-    try:
     if blocking:
-            subprocess.run(cmd, check=True, timeout=5)
+        subprocess.run(cmd, check=True)
     else:
         subprocess.Popen(cmd, creationflags=subprocess.CREATE_NO_WINDOW if system == "Windows" else 0)
-    except (FileNotFoundError, subprocess.TimeoutExpired) as e:
-        logging.warning("Text-to-speech command failed: %s | Error: %s", cmd, e)
 
 
 def log_say(text: str, play_sounds: bool = True, blocking: bool = False):

src/lerobot/utils/video_annotation.py

@@ -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/cameras/test_opencv.py

@@ -20,7 +20,7 @@
 # ```
 
 from pathlib import Path
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import cv2
 import numpy as np
@@ -123,73 +123,6 @@ def test_invalid_width_connect():
         camera.connect(warmup=False)
 
 
-def test_connect_cleans_up_after_settings_failure_and_allows_retry():
-    config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=0)
-    camera = OpenCVCamera(config)
-    opened_captures = []
-
-    def fail_settings():
-        opened_captures.append(camera.videocapture)
-        raise RuntimeError("settings failed")
-
-    with (
-        patch.object(camera, "_configure_capture_settings", side_effect=fail_settings),
-        pytest.raises(RuntimeError, match="settings failed"),
-    ):
-        camera.connect(warmup=False)
-
-    assert camera.videocapture is None
-    assert camera.thread is None
-    assert not camera.is_connected
-    assert opened_captures[0] is not None
-    assert not opened_captures[0].isOpened()
-
-    camera.connect(warmup=False)
-    assert camera.is_connected
-    camera.disconnect()
-
-
-def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
-    config = OpenCVCameraConfig(index_or_path=DEFAULT_PNG_FILE_PATH, warmup_s=1)
-    camera = OpenCVCamera(config)
-    read_threads = []
-
-    def fail_warmup(*_args, **_kwargs):
-        read_threads.append(camera.thread)
-        raise TimeoutError("no frame")
-
-    with (
-        patch.object(camera, "async_read", side_effect=fail_warmup),
-        pytest.raises(TimeoutError, match="no frame"),
-    ):
-        camera.connect()
-
-    assert camera.videocapture is None
-    assert camera.thread is None
-    assert not camera.is_connected
-    assert read_threads[0] is not None
-    assert not read_threads[0].is_alive()
-
-    camera.connect(warmup=False)
-    assert camera.is_connected
-    camera.disconnect()
-
-
-def test_find_cameras_releases_unopened_handles():
-    module_path = OpenCVCamera.__module__
-    unopened_capture = MagicMock()
-    unopened_capture.isOpened.return_value = False
-
-    with (
-        patch(f"{module_path}.platform.system", return_value="Darwin"),
-        patch(f"{module_path}.MAX_OPENCV_INDEX", 1),
-        patch(f"{module_path}.cv2.VideoCapture", return_value=unopened_capture),
-    ):
-        assert OpenCVCamera.find_cameras() == []
-
-    unopened_capture.release.assert_called_once_with()
-
-
 @pytest.mark.parametrize("index_or_path", TEST_IMAGE_PATHS, ids=TEST_IMAGE_SIZES)
 def test_read(index_or_path):
     config = OpenCVCameraConfig(index_or_path=index_or_path, warmup_s=0)

tests/cameras/test_realsense.py

@@ -20,7 +20,7 @@
 # ```
 
 from pathlib import Path
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -30,8 +30,6 @@ from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnected
 
 pytest.importorskip("pyrealsense2")
 
-import pyrealsense2 as rs
-
 from lerobot.cameras.realsense import RealSenseCamera, RealSenseCameraConfig
 
 TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "cameras"
@@ -63,17 +61,6 @@ def test_abc_implementation():
     _ = RealSenseCamera(config)
 
 
-@pytest.mark.parametrize("option", ["exposure", "gain", "white_balance"])
-def test_manual_color_option_requires_rgb(option):
-    with pytest.raises(ValueError, match="use_rgb=True"):
-        RealSenseCameraConfig(
-            serial_number_or_name="042",
-            use_rgb=False,
-            use_depth=True,
-            **{option: 100},
-        )
-
-
 def test_connect():
     config = RealSenseCameraConfig(serial_number_or_name="042", warmup_s=0)
 
@@ -96,27 +83,6 @@ def test_connect_invalid_camera_path(patch_realsense):
         camera.connect(warmup=False)
 
 
-def test_connect_cleans_up_when_sensor_configuration_fails():
-    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
-    camera = RealSenseCamera(config)
-    pipeline = MagicMock()
-    pipeline.start.return_value = MagicMock()
-
-    with (
-        patch("lerobot.cameras.realsense.camera_realsense.rs.pipeline", return_value=pipeline),
-        patch.object(camera, "_configure_rs_pipeline_config"),
-        patch.object(camera, "_configure_capture_settings"),
-        patch.object(camera, "_configure_sensor_options", side_effect=ValueError("invalid exposure")),
-        pytest.raises(ValueError, match="invalid exposure"),
-    ):
-        camera.connect(warmup=False)
-
-    pipeline.stop.assert_called_once_with()
-    assert camera.rs_pipeline is None
-    assert camera.rs_profile is None
-    assert not camera.is_connected
-
-
 def test_invalid_width_connect():
     config = RealSenseCameraConfig(serial_number_or_name="042", width=99999, height=480, fps=30)
     camera = RealSenseCamera(config)
@@ -125,33 +91,6 @@ def test_invalid_width_connect():
         camera.connect(warmup=False)
 
 
-def test_connect_cleans_up_after_warmup_failure_and_allows_retry():
-    config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30)
-    camera = RealSenseCamera(config)
-    read_threads = []
-
-    def fail_warmup(*_args, **_kwargs):
-        read_threads.append(camera.thread)
-        raise TimeoutError("no frame")
-
-    with (
-        patch.object(camera, "async_read", side_effect=fail_warmup),
-        pytest.raises(TimeoutError, match="no frame"),
-    ):
-        camera.connect()
-
-    assert camera.rs_pipeline is None
-    assert camera.rs_profile is None
-    assert camera.thread is None
-    assert not camera.is_connected
-    assert read_threads[0] is not None
-    assert not read_threads[0].is_alive()
-
-    camera.connect(warmup=False)
-    assert camera.is_connected
-    camera.disconnect()
-
-
 def test_read():
     config = RealSenseCameraConfig(serial_number_or_name="042", width=640, height=480, fps=30, warmup_s=0)
     with RealSenseCamera(config) as camera:
@@ -289,203 +228,6 @@ def test_read_latest_too_old():
             _ = camera.read_latest(max_age_ms=0)  # immediately too old
 
 
-def _make_mock_sensor(name: str, supported_options: set | None = None) -> MagicMock:
-    """Build a fake rs.sensor that reports a name and a configurable supported-options set."""
-    supported = supported_options if supported_options is not None else set()
-    sensor = MagicMock()
-    sensor.get_info.return_value = name
-    sensor.supports.side_effect = lambda opt: opt in supported
-    return sensor
-
-
-def _attach_mock_color_sensor(camera: RealSenseCamera, sensor: MagicMock) -> None:
-    """Wire camera.rs_profile so _get_color_sensor finds the given sensor."""
-    profile = MagicMock()
-    device = MagicMock()
-    device.query_sensors.return_value = [sensor]
-    profile.get_device.return_value = device
-    camera.rs_profile = profile
-
-
-def test_get_color_sensor_prefers_rgb_camera():
-    config = RealSenseCameraConfig(serial_number_or_name="042")
-    camera = RealSenseCamera(config)
-
-    rgb = _make_mock_sensor("RGB Camera")
-    stereo = _make_mock_sensor("Stereo Module")
-    profile = MagicMock()
-    device = MagicMock()
-    device.query_sensors.return_value = [stereo, rgb]
-    profile.get_device.return_value = device
-    camera.rs_profile = profile
-
-    assert camera._get_color_sensor() is rgb
-
-
-def test_get_color_sensor_falls_back_to_stereo_module():
-    """D405 has no separate RGB module; color comes from Stereo Module."""
-    config = RealSenseCameraConfig(serial_number_or_name="042")
-    camera = RealSenseCamera(config)
-
-    stereo = _make_mock_sensor("Stereo Module")
-    _attach_mock_color_sensor(camera, stereo)
-
-    assert camera._get_color_sensor() is stereo
-
-
-def test_get_color_sensor_raises_with_available_sensors():
-    config = RealSenseCameraConfig(serial_number_or_name="042")
-    camera = RealSenseCamera(config)
-
-    other = _make_mock_sensor("Motion Module")
-    _attach_mock_color_sensor(camera, other)
-
-    with pytest.raises(RuntimeError, match="Motion Module"):
-        camera._get_color_sensor()
-
-
-def test_configure_sensor_options_skipped_when_none():
-    config = RealSenseCameraConfig(serial_number_or_name="042")
-    camera = RealSenseCamera(config)
-
-    with patch.object(RealSenseCamera, "_get_color_sensor") as mock_get:
-        camera._configure_sensor_options()
-        mock_get.assert_not_called()
-
-
-def test_configure_sensor_options_applies_all_values():
-    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120, gain=64, white_balance=4600)
-    camera = RealSenseCamera(config)
-
-    sensor = _make_mock_sensor(
-        "RGB Camera",
-        supported_options={
-            rs.option.enable_auto_exposure,
-            rs.option.exposure,
-            rs.option.gain,
-            rs.option.enable_auto_white_balance,
-            rs.option.white_balance,
-        },
-    )
-    _attach_mock_color_sensor(camera, sensor)
-
-    camera._configure_sensor_options()
-
-    sensor.set_option.assert_any_call(rs.option.enable_auto_exposure, 0)
-    sensor.set_option.assert_any_call(rs.option.exposure, 120)
-    sensor.set_option.assert_any_call(rs.option.gain, 64)
-    sensor.set_option.assert_any_call(rs.option.enable_auto_white_balance, 0)
-    sensor.set_option.assert_any_call(rs.option.white_balance, 4600)
-
-
-@pytest.mark.parametrize(
-    ("config_field", "option", "label"),
-    [
-        ("exposure", rs.option.exposure, "exposure"),
-        ("gain", rs.option.gain, "gain"),
-        ("white_balance", rs.option.white_balance, "white balance"),
-    ],
-)
-def test_configure_sensor_options_raises_when_requested_option_is_unsupported(config_field, option, label):
-    config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: 100})
-    camera = RealSenseCamera(config)
-
-    sensor = _make_mock_sensor("RGB Camera", supported_options=set())
-    _attach_mock_color_sensor(camera, sensor)
-
-    with pytest.raises(ValueError, match=label):
-        camera._configure_sensor_options()
-
-    sensor.supports.assert_any_call(option)
-    sensor.set_option.assert_not_called()
-
-
-@pytest.mark.parametrize(
-    ("config_field", "option", "value"),
-    [
-        ("exposure", rs.option.exposure, 120),
-        ("gain", rs.option.gain, 64),
-    ],
-)
-def test_configure_sensor_options_exposure_or_gain_disables_auto_exposure(config_field, option, value):
-    """white_balance=None should not touch auto white balance."""
-    config = RealSenseCameraConfig(serial_number_or_name="042", **{config_field: value})
-    camera = RealSenseCamera(config)
-
-    sensor = _make_mock_sensor(
-        "RGB Camera",
-        supported_options={rs.option.enable_auto_exposure, option},
-    )
-    _attach_mock_color_sensor(camera, sensor)
-
-    camera._configure_sensor_options()
-
-    calls = [call.args for call in sensor.set_option.call_args_list]
-    assert (rs.option.enable_auto_exposure, 0) in calls
-    assert (option, value) in calls
-    for opt, _ in calls:
-        assert opt != rs.option.enable_auto_white_balance
-        assert opt != rs.option.white_balance
-
-
-def test_configure_sensor_options_warns_when_auto_exposure_control_is_unsupported(caplog):
-    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=120)
-    camera = RealSenseCamera(config)
-
-    sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.exposure})
-    _attach_mock_color_sensor(camera, sensor)
-
-    with caplog.at_level("WARNING"):
-        camera._configure_sensor_options()
-
-    sensor.set_option.assert_called_once_with(rs.option.exposure, 120)
-    assert "does not support disabling auto-exposure" in caplog.text
-
-
-def test_configure_sensor_options_warns_when_auto_white_balance_control_is_unsupported(caplog):
-    config = RealSenseCameraConfig(serial_number_or_name="042", white_balance=4600)
-    camera = RealSenseCamera(config)
-
-    sensor = _make_mock_sensor("RGB Camera", supported_options={rs.option.white_balance})
-    _attach_mock_color_sensor(camera, sensor)
-
-    with caplog.at_level("WARNING"):
-        camera._configure_sensor_options()
-
-    sensor.set_option.assert_called_once_with(rs.option.white_balance, 4600)
-    assert "does not support disabling auto white balance" in caplog.text
-
-
-def test_configure_sensor_options_out_of_range_raises_value_error():
-    """set_option errors should be re-raised as ValueError with range diagnostics."""
-    config = RealSenseCameraConfig(serial_number_or_name="042", exposure=999999)
-    camera = RealSenseCamera(config)
-
-    sensor = _make_mock_sensor(
-        "RGB Camera",
-        supported_options={rs.option.enable_auto_exposure, rs.option.exposure},
-    )
-
-    def fake_set_option(option, value):
-        if option == rs.option.exposure:
-            raise RuntimeError("value out of range")
-
-    sensor.set_option.side_effect = fake_set_option
-
-    option_range = MagicMock(min=1, max=10000, step=1, default=156)
-    sensor.get_option_range.return_value = option_range
-
-    _attach_mock_color_sensor(camera, sensor)
-
-    with pytest.raises(ValueError, match="exposure") as exc_info:
-        camera._configure_sensor_options()
-
-    msg = str(exc_info.value)
-    assert "999999" in msg
-    assert "min=1" in msg
-    assert "max=10000" in msg
-
-
 @pytest.mark.parametrize(
     "rotation",
     [

tests/configs/test_recipe.py

@@ -29,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_lerobot_dataset.py

@@ -482,20 +482,6 @@ def test_add_frame_works_in_write_mode(tmp_path):
 # ── Resume mode ──────────────────────────────────────────────────────
 
 
-def test_resume_freshly_created_empty_dataset(tmp_path):
-    """resume() accepts a local dataset created before any episode was recorded."""
-    root = tmp_path / "resume_empty_ds"
-    LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=root)
-
-    resumed = LeRobotDataset.resume(repo_id=DUMMY_REPO_ID, root=root)
-
-    assert isinstance(resumed.writer, DatasetWriter)
-    assert resumed.meta.total_episodes == 0
-    assert resumed.meta.total_frames == 0
-    assert resumed.meta.tasks is None
-    assert resumed.meta.episodes is None
-
-
 def test_resume_creates_writer(tmp_path):
     """After resume(), writer is a DatasetWriter."""
     root = tmp_path / "resume_ds"

tests/datasets/test_sampler.py

@@ -25,7 +25,7 @@ from datasets import Dataset  # noqa: E402
 from lerobot.datasets.io_utils import (
     hf_transform_to_torch,
 )
-from lerobot.datasets.sampler import EpisodeAwareSampler
+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/envs/test_robocasa_ordering.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+
+from lerobot.envs.robocasa import RoboCasaEnv, convert_action
+
+
+def test_robocasa_action_uses_openpi_checkpoint_order():
+    action = np.arange(12, dtype=np.float32)
+
+    converted = convert_action(action)
+
+    np.testing.assert_array_equal(converted["action.end_effector_position"], [0, 1, 2])
+    np.testing.assert_array_equal(converted["action.end_effector_rotation"], [3, 4, 5])
+    np.testing.assert_array_equal(converted["action.gripper_close"], [6])
+    np.testing.assert_array_equal(converted["action.base_motion"], [7, 8, 9, 10])
+    np.testing.assert_array_equal(converted["action.control_mode"], [11])
+
+
+def test_robocasa_state_uses_openpi_checkpoint_order():
+    env = object.__new__(RoboCasaEnv)
+    env.obs_type = "pixels_agent_pos"
+    env.camera_name = []
+    raw_observation = {
+        "state.end_effector_position_relative": np.arange(0, 3),
+        "state.end_effector_rotation_relative": np.arange(3, 7),
+        "state.base_position": np.arange(7, 10),
+        "state.base_rotation": np.arange(10, 14),
+        "state.gripper_qpos": np.arange(14, 16),
+    }
+
+    observation = env._format_raw_obs(raw_observation)
+
+    np.testing.assert_array_equal(observation["agent_pos"], np.arange(16, dtype=np.float32))

tests/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,96 @@
+#!/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]
+
+
+@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_checkpoint_portability.py

@@ -0,0 +1,152 @@
+#!/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 json
+import shutil
+from dataclasses import asdict
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from lerobot.configs import FeatureType, NormalizationMode, PolicyFeature
+from lerobot.configs.recipe import MessageTurn, TrainingRecipe
+from lerobot.policies import make_pre_post_processors
+from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, NormalizerProcessorStep
+from lerobot.processor.converters import identity_transition
+from lerobot.processor.render_messages_processor import RenderMessagesStep
+from lerobot.utils.constants import ACTION
+
+
+class _ActionTokenizer:
+    def __call__(self, actions):
+        return np.asarray(actions).round().astype(np.int64)
+
+    def save_pretrained(self, path):
+        path.mkdir(parents=True)
+        (path / "processor_config.json").write_text('{"processor_class": "_ActionTokenizer"}\n')
+
+
+class _PaligemmaTokenizer:
+    vocab_size = 4096
+    bos_token_id = 2
+
+    def encode(self, text, **kwargs):
+        return [10, 11] if text == "Action: " else [12]
+
+
+def _make_pipeline(action_tokenizer_path):
+    recipe = TrainingRecipe(
+        messages=[
+            MessageTurn(role="user", content="${task}", stream="high_level"),
+            MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
+        ]
+    )
+    stats = {ACTION: {"min": torch.tensor([-1.0, -2.0]), "max": torch.tensor([1.0, 2.0])}}
+    normalizer = NormalizerProcessorStep(
+        features={ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(2,))},
+        norm_map={FeatureType.ACTION: NormalizationMode.MIN_MAX},
+        stats=stats,
+    )
+    action_tokenizer = ActionTokenizerProcessorStep(
+        action_tokenizer_name=str(action_tokenizer_path),
+        max_action_tokens=16,
+        fast_skip_tokens=128,
+    )
+    return DataProcessorPipeline(
+        [normalizer, RenderMessagesStep(recipe), action_tokenizer],
+        name="policy_preprocessor",
+        to_transition=identity_transition,
+        to_output=identity_transition,
+    )
+
+
+def test_pi052_pipeline_embeds_and_loads_fitted_action_tokenizer(tmp_path, monkeypatch):
+    original_cache = tmp_path / "original_fast_cache"
+    original_cache.mkdir()
+    tokenizer = _ActionTokenizer()
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained",
+        lambda path, **kwargs: tokenizer,
+    )
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained",
+        lambda *args, **kwargs: _PaligemmaTokenizer(),
+    )
+    monkeypatch.setattr(
+        "lerobot.policies.pi052.fit_fast_tokenizer.fit_fast_tokenizer",
+        lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("FAST fitting must not run")),
+    )
+
+    pipeline = _make_pipeline(original_cache)
+    expected_tokens = pipeline.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0]
+    expected_recipe = asdict(pipeline.steps[1].recipe)
+    expected_state = pipeline.steps[0].state_dict()
+    checkpoint = tmp_path / "checkpoint"
+    pipeline.save_pretrained(checkpoint)
+    DataProcessorPipeline(
+        [],
+        name="policy_postprocessor",
+        to_transition=identity_transition,
+        to_output=identity_transition,
+    ).save_pretrained(checkpoint)
+
+    saved_config = json.loads((checkpoint / "policy_preprocessor.json").read_text())
+    tokenizer_step = saved_config["steps"][2]
+    assert tokenizer_step["config"]["action_tokenizer_name"] == "action_tokenizer"
+    assert tokenizer_step["artifacts"] == {"action_tokenizer_name": "action_tokenizer"}
+    assert (checkpoint / "action_tokenizer" / "processor_config.json").is_file()
+
+    shutil.rmtree(original_cache)
+    loaded, _ = make_pre_post_processors(
+        SimpleNamespace(type="pi052", auto_fit_fast_tokenizer=True),
+        pretrained_path=str(checkpoint),
+        dataset_repo_id="org/dataset-that-must-not-be-read",
+    )
+
+    assert asdict(loaded.steps[1].recipe) == expected_recipe
+    for key, tensor in expected_state.items():
+        torch.testing.assert_close(loaded.steps[0].state_dict()[key], tensor)
+    torch.testing.assert_close(
+        loaded.steps[-1]._tokenize_action(torch.tensor([[[0.2, 0.8]]]))[0],
+        expected_tokens,
+    )
+
+
+def test_pi052_pipeline_rejects_missing_fitted_action_tokenizer(tmp_path, monkeypatch):
+    tokenizer = _ActionTokenizer()
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoProcessor.from_pretrained",
+        lambda path, **kwargs: tokenizer,
+    )
+    monkeypatch.setattr(
+        "lerobot.processor.tokenizer_processor.AutoTokenizer.from_pretrained",
+        lambda *args, **kwargs: _PaligemmaTokenizer(),
+    )
+
+    pipeline = _make_pipeline(tmp_path / "original_fast_cache")
+    checkpoint = tmp_path / "checkpoint"
+    pipeline.save_pretrained(checkpoint)
+    shutil.rmtree(checkpoint / "action_tokenizer")
+
+    with pytest.raises(FileNotFoundError, match="Checkpoint artifacts are incomplete"):
+        DataProcessorPipeline.from_pretrained(
+            checkpoint,
+            config_filename="policy_preprocessor.json",
+            to_transition=identity_transition,
+            to_output=identity_transition,
+        )

tests/policies/pi052/test_pi052_fast_action_loss.py

@@ -0,0 +1,162 @@
+#!/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."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+from torch.nn import functional as F  # noqa: N812
+
+pytest.importorskip("transformers")
+pytest.importorskip("liger_kernel")
+
+from lerobot.policies.pi052.modeling_pi052 import PI052Policy, _fast_lin_ce  # noqa: E402
+from lerobot.policies.pi052.processor_pi052 import make_pi052_pre_post_processors  # 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
+
+
+def test_fast_ce_averages_each_action_sample_equally():
+    torch.manual_seed(0)
+    hidden = torch.randn(2, 5, 8)
+    lm_head_weight = torch.eye(8)
+    action_tokens = torch.tensor([[1, 2, 0, 0, 0], [1, 3, 4, 5, 6]])
+    action_code_mask = torch.tensor([[False, True, False, False, False], [False, True, True, True, True]])
+
+    loss = _fast_lin_ce(
+        hidden,
+        lm_head_weight,
+        action_tokens,
+        action_code_mask,
+        predict_actions_t=None,
+        reduction="mean",
+    )
+    per_sample = _fast_lin_ce(
+        hidden,
+        lm_head_weight,
+        action_tokens,
+        action_code_mask,
+        predict_actions_t=None,
+        reduction="none",
+    )
+
+    assert torch.allclose(loss, per_sample.mean())
+
+
+def test_pi052_rejects_fast_loss_without_recipe():
+    config = SimpleNamespace(recipe_path=None, enable_fast_action_loss=True)
+
+    with pytest.raises(ValueError, match="recipe_path"):
+        make_pi052_pre_post_processors(config)
+
+
+def test_pi052_rejects_missing_fast_batch_keys():
+    policy = PI052Policy.__new__(PI052Policy)
+    nn.Module.__init__(policy)
+    policy.config = SimpleNamespace(
+        enable_fast_action_loss=True,
+        fast_action_loss_weight=1.0,
+        flow_loss_weight=0.0,
+        text_loss_weight=1.0,
+    )
+    batch = {
+        "text_labels": torch.tensor([[1, 2]]),
+        "predict_actions": torch.tensor([True]),
+    }
+
+    with pytest.raises(ValueError, match="FAST action loss is enabled"):
+        policy.forward(batch)

tests/policies/pi052/test_pi052_fit_fast_tokenizer.py

@@ -0,0 +1,122 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import pytest
+
+from lerobot.policies.pi052.fit_fast_tokenizer import (
+    _apply_relative_actions,
+    _dataset_signature,
+    _is_global_leader,
+    _normalize_actions,
+    _select_episode_indices,
+    _validate_fast_reconstruction,
+)
+
+
+def test_fast_tokenizer_fit_uses_training_mean_std_normalization():
+    actions = np.array([[[1.0, 7.0], [3.0, 3.0]]], dtype=np.float32)
+    stats = {"mean": [2.0, 5.0], "std": [0.5, 2.0]}
+
+    normalized = _normalize_actions(actions, "MEAN_STD", stats)
+
+    np.testing.assert_allclose(normalized, [[[-2.0, 1.0], [2.0, -1.0]]])
+
+
+def test_fast_tokenizer_fit_quantiles_match_training_without_clipping():
+    actions = np.array([[[-1.0], [3.0]]], dtype=np.float32)
+    stats = {"q01": [0.0], "q99": [2.0]}
+
+    normalized = _normalize_actions(actions, "QUANTILES", stats)
+
+    np.testing.assert_allclose(normalized, [[[-2.0], [2.0]]])
+
+
+def test_fast_tokenizer_cache_signature_tracks_stats_and_episode_selection():
+    kwargs = {
+        "dataset_repo_id": "org/dataset",
+        "base_tokenizer_name": "physical-intelligence/fast",
+        "n_samples": 100,
+        "chunk_size": 20,
+        "normalization_mode": "QUANTILES",
+        "dataset_revision": "main",
+        "episodes": [1, 2, 3],
+        "exclude_episodes": [2],
+        "use_relative_actions": False,
+        "relative_action_mask": None,
+    }
+
+    first = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [1.0]})
+    changed_stats = _dataset_signature(**kwargs, action_stats={"q01": [0.0], "q99": [2.0]})
+    changed_selection = _dataset_signature(
+        **{**kwargs, "exclude_episodes": [2, 3]},
+        action_stats={"q01": [0.0], "q99": [1.0]},
+    )
+
+    assert first != changed_stats
+    assert first != changed_selection
+
+
+def test_fast_tokenizer_uses_only_global_rank_zero(monkeypatch):
+    monkeypatch.setenv("RANK", "8")
+    monkeypatch.setenv("LOCAL_RANK", "0")
+    assert not _is_global_leader()
+
+    monkeypatch.setenv("RANK", "0")
+    assert _is_global_leader()
+
+
+def test_fast_tokenizer_episode_selection_applies_allowlist_and_exclusions():
+    selected = _select_episode_indices([0, 1, 2, 3], episodes=[1, 2, 3], exclude_episodes=[2])
+
+    assert selected == [1, 3]
+
+
+def test_fast_tokenizer_relative_actions_match_training_transform():
+    actions = np.array([[[2.0, 10.0], [3.0, 11.0]]], dtype=np.float32)
+    states = np.array([[1.0, 4.0]], dtype=np.float32)
+
+    relative = _apply_relative_actions(actions, states, [True, False])
+
+    np.testing.assert_allclose(relative, [[[1.0, 10.0], [2.0, 11.0]]])
+
+
+class _RoundTripTokenizer:
+    def __init__(self, offset: float = 0.0):
+        self.offset = offset
+
+    def __call__(self, actions):
+        return actions
+
+    def decode(self, tokens):
+        return tokens + self.offset
+
+
+def test_fast_tokenizer_reconstruction_validation_reports_error():
+    actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
+
+    report, decoded = _validate_fast_reconstruction(_RoundTripTokenizer(0.05), actions, 0.1, 0.1)
+
+    np.testing.assert_allclose(decoded, actions + 0.05)
+    assert report["reconstruction_rmse"] == pytest.approx(0.05)
+    assert report["max_dim_rmse"] == pytest.approx(0.05)
+
+
+def test_fast_tokenizer_reconstruction_validation_rejects_large_error():
+    actions = np.arange(24, dtype=np.float32).reshape(2, 3, 4) / 24
+
+    with pytest.raises(RuntimeError, match="exceeds the configured limit"):
+        _validate_fast_reconstruction(_RoundTripTokenizer(0.25), actions, 0.1, 0.2)

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_joint_sequence.py

@@ -0,0 +1,148 @@
+#!/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 joint-sequence (paper-style) subtask conditioning.
+
+Joint recipes train the subtask text and the action losses in one sequence,
+with the supervised subtask span attended causally. At inference the same
+layout is rebuilt around the *generated* subtask, so these tests pin:
+
+- the inference-side encoder produces the same token ids and target positions
+  as the training-time tokenizer step for the same messages;
+- OR-ing causal marks into a prefix reproduces the training-time attention
+  pattern (prompt cannot see the subtask; subtask is causal over itself);
+- the joint recipe file stays a valid message recipe;
+- the FAST id mapping with the default ``fast_skip_tokens`` stays clear of
+  PaliGemma's ``<loc>`` range so VQA and FAST supervision never collide.
+"""
+
+from pathlib import Path
+
+import torch
+
+from lerobot.configs.recipe import TrainingRecipe
+from lerobot.policies.pi052.text_processor_pi052 import (
+    PI052TextTokenizerStep,
+    encode_prompt_with_targets,
+)
+
+
+class _CharTokenizer:
+    """Char-level stub: 1 char = 1 token, so offsets are trivially aligned."""
+
+    pad_token_id = 0
+    eos_token = "\x1f"  # unit separator — a 1-char "EOS" for testing
+
+    def __call__(self, text, max_length=None, padding=None, return_tensors=None, **kwargs):
+        limit = max_length if max_length is not None else len(text)
+        ids = [ord(c) % 251 + 1 for c in text[:limit]]
+        offsets = [(i, i + 1) for i in range(len(ids))]
+        attention = [1] * len(ids)
+        if padding == "max_length" and max_length is not None 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),
+        }
+
+
+_MESSAGES = [
+    {"role": "user", "content": "fold the towel"},
+    {"role": "assistant", "content": "grab the near corner"},
+]
+
+
+def test_encode_prompt_with_targets_matches_training_labels():
+    tokenizer = _CharTokenizer()
+
+    step = PI052TextTokenizerStep(max_length=120)
+    step._tokenizer = tokenizer
+    train_ids, train_attn, labels, predict_actions, _prompt = step._encode_messages(
+        tokenizer,
+        [dict(m) for m in _MESSAGES],
+        message_streams=["low_level", "low_level"],
+        target_indices=[1],
+        complementary={},
+    )
+    assert bool(predict_actions)
+
+    ids, attn, marks = encode_prompt_with_targets(tokenizer, [dict(m) for m in _MESSAGES], [1])
+
+    n = int(attn.sum())
+    assert n == int(train_attn.sum())
+    assert torch.equal(ids[0, :n], train_ids[:n])
+    # Causal marks at inference must cover exactly the supervised label span.
+    assert torch.equal(marks[0, :n], labels[:n] != -100)
+    assert marks.any(), "the assistant target span must be marked"
+    # The user turn must stay unmarked (bidirectional prompt).
+    user_len = len("User: fold the towel\n")
+    assert not marks[0, :user_len].any()
+
+
+def test_apply_causal_language_marks_reproduces_training_mask():
+    from lerobot.policies.pi05.modeling_pi05 import make_att_2d_masks
+    from lerobot.policies.pi052.modeling_pi052 import _apply_causal_language_marks
+
+    n_img, n_lang = 4, 8
+    prefix_len = n_img + n_lang
+    pad = torch.ones((1, prefix_len), dtype=torch.bool)
+    att = torch.zeros((1, prefix_len), dtype=torch.bool)
+    # Subtask span = language positions 5..7 (prefix positions 9..11).
+    marks = torch.zeros((1, n_lang), dtype=torch.bool)
+    marks[0, 5:8] = True
+
+    att_marked = _apply_causal_language_marks(att, marks)
+    att_2d = make_att_2d_masks(pad, att_marked)[0]
+
+    subtask = [n_img + 5, n_img + 6, n_img + 7]
+    # Prompt and images never see the subtask.
+    for q in range(n_img + 5):
+        for k in subtask:
+            assert not att_2d[q, k], f"prompt position {q} must not attend subtask position {k}"
+    # Subtask tokens see the full prompt and earlier subtask tokens only.
+    for qi, q in enumerate(subtask):
+        for k in range(n_img + 5):
+            assert att_2d[q, k]
+        for ki, k in enumerate(subtask):
+            assert bool(att_2d[q, k]) == (ki <= qi)
+
+
+def test_joint_recipe_is_a_valid_message_recipe():
+    recipe_path = Path(__file__).parents[3] / "src" / "lerobot" / "configs" / "recipes" / "subtask_joint.yaml"
+    recipe = TrainingRecipe.from_yaml(recipe_path)
+    assert recipe.messages is not None and len(recipe.messages) == 2
+    assert all(turn.stream == "low_level" for turn in recipe.messages)
+    assert not recipe.messages[0].target
+    assert recipe.messages[1].target
+    assert recipe.messages[1].if_present == "subtask"
+
+
+def test_default_fast_mapping_clears_loc_and_seg_ranges():
+    from lerobot.policies.pi052.configuration_pi052 import PI052Config
+    from lerobot.policies.pi052.modeling_pi052 import _FAST_ACTION_VOCAB_SIZE
+
+    skip = PI052Config.__dataclass_fields__["fast_skip_tokens"].default
+    assert skip == 1152
+
+    paligemma_vocab = 257152
+    fast_ids = paligemma_vocab - 1 - skip - torch.arange(_FAST_ACTION_VOCAB_SIZE)
+    # Below the <loc> range [256000, 257024) and the <seg> range [257024, 257152).
+    assert int(fast_ids.max()) < 256000
+    assert int(fast_ids.min()) >= 0

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