feat(flow-matching): cover per-policy sampling divergences

Extend the shared samplers to represent the non-openpi conventions Martino
catalogued, without changing existing openpi behavior:

- sample_noise: add distribution="uniform" for evo1's rand*2-1 noise.
- sample_time_beta: add complement flag (groot/wall_x forward t=(1-beta)*s)
  and optional clamp_min/clamp_max (evo1 Beta(2,2) clamped to [0.02, 0.98]);
  scale/offset now default to the identity mapping.
- sample_beta: build the Beta distribution via a cached, CPU-side helper
  (groot convention) so it is constructed once per (alpha, beta).

Adds tests for uniform noise, complement/clamp timesteps, and caching.

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), [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) |
+| **VLAs Models**            | [Pi0](./docs/source/pi0.mdx), [Pi0Fast](./docs/source/pi0fast.mdx), [Pi0.5](./docs/source/pi05.mdx), [GR00T N1.7](./docs/source/policy_groot_README.md), [SmolVLA](./docs/source/policy_smolvla_README.md), [XVLA](./docs/source/xvla.mdx), [EO-1](./docs/source/eo1.mdx), [MolmoAct2](./docs/source/molmoact2.mdx), [WALL-OSS](./docs/source/walloss.mdx), [EVO1](./docs/source/evo1.mdx) |
 | **World Models**           | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx)                                                                                                                                                                                                                                                                   |
 | **Reward Models**          | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx)                                                                                                                                                                                                                                                                         |
 

docs/source/_toctree.yml

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

docs/source/bring_your_own_policies.mdx

@@ -151,12 +151,12 @@ class MyPolicy(PreTrainedPolicy):
 The methods called by the train/eval loops:
 
 | Method                                                            | Used by           | What it does                                                                                                                                                                                                                                                                                 |
-| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | `reset() -> None`                                                 | `lerobot-eval`    | Clear per-episode state at the start of each episode.                                                                                                                                                                                                                                        |
 | `select_action(batch, **kwargs) -> Tensor`                        | `lerobot-eval`    | Return the next action `(B, action_dim)`. Called every step.                                                                                                                                                                                                                                 |
 | `predict_action_chunk(batch, **kwargs) -> Tensor`                 | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk.                                                                                                                                       |
 | `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train`   | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting.                                                                                                                                                                                         |
-| `get_optim_params() -> dict`                                      | the optimizer     | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). |
+| `get_optim_params() -> dict`                                      | the optimizer     | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
 | `update() -> None` _(optional)_                                   | `lerobot-train`   | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this).                                                                                                                                                                                   |
 
 Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
@@ -189,162 +189,6 @@ def make_my_policy_pre_post_processors(
 
 ---
 
-## Adding high- and low-level language control
-
-The policy API above is sufficient for training and standard evaluation. To use a language-conditioned policy with interactive `lerobot-rollout`, also register a runtime adapter. The adapter keeps policy-specific prompting and tokenization out of the generic control loop.
-
-The runtime supports two policy shapes:
-
-| Policy shape     | Behavior                                                                | Adapter                                        |
-| ---------------- | ----------------------------------------------------------------------- | ---------------------------------------------- |
-| Low-level / flat | The operator's task or subtask directly conditions action prediction.   | Reuse `DirectTaskPolicyAdapter`.               |
-| High + low level | The policy generates subtasks or memory, then conditions actions on it. | Subclass `BaseLanguageAdapter`, as PI052 does. |
-
-During a rollout, `RuntimeState` stores the high-level task and the active language context:
-
-```text
-task ──> adapter.generate_text("subtask", ...) ──> state.language_context["subtask"]
-                                                       │
-observation ──> processors ──> adapter.select_action() ─┴─> action chunk ──> robot
-```
-
-The generic runtime handles generation frequency, pause/resume, prompt replacement, action queues, and dispatch. The adapter only translates between that runtime contract and your policy.
-
-### Low-level policies
-
-If your policy already consumes the live task through its normal preprocessor and implements `predict_action_chunk`, register the shared direct adapter. PI0.5 and MolmoAct2 use this path:
-
-```python
-# src/lerobot/runtime/registry.py
-_ADAPTERS = {
-    # ...
-    "my_policy": "lerobot.runtime.adapter:DirectTaskPolicyAdapter",
-}
-```
-
-Run it with direct-subtask mode so the operator supplies the instruction used by the action policy:
-
-```bash
-lerobot-rollout \
-    --language \
-    --policy.path=user/my_policy_checkpoint \
-    --robot.type=so101_follower \
-    --robot.port=/dev/ttyACM0 \
-    --direct_subtask
-```
-
-The rollout context builds the observation batch with the current instruction before `DirectTaskPolicyAdapter` calls `policy.predict_action_chunk(observation)`. No text-generation method is required.
-
-### Hierarchical policies
-
-For a policy that generates language and actions, subclass [`BaseLanguageAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/runtime/adapter.py) and implement two methods:
-
-- `generate_text(kind, observation, state, user_text=None) -> str` generates a `subtask`, `memory`, or interjection response.
-- `select_action(observation, state)` builds the low-level prompt from the active context and returns an action chunk.
-
-This abbreviated adapter follows [`PI052PolicyAdapter`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py):
-
-```python
-# inference/my_policy_adapter.py
-from typing import Any
-
-from lerobot.runtime import RuntimeState
-from lerobot.runtime.adapter import BaseLanguageAdapter
-from lerobot.utils.constants import (
-    OBS_LANGUAGE_ATTENTION_MASK,
-    OBS_LANGUAGE_TOKENS,
-)
-
-
-class MyPolicyAdapter(BaseLanguageAdapter):
-    def select_action(self, observation: dict[str, Any], state: RuntimeState):
-        instruction = state.language_context.get("subtask") or state.task or ""
-        tokens, attention_mask = tokenize_instruction(instruction)
-
-        batch = dict(observation)
-        batch[OBS_LANGUAGE_TOKENS] = tokens
-        batch[OBS_LANGUAGE_ATTENTION_MASK] = attention_mask
-        return self.policy.predict_action_chunk(batch)
-
-    def generate_text(
-        self,
-        kind: str,
-        observation: dict[str, Any] | None,
-        state: RuntimeState,
-        user_text: str | None = None,
-    ) -> str:
-        messages = self.build_messages(kind, state, user_text)
-        batch, tokenizer = tokenize_messages(messages, observation)
-        return self.policy.select_message(
-            batch,
-            tokenizer=tokenizer,
-            min_new_tokens=self.gen.min_new_tokens,
-            temperature=self.gen.temperature,
-            top_p=self.gen.top_p,
-        )
-
-    def build_messages(
-        self, kind: str, state: RuntimeState, user_text: str | None
-    ) -> list[dict[str, str]]:
-        if kind == "subtask":
-            return [{"role": "user", "content": state.task or ""}]
-        if kind == "memory":
-            return [
-                {"role": "user", "content": state.task or ""},
-                {
-                    "role": "user",
-                    "content": f"Completed subtask: {state.extra.get('prior_subtask', '')}",
-                },
-            ]
-        if kind == "interjection":
-            return [
-                {"role": "user", "content": state.task or ""},
-                {"role": "user", "content": user_text or ""},
-            ]
-        raise ValueError(f"Unsupported text kind: {kind}")
-```
-
-`tokenize_instruction` and `tokenize_messages` are policy-specific helpers. They must reproduce the prompt format used during training; PI052, for example, adds the discretized robot state to its low-level subtask prompt and uses the same PaliGemma formatting for `select_message`.
-
-`BaseLanguageAdapter` provides the default hierarchy: regenerate a subtask at action-chunk boundaries, update memory when the subtask changes, and handle user interjections. Override `_regenerate_context` only if your policy uses a different hierarchy.
-
-Register the adapter with a lazy import so importing LeRobot does not load the model or its optional dependencies:
-
-```python
-# src/lerobot/runtime/registry.py
-_ADAPTERS = {
-    # ...
-    "my_policy": "lerobot.policies.my_policy.inference.my_policy_adapter:MyPolicyAdapter",
-}
-```
-
-The key must match the policy's registered type. Once registered, the same checkpoint works through the shared entry point:
-
-```bash
-lerobot-rollout \
-    --language \
-    --policy.path=user/my_hierarchical_checkpoint \
-    --robot.type=so101_follower \
-    --robot.port=/dev/ttyACM0 \
-    --task="put the cup in the sink"
-```
-
-For RoboCasa-compatible policies, replace the robot arguments with `--sim --sim.task=<task>`. Without `--direct_subtask`, the adapter generates the low-level subtask; with it, the operator bypasses high-level generation and supplies each subtask.
-
-### Keep training and deployment aligned
-
-The adapter is intentionally small, but its prompts are part of the model contract:
-
-- Use the same tokenizer, role formatting, special tokens, image ordering, and state encoding as training.
-- Condition `select_action` on `state.language_context["subtask"]`, falling back to `state.task` for direct or not-yet-generated prompts.
-- Return a full action chunk from `select_action`; the runtime handles control-rate dispatch.
-- Keep optional model dependencies inside lazy imports.
-- Test adapter selection, generated-message routing, action-batch construction, and direct-subtask behavior with a lightweight fake policy.
-
-PI052 is the complete in-tree reference: its [processor](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/processor_pi052.py) renders the training recipe, its [policy](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/modeling_pi052.py) exposes text and action generation, and its [adapter](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi052/inference/pi052_adapter.py) reconstructs those same prompts at deployment.
-
----
-
 ## Path A: Out-of-tree plugin
 
 The fastest way to ship a policy: package it as a standalone Python distribution and install it alongside LeRobot. No PR required, you own the release cycle, and you can publish to PyPI under your own namespace.
@@ -451,12 +295,10 @@ The file names are load-bearing: the factory does lazy imports by name, and the
 
 ### Wiring
 
-Four places need to know about your policy. All by name.
+Two places need to know about your policy. All by name.
 
-1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
-2. **`factory.py:get_policy_class`** — add a branch returning `MyPolicy` from a lazy import.
-3. **`factory.py:make_policy_config`** and **`factory.py:make_pre_post_processors`** — same idea, two more branches.
-4. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
+1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
+2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
 
 Mirror an existing policy that's structurally similar to yours; the diff is small.
 
@@ -488,6 +330,10 @@ This way:
 
 Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything.
 
+### Avoid copying a modeling file — subclass it
+
+If your policy needs to modify a backbone that already exists in `transformers` (custom conditioning, extra inputs, a swapped sub-module), **do not vendor a copy of its `modeling_*.py`**. Instead, subclass the smallest upstream unit and override only what changes. [`pi_gemma.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi_gemma.py) is the canonical reference: it injects AdaRMS conditioning into PaliGemma/Gemma in ~370 lines by subclassing `GemmaModel`/`PaliGemmaModel` and overriding the decoder-layer forward, instead of forking the ~2,000-line modeling file. Model surgery on a _loaded_ native model is also fine (layer truncation, tokenizer expansion, hidden-state capture — see `evo1/internvl3_embedder.py`, `eo1/modeling_eo1.py`, `groot/groot_n1_7.py` for working examples). Reviewers will ask for this pattern when a PR arrives with a copied modeling file; the only accepted exception is a model that does not exist in `transformers` at all.
+
 ### Benchmarks and a published checkpoint
 
 A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
@@ -523,7 +369,7 @@ If your policy is real-robot-only and no sim benchmark applies, swap the sim eva
 The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for:
 
 - [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class.
-- [ ] `factory.py` and `policies/__init__.py` are wired (lazy imports for modeling).
+- [ ] `policies/__init__.py` re-exports the config (this registers the policy; the factory resolves modeling/processor by naming convention).
 - [ ] `make_my_policy_pre_post_processors` follows the naming convention.
 - [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
 - [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.

docs/source/inference.mdx

@@ -1,6 +1,6 @@
 # Policy Deployment (lerobot-rollout)
 
-`lerobot-rollout` is the single CLI for deploying trained policies on real robots or in an interactive simulator. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording, language-driven control, and human-in-the-loop data collection.
+`lerobot-rollout` is the single CLI for deploying trained policies on real robots. It supports multiple execution strategies and inference backends, from quick evaluation to continuous recording and human-in-the-loop data collection.
 
 ## Quick Start
 
@@ -197,52 +197,6 @@ Teleop is optional — if omitted the robot holds its position during the reset
 
 ---
 
-## Interactive language control
-
-Language-conditioned policies can expose a high-level text head in addition to
-their action head. Add `--language` to open-prompt one of these policies on a
-real robot. Language-only flags such as `--direct_subtask` select this mode
-automatically.
-
-MolmoAct2 has no high-level planner, so use direct-subtask mode and type each
-next low-level instruction yourself:
-
-```bash
-lerobot-rollout \
-    --policy.path=lerobot/MolmoAct2-SO100_101-LeRobot \
-    --policy.device=cuda \
-    --robot.type=so101_follower \
-    --robot.port=/dev/ttyACM1 \
-    --robot.cameras='{"cam0":{"type":"opencv","index_or_path":"/dev/video0","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200},"cam1":{"type":"opencv","index_or_path":"/dev/video2","width":640,"height":480,"fps":30,"fourcc":"MJPG","backend":200}}' \
-    --direct_subtask \
-    --robot.max_relative_target='{"shoulder_pan":5,"shoulder_lift":5,"elbow_flex":5,"wrist_flex":5,"wrist_roll":5,"gripper":5}'
-```
-
-The robot starts paused. Type a subtask, then use `/resume` and `/pause` to
-control action dispatch. Check the workspace and motion limits before resuming.
-Without `--direct_subtask`, a policy such as PI052 generates its active subtask
-from the high-level `--task` itself.
-
-RoboCasa uses the same runtime and processor path. `--sim` selects it
-automatically, so no robot configuration is needed:
-
-```bash
-MUJOCO_GL=egl lerobot-rollout \
-    --policy.path=lerobot/pi052_robocasa \
-    --sim --sim.task=CloseFridge --sim.split=pretrain \
-    --task="close the fridge" \
-    --disable_memory \
-    --sim.render_size=384 \
-    --sim.views=robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right \
-    --mode=action --ctrl_hz=20
-```
-
-Open `http://localhost:8010` for the live simulator view. Add
-`--sim.direct_subtask` to bypass the language planner and make each typed prompt
-the action policy's current subtask.
-
----
-
 ## Inference Backends
 
 Select a backend with `--inference.type=<name>`. All strategies work with both backends.

docs/source/language_and_recipes.mdx

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

docs/source/pi052.mdx

@@ -1,274 +0,0 @@
-# π₀.₅ with language supervision (Pi052)
-
-Pi052 extends [Pi05](./pi05) with a trainable PaliGemma language head and a
-runtime that alternates language generation with action generation. A single
-checkpoint can predict a low-level subtask, optionally update memory or answer
-visual questions, and condition its flow-matching action expert on that text.
-
-Use Pi05 when you only need task-conditioned actions. Use Pi052 when the policy
-must generate or consume intermediate language during a rollout.
-
-## How Pi052 differs from Pi05
-
-| Capability          | Pi05                                                   | Pi052                                                                             |
-| ------------------- | ------------------------------------------------------ | --------------------------------------------------------------------------------- |
-| Action model        | PaliGemma vision-language prefix + Gemma action expert | Same base architecture                                                            |
-| Language head       | Not trained for runtime generation                     | Re-enabled and trained with text cross-entropy                                    |
-| Action conditioning | Episode task                                           | Active low-level subtask plus normalized robot state                              |
-| Training targets    | Flow-matching actions                                  | Flow actions, recipe-selected text, and optional FAST action tokens               |
-| Dataset requirement | Standard images, state, actions, and task              | The same fields plus language annotations for every language capability you train |
-| Rollout             | Direct task-to-action policy                           | Hierarchical task → subtask → action loop, with optional memory and VQA           |
-
-Pi052 can initialize from a Pi05 checkpoint. The policy architecture remains
-compatible, while Pi052 builds its own processors so recipe labels and FAST
-labels are not silently replaced by the Pi05 processor stack.
-
-## Install
-
-Install LeRobot with the PI dependencies:
-
-```bash
-git clone https://github.com/huggingface/lerobot.git
-cd lerobot
-python -m venv .venv
-source .venv/bin/activate
-pip install -e ".[pi]"
-```
-
-The `pi` extra includes the PaliGemma/FAST dependencies. Install
-`liger-kernel` for the supported fused training kernels; optional FlashRT
-backends also require the Hugging Face `kernels` package and a supported CUDA
-GPU.
-
-## Prepare language-annotated data
-
-Pi052 does not infer supervised subtasks from a normal LeRobot dataset during
-training. The dataset must contain the language targets used by the selected
-recipe in the optional `language_persistent` and `language_events` columns.
-
-At minimum, annotate a continuous `subtask` timeline so each training frame has
-an active low-level instruction. Add `memory`, VQA, interjections, and speech
-annotations only if the recipe trains those capabilities.
-
-The provided recipes are:
-
-| Recipe                                | Required annotations                                                    | Trains                                             |
-| ------------------------------------- | ----------------------------------------------------------------------- | -------------------------------------------------- |
-| `recipes/subtask.yaml`                | `subtask`                                                               | Subtask prediction and subtask-conditioned actions |
-| `recipes/subtask_joint.yaml`          | `subtask`                                                               | Paper-style joint sequence: subtask text and actions in one sample |
-| `recipes/subtask_mem.yaml`            | `subtask`, `memory`                                                     | Subtasks, actions, and memory updates              |
-| `recipes/subtask_mem_vqa_speech.yaml` | `subtask`, `memory`, `vqa`; interjection/speech rows for those branches | Subtasks, actions, memory, VQA, and spoken replies |
-
-The blend recipes factorize training into separate high-level (task → subtask)
-and low-level (subtask → actions) samples, matching how inference decomposes
-π(a|o, subtask)·π(subtask|o, task). `recipes/subtask_joint.yaml` instead uses
-the π0.5 paper's single-sequence layout — the supervised subtask span is
-attended causally and conditions the FAST and flow losses in the same forward.
-Checkpoints trained with the joint recipe must set
-`--policy.joint_subtask_conditioning=true` at inference so the flow prefix
-rebuilds the same layout (task turn with state, then the generated subtask as a
-causal assistant turn); leave it `false` for the blend recipes.
-
-Use `lerobot-annotate` to generate these columns. The repository includes a
-Hugging Face Jobs launcher that you can edit for your source and destination
-datasets. For a local annotation run, first install
-`pip install -e ".[annotations]"`:
-
-```bash
-HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
-```
-
-Before a long training run, inspect several episodes and verify that subtasks
-are temporally correct and cover the full demonstration. See
-[Annotation Pipeline](./annotation_pipeline) for generation and validation, and
-[Language Columns and Recipes](./language_and_recipes) for the schema and
-recipe resolver.
-
-<Tip>
-  If a dataset has no language columns, recipe rendering becomes a no-op and
-  Pi052 falls back to the plain Pi05 prompt path. This is useful for
-  compatibility but does not train the language planner.
-</Tip>
-
-## Train Pi052
-
-This example initializes Pi052 from the public Pi05 base checkpoint and trains
-the default subtask-and-memory recipe:
-
-```bash
-lerobot-train \
-    --dataset.repo_id=${HF_USER}/my_language_annotated_dataset \
-    --policy.type=pi052 \
-    --policy.pretrained_path=lerobot/pi05_base \
-    --policy.recipe_path=recipes/subtask_mem.yaml \
-    --policy.dtype=bfloat16 \
-    --policy.device=cuda \
-    --policy.freeze_vision_encoder=false \
-    --policy.gradient_checkpointing=true \
-    --batch_size=8 \
-    --steps=30000 \
-    --output_dir=outputs/pi052 \
-    --job_name=pi052 \
-    --wandb.enable=true
-```
-
-For subtask-only data, change the recipe to `recipes/subtask.yaml` and disable
-memory during rollout. Start with a small run and confirm that W&B examples show
-the expected prompt, text target, and action endpoints before scaling up.
-
-### Main training controls
-
-| Option                           |                    Default | Purpose                                                        |
-| -------------------------------- | -------------------------: | -------------------------------------------------------------- |
-| `policy.recipe_path`             | `recipes/subtask_mem.yaml` | Selects the language/action objective mixture                  |
-| `policy.text_loss_weight`        |                      `1.0` | Language-head cross-entropy weight; `0` disables text training |
-| `policy.flow_loss_weight`        |                     `10.0` | Continuous action flow-loss weight                             |
-| `policy.enable_fast_action_loss` |                     `true` | Adds discrete FAST action-token supervision                    |
-| `policy.fast_action_loss_weight` |                      `1.0` | FAST cross-entropy weight                                      |
-| `policy.knowledge_insulation`    |                     `true` | Blocks action-loss gradients through the VLM K/V path          |
-| `policy.flow_num_repeats`        |                        `5` | Reuses one VLM prefix for independent denoising targets        |
-| `policy.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,21 +110,15 @@ lerobot-train \
 ### Key Training Parameters
 
 | Parameter                              | Description                                        | Default                         |
-| --------------------------------------- | -------------------------------------------------- | ------------------------------- |
+| -------------------------------------- | -------------------------------------------------- | ------------------------------- |
 | `--policy.gradient_checkpointing=true` | Reduces memory usage significantly during training | `false`                         |
 | `--policy.dtype=bfloat16`              | Use mixed precision training for efficiency        | `float32`                       |
 | `--policy.chunk_size`                  | Number of action steps to predict (action horizon) | `50`                            |
 | `--policy.n_action_steps`              | Number of 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

examples/annotations/run_hf_job.py

@@ -46,8 +46,11 @@ CMD = (
     "apt-get update -qq && apt-get install -y -qq git ffmpeg && "
     "pip install --no-deps "
     "'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
+    # Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 /
+    # draccus 0.11, which break lerobot at import time.
     "pip install --upgrade-strategy only-if-needed "
-    "datasets pyarrow av jsonlines draccus gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
+    "'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
+    "'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
     "openai && "
     "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
     "export VLLM_VIDEO_BACKEND=pyav && "

pyproject.toml

@@ -150,7 +150,6 @@ pygame-dep = ["pygame>=2.5.1,<2.7.0"]
 # There is no cmeel-urdfdom 5.x; <5 selects the 4.x ABI the placo/pin wheels are built against.
 placo-dep = ["placo>=0.9.6,<0.9.16", "cmeel-urdfdom>=4,<5", "cmeel-tinyxml2<11"]
 transformers-dep = ["transformers>=5.4.0,<5.6.0"]
-sentencepiece-dep = ["sentencepiece>=0.2.0,<0.3.0"] # FAST action tokenizer backend (pi052, pi0_fast)
 grpcio-dep = ["grpcio>=1.73.1,<2.0.0", "protobuf>=6.31.1,<8.0.0"]
 accelerate-dep = ["accelerate>=1.14.0,<2.0.0"]
 can-dep = ["python-can>=4.2.0,<5.0.0"]
@@ -213,7 +212,7 @@ wallx = [
     "torchdiffeq>=0.2.4,<0.3.0",
     "lerobot[qwen-vl-utils-dep]",
 ]
-pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]", "lerobot[sentencepiece-dep]"]
+pi = ["lerobot[transformers-dep]", "lerobot[scipy-dep]"]
 molmoact2 = ["lerobot[transformers-dep]", "lerobot[peft-dep]", "lerobot[scipy-dep]"]
 smolvla = ["lerobot[transformers-dep]", "num2words>=0.5.14,<0.6.0", "lerobot[accelerate-dep]"]
 multi_task_dit = ["lerobot[transformers-dep]", "lerobot[diffusers-dep]"]
@@ -414,8 +413,6 @@ ignore = [
 "__init__.py" = ["F401", "F403", "E402"]
 # E402: conditional-import guards (TYPE_CHECKING / is_package_available) must precede the imports they protect
 "src/lerobot/scripts/convert_dataset_v21_to_v30.py" = ["E402"]
-"src/lerobot/policies/wall_x/**" = ["N801", "N812", "SIM102", "SIM108", "SIM210", "SIM211", "B006", "B007", "SIM118"] # Supprese these as they are coming from original Qwen2_5_vl code TODO(pepijn): refactor original
-
 [tool.ruff.lint.isort]
 combine-as-imports = true
 known-first-party = ["lerobot"]

scripts/backfill_pi052_action_tokenizer.py

@@ -1,227 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Audit or backfill checkpoint-local FAST artifacts for PI052 model repositories."""
-
-from __future__ import annotations
-
-import argparse
-import hashlib
-import io
-import json
-from pathlib import Path, PurePosixPath
-from typing import Any
-
-from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
-
-DEFAULT_REPOSITORIES = (
-    "pepijn223/pi052_atomic4_01_baseline",
-    "pepijn223/pi052_atomic4_02_lr_1e5",
-    "pepijn223/pi052_atomic4_03_recipe_50_50",
-    "pepijn223/pi052_atomic4_04_flow_weight_10",
-    "pepijn223/pi052_atomic4_05_flow_repeat_1",
-    "pepijn223/pi052_atomic4_06_ki_off",
-)
-CHECKPOINT_DIRECTORIES = (
-    "",
-    "checkpoints/003000/pretrained_model",
-    "checkpoints/006000/pretrained_model",
-    "checkpoints/009000/pretrained_model",
-    "checkpoints/012000/pretrained_model",
-)
-TOKENIZER_DIRECTORY = "action_tokenizer"
-
-
-def artifact_fingerprint(files: list[tuple[str, bytes]]) -> str:
-    digest = hashlib.sha256()
-    for relative_path, content in sorted(files):
-        encoded_path = relative_path.encode()
-        digest.update(len(encoded_path).to_bytes(8, "big"))
-        digest.update(encoded_path)
-        digest.update(len(content).to_bytes(8, "big"))
-        digest.update(content)
-    return digest.hexdigest()
-
-
-def tokenizer_files(tokenizer_path: Path) -> list[tuple[str, Path]]:
-    return [
-        (path.relative_to(tokenizer_path).as_posix(), path)
-        for path in sorted(tokenizer_path.rglob("*"))
-        if path.is_file()
-    ]
-
-
-def _repo_path(directory: str, filename: str) -> str:
-    return (PurePosixPath(directory) / filename).as_posix() if directory else filename
-
-
-def _download_json(repo_id: str, path_in_repo: str, revision: str | None = None) -> dict[str, Any]:
-    path = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
-    return json.loads(Path(path).read_text())
-
-
-def make_portable_preprocessor(config: dict[str, Any]) -> dict[str, Any]:
-    config = json.loads(json.dumps(config))
-    action_steps = [
-        step for step in config["steps"] if step.get("registry_name") == "action_tokenizer_processor"
-    ]
-    if len(action_steps) != 1:
-        raise ValueError(f"Expected one action tokenizer step, found {len(action_steps)}")
-    action_step = action_steps[0]
-    action_step["config"]["action_tokenizer_name"] = TOKENIZER_DIRECTORY
-    action_step["artifacts"] = {"action_tokenizer_name": TOKENIZER_DIRECTORY}
-
-    recipe_steps = [
-        step for step in config["steps"] if step.get("registry_name") == "render_messages_processor"
-    ]
-    if len(recipe_steps) != 1 or not recipe_steps[0].get("config", {}).get("recipe"):
-        raise ValueError("PI052 preprocessor does not contain an embedded training recipe")
-    return config
-
-
-def _json_operation(path_in_repo: str, content: dict[str, Any]) -> CommitOperationAdd:
-    serialized = (json.dumps(content, indent=2) + "\n").encode()
-    return CommitOperationAdd(path_in_repo=path_in_repo, path_or_fileobj=io.BytesIO(serialized))
-
-
-def prepare_operations(
-    repo_id: str,
-    tokenizer_path: Path,
-    revision: str | None = None,
-) -> list[CommitOperationAdd]:
-    operations: list[CommitOperationAdd] = []
-    files = tokenizer_files(tokenizer_path)
-    for directory in CHECKPOINT_DIRECTORIES:
-        preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
-        operations.append(
-            _json_operation(
-                preprocessor_path,
-                make_portable_preprocessor(_download_json(repo_id, preprocessor_path, revision)),
-            )
-        )
-        for relative_path, local_path in files:
-            operations.append(
-                CommitOperationAdd(
-                    path_in_repo=_repo_path(
-                        directory,
-                        f"{TOKENIZER_DIRECTORY}/{relative_path}",
-                    ),
-                    path_or_fileobj=str(local_path),
-                )
-            )
-    return operations
-
-
-def audit_repository(
-    api: HfApi,
-    repo_id: str,
-    expected_tokenizer_fingerprint: str,
-    revision: str | None = None,
-) -> None:
-    info = api.model_info(repo_id, revision=revision)
-    repository_files = {sibling.rfilename for sibling in info.siblings or []}
-
-    for directory in CHECKPOINT_DIRECTORIES:
-        preprocessor_path = _repo_path(directory, "policy_preprocessor.json")
-        policy_config_path = _repo_path(directory, "config.json")
-        postprocessor_path = _repo_path(directory, "policy_postprocessor.json")
-        for required_path in (preprocessor_path, policy_config_path, postprocessor_path):
-            if required_path not in repository_files:
-                raise FileNotFoundError(f"{repo_id}@{revision or 'main'} is missing {required_path}")
-
-        preprocessor = _download_json(repo_id, preprocessor_path, revision)
-        portable_preprocessor = make_portable_preprocessor(preprocessor)
-        if preprocessor != portable_preprocessor:
-            raise ValueError(f"{repo_id}:{preprocessor_path} is not portable")
-
-        normalizer_steps = [
-            step for step in preprocessor["steps"] if step.get("registry_name") == "normalizer_processor"
-        ]
-        if len(normalizer_steps) != 1 or "state_file" not in normalizer_steps[0]:
-            raise ValueError(f"{repo_id}:{preprocessor_path} is missing normalizer state metadata")
-        normalizer_path = _repo_path(directory, normalizer_steps[0]["state_file"])
-        if normalizer_path not in repository_files:
-            raise FileNotFoundError(f"{repo_id} is missing {normalizer_path}")
-
-        remote_tokenizer_files: list[tuple[str, bytes]] = []
-        for relative_path in _tokenizer_relative_paths(repository_files, directory):
-            path_in_repo = _repo_path(directory, f"{TOKENIZER_DIRECTORY}/{relative_path}")
-            downloaded = hf_hub_download(repo_id, path_in_repo, repo_type="model", revision=revision)
-            remote_tokenizer_files.append((relative_path, Path(downloaded).read_bytes()))
-        fingerprint = artifact_fingerprint(remote_tokenizer_files)
-        if fingerprint != expected_tokenizer_fingerprint:
-            raise ValueError(
-                f"{repo_id}:{_repo_path(directory, TOKENIZER_DIRECTORY)} fingerprint "
-                f"{fingerprint} != {expected_tokenizer_fingerprint}"
-            )
-
-
-def _tokenizer_relative_paths(repository_files: set[str], directory: str) -> list[str]:
-    prefix = _repo_path(directory, TOKENIZER_DIRECTORY).rstrip("/") + "/"
-    paths = sorted(path.removeprefix(prefix) for path in repository_files if path.startswith(prefix))
-    if not paths:
-        raise FileNotFoundError(f"Missing tokenizer artifact directory {prefix.rstrip('/')}")
-    return paths
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("--tokenizer-path", type=Path, required=True)
-    parser.add_argument("--repo-id", action="append", dest="repo_ids")
-    parser.add_argument("--revision")
-    parser.add_argument("--dry-run", action="store_true")
-    parser.add_argument("--audit-only", action="store_true")
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    tokenizer_path = args.tokenizer_path.resolve()
-    if not tokenizer_path.is_dir():
-        raise FileNotFoundError(f"Tokenizer directory does not exist: {tokenizer_path}")
-
-    files = tokenizer_files(tokenizer_path)
-    fingerprint = artifact_fingerprint([(relative_path, path.read_bytes()) for relative_path, path in files])
-    api = HfApi()
-    repositories = tuple(args.repo_ids or DEFAULT_REPOSITORIES)
-    print(f"Tokenizer fingerprint: {fingerprint}")
-
-    for repo_id in repositories:
-        if args.audit_only:
-            audit_repository(api, repo_id, fingerprint, args.revision)
-            print(f"AUDIT OK {repo_id}@{args.revision or 'main'}")
-            continue
-
-        operations = prepare_operations(repo_id, tokenizer_path, args.revision)
-        if args.dry_run:
-            print(f"DRY RUN {repo_id}: {len(operations)} files")
-            for operation in operations:
-                print(f"  {operation.path_in_repo}")
-            continue
-
-        commit = api.create_commit(
-            repo_id=repo_id,
-            repo_type="model",
-            operations=operations,
-            commit_message="Embed fitted FAST tokenizer for portable PI052 checkpoints",
-            revision=args.revision,
-        )
-        audit_repository(api, repo_id, fingerprint, commit.oid)
-        print(f"BACKFILLED {repo_id}@{commit.oid}")
-
-
-if __name__ == "__main__":
-    main()

src/lerobot/configs/default.py

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

src/lerobot/configs/policies.py

@@ -205,24 +205,30 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
                     f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
                 ) from e
 
-        # HACK: Parse the original config to get the config subclass, so that we can
-        # apply cli overrides.
-        # This is very ugly, ideally we'd like to be able to do that natively with draccus
-        # something like --policy.path (in addition to --policy.type)
-        with draccus.config_type("json"):
-            orig_config = draccus.parse(cls, config_file, args=[])
-
         if config_file is None:
             raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
 
         with open(config_file) as f:
             config = json.load(f)
 
-        config.pop("type")
+        # Resolve the concrete config subclass from the serialized "type" tag, then parse
+        # the config (with CLI overrides) directly for that class. The "type" key is
+        # stripped because draccus only consumes it when parsing the registry base class.
+        policy_type = config.pop("type", None)
+        if policy_type is None:
+            raise ValueError(f"Missing 'type' field in {CONFIG_NAME} of {model_id}")
+        try:
+            config_cls = cls.get_choice_class(policy_type)
+        except Exception as e:
+            raise ValueError(
+                f"Policy type '{policy_type}' (from {CONFIG_NAME} of {model_id}) is not registered. "
+                f"Available policy types: {cls.get_known_choices()}"
+            ) from e
+
         with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
             json.dump(config, f)
             config_file = f.name
 
         cli_overrides = policy_kwargs.pop("cli_overrides", [])
         with draccus.config_type("json"):
-            return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)
+            return draccus.parse(config_cls, config_file, args=cli_overrides)

src/lerobot/configs/recipe.py

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

src/lerobot/configs/recipes/subtask.yaml

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

src/lerobot/configs/recipes/subtask_joint.yaml

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

src/lerobot/configs/recipes/subtask_mem.yaml

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

src/lerobot/configs/recipes/subtask_mem_vqa_speech.yaml

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

src/lerobot/datasets/dataset_reader.py

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

src/lerobot/datasets/factory.py

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

src/lerobot/datasets/language_render.py

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

src/lerobot/envs/configs.py

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

src/lerobot/envs/robocasa.py

@@ -33,8 +33,8 @@ logger = logging.getLogger(__name__)
 
 # Dimensions for the flat action/state vectors used by the LeRobot wrapper.
 # These correspond to the PandaOmron robot in RoboCasa365.
-OBS_STATE_DIM = 16  # ee_pos_rel(3) + ee_quat_rel(4) + base_pos(3) + base_quat(4) + gripper_qpos(2)
-ACTION_DIM = 12  # ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
+OBS_STATE_DIM = 16  # base_pos(3) + base_quat(4) + ee_pos_rel(3) + ee_quat_rel(4) + gripper_qpos(2)
+ACTION_DIM = 12  # base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
 ACTION_LOW = -1.0
 ACTION_HIGH = 1.0
 
@@ -101,15 +101,14 @@ def _resolve_tasks(task: str) -> tuple[list[str], str | None]:
 def convert_action(flat_action: np.ndarray) -> dict[str, Any]:
     """Split a flat (12,) action vector into a RoboCasa action dict.
 
-    Layout (openpi / robocasa.utils.env_utils.convert_action order):
-        ee_pos(3) + ee_rot(3) + gripper(1) + base_motion(4) + control_mode(1)
+    Layout: base_motion(4) + control_mode(1) + ee_pos(3) + ee_rot(3) + gripper(1)
     """
     return {
-        "action.end_effector_position": flat_action[0:3],
-        "action.end_effector_rotation": flat_action[3:6],
-        "action.gripper_close": flat_action[6:7],
-        "action.base_motion": flat_action[7:11],
-        "action.control_mode": flat_action[11:12],
+        "action.base_motion": flat_action[0:4],
+        "action.control_mode": flat_action[4:5],
+        "action.end_effector_position": flat_action[5:8],
+        "action.end_effector_rotation": flat_action[8:11],
+        "action.gripper_close": flat_action[11:12],
     }
 
 
@@ -137,16 +136,9 @@ class RoboCasaEnv(gym.Env):
         episode_length: int | None = None,
         obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
         episode_index: int = 0,
-        terminate_on_success: bool = True,
-        horizon: int | None = None,
     ):
         super().__init__()
         self.task = task
-        # When False, a task-success does NOT end/reset the episode — used by the
-        # interactive sim so one kitchen persists across sequential prompts.
-        self.terminate_on_success = terminate_on_success
-        # Underlying robosuite horizon (steps before truncation). None -> default.
-        self.horizon = horizon
         self.obs_type = obs_type
         self.render_mode = render_mode
         self.observation_width = observation_width
@@ -218,16 +210,12 @@ class RoboCasaEnv(gym.Env):
         # (only None/"all"/"pretrain"/"target" are valid). Always pass a
         # valid value so we don't hit that default. Extra kwargs are
         # forwarded to the underlying kitchen env via create_env/robosuite.make.
-        extra_kwargs: dict[str, Any] = {}
-        if self.horizon is not None:
-            extra_kwargs["horizon"] = int(self.horizon)
         self._env = RoboCasaGymEnv(
             env_name=self.task,
             camera_widths=self.observation_width,
             camera_heights=self.observation_height,
             split=self.split if self.split is not None else "all",
             obj_registries=self.obj_registries,
-            **extra_kwargs,
         )
 
         ep_meta = self._env.env.get_ep_meta()
@@ -242,14 +230,12 @@ class RoboCasaEnv(gym.Env):
             return {"pixels": images}
 
         # `state.*` keys come from PandaOmronKeyConverter inside the wrapper.
-        # openpi state order: ee first, then base, then gripper (matches the
-        # openpi robocasa pipeline / examples/robocasa/main.py state layout).
         agent_pos = np.concatenate(
             [
-                raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
-                raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
                 raw_obs.get("state.base_position", np.zeros(3)),
                 raw_obs.get("state.base_rotation", np.zeros(4)),
+                raw_obs.get("state.end_effector_position_relative", np.zeros(3)),
+                raw_obs.get("state.end_effector_rotation_relative", np.zeros(4)),
                 raw_obs.get("state.gripper_qpos", np.zeros(2)),
             ],
             axis=-1,
@@ -294,7 +280,7 @@ class RoboCasaEnv(gym.Env):
         raw_obs, reward, done, truncated, info = self._env.step(action_dict)
 
         is_success = bool(info.get("success", False))
-        terminated = done or (is_success and self.terminate_on_success)
+        terminated = done or is_success
         info.update({"task": self.task, "done": done, "is_success": is_success})
 
         observation = self._format_raw_obs(raw_obs)
@@ -327,8 +313,6 @@ def _make_env_fns(
     split: str | None,
     episode_length: int | None,
     obj_registries: Sequence[str],
-    terminate_on_success: bool = True,
-    horizon: int | None = None,
 ) -> list[Callable[[], RoboCasaEnv]]:
     """Build n_envs factory callables for a single task.
 
@@ -351,8 +335,6 @@ def _make_env_fns(
             episode_length=episode_length,
             obj_registries=obj_registries,
             episode_index=episode_index,
-            terminate_on_success=terminate_on_success,
-            horizon=horizon,
         )
 
     return [partial(_make_env, i) for i in range(n_envs)]
@@ -366,8 +348,6 @@ def create_robocasa_envs(
     env_cls: Callable[[Sequence[Callable[[], Any]]], Any] | None = None,
     episode_length: int | None = None,
     obj_registries: Sequence[str] = DEFAULT_OBJ_REGISTRIES,
-    terminate_on_success: bool = True,
-    horizon: int | None = None,
 ) -> dict[str, dict[int, Any]]:
     """Create vectorized RoboCasa365 environments with a consistent return shape.
 
@@ -429,8 +409,6 @@ def create_robocasa_envs(
             split=split,
             episode_length=episode_length,
             obj_registries=obj_registries,
-            terminate_on_success=terminate_on_success,
-            horizon=horizon,
         )
 
         if is_async:

src/lerobot/optim/optimizers.py

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

src/lerobot/policies/__init__.py

@@ -28,11 +28,11 @@ 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
 from .utils import make_robot_action, prepare_observation_for_inference
+from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig as VLAJEPAConfig
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
 from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
 from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
@@ -56,9 +56,9 @@ __all__ = [
     "PI0Config",
     "PI0FastConfig",
     "PI05Config",
-    "PI052Config",
     "SmolVLAConfig",
     "TDMPCConfig",
+    "VLAJEPAConfig",
     "VQBeTConfig",
     "WallXConfig",
     "XVLAConfig",

src/lerobot/policies/act/processor_act.py

@@ -18,17 +18,10 @@ from typing import Any
 import torch
 
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_pre_post_processors,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_act import ACTConfig
 
@@ -54,34 +47,4 @@ def make_act_pre_post_processors(
         tuple[PolicyProcessorPipeline[dict[str, Any], dict[str, Any]], PolicyProcessorPipeline[PolicyAction, PolicyAction]]: A tuple containing the
         pre-processor pipeline and the post-processor pipeline.
     """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-            device=config.device,
-        ),
-    ]
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        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,
-        ),
-    )
+    return make_default_pre_post_processors(config, dataset_stats, normalizer_device=config.device)

src/lerobot/policies/common/flow_matching.py

@@ -0,0 +1,161 @@
+#!/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.
+
+"""Flow-matching sampling primitives shared across policies.
+
+Canonical versions of the beta-distributed timestep sampler and the forward-Euler
+denoising loop (with its real-time-chunking hook) that the openpi-derived policies
+(pi0, pi05, smolvla, eo1) historically each carried a copy of. All functions are
+stateless; adopting them does not affect checkpoints.
+"""
+
+from collections.abc import Callable
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+import torch
+from torch import Tensor
+
+if TYPE_CHECKING:
+    from lerobot.policies.rtc.modeling_rtc import RTCProcessor
+
+
+@lru_cache(maxsize=None)
+def _beta_distribution(alpha: float, beta: float) -> "torch.distributions.Beta":
+    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so build on CPU.
+    # Cached (groot convention) so the distribution object is constructed once per (alpha, beta).
+    alpha_t = torch.tensor(alpha, dtype=torch.float32)
+    beta_t = torch.tensor(beta, dtype=torch.float32)
+    return torch.distributions.Beta(alpha_t, beta_t)
+
+
+def sample_beta(alpha: float, beta: float, bsize: int, device) -> Tensor:  # see openpi (exact copy)
+    return _beta_distribution(alpha, beta).sample((bsize,)).to(device)
+
+
+def sample_noise(shape, device, *, distribution: str = "normal") -> Tensor:
+    """Float32 flow-matching noise sample.
+
+    ``distribution="normal"`` (default, openpi: pi0/pi05/eo1/smolvla/groot/wall_x) draws
+    standard-normal noise. ``distribution="uniform"`` (evo1) draws uniformly from
+    ``[-1, 1)`` via ``rand * 2 - 1``.
+    """
+    if distribution == "normal":
+        return torch.normal(
+            mean=0.0,
+            std=1.0,
+            size=shape,
+            dtype=torch.float32,
+            device=device,
+        )
+    if distribution == "uniform":
+        return torch.rand(shape, dtype=torch.float32, device=device) * 2 - 1
+    raise ValueError(f"Unknown noise distribution: {distribution!r} (expected 'normal' or 'uniform')")
+
+
+def sample_time_beta(
+    bsize: int,
+    device,
+    *,
+    alpha: float,
+    beta: float,
+    scale: float = 1.0,
+    offset: float = 0.0,
+    complement: bool = False,
+    clamp_min: float | None = None,
+    clamp_max: float | None = None,
+) -> Tensor:
+    """Beta-distributed flow-matching timesteps.
+
+    Computes ``t = f(Beta(alpha, beta)) * scale + offset`` where ``f`` is the identity by
+    default or ``1 - x`` when ``complement=True``, then optionally clamps to
+    ``[clamp_min, clamp_max]``. This covers the known per-policy conventions:
+
+    * openpi backward (pi0/pi05/eo1/smolvla): ``scale=0.999, offset=0.001``.
+    * forward (groot/wall_x): ``complement=True, scale=0.999`` giving ``(1 - beta) * 0.999``.
+    * evo1: ``alpha=beta=2, clamp_min=0.02, clamp_max=0.98``.
+    """
+    time_beta = sample_beta(alpha, beta, bsize, device)
+    if complement:
+        time_beta = 1.0 - time_beta
+    time = time_beta * scale + offset
+    if clamp_min is not None or clamp_max is not None:
+        time = time.clamp(min=clamp_min, max=clamp_max)
+    return time.to(dtype=torch.float32, device=device)
+
+
+def euler_integrate(
+    denoise_fn: Callable[[Tensor, Tensor], Tensor],
+    noise: Tensor,
+    num_steps: int,
+    *,
+    rtc_processor: "RTCProcessor | None" = None,
+    rtc_enabled: bool = False,
+    inference_delay: int | None = None,
+    prev_chunk_left_over: Tensor | None = None,
+    execution_horizon: int | None = None,
+) -> Tensor:
+    """Forward-Euler integration of a velocity field from t=1 (noise) to t=0 (actions).
+
+    This is the openpi sampling loop: ``dt = -1/num_steps``, ``time = 1.0 + step*dt``,
+    ``x_t <- x_t + dt * v_t``, with the optional real-time-chunking (RTC) guidance hook
+    wrapping the velocity computation and debug tracking after each step.
+
+    Args:
+        denoise_fn: Computes the velocity ``v_t`` from ``(x_t, time_tensor)`` where
+            ``time_tensor`` is a float32 tensor of shape ``(batch_size,)``. The returned
+            velocity must have the same shape and dtype as ``x_t``.
+        noise: Initial sample ``x_1`` of shape ``(batch_size, ...)``.
+        num_steps: Number of Euler steps.
+        rtc_processor: Optional RTC processor. Debug tracking fires whenever it is set and
+            has debugging enabled, even if RTC guidance itself is disabled (this mirrors
+            the historical per-policy loops).
+        rtc_enabled: Whether to route the velocity computation through
+            ``rtc_processor.denoise_step`` (requires ``rtc_processor``).
+        inference_delay: RTC guidance parameter, forwarded verbatim.
+        prev_chunk_left_over: RTC guidance parameter, forwarded verbatim.
+        execution_horizon: RTC guidance parameter, forwarded verbatim.
+    """
+    bsize = noise.shape[0]
+    device = noise.device
+
+    dt = -1.0 / num_steps
+    x_t = noise
+    for step in range(num_steps):
+        time = 1.0 + step * dt
+        time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+
+        def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
+            return denoise_fn(input_x_t, current_timestep)
+
+        if rtc_enabled:
+            v_t = 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 rtc_processor is not None and rtc_processor.is_debug_enabled():
+            rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
+
+    return x_t

src/lerobot/policies/common/vla_utils.py

@@ -0,0 +1,243 @@
+#!/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.
+
+"""Helpers shared by the openpi-derived VLA policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla).
+
+These are the canonical versions of functions that historically were copy-pasted per
+policy. They are pure (no parameters, no module state), so importing them from here
+instead of a policy-local copy has no effect on checkpoints.
+"""
+
+import math
+from typing import TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F  # noqa: N812
+from torch import Tensor
+
+from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
+from lerobot.utils.device_utils import get_safe_dtype
+from lerobot.utils.import_utils import _transformers_available, require_package
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers import DynamicCache
+else:
+    DynamicCache = None
+
+
+def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
+    time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
+) -> Tensor:
+    """Computes sine-cosine positional embedding vectors for scalar positions."""
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+
+    if time.ndim != 1:
+        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
+
+    dtype = get_safe_dtype(torch.float64, device.type)
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+
+    # Compute the outer product
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None]
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+
+
+def make_att_2d_masks(pad_masks: Tensor, att_masks: Tensor) -> Tensor:  # see openpi (exact copy)
+    """Copied from big_vision.
+
+    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
+    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
+    setup several types of attention, for example:
+
+      [[1 1 1 1 1 1]]: pure causal attention.
+
+      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
+          themselves and the last 3 tokens have a causal attention. The first
+          entry could also be a 1 without changing behaviour.
+
+      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
+          block can attend all previous blocks and all tokens on the same block.
+
+    Args:
+      input_mask: bool[B, N] true if its part of the input, false if padding.
+      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
+        it and 0 where it shares the same attention mask as the previous token.
+    """
+    if att_masks.ndim != 2:
+        raise ValueError(att_masks.ndim)
+    if pad_masks.ndim != 2:
+        raise ValueError(pad_masks.ndim)
+
+    cumsum = torch.cumsum(att_masks, dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
+    return att_2d_masks & pad_2d_masks
+
+
+def prepare_attention_masks_4d(att_2d_masks: Tensor, dtype: torch.dtype | None = None) -> Tensor:
+    """Expand boolean 2D attention masks to the additive 4D layout expected by transformers.
+
+    Valid positions become 0.0 and masked positions the large negative openpi constant.
+    """
+    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 clone_past_key_values(past_key_values):
+    """Clone the DynamicCache returned by prefix prefill for compiled denoising."""
+    if DynamicCache is None:
+        require_package("transformers", extra="transformers-dep")
+
+    return DynamicCache(
+        tuple(
+            (keys.clone(), values.clone(), sliding_window) for keys, values, sliding_window in past_key_values
+        )
+    )
+
+
+def pad_vector(vector: Tensor, new_dim: int, *, truncate: bool = False) -> Tensor:
+    """Pad the last dimension of a vector to new_dim with zeros.
+
+    Can be (batch_size x sequence_length x features_dimension)
+    or (batch_size x features_dimension)
+
+    With ``truncate=False`` (openpi behavior), vectors whose last dimension is already
+    >= new_dim are returned unchanged. With ``truncate=True`` (xVLA behavior), the last
+    dimension is truncated to exactly ``new_dim`` (which may be 0).
+    """
+    if vector.shape[-1] == new_dim:
+        return vector
+    if not truncate:
+        if vector.shape[-1] >= new_dim:
+            return vector
+        return F.pad(vector, (0, new_dim - vector.shape[-1]))
+    shape = list(vector.shape)
+    current_dim = shape[-1]
+    shape[-1] = new_dim
+    new_vector = vector.new_zeros(*shape)
+    length = min(current_dim, new_dim)
+    new_vector[..., :length] = vector[..., :length]
+    return new_vector
+
+
+def resize_with_pad_torch(  # see openpi `resize_with_pad_torch` (exact copy)
+    images: torch.Tensor,
+    height: int,
+    width: int,
+    mode: str = "bilinear",
+) -> torch.Tensor:
+    """PyTorch version of resize_with_pad. Resizes an image to a target height and width without distortion
+    by padding with black. If the image is float32, it must be in the range [-1, 1].
+
+    Padding is centered (openpi convention). For the top-left-padding variant used by
+    smolvla/xvla, see :func:`resize_with_pad`.
+
+    Args:
+        images: Tensor of shape [*b, h, w, c] or [*b, c, h, w]
+        height: Target height
+        width: Target width
+        mode: Interpolation mode ('bilinear', 'nearest', etc.)
+
+    Returns:
+        Resized and padded tensor with same shape format as input
+    """
+    # Check if input is in channels-last format [*b, h, w, c] or channels-first [*b, c, h, w]
+    if images.shape[-1] <= 4:  # Assume channels-last format
+        channels_last = True
+        if images.dim() == 3:
+            images = images.unsqueeze(0)  # Add batch dimension
+        images = images.permute(0, 3, 1, 2)  # [b, h, w, c] -> [b, c, h, w]
+    else:
+        channels_last = False
+        if images.dim() == 3:
+            images = images.unsqueeze(0)  # Add batch dimension
+
+    batch_size, channels, cur_height, cur_width = images.shape
+
+    # Calculate resize ratio
+    ratio = max(cur_width / width, cur_height / height)
+    resized_height = int(cur_height / ratio)
+    resized_width = int(cur_width / ratio)
+
+    # Resize
+    resized_images = F.interpolate(
+        images,
+        size=(resized_height, resized_width),
+        mode=mode,
+        align_corners=False if mode == "bilinear" else None,
+    )
+
+    # Handle dtype-specific clipping
+    if images.dtype == torch.uint8:
+        resized_images = torch.round(resized_images).clamp(0, 255).to(torch.uint8)
+    elif images.dtype == torch.float32:
+        resized_images = resized_images.clamp(0.0, 1.0)
+    else:
+        raise ValueError(f"Unsupported image dtype: {images.dtype}")
+
+    # Calculate padding
+    pad_h0, remainder_h = divmod(height - resized_height, 2)
+    pad_h1 = pad_h0 + remainder_h
+    pad_w0, remainder_w = divmod(width - resized_width, 2)
+    pad_w1 = pad_w0 + remainder_w
+
+    # Pad
+    constant_value = 0 if images.dtype == torch.uint8 else 0.0
+    padded_images = F.pad(
+        resized_images,
+        (pad_w0, pad_w1, pad_h0, pad_h1),  # left, right, top, bottom
+        mode="constant",
+        value=constant_value,
+    )
+
+    # Convert back to original format if needed
+    if channels_last:
+        padded_images = padded_images.permute(0, 2, 3, 1)  # [b, c, h, w] -> [b, h, w, c]
+
+    return padded_images
+
+
+def resize_with_pad(img: torch.Tensor, height: int, width: int, *, pad_value: float) -> torch.Tensor:
+    """Resize a (b, c, h, w) image without distortion, padding on the LEFT and TOP.
+
+    This is the smolvla/xvla convention. For the centered-padding openpi variant, see
+    :func:`resize_with_pad_torch`. ``pad_value`` is keyword-only on purpose: callers
+    historically used different values (0, -1) and must state their choice explicitly.
+    """
+    if img.ndim != 4:
+        raise ValueError(f"(b,c,h,w) expected, but got {img.shape}")
+
+    current_height, current_width = img.shape[2:]
+    if current_height == height and current_width == width:
+        return img
+
+    ratio = max(current_width / width, current_height / height)
+    resized_height = int(current_height / ratio)
+    resized_width = int(current_width / ratio)
+    resized_img = F.interpolate(
+        img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
+    )
+
+    pad_height = max(0, height - resized_height)
+    pad_width = max(0, width - resized_width)
+    padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
+    return padded_img

src/lerobot/policies/diffusion/processor_diffusion.py

@@ -19,17 +19,10 @@ from typing import Any
 import torch
 
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_pre_post_processors,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_diffusion import DiffusionConfig
 
@@ -63,32 +56,4 @@ def make_diffusion_pre_post_processors(
     Returns:
         A tuple containing the configured pre-processor and post-processor pipelines.
     """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        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,
-        ),
-    )
+    return make_default_pre_post_processors(config, dataset_stats)

src/lerobot/policies/eo1/processor_eo1.py

@@ -23,24 +23,16 @@ import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
     ComplementaryDataProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.types import TransitionKey
-from lerobot.utils.constants import (
-    OBS_STATE,
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
-)
+from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.import_utils import _transformers_available, require_package
 
 from .configuration_eo1 import EO1Config
@@ -242,14 +234,12 @@ def make_eo1_pre_post_processors(
 ]:
     """Build pre/post processor pipelines for EO1."""
 
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     input_steps: list[ProcessorStep] = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        steps.rename_observations,
+        steps.add_batch_dim,
+        steps.normalize,
         EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size),
         EO1QwenProcessorStep(
             processor_name=config.vlm_base,
@@ -257,27 +247,12 @@ def make_eo1_pre_post_processors(
             image_max_pixels=config.image_max_pixels,
             use_fast_processor=config.use_fast_processor,
         ),
-        DeviceProcessorStep(device=config.device),
+        steps.to_device,
     ]
 
     output_steps: list[ProcessorStep] = [
-        UnnormalizerProcessorStep(
-            features=config.output_features,
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        DeviceProcessorStep(device="cpu"),
+        steps.unnormalize,
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/factory.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import importlib
+import inspect
 import logging
 from typing import TYPE_CHECKING, Any, TypedDict, Unpack
 
@@ -44,26 +45,10 @@ from lerobot.utils.constants import (
 )
 from lerobot.utils.feature_utils import dataset_to_policy_features
 
-from .act.configuration_act import ACTConfig
-from .diffusion.configuration_diffusion import DiffusionConfig
-from .eo1.configuration_eo1 import EO1Config
 from .evo1.configuration_evo1 import Evo1Config
-from .fastwam.configuration_fastwam import FastWAMConfig
-from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
 from .groot.configuration_groot import GrootConfig
-from .lingbot_va.configuration_lingbot_va import LingBotVAConfig
-from .molmoact2.configuration_molmoact2 import MolmoAct2Config
-from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
-from .pi0.configuration_pi0 import PI0Config
-from .pi05.configuration_pi05 import PI05Config
 from .pretrained import PreTrainedPolicy
-from .smolvla.configuration_smolvla import SmolVLAConfig
-from .tdmpc.configuration_tdmpc import TDMPCConfig
 from .utils import validate_visual_features_consistency
-from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
-from .vqbet.configuration_vqbet import VQBeTConfig
-from .wall_x.configuration_wall_x import WallXConfig
-from .xvla.configuration_xvla import XVLAConfig
 
 
 def _reconnect_relative_absolute_steps(
@@ -88,104 +73,23 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
     """
     Retrieves a policy class by its registered name.
 
-    This function uses dynamic imports to avoid loading all policy classes into memory
-    at once, improving startup time and reducing dependencies.
+    Resolution is convention-based: the draccus-registered config class of ``name`` is
+    looked up, its ``configuration_*`` module path is rewritten to ``modeling_*``, and
+    the ``<X>Policy`` class is imported from there. The modeling module is only imported
+    at call time, keeping heavy optional dependencies lazy. This works for both built-in
+    policies and third-party lerobot plugins (anything registered via
+    ``@PreTrainedConfig.register_subclass``).
 
     Args:
-        name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
-            "multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
-            "molmoact2", "eo1", "evo1".
+        name: The registered name of the policy (e.g. "act", "diffusion", "pi0").
     Returns:
         The policy class corresponding to the given name.
 
     Raises:
-        NotImplementedError: If the policy name is not recognized.
+        ValueError: If the policy name is not registered.
+        ImportError: If the policy's optional dependencies are not installed.
     """
-    if name == "tdmpc":
-        from .tdmpc.modeling_tdmpc import TDMPCPolicy
-
-        return TDMPCPolicy
-    elif name == "diffusion":
-        from .diffusion.modeling_diffusion import DiffusionPolicy
-
-        return DiffusionPolicy
-    elif name == "act":
-        from .act.modeling_act import ACTPolicy
-
-        return ACTPolicy
-    elif name == "multi_task_dit":
-        from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
-
-        return MultiTaskDiTPolicy
-    elif name == "vqbet":
-        from .vqbet.modeling_vqbet import VQBeTPolicy
-
-        return VQBeTPolicy
-    elif name == "pi0":
-        from .pi0.modeling_pi0 import PI0Policy
-
-        return PI0Policy
-    elif name == "pi0_fast":
-        from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
-
-        return PI0FastPolicy
-    elif name == "pi05":
-        from .pi05.modeling_pi05 import PI05Policy
-
-        return PI05Policy
-    elif name == "pi052":
-        from .pi052.modeling_pi052 import PI052Policy
-
-        return PI052Policy
-    elif name == "gaussian_actor":
-        from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
-
-        return GaussianActorPolicy
-    elif name == "smolvla":
-        from .smolvla.modeling_smolvla import SmolVLAPolicy
-
-        return SmolVLAPolicy
-    elif name == "groot":
-        from .groot.modeling_groot import GrootPolicy
-
-        return GrootPolicy
-    elif name == "xvla":
-        from .xvla.modeling_xvla import XVLAPolicy
-
-        return XVLAPolicy
-    elif name == "wall_x":
-        from .wall_x.modeling_wall_x import WallXPolicy
-
-        return WallXPolicy
-    elif name == "eo1":
-        from .eo1.modeling_eo1 import EO1Policy
-
-        return EO1Policy
-    elif name == "molmoact2":
-        from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
-
-        return MolmoAct2Policy
-    elif name == "vla_jepa":
-        from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
-
-        return VLAJEPAPolicy
-    elif name == "lingbot_va":
-        from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
-
-        return LingBotVAPolicy
-    elif name == "fastwam":
-        from .fastwam.modeling_fastwam import FastWAMPolicy
-
-        return FastWAMPolicy
-    elif name == "evo1":
-        from .evo1.modeling_evo1 import Evo1Policy
-
-        return Evo1Policy
-    else:
-        try:
     return _get_policy_cls_from_policy_name(name=name)
-        except Exception as e:
-            raise ValueError(f"Policy type '{name}' is not available.") from e
 
 
 def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
@@ -196,9 +100,8 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
     mapping a string identifier to the corresponding config class.
 
     Args:
-        policy_type: The type of the policy. Supported types include "tdmpc",
-                     "multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "pi052",
-                     "gaussian_actor", "smolvla", "wall_x", "molmoact2", "eo1", "evo1".
+        policy_type: The registered type of the policy (any name registered via
+                     ``@PreTrainedConfig.register_subclass``, e.g. "act", "diffusion", "pi0").
         **kwargs: Keyword arguments to be passed to the configuration class constructor.
 
     Returns:
@@ -207,52 +110,11 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
     Raises:
         ValueError: If the `policy_type` is not recognized.
     """
-    if policy_type == "tdmpc":
-        return TDMPCConfig(**kwargs)
-    elif policy_type == "diffusion":
-        return DiffusionConfig(**kwargs)
-    elif policy_type == "act":
-        return ACTConfig(**kwargs)
-    elif policy_type == "multi_task_dit":
-        return MultiTaskDiTConfig(**kwargs)
-    elif policy_type == "vqbet":
-        return VQBeTConfig(**kwargs)
-    elif policy_type == "pi0":
-        return PI0Config(**kwargs)
-    elif policy_type == "pi05":
-        return PI05Config(**kwargs)
-    elif policy_type == "pi052":
-        from .pi052.configuration_pi052 import PI052Config
-
-        return PI052Config(**kwargs)
-    elif policy_type == "gaussian_actor":
-        return GaussianActorConfig(**kwargs)
-    elif policy_type == "smolvla":
-        return SmolVLAConfig(**kwargs)
-    elif policy_type == "groot":
-        return GrootConfig(**kwargs)
-    elif policy_type == "xvla":
-        return XVLAConfig(**kwargs)
-    elif policy_type == "wall_x":
-        return WallXConfig(**kwargs)
-    elif policy_type == "eo1":
-        return EO1Config(**kwargs)
-    elif policy_type == "molmoact2":
-        return MolmoAct2Config(**kwargs)
-    elif policy_type == "vla_jepa":
-        return VLAJEPAConfig(**kwargs)
-    elif policy_type == "lingbot_va":
-        return LingBotVAConfig(**kwargs)
-    elif policy_type == "fastwam":
-        return FastWAMConfig(**kwargs)
-    elif policy_type == "evo1":
-        return Evo1Config(**kwargs)
-    else:
     try:
         config_cls = PreTrainedConfig.get_choice_class(policy_type)
-            return config_cls(**kwargs)
     except Exception as e:
         raise ValueError(f"Policy type '{policy_type}' is not available.") from e
+    return config_cls(**kwargs)
 
 
 class ProcessorConfigKwargs(TypedDict, total=False):
@@ -275,12 +137,6 @@ class ProcessorConfigKwargs(TypedDict, total=False):
     preprocessor_overrides: dict[str, Any] | None
     postprocessor_overrides: dict[str, Any] | None
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None
-    # Dataset repo used for optional processor fitting; omit it to use universal tokenizers.
-    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
 
 
@@ -312,13 +168,9 @@ def make_pre_post_processors(
         A tuple containing the input (pre-processor) and output (post-processor) pipelines.
 
     Raises:
-        NotImplementedError: If a processor factory is not implemented for the given
-            policy configuration type.
+        ValueError: If no processor factory exists for the given policy configuration type.
     """
     if pretrained_path:
-        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
 
@@ -368,195 +220,14 @@ def make_pre_post_processors(
             )
         return preprocessor, postprocessor
 
-    # Create a new processor based on policy type
-    if isinstance(policy_cfg, TDMPCConfig):
-        from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
-
-        processors = make_tdmpc_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, DiffusionConfig):
-        from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
-
-        processors = make_diffusion_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, ACTConfig):
-        from .act.processor_act import make_act_pre_post_processors
-
-        processors = make_act_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, MultiTaskDiTConfig):
-        from .multi_task_dit.processor_multi_task_dit import (
-            make_multi_task_dit_pre_post_processors,
-        )
-
-        processors = make_multi_task_dit_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, VQBeTConfig):
-        from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
-
-        processors = make_vqbet_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, PI0Config):
-        from .pi0.processor_pi0 import make_pi0_pre_post_processors
-
-        processors = make_pi0_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif policy_cfg.type == "pi0_fast":
-        from .pi0_fast.processor_pi0_fast import make_pi0_fast_pre_post_processors
-
-        processors = make_pi0_fast_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            dataset_repo_id=kwargs.get("dataset_repo_id"),
-            dataset_root=kwargs.get("dataset_root"),
-            dataset_revision=kwargs.get("dataset_revision"),
-            episodes=kwargs.get("dataset_episodes"),
-            exclude_episodes=kwargs.get("dataset_exclude_episodes"),
-        )
-
-    elif policy_cfg.type == "pi052":
-        # PI052 must precede PI05 because its config subclasses PI05Config.
-        from .pi052.processor_pi052 import make_pi052_pre_post_processors
-
-        processors = make_pi052_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            # Without a dataset repo, FAST auto-fit falls back to the universal tokenizer.
-            dataset_repo_id=kwargs.get("dataset_repo_id"),
-            dataset_root=kwargs.get("dataset_root"),
-            dataset_revision=kwargs.get("dataset_revision"),
-            episodes=kwargs.get("dataset_episodes"),
-            exclude_episodes=kwargs.get("dataset_exclude_episodes"),
-        )
-
-    elif isinstance(policy_cfg, PI05Config):
-        from .pi05.processor_pi05 import make_pi05_pre_post_processors
-
-        processors = make_pi05_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, GaussianActorConfig):
-        from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
-
-        processors = make_gaussian_actor_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, SmolVLAConfig):
-        from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
-
-        processors = make_smolvla_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, GrootConfig):
-        from .groot.processor_groot import make_groot_pre_post_processors
-
-        processors = make_groot_pre_post_processors(
+    # Create new processors from the policy config, resolving the per-policy factory
+    # function by naming convention (lazy import keeps optional dependencies optional).
+    return _make_processors_from_policy_config(
         config=policy_cfg,
         dataset_stats=kwargs.get("dataset_stats"),
         dataset_meta=kwargs.get("dataset_meta"),
     )
 
-    elif isinstance(policy_cfg, XVLAConfig):
-        from .xvla.processor_xvla import (
-            make_xvla_pre_post_processors,
-        )
-
-        processors = make_xvla_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, WallXConfig):
-        from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
-
-        processors = make_wall_x_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, EO1Config):
-        from .eo1.processor_eo1 import make_eo1_pre_post_processors
-
-        processors = make_eo1_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-    elif isinstance(policy_cfg, Evo1Config):
-        from .evo1.processor_evo1 import make_evo1_pre_post_processors
-
-        processors = make_evo1_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, MolmoAct2Config):
-        from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
-
-        processors = make_molmoact2_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            dataset_meta=kwargs.get("dataset_meta"),
-        )
-
-    elif isinstance(policy_cfg, VLAJEPAConfig):
-        from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
-
-        processors = make_vla_jepa_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, LingBotVAConfig):
-        from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
-
-        processors = make_lingbot_va_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, FastWAMConfig):
-        from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
-
-        processors = make_fastwam_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    else:
-        try:
-            processors = _make_processors_from_policy_config(
-                config=policy_cfg,
-                dataset_stats=kwargs.get("dataset_stats"),
-            )
-        except Exception as e:
-            raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
-
-    return processors
-
 
 def make_policy(
     cfg: PreTrainedConfig,
@@ -699,10 +370,12 @@ def make_policy(
     return policy
 
 
-def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
+def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
     """Get policy class from its registered name using dynamic imports.
 
-    This is used as a helper function to import policies from 3rd party lerobot plugins.
+    Works for built-in policies and 3rd party lerobot plugins alike: the config class
+    registered under ``name`` is resolved via the draccus ChoiceRegistry, and the policy
+    class is imported from the sibling ``modeling_*`` module by naming convention.
 
     Args:
         name: The name of the policy.
@@ -728,22 +401,39 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
         "configuration_", "modeling_"
     )  # e.g., configuration_diffusion -> modeling_diffusion
 
+    try:
         module = importlib.import_module(module_path)
-    policy_cls = getattr(module, cls_name)
+    except ModuleNotFoundError as e:
+        if e.name == module_path:
+            # The modeling_* module itself does not exist for this policy type. A missing
+            # optional dependency inside an existing module propagates unchanged instead,
+            # so its actionable install hint stays visible.
+            raise ValueError(f"Policy class for '{name}' is not implemented.") from e
+        raise
+    policy_cls = getattr(module, cls_name, None)
+    if policy_cls is None:
+        raise ValueError(
+            f"Policy class '{cls_name}' not found in '{module_path}'. "
+            f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
+        )
     return policy_cls
 
 
 def _make_processors_from_policy_config(
     config: PreTrainedConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    dataset_meta: Any | None = None,
 ) -> tuple[Any, Any]:
     """Create pre- and post-processors from a policy configuration using dynamic imports.
 
-    This is used as a helper function to import processor factories from 3rd party lerobot plugins.
+    Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
+    by naming convention. Works for built-in policies and 3rd party lerobot plugins.
 
     Args:
         config: The policy configuration object.
         dataset_stats: Dataset statistics for normalization.
+        dataset_meta: Dataset metadata, forwarded only to factories that declare a
+            ``dataset_meta`` parameter (e.g. groot, molmoact2).
     Returns:
         A tuple containing the input (pre-processor) and output (post-processor) pipelines.
     """
@@ -756,6 +446,19 @@ def _make_processors_from_policy_config(
     logging.debug(
         f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
     )
+    try:
         module = importlib.import_module(module_path)
-    function = getattr(module, function_name)
-    return function(config, dataset_stats=dataset_stats)
+    except ModuleNotFoundError as e:
+        if e.name == module_path:
+            # The processor_* module itself does not exist for this policy type. A missing
+            # optional dependency inside an existing module propagates unchanged instead,
+            # so its actionable install hint stays visible.
+            raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
+        raise
+    function = getattr(module, function_name, None)
+    if function is None:
+        raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
+    call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
+    if "dataset_meta" in inspect.signature(function).parameters:
+        call_kwargs["dataset_meta"] = dataset_meta
+    return function(config, **call_kwargs)

src/lerobot/policies/fastwam/processor_fastwam.py

@@ -22,20 +22,11 @@ import torch
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
     ActionProcessorStep,
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
-)
-from lerobot.utils.constants import (
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
 
 from .configuration_fastwam import FastWAMConfig
@@ -105,38 +96,20 @@ def make_fastwam_pre_post_processors(
     # anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
     # checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
 
+    steps = make_default_policy_processor_steps(config, normalization_stats, normalizer_device=config.device)
+
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=normalization_stats,
-            device=config.device,
-        ),
+        steps.rename_observations,
+        steps.add_batch_dim,
+        steps.to_device,
+        steps.normalize,
     ]
     output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features,
-            norm_map=config.normalization_mapping,
-            stats=normalization_stats,
-        ),
+        steps.unnormalize,
     ]
     if config.toggle_action_dimensions:
         output_steps.append(
             FastWAMActionToggleProcessorStep(toggle_dimensions=config.toggle_action_dimensions)
         )
-    output_steps.append(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,
-        ),
-    )
+    output_steps.append(steps.to_cpu)
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py

@@ -20,17 +20,10 @@ from typing import Any
 import torch
 
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_pre_post_processors,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_gaussian_actor import GaussianActorConfig
 
@@ -62,33 +55,4 @@ def make_gaussian_actor_pre_post_processors(
     Returns:
         A tuple containing the configured pre-processor and post-processor pipelines.
     """
-
-    # Add remaining processors
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        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,
-        ),
-    )
+    return make_default_pre_post_processors(config, dataset_stats)

src/lerobot/policies/lingbot_va/processor_lingbot_va.py

@@ -25,19 +25,12 @@ import torch
 
 from lerobot.configs.types import FeatureType, NormalizationMode
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
-    RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
-)
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import (
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
 
 from .configuration_lingbot_va import LingBotVAConfig
@@ -52,15 +45,13 @@ def make_lingbot_va_pre_post_processors(
 ]:
     """Build the pre/post processor pipelines for LingBot-VA."""
 
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     input_steps: list[ProcessorStep] = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        DeviceProcessorStep(device=config.device),
+        steps.rename_observations,
+        steps.add_batch_dim,
+        steps.normalize,
+        steps.to_device,
     ]
 
     # Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
@@ -70,18 +61,7 @@ def make_lingbot_va_pre_post_processors(
             norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
             stats=dataset_stats,
         ),
-        DeviceProcessorStep(device="cpu"),
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py

@@ -19,18 +19,12 @@ from typing import Any
 import torch
 
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
     TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_multi_task_dit import MultiTaskDiTConfig
 
@@ -66,9 +60,11 @@ def make_multi_task_dit_pre_post_processors(
         A tuple containing the configured pre-processor and post-processor pipelines.
     """
 
+    steps = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=config.device)
+
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
+        steps.rename_observations,
+        steps.add_batch_dim,
         TokenizerProcessorStep(
             tokenizer_name=config.text_encoder_name,
             padding=config.tokenizer_padding,
@@ -76,32 +72,12 @@ def make_multi_task_dit_pre_post_processors(
             max_length=config.tokenizer_max_length,
             truncation=config.tokenizer_truncation,
         ),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-            device=config.device,
-        ),
+        steps.to_device,
+        steps.normalize,
     ]
     output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features,
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        DeviceProcessorStep(device="cpu"),
+        steps.unnormalize,
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/pi0/processor_pi0.py

@@ -21,22 +21,16 @@ import torch
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
-    AddBatchDimensionProcessorStep,
     ComplementaryDataProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
     RelativeActionsProcessorStep,
-    RenameObservationsProcessorStep,
     TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_pi0 import PI0Config
 
@@ -136,10 +130,12 @@ def make_pi0_pre_post_processors(
         action_names=getattr(config, "action_feature_names", None),
     )
 
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     # OpenPI order: raw → relative → normalize → model → unnormalize → absolute
     input_steps: list[ProcessorStep] = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
+        steps.rename_observations,  # To mimic the same processor as pretrained one
+        steps.add_batch_dim,
         Pi0NewLineProcessor(),  # Add newlines before tokenization for PaliGemma
         TokenizerProcessorStep(
             tokenizer_name="google/paligemma-3b-pt-224",
@@ -147,32 +143,15 @@ def make_pi0_pre_post_processors(
             padding_side="right",
             padding="max_length",
         ),
-        DeviceProcessorStep(device=config.device),
+        steps.to_device,
         relative_step,
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        steps.normalize,
     ]
 
     output_steps: list[ProcessorStep] = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
+        steps.unnormalize,
         AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
-        DeviceProcessorStep(device="cpu"),
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/pi05/modeling_pi05.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import builtins
-import json
 import logging
 import math
 from collections import deque
@@ -24,7 +23,6 @@ from typing import TYPE_CHECKING, Literal, TypedDict, Unpack
 
 import torch
 import torch.nn.functional as F  # noqa: N812
-from safetensors.torch import load_file
 from torch import Tensor, nn
 
 from lerobot.utils.import_utils import _transformers_available, require_package
@@ -34,7 +32,6 @@ if TYPE_CHECKING or _transformers_available:
     from transformers.cache_utils import DynamicCache
     from transformers.models.auto import CONFIG_MAPPING
     from transformers.models.gemma import modeling_gemma
-    from transformers.utils import cached_file
 
     from ..pi_gemma import (
         PaliGemmaForConditionalGenerationWithPiGemma,
@@ -50,7 +47,6 @@ else:
     _gated_residual = None
     layernorm_forward = None
     PaliGemmaForConditionalGenerationWithPiGemma = None
-    cached_file = None
 from lerobot.configs import PreTrainedConfig
 from lerobot.utils.constants import (
     ACTION,
@@ -70,84 +66,6 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
-_SAFETENSORS_FILE = "model.safetensors"
-_SAFETENSORS_INDEX = "model.safetensors.index.json"
-
-
-def _resolve_weight_files(
-    pretrained_name_or_path: str | Path,
-    *,
-    force_download: bool,
-    resume_download: bool | None,
-    proxies: dict | None,
-    token: str | bool | None,
-    cache_dir: str | Path | None,
-    local_files_only: bool,
-    revision: str | None,
-) -> list[Path]:
-    model_id = str(pretrained_name_or_path)
-    local_dir = Path(model_id)
-    load_kwargs = {
-        "revision": revision,
-        "cache_dir": cache_dir,
-        "force_download": force_download,
-        "resume_download": resume_download,
-        "proxies": proxies,
-        "token": token,
-        "local_files_only": local_files_only,
-    }
-
-    if local_dir.is_dir():
-        index_path = local_dir / _SAFETENSORS_INDEX
-        single_path = local_dir / _SAFETENSORS_FILE
-    else:
-        resolved_index = cached_file(
-            model_id,
-            _SAFETENSORS_INDEX,
-            _raise_exceptions_for_missing_entries=False,
-            **load_kwargs,
-        )
-        index_path = Path(resolved_index) if resolved_index is not None else None
-        single_path = None
-        if index_path is None:
-            resolved_file = cached_file(model_id, _SAFETENSORS_FILE, **load_kwargs)
-            single_path = Path(resolved_file) if resolved_file is not None else None
-
-    if index_path is None or not index_path.is_file():
-        if single_path is None or not single_path.is_file():
-            raise FileNotFoundError(f"No {_SAFETENSORS_FILE} found in {model_id!r}.")
-        return [single_path]
-
-    index = json.loads(index_path.read_text())
-    shard_names = sorted(set(index.get("weight_map", {}).values()))
-    if not shard_names:
-        raise ValueError(f"Invalid safetensors index without a weight_map: {index_path}")
-    if local_dir.is_dir():
-        files = [local_dir / name for name in shard_names]
-    else:
-        files = []
-        for name in shard_names:
-            resolved_file = cached_file(model_id, name, **load_kwargs)
-            if resolved_file is None:
-                raise FileNotFoundError(f"Checkpoint shard {name!r} not found in {model_id!r}.")
-            files.append(Path(resolved_file))
-    missing = [str(path) for path in files if not path.is_file()]
-    if missing:
-        raise FileNotFoundError(f"Missing checkpoint shards: {missing}")
-    return files
-
-
-def _load_weight_files(files: list[Path]) -> dict[str, Tensor]:
-    state_dict: dict[str, Tensor] = {}
-    for path in files:
-        shard = load_file(path)
-        overlap = state_dict.keys() & shard.keys()
-        if overlap:
-            raise ValueError(f"Duplicate checkpoint keys in {path}: {sorted(overlap)[:5]}")
-        state_dict.update(shard)
-    return state_dict
-
-
 def get_safe_dtype(target_dtype, device_type):
     """Get a safe dtype for the given device type."""
     if device_type == "mps" and target_dtype == torch.float64:
@@ -645,12 +563,6 @@ class PaliGemmaWithExpertModel(
 class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
     """Core PI05 PyTorch model."""
 
-    use_hf_vision_checkpointing_api = False
-    checkpoint_vision_embeddings = True
-    use_typed_attention_masks = False
-    use_on_device_suffix_mask = False
-    precompute_denoise_times = False
-
     def __init__(self, config: PI05Config, rtc_processor: RTCProcessor | None = None):
         super().__init__()
         self.config = config
@@ -694,11 +606,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         """Enable gradient checkpointing for memory optimization."""
         self.gradient_checkpointing_enabled = True
         self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = True
-        vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
-        if self.use_hf_vision_checkpointing_api:
-            vision_tower.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
-        else:
-            vision_tower.gradient_checkpointing = True
+        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = True
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = True
         logging.info("Enabled gradient checkpointing for PI05Pytorch model")
 
@@ -706,11 +614,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         """Disable gradient checkpointing."""
         self.gradient_checkpointing_enabled = False
         self.paligemma_with_expert.paligemma.model.language_model.gradient_checkpointing = False
-        vision_tower = self.paligemma_with_expert.paligemma.model.vision_tower
-        if self.use_hf_vision_checkpointing_api:
-            vision_tower.gradient_checkpointing_disable()
-        else:
-            vision_tower.gradient_checkpointing = False
+        self.paligemma_with_expert.paligemma.model.vision_tower.gradient_checkpointing = False
         self.paligemma_with_expert.gemma_expert.model.gradient_checkpointing = False
         logging.info("Disabled gradient checkpointing for PI05Pytorch model")
 
@@ -725,13 +629,10 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             )
         return func(*args, **kwargs)
 
-    def _prepare_attention_masks_4d(self, att_2d_masks, dtype=None):
+    def _prepare_attention_masks_4d(self, att_2d_masks):
         """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
+        return torch.where(att_2d_masks_4d, 0.0, OPENPI_ATTENTION_MASK_VALUE)
 
     def sample_noise(self, shape, device):
         return torch.normal(
@@ -757,16 +658,13 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         pad_masks = []
         att_masks = []
 
-        if self.checkpoint_vision_embeddings:
+        # Process images
+        for img, img_mask in zip(images, img_masks, strict=True):
 
-            def embed_image(img):
-                return self._apply_checkpoint(self.paligemma_with_expert.embed_image, img)
+            def image_embed_func(img):
+                return self.paligemma_with_expert.embed_image(img)
 
-            img_embs = [embed_image(img) for img in images]
-        else:
-            img_embs = [self.paligemma_with_expert.embed_image(img) for img in images]
-
-        for img_emb, img_mask in zip(img_embs, img_masks, strict=True):
+            img_emb = self._apply_checkpoint(image_embed_func, img)
             bsize, num_img_embs = img_emb.shape[:2]
 
             embs.append(img_emb)
@@ -836,12 +734,6 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         embs = torch.cat(embs, dim=1)
         pad_masks = torch.cat(pad_masks, dim=1)
-        if self.use_on_device_suffix_mask:
-            n = len(att_masks)
-            att_masks = torch.zeros(n, dtype=embs.dtype, device=embs.device)
-            att_masks[0] = 1
-            att_masks = att_masks[None, :].expand(bsize, n)
-        else:
         att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
         att_masks = att_masks[None, :].expand(bsize, len(att_masks))
 
@@ -927,8 +819,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
         prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1
 
-        mask_dtype = prefix_embs.dtype if self.use_typed_attention_masks else None
-        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks, dtype=mask_dtype)
+        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(prefix_att_2d_masks)
         self.paligemma_with_expert.paligemma.model.language_model.config._attn_implementation = "eager"  # noqa: SLF001
 
         _, past_key_values = self.paligemma_with_expert.forward(
@@ -941,19 +832,10 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         dt = -1.0 / num_steps
 
-        times = None
-        if self.precompute_denoise_times:
-            times = torch.tensor(
-                [1.0 + step * dt for step in range(num_steps)], dtype=torch.float32, device=device
-            )
-
         x_t = noise
         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(
@@ -1031,9 +913,6 @@ class PI05Policy(PreTrainedPolicy):
 
     config_class = PI05Config
     name = "pi05"
-    model_class = PI05Pytorch
-    eval_after_pretrained_load = False
-    show_openpi_disclaimer = True
 
     def __init__(
         self,
@@ -1051,7 +930,7 @@ class PI05Policy(PreTrainedPolicy):
 
         # Initialize the core PI05 model
         self.init_rtc_processor()
-        self.model = self.model_class(config, rtc_processor=self.rtc_processor)
+        self.model = PI05Pytorch(config, rtc_processor=self.rtc_processor)
 
         # Enable gradient checkpointing if requested
         if config.gradient_checkpointing:
@@ -1077,8 +956,7 @@ class PI05Policy(PreTrainedPolicy):
         strict: bool = True,
         **kwargs,
     ) -> T:
-        """Load PI05-compatible single-file or sharded safetensors checkpoints."""
-        if cls.show_openpi_disclaimer:
+        """Override the from_pretrained method to handle key remapping and display important disclaimer."""
         print(
             "The PI05 model is a direct port of the OpenPI implementation. \n"
             "This implementation follows the original OpenPI structure for compatibility. \n"
@@ -1087,6 +965,7 @@ 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,
@@ -1100,34 +979,84 @@ class PI05Policy(PreTrainedPolicy):
                 **kwargs,
             )
 
+        # Initialize model without loading weights
+        # Check if dataset_stats were provided in kwargs
         model = cls(config, **kwargs)
-        files = _resolve_weight_files(
+
+        # Load state dict (expects keys with "model." prefix)
+        try:
+            print(f"Loading model from: {pretrained_name_or_path}")
+            try:
+                from transformers.utils import cached_file
+
+                resolved_file = cached_file(
                     pretrained_name_or_path,
-            force_download=force_download,
-            resume_download=resume_download,
-            proxies=proxies,
-            token=token,
-            cache_dir=cache_dir,
-            local_files_only=local_files_only,
-            revision=revision,
+                    "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),
                 )
-        fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), model.config)
-        remapped_state_dict = {
-            key if key.startswith("model.") else f"model.{key}": value
-            for key, value in fixed_state_dict.items()
-        }
-        remapped_state_dict = model._prepare_pretrained_state_dict(remapped_state_dict)
-        missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
-        if missing_keys:
-            logging.warning("Missing %s checkpoint keys: %s", cls.name, missing_keys)
-        if unexpected_keys:
-            logging.warning("Unexpected %s checkpoint keys: %s", cls.name, unexpected_keys)
-        if model.eval_after_pretrained_load:
-            model.eval()
+                from safetensors.torch import load_file
+
+                original_state_dict = load_file(resolved_file)
+                print("✓ Loaded state dict from model.safetensors")
+            except Exception as e:
+                print(f"Could not load state dict from remote files: {e}")
+                print("Returning model without loading pretrained weights")
                 return model
 
-    def _prepare_pretrained_state_dict(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
-        return state_dict
+            # First, fix any key differences (see openpi model.py, _fix_pytorch_state_dict_keys)
+            fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)
+
+            # Then add "model." prefix for all keys that don't already have it
+            remapped_state_dict = {}
+            remap_count = 0
+
+            for key, value in fixed_state_dict.items():
+                if not key.startswith("model."):
+                    new_key = f"model.{key}"
+                    remapped_state_dict[new_key] = value
+                    remap_count += 1
+                else:
+                    remapped_state_dict[key] = value
+
+            if remap_count > 0:
+                print(f"Remapped {remap_count} state dict keys")
+
+            # Load the remapped state dict into the model
+            missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
+
+            if missing_keys:
+                print(f"Missing keys when loading state dict: {len(missing_keys)} keys")
+                if len(missing_keys) <= 5:
+                    for key in missing_keys:
+                        print(f"  - {key}")
+                else:
+                    for key in missing_keys[:5]:
+                        print(f"  - {key}")
+                    print(f"  ... and {len(missing_keys) - 5} more")
+
+            if unexpected_keys:
+                print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
+                if len(unexpected_keys) <= 5:
+                    for key in unexpected_keys:
+                        print(f"  - {key}")
+                else:
+                    for key in unexpected_keys[:5]:
+                        print(f"  - {key}")
+                    print(f"  ... and {len(unexpected_keys) - 5} more")
+
+            if not missing_keys and not unexpected_keys:
+                print("All keys loaded successfully!")
+
+        except Exception as e:
+            print(f"Warning: Could not load state dict: {e}")
+
+        return model
 
     def _fix_pytorch_state_dict_keys(
         self, state_dict, model_config
@@ -1299,16 +1228,12 @@ class PI05Policy(PreTrainedPolicy):
 
         # Action queue logic for n_action_steps > 1
         if len(self._action_queue) == 0:
-            action_batch = self._prepare_action_batch(batch)
-            actions = self.predict_action_chunk(action_batch)[:, : self.config.n_action_steps]
+            actions = self.predict_action_chunk(batch)[:, : self.config.n_action_steps]
             # Transpose to get shape (n_action_steps, batch_size, action_dim)
             self._action_queue.extend(actions.transpose(0, 1))
 
         return self._action_queue.popleft()
 
-    def _prepare_action_batch(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
-        return batch
-
     @torch.no_grad()
     def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
         """Predict a chunk of actions given environment observations."""

src/lerobot/policies/pi05/processor_pi05.py

@@ -24,26 +24,17 @@ import torch
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
     RelativeActionsProcessorStep,
-    RenameObservationsProcessorStep,
     TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
 from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import (
-    OBS_STATE,
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
-)
+from lerobot.utils.constants import OBS_STATE
 
 from .configuration_pi05 import PI05Config
 
@@ -135,18 +126,16 @@ def make_pi05_pre_post_processors(
         action_names=getattr(config, "action_feature_names", None),
     )
 
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     # OpenPI order: raw → relative → normalize → model → unnormalize → absolute
     input_steps: list[ProcessorStep] = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
+        steps.rename_observations,  # To mimic the same processor as pretrained one
+        steps.add_batch_dim,
         relative_step,
         # NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
         # because the tokenizer step expects normalized state in [-1, 1] range for discretization
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        steps.normalize,
         Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
         TokenizerProcessorStep(
             tokenizer_name="google/paligemma-3b-pt-224",
@@ -154,26 +143,13 @@ def make_pi05_pre_post_processors(
             padding_side="right",
             padding="max_length",
         ),
-        DeviceProcessorStep(device=config.device),
+        steps.to_device,
     ]
 
     output_steps: list[ProcessorStep] = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
+        steps.unnormalize,
         AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
-        DeviceProcessorStep(device="cpu"),
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/pi052/__init__.py

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

src/lerobot/policies/pi052/configuration_pi052.py

@@ -1,195 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""PI0.5 with hierarchical text generation and flow-matched actions."""
-
-from dataclasses import dataclass
-
-from lerobot.configs import PreTrainedConfig
-from lerobot.optim.optimizers import AdamWConfig
-
-from ..pi05.configuration_pi05 import PI05Config
-
-
-@PreTrainedConfig.register_subclass("pi052")
-@dataclass
-class PI052Config(PI05Config):
-    """PI0.5 configuration for recipe-driven text and action supervision."""
-
-    # Recipe / language stack ---------------------------------------------
-    recipe_path: str | None = "recipes/subtask_mem.yaml"
-    """Recipe path relative to ``src/lerobot/configs/``, or ``None`` for the plain PI0.5 prompt."""
-
-    apply_chat_template: bool = False
-    """Whether to apply a tokenizer chat template.
-
-    PaliGemma defaults to plain recipe-rendered prefixes because it is not chat-pretrained.
-    """
-
-    # Balance frequent recipe text supervision against the paper's α=10 flow weight.
-    text_loss_weight: float = 1.0
-    """LM-head cross-entropy weight; ``0`` disables text training."""
-
-    flow_loss_weight: float = 10.0
-    """Weight on action-expert flow matching relative to text supervision."""
-
-    # Backbone training ---------------------------------------------------
-    unfreeze_lm_head: bool = True
-    """Keep PaliGemma's language head trainable for hierarchical inference."""
-
-    # Optional context dropout improves tolerance to missing or stale language state.
-    plan_dropout_prob: float = 0.0
-    memory_dropout_prob: float = 0.0
-    subtask_dropout_prob: float = 0.0
-
-    # FAST adds discrete-action CE to the text and flow objectives from paper §III.B-C.
-    enable_fast_action_loss: bool = True
-    """Add FAST-tokenized action cross-entropy to text CE and flow matching."""
-
-    action_tokenizer_name: str = "physical-intelligence/fast"
-    """HF identifier for the FAST action tokenizer."""
-
-    max_action_tokens: int = 256
-    """Maximum number of FAST tokens per action chunk."""
-
-    fast_skip_tokens: int = 1152
-    """Number of top-of-vocab tokens the FAST id mapping skips.
-
-    1152 skips PaliGemma's 128 ``<seg>`` and 1024 ``<loc>`` special tokens so
-    FAST codes land in plain-text ids below 256000 and never collide with the
-    ``<loc>`` targets used for VQA. openpi's pi0-FAST convention is 128 (FAST
-    occupies the ``<loc>`` range); use 128 only to stay weight-compatible with
-    checkpoints trained that way."""
-
-    fast_action_loss_weight: float = 1.0
-    """Weight on FAST action-token CE relative to continuous-flow supervision."""
-
-    subtask_replan_steps: int = 0
-    """Environment steps between subtask generations during evaluation.
-
-    Non-positive values regenerate each action chunk while still refreshing the action prompt every chunk.
-    """
-
-    joint_subtask_conditioning: bool = False
-    """Condition low-level action inference on the task plus the generated subtask.
-
-    Matches paper-style joint-sequence recipes (``recipes/subtask_joint.yaml``)
-    where one sample supervises the subtask text and conditions the action
-    losses on it: the inference prefix becomes
-    ``User: {task}, State: ...;\\nAssistant: {subtask}<eos>`` with the subtask
-    span attended causally, exactly as trained. Leave ``False`` for the blend
-    recipes, whose low-level samples use ``User: {subtask}, State: ...;``."""
-
-    auto_fit_fast_tokenizer: bool = False
-    """Fit and cache a dataset-specific FAST tokenizer before training.
-
-    Disabled by default to avoid the extra dataset pass and use the universal tokenizer.
-    """
-
-    fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
-    """Where fitted FAST tokenizers are stored. ``~`` expands."""
-
-    fast_tokenizer_fit_samples: int = 1024
-    """Number of action chunks sampled when fitting FAST."""
-
-    fast_tokenizer_validation_samples: int = 256
-    """Held-out action chunks used to validate tokenizer reconstruction."""
-
-    fast_tokenizer_max_reconstruction_rmse: float = 0.10
-    """Maximum normalized RMSE allowed across held-out action chunks."""
-
-    fast_tokenizer_max_dim_rmse: float = 0.20
-    """Maximum normalized RMSE allowed for any nonconstant action dimension."""
-
-    # Knowledge insulation detaches VLM K/V from action-loss gradients (paper §III.B).
-    knowledge_insulation: bool = True
-    """Block action-loss gradients through VLM keys and values."""
-
-    # Optional training backends. Defaults preserve the eager/SDPA path.
-    use_flashrt_adarms: bool = False
-    """Use FlashRT adaptive RMSNorm kernels when available."""
-
-    use_compiled_text_ce: bool = False
-    """Compile the materialized-logits text and FAST CE path."""
-
-    use_compiled_vision: bool = False
-    """Compile the SigLIP tower for no-grad flow and inference passes."""
-
-    use_flex_attention: bool = False
-    """Use FlexAttention for amortized KI, with SDPA fallback where unsupported."""
-
-    use_manual_attention: bool = False
-    """Use materialized-logits attention for explicitly profiled KI shapes."""
-
-    manual_attention_scope: str = "all"
-    """Apply manual attention to all KI queries or only action queries."""
-
-    # Scale language-head updates relative to the base optimizer schedule.
-    lm_head_lr_scale: float = 1.0
-
-    # Scale backbone and action-expert optimizer groups independently.
-    backbone_lr_scale: float = 1.0
-    action_expert_lr_scale: float = 1.0
-
-    # Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow.
-    flow_num_repeats: int = 5
-
-    # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
-    text_ce_z_loss_weight: float = 1e-4
-
-    use_flashrt_fp8_mlp: bool = False
-    """Enable calibrated FlashRT FP8 kernels for Gemma and SigLIP MLPs.
-
-    Apply after loading with ``PI052Policy.apply_flashrt_fp8_mlp``; unavailable kernels keep BF16.
-    """
-
-    # Keep serialized PI052 AdamW options local because PI05Config lacks them.
-    optimizer_foreach: bool | None = False
-    optimizer_fused: bool | None = True
-
-    def get_optimizer_preset(self) -> AdamWConfig:
-        return AdamWConfig(
-            lr=self.optimizer_lr,
-            betas=self.optimizer_betas,
-            eps=self.optimizer_eps,
-            weight_decay=self.optimizer_weight_decay,
-            grad_clip_norm=self.optimizer_grad_clip_norm,
-            foreach=self.optimizer_foreach,
-            fused=self.optimizer_fused,
-        )
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        if self.enable_fast_action_loss and not self.recipe_path:
-            raise ValueError("PI052 FAST action loss requires recipe_path to build action supervision.")
-        if self.text_loss_weight > 0 and self.unfreeze_lm_head:
-            self.train_expert_only = False
-        if self.flow_num_repeats < 1:
-            raise ValueError(f"flow_num_repeats must be >= 1, got {self.flow_num_repeats}")
-        if self.fast_tokenizer_validation_samples < 1:
-            raise ValueError("fast_tokenizer_validation_samples must be >= 1")
-        if self.fast_tokenizer_max_reconstruction_rmse <= 0 or self.fast_tokenizer_max_dim_rmse <= 0:
-            raise ValueError("FAST tokenizer reconstruction thresholds must be positive")
-        if self.manual_attention_scope not in {"all", "action"}:
-            raise ValueError(
-                f"manual_attention_scope must be 'all' or 'action', got {self.manual_attention_scope!r}"
-            )
-        if self.use_flex_attention and self.use_manual_attention:
-            raise ValueError("use_flex_attention and use_manual_attention are mutually exclusive")
-        if self.use_flex_attention and self.flow_num_repeats == 1:
-            raise ValueError("use_flex_attention requires flow_num_repeats > 1")
-        if not self.knowledge_insulation and (
-            self.use_flex_attention or self.use_manual_attention or self.use_flashrt_adarms
-        ):
-            raise ValueError("KI attention and AdaRMS optimizations require knowledge_insulation=True")

src/lerobot/policies/pi052/fit_fast_tokenizer.py

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

src/lerobot/policies/pi052/flashrt_fp8.py

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

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

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

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

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

src/lerobot/policies/pi052/processor_pi052.py

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

src/lerobot/policies/pi052/text_processor_pi052.py

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

src/lerobot/policies/pi0_fast/configuration_pi0_fast.py

@@ -61,9 +61,6 @@ class PI0FastConfig(PreTrainedConfig):
     tokenizer_max_length: int = 200  # see openpi `__post_init__`
     text_tokenizer_name: str = "google/paligemma-3b-pt-224"
     action_tokenizer_name: str = "lerobot/fast-action-tokenizer"
-    auto_fit_fast_tokenizer: bool = False
-    fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
-    fast_tokenizer_fit_samples: int = 1024
     temperature: float = 0.0
     max_decoding_steps: int = 256
     fast_skip_tokens: int = 128

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -25,26 +25,17 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     ActionTokenizerProcessorStep,
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
     RelativeActionsProcessorStep,
-    RenameObservationsProcessorStep,
     TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
 from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import (
-    OBS_STATE,
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
-)
+from lerobot.utils.constants import OBS_STATE
 
 from .configuration_pi0_fast import PI0FastConfig
 
@@ -101,11 +92,6 @@ class Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(ProcessorStep):
 def make_pi0_fast_pre_post_processors(
     config: PI0FastConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    dataset_repo_id: str | None = None,
-    dataset_root: str | None = None,
-    dataset_revision: str | None = None,
-    episodes: list[int] | None = None,
-    exclude_episodes: list[int] | None = None,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -140,6 +126,8 @@ def make_pi0_fast_pre_post_processors(
         action_names=getattr(config, "action_feature_names", None),
     )
 
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     # Pi0Fast order: relative → normalize → tokenize → model → unnormalize → absolute
     # This matches pi0/pi0.5: RelativeActionsProcessorStep runs first on raw absolute actions,
     # caching the raw state. NormalizerProcessorStep then normalizes the raw relative actions,
@@ -148,27 +136,11 @@ 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] = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
+        steps.rename_observations,  # To mimic the same processor as pretrained one
+        steps.add_batch_dim,
         relative_step,
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        steps.normalize,
         Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(max_state_dim=config.max_state_dim),
         TokenizerProcessorStep(
             tokenizer_name=config.text_tokenizer_name,
@@ -177,31 +149,18 @@ def make_pi0_fast_pre_post_processors(
             padding="max_length",
         ),
         ActionTokenizerProcessorStep(
-            action_tokenizer_name=action_tokenizer_path,
+            action_tokenizer_name=config.action_tokenizer_name,
             max_action_tokens=config.max_action_tokens,
             fast_skip_tokens=config.fast_skip_tokens,
             paligemma_tokenizer_name=config.text_tokenizer_name,
         ),
-        DeviceProcessorStep(device=config.device),
+        steps.to_device,
     ]
 
     output_steps: list[ProcessorStep] = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
+        steps.unnormalize,
         AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
-        DeviceProcessorStep(device="cpu"),
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/pi_gemma.py

@@ -14,27 +14,18 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Literal
+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
 
-# Default PaliGemma SigLIP input resolution. Mirrors
-# ``pi05.configuration_pi05.DEFAULT_IMAGE_SIZE``; duplicated as a plain constant
-# to avoid importing the pi05 package here (which would create an import cycle:
-# pi_gemma -> pi05.__init__ -> modeling_pi05 -> pi_gemma).
-DEFAULT_IMAGE_SIZE = 224
-
 if TYPE_CHECKING or _transformers_available:
     from transformers.cache_utils import DynamicCache
     from transformers.masking_utils import create_causal_mask
     from transformers.modeling_layers import GradientCheckpointingLayer
     from transformers.modeling_outputs import BaseModelOutputWithPast
-    from transformers.models.auto import CONFIG_MAPPING
-    from transformers.models.gemma import modeling_gemma
     from transformers.models.gemma.modeling_gemma import (
         GemmaAttention,
         GemmaConfig,
@@ -58,8 +49,6 @@ else:
     GradientCheckpointingLayer = None
     BaseModelOutputWithPast = None
     create_causal_mask = None
-    CONFIG_MAPPING = None
-    modeling_gemma = None
 
 
 def _gated_residual(
@@ -132,10 +121,7 @@ class PiGemmaRMSNorm(nn.Module):
         if cond.shape[-1] != self.cond_dim:
             raise ValueError(f"Expected cond dim {self.cond_dim}, got {cond.shape[-1]}")
         modulation = self.dense(cond)
-        # Per-sample cond (B, cond_dim) → broadcast over the sequence. A
-        # per-token cond (B, T, cond_dim) is already aligned with x and must
-        # not be unsqueezed (used by pi052's amortized K_repeat path).
-        if len(x.shape) == 3 and modulation.dim() == 2:
+        if len(x.shape) == 3:
             modulation = modulation.unsqueeze(1)
         scale, shift, gate = modulation.chunk(3, dim=-1)
         normed = normed * (1 + scale.float()) + shift.float()
@@ -289,8 +275,6 @@ class PiGemmaModel(GemmaModel):  # type: ignore[misc]
         # Convert to bfloat16 if the first layer uses bfloat16
         if len(self.layers) > 0 and self.layers[0].self_attn.q_proj.weight.dtype == torch.bfloat16:
             hidden_states = hidden_states.to(torch.bfloat16)
-        if causal_mask is not None and torch.is_floating_point(causal_mask):
-            causal_mask = causal_mask.to(dtype=hidden_states.dtype)
 
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -383,374 +367,3 @@ __all__ = [
     "PaliGemmaModelWithPiGemma",
     "PaliGemmaForConditionalGenerationWithPiGemma",
 ]
-
-
-# PI0.5 / PI052 dual-expert backbone: generic PaliGemma + Gemma action-expert
-# transformer machinery used by the pi052 policy. GemmaVariantConfig is openpi's
-# width/depth variant config (renamed from GemmaConfig to avoid clashing with
-# transformers' GemmaConfig).
-
-
-def sdpa_attention_forward(
-    module,
-    query: torch.Tensor,
-    key: torch.Tensor,
-    value: torch.Tensor,
-    attention_mask: torch.Tensor | None,
-    scaling: float,
-    dropout: float = 0.0,
-):
-    """Drop-in for ``modeling_gemma.eager_attention_forward`` using
-    ``torch.nn.functional.scaled_dot_product_attention``.
-
-    PyTorch SDPA picks the memory-efficient kernel for arbitrary additive
-    bias masks (the FA backend only accepts causal/sliding-window). On
-    H100 that is ~1.3-1.7x faster and uses ~30-40% less attention memory
-    than the eager softmax(QK^T)+matmul path. Mirrors eager's signature
-    and output shape (``(B, Lq, H, D)``) so call sites are unchanged.
-    """
-    n_rep = module.num_key_value_groups
-    if n_rep > 1:
-        key = key.repeat_interleave(n_rep, dim=1)
-        value = value.repeat_interleave(n_rep, dim=1)
-    if attention_mask is not None and attention_mask.dtype != query.dtype:
-        attention_mask = attention_mask.to(dtype=query.dtype)
-    attn_output = F.scaled_dot_product_attention(
-        query,
-        key,
-        value,
-        attn_mask=attention_mask,
-        dropout_p=dropout if module.training else 0.0,
-        is_causal=False,
-        scale=scaling,
-    )
-    return attn_output.transpose(1, 2).contiguous(), None
-
-
-# Define the complete layer computation function for gradient checkpointing
-def compute_layer_complete(
-    layer_idx, inputs_embeds, attention_mask, position_ids, adarms_cond, paligemma, gemma_expert
-):
-    models = [paligemma.model.language_model, gemma_expert.model]
-    query_states = []
-    key_states = []
-    value_states = []
-    gates = []
-    for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
-        hidden_states, gate = layernorm_forward(layer.input_layernorm, hidden_states, adarms_cond[i])
-        gates.append(gate)
-        input_shape = hidden_states.shape[:-1]
-        hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
-        query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-        key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-        value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
-        query_states.append(query_state)
-        key_states.append(key_state)
-        value_states.append(value_state)
-    # Concatenate and process attention
-    query_states = torch.cat(query_states, dim=2)
-    key_states = torch.cat(key_states, dim=2)
-    value_states = torch.cat(value_states, dim=2)
-    dummy_tensor = torch.zeros(
-        query_states.shape[0],
-        query_states.shape[2],
-        query_states.shape[-1],
-        device=query_states.device,
-        dtype=query_states.dtype,
-    )
-    cos, sin = paligemma.model.language_model.rotary_emb(dummy_tensor, position_ids)
-    query_states, key_states = modeling_gemma.apply_rotary_pos_emb(
-        query_states, key_states, cos, sin, unsqueeze_dim=1
-    )
-    batch_size = query_states.shape[0]
-    scaling = paligemma.model.language_model.layers[layer_idx].self_attn.scaling
-    att_output, _ = sdpa_attention_forward(
-        paligemma.model.language_model.layers[layer_idx].self_attn,
-        query_states,
-        key_states,
-        value_states,
-        attention_mask,
-        scaling,
-    )
-    # Get head_dim from the current layer, not from the model
-    head_dim = paligemma.model.language_model.layers[layer_idx].self_attn.head_dim
-    att_output = att_output.reshape(batch_size, -1, 1 * 8 * head_dim)
-    # Process layer outputs
-    outputs_embeds = []
-    start_pos = 0
-    for i, hidden_states in enumerate(inputs_embeds):
-        layer = models[i].layers[layer_idx]
-        end_pos = start_pos + hidden_states.shape[1]
-        if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
-            att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
-        out_emb = layer.self_attn.o_proj(att_output[:, start_pos:end_pos])
-        # first residual
-        out_emb = _gated_residual(hidden_states, out_emb, gates[i])
-        after_first_residual = out_emb.clone()
-        out_emb, gate = layernorm_forward(layer.post_attention_layernorm, out_emb, adarms_cond[i])
-        # Convert to bfloat16 if the next layer (mlp) uses bfloat16
-        if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
-            out_emb = out_emb.to(dtype=torch.bfloat16)
-        out_emb = layer.mlp(out_emb)
-        # second residual
-        out_emb = _gated_residual(after_first_residual, out_emb, gate)
-        outputs_embeds.append(out_emb)
-        start_pos = end_pos
-    return outputs_embeds
-
-
-class GemmaVariantConfig:  # see openpi `gemma.py: Config`
-    """Configuration for Gemma model variants."""
-
-    def __init__(self, width, depth, mlp_dim, num_heads, num_kv_heads, head_dim):
-        self.width = width
-        self.depth = depth
-        self.mlp_dim = mlp_dim
-        self.num_heads = num_heads
-        self.num_kv_heads = num_kv_heads
-        self.head_dim = head_dim
-
-
-def get_gemma_config(variant: str) -> GemmaVariantConfig:  # see openpi `gemma.py: get_config`
-    """Returns config for specified gemma variant."""
-    if variant == "gemma_300m":
-        return GemmaVariantConfig(
-            width=1024,
-            depth=18,
-            mlp_dim=4096,
-            num_heads=8,
-            num_kv_heads=1,
-            head_dim=256,
-        )
-    elif variant == "gemma_2b":
-        return GemmaVariantConfig(
-            width=2048,
-            depth=18,
-            mlp_dim=16_384,
-            num_heads=8,
-            num_kv_heads=1,
-            head_dim=256,
-        )
-    else:
-        raise ValueError(f"Unknown variant: {variant}")
-
-
-class PaliGemmaWithExpertModel(
-    nn.Module
-):  # see openpi `gemma_pytorch.py: PaliGemmaWithExpertModel` this class is almost a exact copy of PaliGemmaWithExpertModel in openpi
-    """PaliGemma model with action expert for PI05."""
-
-    def __init__(
-        self,
-        vlm_config,
-        action_expert_config,
-        use_adarms=None,
-        precision: Literal["bfloat16", "float32"] = "bfloat16",
-        image_size: int = DEFAULT_IMAGE_SIZE,
-        freeze_vision_encoder: bool = False,
-        train_expert_only: bool = False,
-    ):
-        if use_adarms is None:
-            use_adarms = [False, False]
-        super().__init__()
-        self.freeze_vision_encoder = freeze_vision_encoder
-        self.train_expert_only = train_expert_only
-
-        vlm_config_hf = CONFIG_MAPPING["paligemma"]()
-        vlm_config_hf._vocab_size = 257152  # noqa: SLF001
-        vlm_config_hf.image_token_index = 257152
-        vlm_config_hf.text_config.hidden_size = vlm_config.width
-        vlm_config_hf.text_config.intermediate_size = vlm_config.mlp_dim
-        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_heads
-        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
-        vlm_config_hf.text_config.num_hidden_layers = vlm_config.depth
-        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_kv_heads
-        vlm_config_hf.text_config.hidden_activation = "gelu_pytorch_tanh"
-        vlm_config_hf.text_config.dtype = "float32"
-        vlm_config_hf.text_config.vocab_size = 257152
-        vlm_config_hf.text_config.use_adarms = use_adarms[0]
-        vlm_config_hf.text_config.adarms_cond_dim = vlm_config.width if use_adarms[0] else None
-        vlm_config_hf.vision_config.image_size = image_size
-        vlm_config_hf.vision_config.intermediate_size = 4304
-        vlm_config_hf.vision_config.projection_dim = 2048
-        vlm_config_hf.vision_config.projector_hidden_act = "gelu_fast"
-        vlm_config_hf.vision_config.dtype = "float32"
-
-        action_expert_config_hf = CONFIG_MAPPING["gemma"](
-            head_dim=action_expert_config.head_dim,
-            hidden_size=action_expert_config.width,
-            intermediate_size=action_expert_config.mlp_dim,
-            num_attention_heads=action_expert_config.num_heads,
-            num_hidden_layers=action_expert_config.depth,
-            num_key_value_heads=action_expert_config.num_kv_heads,
-            vocab_size=257152,
-            hidden_activation="gelu_pytorch_tanh",
-            dtype="float32",
-            use_adarms=use_adarms[1],
-            adarms_cond_dim=action_expert_config.width if use_adarms[1] else None,
-        )
-
-        self.paligemma = PaliGemmaForConditionalGenerationWithPiGemma(config=vlm_config_hf)
-        self.gemma_expert = PiGemmaForCausalLM(config=action_expert_config_hf)
-        self.gemma_expert.model.embed_tokens = None
-
-        self.to_bfloat16_for_selected_params(precision)
-        self._set_requires_grad()
-
-    def to_bfloat16_for_selected_params(self, precision: Literal["bfloat16", "float32"] = "bfloat16"):
-        if precision == "bfloat16":
-            self.to(dtype=torch.bfloat16)
-        elif precision == "float32":
-            self.to(dtype=torch.float32)
-            return
-        else:
-            raise ValueError(f"Invalid precision: {precision}")
-
-        # Keep full vision path in float32 so we never toggle (toggle causes optimizer
-        # "same dtype" error). Saves memory vs full float32; more memory than only 3 params.
-        params_to_keep_float32 = [
-            "vision_tower",
-            "multi_modal_projector",
-            "lm_head",
-            "input_layernorm",
-            "post_attention_layernorm",
-            "model.norm",
-        ]
-
-        for name, param in self.named_parameters():
-            if any(selector in name for selector in params_to_keep_float32):
-                param.data = param.data.to(dtype=torch.float32)
-
-    def _set_requires_grad(self):
-        if self.freeze_vision_encoder:
-            self.paligemma.model.vision_tower.eval()
-            for param in self.paligemma.model.vision_tower.parameters():
-                param.requires_grad = False
-        if self.train_expert_only:
-            self.paligemma.eval()
-            for param in self.paligemma.parameters():
-                param.requires_grad = False
-
-    def train(self, mode: bool = True):
-        super().train(mode)
-        if self.freeze_vision_encoder:
-            self.paligemma.model.vision_tower.eval()
-        if self.train_expert_only:
-            self.paligemma.eval()
-
-    def embed_image(self, image: torch.Tensor):
-        # Vision tower and multi_modal_projector are kept in float32 (params_to_keep_float32).
-        out_dtype = image.dtype
-        if image.dtype != torch.float32:
-            image = image.to(torch.float32)
-        image_outputs = self.paligemma.model.get_image_features(image)
-        # OpenPI / big_vision convention: image (soft) tokens are NOT scaled by the
-        # Gemma embedder normalizer (sqrt(hidden_size)) — only text tokens are. lerobot/pi05_base
-        # was trained in this regime, so scaling image features here over-scales them ~45x and
-        # breaks the pretrained vision-language alignment. Keep image features un-normalized.
-        features = image_outputs.pooler_output
-        if features.dtype != out_dtype:
-            features = features.to(out_dtype)
-        return features
-
-    def embed_language_tokens(self, tokens: torch.Tensor):
-        return self.paligemma.model.language_model.embed_tokens(tokens)
-
-    def forward(
-        self,
-        attention_mask: torch.Tensor | None = None,
-        position_ids: torch.LongTensor | None = None,
-        past_key_values: list[torch.FloatTensor] | None = None,
-        inputs_embeds: list[torch.FloatTensor] | None = None,
-        use_cache: bool | None = None,
-        adarms_cond: list[torch.Tensor] | None = None,
-    ):
-        if adarms_cond is None:
-            adarms_cond = [None, None]
-        if inputs_embeds[1] is None:
-            prefix_output = self.paligemma.model.language_model.forward(
-                inputs_embeds=inputs_embeds[0],
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                use_cache=use_cache,
-                adarms_cond=adarms_cond[0] if adarms_cond is not None else None,
-            )
-            prefix_past_key_values = prefix_output.past_key_values
-            prefix_output = prefix_output.last_hidden_state
-            suffix_output = None
-        elif inputs_embeds[0] is None:
-            suffix_output = self.gemma_expert.model.forward(
-                inputs_embeds=inputs_embeds[1],
-                attention_mask=attention_mask,
-                position_ids=position_ids,
-                past_key_values=past_key_values,
-                use_cache=use_cache,
-                adarms_cond=adarms_cond[1] if adarms_cond is not None else None,
-            )
-            suffix_output = suffix_output.last_hidden_state
-            prefix_output = None
-            prefix_past_key_values = None
-        else:
-            models = [self.paligemma.model.language_model, self.gemma_expert.model]
-            num_layers = self.paligemma.config.text_config.num_hidden_layers
-
-            # Check if gradient checkpointing is enabled for any of the models
-            use_gradient_checkpointing = (
-                hasattr(self.gemma_expert.model, "gradient_checkpointing")
-                and self.gemma_expert.model.gradient_checkpointing
-                and self.training
-            ) or (hasattr(self, "gradient_checkpointing") and self.gradient_checkpointing and self.training)
-
-            # Process all layers with gradient checkpointing if enabled
-            for layer_idx in range(num_layers):
-                if use_gradient_checkpointing:
-                    inputs_embeds = torch.utils.checkpoint.checkpoint(
-                        compute_layer_complete,
-                        layer_idx,
-                        inputs_embeds,
-                        attention_mask,
-                        position_ids,
-                        adarms_cond,
-                        use_reentrant=False,
-                        preserve_rng_state=False,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
-                    )
-                else:
-                    inputs_embeds = compute_layer_complete(
-                        layer_idx,
-                        inputs_embeds,
-                        attention_mask,
-                        position_ids,
-                        adarms_cond,
-                        paligemma=self.paligemma,
-                        gemma_expert=self.gemma_expert,
-                    )
-
-            # final norm
-            def compute_final_norms(inputs_embeds, adarms_cond):
-                outputs_embeds = []
-                for i, hidden_states in enumerate(inputs_embeds):
-                    out_emb, _ = layernorm_forward(models[i].norm, hidden_states, adarms_cond[i])
-                    outputs_embeds.append(out_emb)
-                return outputs_embeds
-
-            # Apply gradient checkpointing to final norm if enabled
-            if use_gradient_checkpointing:
-                outputs_embeds = torch.utils.checkpoint.checkpoint(
-                    compute_final_norms,
-                    inputs_embeds,
-                    adarms_cond,
-                    use_reentrant=False,
-                    preserve_rng_state=False,
-                )
-            else:
-                outputs_embeds = compute_final_norms(inputs_embeds, adarms_cond)
-
-            prefix_output = outputs_embeds[0]
-            suffix_output = outputs_embeds[1]
-            prefix_past_key_values = None
-
-        return [prefix_output, suffix_output], prefix_past_key_values

src/lerobot/policies/pretrained.py

@@ -23,8 +23,6 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
 
-import packaging
-import safetensors
 from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
@@ -34,6 +32,7 @@ from torch import Tensor, nn
 from lerobot.__version__ import __version__
 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 .utils import log_model_loading_keys
@@ -221,26 +220,10 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        # Create base kwargs
-        kwargs = {"strict": strict}
-
-        # Add device parameter for newer versions that support it
-        if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
-            kwargs["device"] = map_location
-
-        # Load the model with appropriate kwargs
-        missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
-        log_model_loading_keys(missing_keys, unexpected_keys)
-
-        # For older versions, manually move to device if needed
-        if "device" not in kwargs and map_location != "cpu":
-            logging.warning(
-                "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
-                " This means that the model is loaded on 'cpu' first and then copied to the device."
-                " This leads to a slower loading time."
-                " Please update safetensors to version 0.4.3 or above for improved performance."
+        missing_keys, unexpected_keys = load_model_as_safetensor(
+            model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
         )
-            model.to(map_location)
+        log_model_loading_keys(missing_keys, unexpected_keys)
         return model
 
     @abc.abstractmethod

src/lerobot/policies/smolvla/processor_smolvla.py

@@ -19,19 +19,13 @@ from typing import Any
 import torch
 
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
     NewLineTaskProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
     TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_smolvla import SmolVLAConfig
 
@@ -66,9 +60,11 @@ def make_smolvla_pre_post_processors(
         A tuple containing the configured pre-processor and post-processor pipelines.
     """
 
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
-        AddBatchDimensionProcessorStep(),
+        steps.rename_observations,  # To mimic the same processor as pretrained one
+        steps.add_batch_dim,
         NewLineTaskProcessorStep(),
         TokenizerProcessorStep(
             tokenizer_name=config.vlm_model_name,
@@ -76,28 +72,11 @@ def make_smolvla_pre_post_processors(
             padding_side="right",
             max_length=config.tokenizer_max_length,
         ),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        steps.to_device,
+        steps.normalize,
     ]
     output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        DeviceProcessorStep(device="cpu"),
+        steps.unnormalize,
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/tdmpc/processor_tdmpc.py

@@ -19,17 +19,10 @@ from typing import Any
 import torch
 
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_pre_post_processors,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_tdmpc import TDMPCConfig
 
@@ -61,32 +54,4 @@ def make_tdmpc_pre_post_processors(
     Returns:
         A tuple containing the configured pre-processor and post-processor pipelines.
     """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        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,
-        ),
-    )
+    return make_default_pre_post_processors(config, dataset_stats)

src/lerobot/policies/vla_jepa/processor_vla_jepa.py

@@ -20,20 +20,16 @@ import torch
 
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
     EnvTransition,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
-from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 
 @ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
@@ -112,15 +108,12 @@ def make_vla_jepa_pre_post_processors(
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
 ]:
     features = {**config.input_features, **config.output_features}
+    steps = make_default_policy_processor_steps(config, dataset_stats)
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features=features,
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
+        steps.rename_observations,
+        steps.add_batch_dim,
+        steps.to_device,
+        steps.normalize,
     ]
     output_steps: list[ProcessorStep] = []
     if config.clip_normalized_actions:
@@ -129,6 +122,8 @@ def make_vla_jepa_pre_post_processors(
         output_steps.append(
             PreSnapGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
         )
+    # NOTE: unlike the default policy unnormalizer (output features only), VLA-JEPA
+    # unnormalizes over BOTH input and output features.
     output_steps.append(
         UnnormalizerProcessorStep(
             features=features,
@@ -140,16 +135,5 @@ def make_vla_jepa_pre_post_processors(
         output_steps.append(
             BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)
         )
-    output_steps.append(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,
-        ),
-    )
+    output_steps.append(steps.to_cpu)
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)

src/lerobot/policies/vqbet/processor_vqbet.py

@@ -20,17 +20,10 @@ from typing import Any
 import torch
 
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_pre_post_processors,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_vqbet import VQBeTConfig
 
@@ -62,32 +55,4 @@ def make_vqbet_pre_post_processors(
     Returns:
         A tuple containing the configured pre-processor and post-processor pipelines.
     """
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),  # Let the possibility to the user to rename the keys
-        AddBatchDimensionProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-    ]
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        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,
-        ),
-    )
+    return make_default_pre_post_processors(config, dataset_stats)

src/lerobot/policies/wall_x/configuration_wall_x.py

@@ -58,10 +58,14 @@ class WallXConfig(PreTrainedConfig):
     # Action prediction mode: "diffusion" or "fast"
     prediction_mode: str = "diffusion"
 
-    # Attention Implementation, options: "eager", "flash_attention_2", "sdpa"
-    # NOTE: flash-attn==2.7.4.post1 is required for flash_attention_2 implementation
+    # Wall-X's bidirectional action-token islands currently require eager attention.
     attn_implementation: str = "eager"
 
+    # Vision attention is independent from the text action-token mask. ``auto`` uses
+    # PyTorch's packed variable-length attention when the runtime supports it and
+    # otherwise falls back to the native per-chunk SDPA implementation.
+    vision_attn_implementation: str = "auto"
+
     # ==================== Optimizer Presets ====================
     optimizer_lr: float = 2e-5
     optimizer_betas: tuple[float, float] = (0.9, 0.95)
@@ -86,6 +90,18 @@ class WallXConfig(PreTrainedConfig):
         if self.prediction_mode not in ["diffusion", "fast"]:
             raise ValueError(f"prediction_mode must be 'diffusion' or 'fast', got {self.prediction_mode}")
 
+        if self.attn_implementation != "eager":
+            raise ValueError(
+                "Wall-X currently supports only attn_implementation='eager' because its "
+                "bidirectional action-token islands require an explicit attention mask."
+            )
+
+        if self.vision_attn_implementation not in {"auto", "sdpa", "varlen"}:
+            raise ValueError(
+                "vision_attn_implementation must be one of 'auto', 'sdpa', or 'varlen', got "
+                f"{self.vision_attn_implementation!r}"
+            )
+
         # Assign use_fast_tokenizer based on prediction_mode
         if self.prediction_mode == "fast":
             self.use_fast_tokenizer = True

src/lerobot/policies/wall_x/modeling_wall_x.py

@@ -43,11 +43,14 @@ from typing import TYPE_CHECKING, Any
 import numpy as np
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
-from PIL import Image
+import torch.nn.functional as functional
+from safetensors import SafetensorError
+from safetensors.torch import load_file
 from torch import Tensor
 from torch.distributions import Beta
 from torch.nn import CrossEntropyLoss
+from torchvision.transforms import InterpolationMode
+from torchvision.transforms.v2 import functional as tv_functional
 
 from lerobot.utils.constants import ACTION, OBS_STATE
 from lerobot.utils.import_utils import (
@@ -74,17 +77,17 @@ if TYPE_CHECKING or _wallx_deps_available:
     from qwen_vl_utils.vision_process import smart_resize
     from torchdiffeq import odeint
     from transformers import AutoProcessor, BatchFeature
-    from transformers.cache_utils import StaticCache
     from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+        Qwen2_5_VisionTransformerPretrainedModel,
         Qwen2_5_VLForConditionalGeneration,
     )
-    from transformers.utils import is_torchdynamo_compiling
+    from transformers.utils import cached_file, is_torchdynamo_compiling
 
-    from .qwen_model.configuration_qwen2_5_vl import Qwen2_5_VLConfig
-    from .qwen_model.qwen2_5_vl_moe import (
-        Qwen2_5_VisionTransformerPretrainedModel,
+    from .qwen_model import (
         Qwen2_5_VLACausalLMOutputWithPast,
+        Qwen2_5_VLConfig,
         Qwen2_5_VLMoEModel,
+        configure_wall_x_vision_attention,
     )
 else:
     LoraConfig = None
@@ -93,13 +96,14 @@ else:
     odeint = None
     AutoProcessor = None
     BatchFeature = None
-    StaticCache = None
     Qwen2_5_VLForConditionalGeneration = None
+    cached_file = None
     is_torchdynamo_compiling = None
     Qwen2_5_VLConfig = None
     Qwen2_5_VisionTransformerPretrainedModel = None
     Qwen2_5_VLACausalLMOutputWithPast = None
     Qwen2_5_VLMoEModel = None
+    configure_wall_x_vision_attention = None
 
 from .utils import (
     get_wallx_normal_text,
@@ -111,6 +115,75 @@ from .utils import (
 logger = logging.getLogger(__name__)
 
 
+def _wall_x_resize_dimensions(height: int, width: int) -> tuple[int, int, int, int]:
+    """Return the intermediate and final Wall-X resize dimensions as ``(H, W, H, W)``."""
+    if RESOLUTION == -1:
+        intermediate_height, intermediate_width = height, width
+    elif width > height:
+        intermediate_width = RESOLUTION
+        intermediate_height = int(RESOLUTION * height / width)
+    else:
+        intermediate_height = RESOLUTION
+        intermediate_width = int(RESOLUTION * width / height)
+
+    resized_height, resized_width = smart_resize(
+        intermediate_height,
+        intermediate_width,
+        factor=IMAGE_FACTOR,
+        min_pixels=MIN_PIXELS,
+        max_pixels=MAX_PIXELS,
+    )
+    return intermediate_height, intermediate_width, resized_height, resized_width
+
+
+def _resize_wall_x_image_batch(images: Tensor) -> tuple[Tensor, tuple[int, int, int, int]]:
+    """Quantize and resize a BCHW camera batch without leaving its current device."""
+    if images.ndim != 4:
+        raise ValueError(f"Wall-X images must be BCHW tensors, got shape {tuple(images.shape)}")
+
+    original_height, original_width = images.shape[-2:]
+    intermediate_height, intermediate_width, resized_height, resized_width = _wall_x_resize_dimensions(
+        original_height, original_width
+    )
+
+    if images.is_floating_point():
+        # Match the previous PIL path, which quantized via `(image * 255).to(torch.uint8)`.
+        images = (images * 255).to(torch.uint8)
+    elif images.dtype != torch.uint8:
+        raise TypeError(f"Wall-X images must be floating point or uint8, got {images.dtype}")
+
+    if images.shape[-2:] != (intermediate_height, intermediate_width):
+        images = tv_functional.resize(
+            images,
+            [intermediate_height, intermediate_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        )
+    if images.shape[-2:] != (resized_height, resized_width):
+        images = tv_functional.resize(
+            images,
+            [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=True,
+        )
+
+    return images, (original_height, original_width, resized_height, resized_width)
+
+
+def _prepare_wall_x_image_inputs(
+    batch: dict[str, Any], img_keys: list[str]
+) -> tuple[list[list[Tensor]], dict[str, tuple[int, int, int, int]]]:
+    """Resize each camera as a batch, then restore sample-major/camera-minor ordering."""
+    resized_by_key: dict[str, Tensor] = {}
+    dimensions_by_key: dict[str, tuple[int, int, int, int]] = {}
+    for key in img_keys:
+        resized_by_key[key], dimensions_by_key[key] = _resize_wall_x_image_batch(batch[key])
+
+    batch_size = batch[img_keys[0]].shape[0]
+    image_inputs = [[resized_by_key[key][i] for key in img_keys] for i in range(batch_size)]
+    return image_inputs, dimensions_by_key
+
+
 class SinusoidalPosEmb(nn.Module):
     """Sinusoidal positional embedding for diffusion timesteps."""
 
@@ -246,7 +319,7 @@ class ActionHead(nn.Module):
         flow = flow.to(torch.float32)
 
         action_pred = self.action_proj_back(action_hidden_states)
-        loss = F.mse_loss(action_pred, flow, reduction="none")
+        loss = functional.mse_loss(action_pred, flow, reduction="none")
 
         if dof_mask is not None:
             dof_mask = dof_mask.reshape(-1, dof_mask.shape[-1]).to(torch.float32)
@@ -254,7 +327,7 @@ class ActionHead(nn.Module):
 
         return loss
 
-    def proprioception_proj(self, proprioception, dof_mask=None, use_history=False):
+    def proprioception_proj(self, proprioception, dof_mask=None):
         """Project proprioceptive data to hidden space."""
         # Ensure proper device and dtype alignment
         proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
@@ -264,9 +337,6 @@ class ActionHead(nn.Module):
         if dof_mask is not None:
             # Concatenate proprioception with DOF mask
             # TODO: Use variable-based dimension checking for better flexibility
-            if use_history:
-                proprioception = torch.cat([proprioception, dof_mask], dim=-1)
-            else:
             proprioception = torch.cat([proprioception, dof_mask], dim=-1)
 
         proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
@@ -281,7 +351,7 @@ class ActionHead(nn.Module):
 _Qwen2_5_VLForAction_Base = Qwen2_5_VLForConditionalGeneration if _wallx_deps_available else nn.Module
 
 
-class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
+class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):  # noqa: N801
     """
     Qwen2.5 Vision-Language Mixture of Experts model for action processing.
 
@@ -305,6 +375,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
         config=None,
         action_tokenizer_path=None,
         attn_implementation: str = "eager",
+        vision_attn_implementation: str = "auto",
         cache_dir: str | PathLike | None = None,
         force_download: bool = False,
         local_files_only: bool = False,
@@ -321,11 +392,14 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             config_path (str, optional): Configuration file path, if None will look for qwen25_config.json in pretrained_model_path
             action_tokenizer_path (str, optional): Action tokenizer path, if None will load from default config
             attn_implementation (str, optional): Attention implementation, if None will load from default config
+            vision_attn_implementation (str, optional): Vision attention backend. ``auto`` uses packed
+                variable-length attention when supported and otherwise falls back to SDPA.
             **kwargs: Additional arguments
 
         Returns:
             Qwen2_5_VLMoEForAction: Loaded model instance
         """
+        Qwen2_5_VLMoEModel._require_eager_attention(attn_implementation)
         if config is None:
             config = cls.config_class.from_pretrained(
                 pretrained_name_or_path,
@@ -339,7 +413,15 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             )
         if attn_implementation is not None:
             config._attn_implementation = attn_implementation
-        processor = AutoProcessor.from_pretrained(pretrained_name_or_path, use_fast=True)
+        processor = AutoProcessor.from_pretrained(
+            pretrained_name_or_path,
+            cache_dir=cache_dir,
+            force_download=force_download,
+            local_files_only=local_files_only,
+            token=token,
+            revision=revision,
+            use_fast=True,
+        )
         if action_tokenizer_path is not None:
             action_tokenizer = AutoProcessor.from_pretrained(action_tokenizer_path, trust_remote_code=True)
             processor.action_processor = action_tokenizer
@@ -351,41 +433,41 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
         config.text_config.pad_token_id = processor.tokenizer.pad_token_id
 
         # Initialize model with configuration and processor
-        model = cls(config, processor=processor, action_tokenizer=action_tokenizer, **kwargs)
+        model = cls(
+            config,
+            processor=processor,
+            action_tokenizer=action_tokenizer,
+            vision_attn_implementation=vision_attn_implementation,
+            **kwargs,
+        )
 
         # Resize token embeddings to match processor tokenizer vocabulary size
         model.resize_token_embeddings(len(processor.tokenizer))
 
-        # Try to load the model.safetensors file
-        print(f"Loading model from: {pretrained_name_or_path}")
+        logger.info("Loading Wall-X model from %s", pretrained_name_or_path)
         try:
-            from transformers.utils import cached_file
-
-            # Try safetensors first
             resolved_file = cached_file(
                 pretrained_name_or_path,
                 "model.safetensors",
-                cache_dir=kwargs.get("cache_dir"),
-                force_download=kwargs.get("force_download", False),
+                cache_dir=cache_dir,
+                force_download=force_download,
                 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),
+                token=token,
+                revision=revision,
+                local_files_only=local_files_only,
             )
-            from safetensors.torch import load_file
-
             sd = 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
+        except (OSError, SafetensorError) as error:
+            raise OSError(
+                f"Failed to load pretrained Wall-X weights from {pretrained_name_or_path!r}"
+            ) from error
+        logger.info("Loaded Wall-X state dict from model.safetensors")
 
         state_dict = {}
         # filter normalizer statistic params
         del_keys = []
-        for key in sd.keys():
+        for key in sd:
             if "action_preprocessor.normalizer" in key:
                 del_keys.append(key)
         for key in del_keys:
@@ -404,6 +486,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
         action_tokenizer=None,
         action_mapper=None,
         flow_loss_weight=1.0,
+        vision_attn_implementation: str = "auto",
     ):
         """
         Initialize the Qwen2.5 VLMoE model for action processing.
@@ -416,10 +499,16 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             action_mapper: Action mapping utility
             flow_loss_weight (float): Weight for flow loss computation
         """
+        Qwen2_5_VLMoEModel._require_eager_attention(config._attn_implementation)
+        config._attn_implementation = "eager"
+        # Text needs eager attention for action-token islands. Vision has no such
+        # constraint, so keep its portable native fallback on SDPA.
+        config.vision_config._attn_implementation = "sdpa"
         super().__init__(config)
 
         # Initialize vision transformer and language model components
         self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)
+        configure_wall_x_vision_attention(self.visual, vision_attn_implementation)
         self.model = Qwen2_5_VLMoEModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
@@ -457,7 +546,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
 
         params_to_keep_float32 = []
 
-        for name, param in self.named_parameters():
+        for name, _param in self.named_parameters():
             if "input_layernorm" in name or "post_attention_layernorm" in name or "model.norm" in name:
                 params_to_keep_float32.append(name)
             if "action_preprocessor" in name:
@@ -491,7 +580,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             "action_token_id": action_token_id,
         }
 
-    def add_lora(self, r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.1):
+    def add_lora(self, r=8, lora_alpha=32, target_modules=None, lora_dropout=0.1):
         """
         Add LoRA (Low-Rank Adaptation) adapters to the model.
 
@@ -501,6 +590,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             target_modules (list): List of module names to apply LoRA to
             lora_dropout (float): Dropout probability for LoRA layers
         """
+        if target_modules is None:
+            target_modules = ["q_proj", "v_proj"]
+
         config = LoraConfig(
             r=r,
             lora_alpha=lora_alpha,
@@ -795,6 +887,9 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
+        if rope_deltas is not None:
+            self.rope_deltas = rope_deltas
+
         # Calculate RoPE position IDs if not provided
         # Note: Cannot calculate rope deltas with 4D attention mask. TODO: Fix this limitation
         if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
@@ -833,7 +928,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             # Process image embeddings
             if pixel_values is not None:
                 pixel_values = pixel_values.type(self.visual.dtype)
-                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
                 mask = input_ids == self.config.image_token_id
                 mask_unsqueezed = mask.unsqueeze(-1)
                 mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
@@ -845,7 +940,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             # Process video embeddings
             if pixel_values_videos is not None:
                 pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
-                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
+                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
                 n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
                 n_video_features = video_embeds.shape[0]
 
@@ -869,7 +964,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
                 proprioception = self.action_preprocessor.proprioception_proj(
                     proprioception,
                     agent_pos_mask,
-                    use_history=proprioception.shape[1] > 1,
                 )
                 mask = input_ids == self.action_token_id_set["propri_token_id"]
                 mask_unsqueezed = mask.unsqueeze(-1)
@@ -919,6 +1013,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             output_attentions=output_attentions,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            cache_position=cache_position,
         )
 
         hidden_states = outputs[0]
@@ -1107,7 +1202,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             # Process image embeddings
             if pixel_values is not None:
                 pixel_values = pixel_values.type(self.visual.dtype)
-                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
+                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw).pooler_output
                 n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
                 n_image_features = image_embeds.shape[0]
 
@@ -1128,7 +1223,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
             # Process video embeddings
             if pixel_values_videos is not None:
                 pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
-                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
+                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw).pooler_output
                 n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
                 n_video_features = video_embeds.shape[0]
 
@@ -1153,7 +1248,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
                 proprio_embed = self.action_preprocessor.proprioception_proj(
                     proprioception,
                     agent_pos_mask,
-                    use_history=proprioception.shape[1] > 1,
                 )
                 proprioception_mask = input_ids == self.action_token_id_set["propri_token_id"]
                 proprio_embed = proprio_embed.to(torch.bfloat16)
@@ -1202,25 +1296,37 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
 
         # Split input sequence for text and fast modes (not needed for diffusion)
         if predict_mode == "text" or predict_mode == "fast":
-            # Look for generation prompt tokens: <|im_start|>assistant
+            generation_prompt = "<|im_start|>assistant\n"
             generation_prompt_ids = torch.tensor(
-                [151644, 77091], device=input_ids.device, dtype=input_ids.dtype
+                self.processor.tokenizer.encode(generation_prompt, add_special_tokens=False),
+                device=input_ids.device,
+                dtype=input_ids.dtype,
             )
-            matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (
-                input_ids[0, 1:] == generation_prompt_ids[1]
+            prompt_length = generation_prompt_ids.numel()
+            if prompt_length == 0:
+                raise ValueError(f"Tokenizer produced no tokens for generation prompt {generation_prompt!r}")
+            if input_ids.shape[1] < prompt_length:
+                matches = torch.empty(0, device=input_ids.device, dtype=torch.bool)
+            else:
+                matches = (
+                    input_ids[0]
+                    .unfold(dimension=0, size=prompt_length, step=1)
+                    .eq(generation_prompt_ids)
+                    .all(dim=-1)
                 )
 
             if matches.any():
                 split_pos = torch.nonzero(matches, as_tuple=True)[0][0].item()
+                prompt_end = split_pos + prompt_length
                 # Extract ground truth output tokens (including newline)
-                gt_output_ids = input_ids[:, split_pos + 3 :]
+                gt_output_ids = input_ids[:, prompt_end:]
                 # Remove output part from input, keeping prompt
-                input_ids = input_ids[:, : split_pos + 3]
-                inputs_embeds = inputs_embeds[:, : split_pos + 3, :]
+                input_ids = input_ids[:, :prompt_end]
+                inputs_embeds = inputs_embeds[:, :prompt_end, :]
                 if attention_mask is not None:
-                    attention_mask = attention_mask[:, : split_pos + 3]
+                    attention_mask = attention_mask[:, :prompt_end]
                 if labels is not None:
-                    labels = labels[:, split_pos + 3 :]
+                    labels = labels[:, prompt_end:]
             else:
                 raise ValueError(
                     "input_ids does not contain the generation prompt tokens <|im_start|>assistant"
@@ -1255,7 +1361,7 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
                 use_cache=True,
                 pad_token_id=self.processor.tokenizer.pad_token_id,
                 temperature=(1.0 if not re_generate else 0.7),  # Higher temperature for regeneration
-                do_sample=(False if not re_generate else True),  # Enable sampling for regeneration
+                do_sample=re_generate,  # Enable sampling for regeneration
             )
 
             # Decode generated and ground truth text
@@ -1524,27 +1630,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base):
         else:
             model_inputs = {"input_ids": input_ids, "inputs_embeds": None}
 
-        # Prepare 4D causal attention mask for static cache
-        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
-            if model_inputs["inputs_embeds"] is not None:
-                batch_size, sequence_length, _ = inputs_embeds.shape
-                device = inputs_embeds.device
-            else:
-                batch_size, sequence_length = input_ids.shape
-                device = input_ids.device
-
-            attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(
-                attention_mask,
-                sequence_length=sequence_length,
-                target_length=past_key_values.get_max_cache_shape(),
-                dtype=self.lm_head.weight.dtype,
-                device=device,
-                cache_position=cache_position,
-                batch_size=batch_size,
-                config=self.config,
-                past_key_values=past_key_values,
-            )
-
         # Assemble all model inputs for generation
         model_inputs.update(
             {
@@ -1749,6 +1834,7 @@ class WallXPolicy(PreTrainedPolicy):
             pretrained_name_or_path=config.pretrained_name_or_path,
             action_tokenizer_path=config.action_tokenizer_path,
             attn_implementation=config.attn_implementation,
+            vision_attn_implementation=config.vision_attn_implementation,
         )
         self.model.to(config.device)
         self.model.to_bfloat16_for_selected_params()
@@ -1768,6 +1854,8 @@ class WallXPolicy(PreTrainedPolicy):
     def preprocess_inputs(
         self,
         batch: dict[str, Any],
+        *,
+        compute_position_ids: bool = False,
     ) -> BatchFeature:
         """
         Convert a batch of LeRobot dataset items to Wall-X model input format.
@@ -1789,50 +1877,21 @@ class WallXPolicy(PreTrainedPolicy):
         # Get batch size from state tensor
         batch_size = batch[OBS_STATE].shape[0]
 
-        # ==================== PROCESS ALL SAMPLES ====================
-        all_image_inputs = []
-        all_texts = []
-
         # Find image keys in batch
         img_keys = [key for key in self.config.image_features if key in batch]
+        if not img_keys:
+            raise ValueError("Wall-X requires at least one image feature in each batch")
+
+        # Resize one camera batch at a time on the tensors' current device. Reassembling
+        # sample-major keeps image_grid_thw aligned with each sample's image placeholders.
+        all_image_inputs, dimensions_by_key = _prepare_wall_x_image_inputs(batch, img_keys)
+        all_texts = []
+
+        # Preserve the existing grounding behavior for multi-camera inputs: the old camera
+        # loop left these values set to the final configured camera's dimensions.
+        orig_height, orig_width, resized_height, resized_width = dimensions_by_key[img_keys[-1]]
 
         for i in range(batch_size):
-            # Vision preprocessing per sample
-            processed_frames = []
-            orig_height, orig_width = None, None
-            resized_height, resized_width = None, None
-
-            for key in img_keys:
-                current_obs = batch[key][i].clone()  # (C, H, W)
-                if current_obs.dim() == 3:
-                    current_obs = current_obs.permute(1, 2, 0)  # (H, W, C)
-
-                img_pil = Image.fromarray((current_obs * 255).to(torch.uint8).cpu().numpy())
-                orig_width, orig_height = img_pil.size
-
-                target_size = RESOLUTION
-                if target_size != -1:
-                    if orig_width > orig_height:
-                        new_width = target_size
-                        new_height = int(target_size * orig_height / orig_width)
-                    else:
-                        new_height = target_size
-                        new_width = int(target_size * orig_width / orig_height)
-                    img_pil = img_pil.resize((new_width, new_height))
-
-                current_width, current_height = img_pil.size
-                resized_height, resized_width = smart_resize(
-                    current_height,
-                    current_width,
-                    factor=IMAGE_FACTOR,
-                    min_pixels=MIN_PIXELS,
-                    max_pixels=MAX_PIXELS,
-                )
-                resized_img = img_pil.resize((resized_width, resized_height))
-                processed_frames.append(resized_img)
-
-            all_image_inputs.append(processed_frames)
-
             # Text preprocessing
             task_text = batch["task"][i] if isinstance(batch["task"], list) else batch["task"]
             instruction_info = {"instruction": task_text}
@@ -1859,8 +1918,8 @@ class WallXPolicy(PreTrainedPolicy):
         agent_pos_mask = (~torch.isnan(agent_pos)).float()
         agent_pos = agent_pos.nan_to_num(nan=0.0)
 
-        if agent_pos.shape[-1] != 20:
-            pad_size = 20 - agent_pos.shape[-1]
+        if agent_pos.shape[-1] < self.config.max_state_dim:
+            pad_size = self.config.max_state_dim - agent_pos.shape[-1]
             agent_pos = torch.cat(
                 [
                     agent_pos,
@@ -1880,6 +1939,10 @@ class WallXPolicy(PreTrainedPolicy):
                 ],
                 dim=-1,
             )
+        elif agent_pos.shape[-1] > self.config.max_state_dim:
+            raise ValueError(
+                f"State dimension {agent_pos.shape[-1]} exceeds max_state_dim {self.config.max_state_dim}"
+            )
 
         # ==================== PROCESS ACTIONS ====================
         action = batch.get(ACTION)  # (batch_size, chunk_size, action_dim)
@@ -1889,8 +1952,8 @@ class WallXPolicy(PreTrainedPolicy):
             dof_mask = (~torch.isnan(action)).float()
             action = action.nan_to_num(nan=0.0)
 
-            if action.shape[-1] != 20:
-                pad_size = 20 - action.shape[-1]
+            if action.shape[-1] < self.config.max_action_dim:
+                pad_size = self.config.max_action_dim - action.shape[-1]
                 action = torch.cat(
                     [action, torch.zeros(action.shape[0], action.shape[1], pad_size, device=action.device)],
                     dim=-1,
@@ -1902,6 +1965,10 @@ class WallXPolicy(PreTrainedPolicy):
                     ],
                     dim=-1,
                 )
+            elif action.shape[-1] > self.config.max_action_dim:
+                raise ValueError(
+                    f"Action dimension {action.shape[-1]} exceeds max_action_dim {self.config.max_action_dim}"
+                )
         else:
             action_dim = self.config.output_features[ACTION].shape[0]
             dof_mask = torch.cat(
@@ -1910,7 +1977,10 @@ class WallXPolicy(PreTrainedPolicy):
                         batch_size, self.config.chunk_size, action_dim, device=batch[OBS_STATE].device
                     ),
                     torch.zeros(
-                        batch_size, self.config.chunk_size, 20 - action_dim, device=batch[OBS_STATE].device
+                        batch_size,
+                        self.config.chunk_size,
+                        self.config.max_action_dim - action_dim,
+                        device=batch[OBS_STATE].device,
                     ),
                 ],
                 dim=-1,
@@ -1930,12 +2000,26 @@ class WallXPolicy(PreTrainedPolicy):
             text=all_texts,
             images=all_image_inputs,
             videos=None,
+            device=batch[OBS_STATE].device,
             padding=True,
             truncation=True,
             return_tensors="pt",
             max_length=TOKENIZER_MAX_LENGTH,
         )
 
+        if compute_position_ids:
+            # Qwen's RoPE indexing uses Python list/scalar conversions. Run it while the
+            # tokenizer and grid metadata are still on CPU, then move the compact result.
+            position_ids, rope_deltas = self.model.get_rope_index(
+                inputs.input_ids,
+                inputs.get("image_grid_thw"),
+                inputs.get("video_grid_thw"),
+                inputs.get("second_per_grid_ts"),
+                inputs.attention_mask,
+            )
+            inputs["position_ids"] = position_ids
+            inputs["rope_deltas"] = rope_deltas
+
         # ==================== ADDITIONAL INPUTS ====================
         action_token_id = self.model.processor.tokenizer.convert_tokens_to_ids("<|action|>")
         moe_token_types = inputs.input_ids == action_token_id
@@ -1952,7 +2036,7 @@ class WallXPolicy(PreTrainedPolicy):
         )
 
         # Move all tensors to the correct device
-        device = self.config.device
+        device = batch[OBS_STATE].device
         for key, value in inputs.items():
             if isinstance(value, torch.Tensor):
                 inputs[key] = value.to(device)
@@ -1972,9 +2056,7 @@ class WallXPolicy(PreTrainedPolicy):
         Returns:
             tuple: (loss, loss_dict)
         """
-        batch = self.preprocess_inputs(
-            batch,
-        )
+        batch = self.preprocess_inputs(batch, compute_position_ids=True)
 
         # Call the underlying model's forward with mode="train"
         outputs = self.model(**batch, mode="train")
@@ -1982,19 +2064,19 @@ class WallXPolicy(PreTrainedPolicy):
         # Extract losses from output
         loss = outputs.loss
         loss_dict = {
-            "loss": loss.item() if loss is not None else 0.0,
+            "loss": loss.detach() if loss is not None else 0.0,
         }
 
         if outputs.flow_loss is not None:
-            loss_dict["flow_loss"] = outputs.flow_loss.item()
+            loss_dict["flow_loss"] = outputs.flow_loss.detach()
         if outputs.cross_entropy_loss is not None:
-            loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.item()
+            loss_dict["cross_entropy_loss"] = outputs.cross_entropy_loss.detach()
 
         # Add channel losses if available
         if outputs.channel_loss_dict is not None:
             for key, value in outputs.channel_loss_dict.items():
                 if isinstance(value, torch.Tensor):
-                    loss_dict[f"channel_{key}"] = value.item()
+                    loss_dict[f"channel_{key}"] = value.detach()
 
         return loss, loss_dict
 

src/lerobot/policies/wall_x/processor_wall_x.py

@@ -20,19 +20,13 @@ import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
     ComplementaryDataProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
 
 from .configuration_wall_x import WallXConfig
 
@@ -65,37 +59,22 @@ def make_wall_x_pre_post_processors(
         A tuple containing the configured pre-processor and post-processor pipelines
     """
 
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
+        steps.rename_observations,
+        steps.add_batch_dim,
         WallXTaskProcessor(),  # Process task description
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        DeviceProcessorStep(device=config.device),
+        steps.normalize,
+        steps.to_device,
     ]
 
     output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        DeviceProcessorStep(device="cpu"),
+        steps.unnormalize,
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
 
 
 @ProcessorStepRegistry.register(name="wall_x_task_processor")

src/lerobot/policies/wall_x/qwen_model/__init__.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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 .configuration_qwen2_5_vl import (
+    Qwen2_5_VLConfig,
+    Qwen2_5_VLTextConfig,
+    Qwen2_5_VLVisionConfig,
+)
+from .qwen2_5_vl_moe import (
+    BlockSparseMLP,
+    Qwen2_5_VLACausalLMOutputWithPast,
+    Qwen2_5_VLDecoderLayer_with_MoE,
+    Qwen2_5_VLMoEModel,
+    SparseMoeBlock,
+)
+from .vision_attention import (
+    WallXVisionAttention,
+    configure_wall_x_vision_attention,
+)
+
+__all__ = [
+    "BlockSparseMLP",
+    "Qwen2_5_VLACausalLMOutputWithPast",
+    "Qwen2_5_VLConfig",
+    "Qwen2_5_VLDecoderLayer_with_MoE",
+    "Qwen2_5_VLMoEModel",
+    "Qwen2_5_VLTextConfig",
+    "Qwen2_5_VLVisionConfig",
+    "SparseMoeBlock",
+    "WallXVisionAttention",
+    "configure_wall_x_vision_attention",
+]

src/lerobot/policies/wall_x/qwen_model/configuration_qwen2_5_vl.py

@@ -1,250 +1,114 @@
-from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_rope_utils import rope_config_validation
+#!/usr/bin/env python
 
+# Copyright 2025 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.
 
-class Qwen2_5_VLVisionConfig(PretrainedConfig):
-    model_type = "qwen2_5_vl"
-    base_config_key = "vision_config"
+"""Wall-X configuration extensions for the native Transformers Qwen2.5-VL config."""
 
-    def __init__(
-        self,
-        depth=32,
-        hidden_size=3584,
-        hidden_act="silu",
-        intermediate_size=3420,
-        num_heads=16,
-        in_channels=3,
-        patch_size=14,
-        spatial_merge_size=2,
-        temporal_patch_size=2,
-        tokens_per_second=4,
-        window_size=112,
-        out_hidden_size=3584,
-        fullatt_block_indexes=[7, 15, 23, 31],
-        initializer_range=0.02,
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
-        self.depth = depth
-        self.hidden_size = hidden_size
-        self.hidden_act = hidden_act
-        self.intermediate_size = intermediate_size
-        self.num_heads = num_heads
-        self.in_channels = in_channels
-        self.patch_size = patch_size
-        self.spatial_merge_size = spatial_merge_size
-        self.temporal_patch_size = temporal_patch_size
-        self.tokens_per_second = tokens_per_second
-        self.window_size = window_size
-        self.fullatt_block_indexes = fullatt_block_indexes
-        self.out_hidden_size = out_hidden_size
-        self.initializer_range = initializer_range
+from huggingface_hub.dataclasses import strict
 
+from lerobot.utils.import_utils import _transformers_available
 
-class Qwen2_5_VLConfig(PretrainedConfig):
-    r"""
-    This is the configuration class to store the configuration of a [`Qwen2_5_VLModel`]. It is used to instantiate a
-    Qwen2-VL model according to the specified arguments, defining the model architecture. Instantiating a configuration
-    with the defaults will yield a similar configuration to that of
-    Qwen2-VL-7B-Instruct [Qwen/Qwen2-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct).
+if TYPE_CHECKING or _transformers_available:
+    from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
+        Qwen2_5_VLConfig as TransformersQwen2_5_VLConfig,
+        Qwen2_5_VLTextConfig as TransformersQwen2_5_VLTextConfig,
+        Qwen2_5_VLVisionConfig,
+    )
+else:
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    @dataclass
+    class _TransformersConfigFallback:
+        """Import-safe stand-in used only when Transformers is unavailable."""
 
+    TransformersQwen2_5_VLConfig = _TransformersConfigFallback
+    TransformersQwen2_5_VLTextConfig = _TransformersConfigFallback
+    Qwen2_5_VLVisionConfig = None
 
-    Args:
-        vocab_size (`int`, *optional*, defaults to 152064):
-            Vocabulary size of the Qwen2_5_VL model. Defines the number of different tokens that can be represented by the
-            `inputs_ids` passed when calling [`Qwen2_5_VLModel`]
-        hidden_size (`int`, *optional*, defaults to 8192):
-            Dimension of the hidden representations.
-        intermediate_size (`int`, *optional*, defaults to 29568):
-            Dimension of the MLP representations.
-        num_hidden_layers (`int`, *optional*, defaults to 80):
-            Number of hidden layers in the Transformer encoder.
-        num_attention_heads (`int`, *optional*, defaults to 64):
-            Number of attention heads for each attention layer in the Transformer encoder.
-        num_key_value_heads (`int`, *optional*, defaults to 8):
-            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
-            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
-            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
-            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
-            by meanpooling all the original heads within that group. For more details checkout [this
-            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to `32`.
-        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
-            The non-linear activation function (function or string) in the decoder.
-        max_position_embeddings (`int`, *optional*, defaults to 32768):
-            The maximum sequence length that this model might ever be used with.
-        initializer_range (`float`, *optional*, defaults to 0.02):
-            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
-        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
-            The epsilon used by the rms normalization layers.
-        use_cache (`bool`, *optional*, defaults to `True`):
-            Whether or not the model should return the last key/values attentions (not used by all models). Only
-            relevant if `config.is_decoder=True`.
-        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
-            Whether the model's input and output word embeddings should be tied.
-        rope_theta (`float`, *optional*, defaults to 1000000.0):
-            The base period of the RoPE embeddings.
-        use_sliding_window (`bool`, *optional*, defaults to `False`):
-            Whether to use sliding window attention.
-        sliding_window (`int`, *optional*, defaults to 4096):
-            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
-        max_window_layers (`int`, *optional*, defaults to 80):
-            The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention.
-        attention_dropout (`float`, *optional*, defaults to 0.0):
-            The dropout ratio for the attention probabilities.
-        vision_config (`Dict`, *optional*):
-            The config for the visual encoder initialization.
-        rope_scaling (`Dict`, *optional*):
-            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
-            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
-            accordingly.
-            Expected contents:
-                `rope_type` (`str`):
-                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
-                    'llama3'], with 'default' being the original RoPE implementation.
-                `factor` (`float`, *optional*):
-                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
-                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
-                    original maximum pre-trained length.
-                `original_max_position_embeddings` (`int`, *optional*):
-                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
-                    pretraining.
-                `attention_factor` (`float`, *optional*):
-                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
-                    computation. If unspecified, it defaults to value recommended by the implementation, using the
-                    `factor` field to infer the suggested value.
-                `beta_fast` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 32.
-                `beta_slow` (`float`, *optional*):
-                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
-                    ramp function. If unspecified, it defaults to 1.
-                `short_factor` (`List[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `long_factor` (`List[float]`, *optional*):
-                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
-                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
-                    size divided by the number of attention heads divided by 2
-                `low_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
-                `high_freq_factor` (`float`, *optional*):
-                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
-
-    ```python
-    >>> from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLConfig
-
-    >>> # Initializing a Qwen2_5_VL style configuration
-    >>> configuration = Qwen2_5_VLConfig()
-
-    >>> # Initializing a model from the Qwen2-VL-7B style configuration
-    >>> model = Qwen2_5_VLForConditionalGeneration(configuration)
-
-    >>> # Accessing the model configuration
-    >>> configuration = model.config
-    ```"""
-
-    model_type = "qwen2_5_vl"
-    sub_configs = {"vision_config": Qwen2_5_VLVisionConfig}
-    keys_to_ignore_at_inference = ["past_key_values"]
-    # Default tensor parallel plan for base model `Qwen2_5_VL`
-    base_model_tp_plan = {
-        "layers.*.self_attn.q_proj": "colwise",
-        "layers.*.self_attn.k_proj": "colwise",
-        "layers.*.self_attn.v_proj": "colwise",
-        "layers.*.self_attn.o_proj": "rowwise",
-        "layers.*.mlp.gate_proj": "colwise",
-        "layers.*.mlp.up_proj": "colwise",
-        "layers.*.mlp.down_proj": "rowwise",
-    }
-    base_model_pp_plan = {
-        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
-        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
-        "norm": (["hidden_states"], ["hidden_states"]),
+# Wall-X checkpoints pre0.6.0 use the legacy, flat Qwen2.5-VL config layout.  The native
+# ``Qwen2_5_VLConfig`` accepts that layout and moves text-model fields into its
+# ``text_config`` sub-config, so only the Wall-X-specific MoE fields need to be
+# declared here.
+_LEGACY_TEXT_ATTRIBUTES = {
+    "attention_dropout",
+    "attention_moe",
+    "dim_inputs",
+    "dof_config",
+    "experts",
+    "hidden_act",
+    "hidden_size",
+    "initializer_range",
+    "intermediate_size",
+    "layer_types",
+    "max_position_embeddings",
+    "max_window_layers",
+    "mlp_moe",
+    "noise_scheduler",
+    "num_attention_heads",
+    "num_experts",
+    "num_hidden_layers",
+    "num_key_value_heads",
+    "pad_token_id",
+    "rms_norm_eps",
+    "sliding_window",
+    "use_cache",
+    "use_sliding_window",
+    "vocab_size",
 }
 
-    def __init__(
-        self,
-        vocab_size=152064,
-        hidden_size=8192,
-        intermediate_size=29568,
-        num_hidden_layers=80,
-        num_attention_heads=64,
-        num_key_value_heads=8,
-        hidden_act="silu",
-        max_position_embeddings=32768,
-        initializer_range=0.02,
-        rms_norm_eps=1e-05,
-        use_cache=True,
-        tie_word_embeddings=False,
-        rope_theta=1000000.0,
-        use_sliding_window=False,
-        sliding_window=4096,
-        max_window_layers=80,
-        attention_dropout=0.0,
-        vision_config=None,
-        rope_scaling=None,
-        num_experts=4,
-        experts=None,
-        dof_config=None,
-        noise_scheduler=None,
-        dim_inputs=(1536, 1536),
-        attention_moe=False,
-        mlp_moe=False,
-        **kwargs,
-    ):
-        if isinstance(vision_config, dict):
-            self.vision_config = self.sub_configs["vision_config"](**vision_config)
-        elif vision_config is None:
-            self.vision_config = self.sub_configs["vision_config"]()
 
-        self.vocab_size = vocab_size
-        self.max_position_embeddings = max_position_embeddings
-        self.hidden_size = hidden_size
-        self.intermediate_size = intermediate_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.use_sliding_window = use_sliding_window
-        self.sliding_window = sliding_window
-        self.max_window_layers = max_window_layers
-        self.layer_types = ["dense"] * num_hidden_layers
+@strict
+class Qwen2_5_VLTextConfig(TransformersQwen2_5_VLTextConfig):  # noqa: N801
+    """Native Qwen2.5-VL text config plus Wall-X's hard-routed MoE settings."""
 
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
+    num_experts: int = 4
+    experts: list[dict] | None = None
+    dof_config: dict | None = None
+    noise_scheduler: dict | None = None
+    dim_inputs: tuple[int, ...] | list[int] = (1536, 1536)
+    attention_moe: bool = False
+    mlp_moe: bool = False
 
-        self.num_key_value_heads = num_key_value_heads
-        self.hidden_act = hidden_act
-        self.initializer_range = initializer_range
-        self.rms_norm_eps = rms_norm_eps
-        self.use_cache = use_cache
-        self.rope_theta = rope_theta
-        self.attention_dropout = attention_dropout
-        self.rope_scaling = rope_scaling
-
-        self.num_experts = num_experts
-        self.experts = experts
-        self.dof_config = dof_config
-        self.noise_scheduler = noise_scheduler
-        self.dim_inputs = tuple(dim_inputs)
-        self.attention_moe = attention_moe
-        self.mlp_moe = mlp_moe
-
-        if self.rope_scaling is not None and "type" in self.rope_scaling:
-            if self.rope_scaling["type"] == "mrope":
-                self.rope_scaling["type"] = "default"
-            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
-        rope_config_validation(self, ignore_keys={"mrope_section"})
-
-        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
-
-    @property
-    def text_config(self):
-        return self
+    def __post_init__(self, **kwargs):
+        self.dim_inputs = tuple(self.dim_inputs)
+        super().__post_init__(**kwargs)
 
 
-__all__ = ["Qwen2_5_VLConfig"]
+@strict
+class Qwen2_5_VLConfig(TransformersQwen2_5_VLConfig):  # noqa: N801
+    """Native composite Qwen2.5-VL config with a Wall-X text sub-config.
+
+    The native composite loader supports both current nested configs and the
+    flat layout used by existing ``wall-oss-flow`` checkpoints.
+    """
+
+    sub_configs = {
+        "vision_config": Qwen2_5_VLVisionConfig,
+        "text_config": Qwen2_5_VLTextConfig,
+    }
+
+    def __getattr__(self, name):
+        """Keep legacy direct access to fields now owned by ``text_config``.
+
+        Wall-X historically used a flat config and accesses fields such as
+        ``hidden_size`` and ``num_experts`` directly. Forwarding unknown
+        attributes preserves that API without duplicating the native config.
+        """
+        text_config = self.__dict__.get("text_config")
+        if name in _LEGACY_TEXT_ATTRIBUTES and text_config is not None and hasattr(text_config, name):
+            return getattr(text_config, name)
+        raise AttributeError(f"{type(self).__name__!s} has no attribute {name!r}")

src/lerobot/policies/wall_x/qwen_model/vision_attention.py

@@ -0,0 +1,208 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+
+"""Wall-X vision attention backends.
+
+Qwen2.5-VL's native non-Flash vision path splits a packed image sequence into
+Python-level chunks before calling attention. Wall-X batches many camera frames,
+so that path launches thousands of tiny attention operations per training step.
+This module keeps the native SDPA path as a portable fallback and adds a packed
+``torch.nn.attention.varlen`` path that consumes Qwen's existing ``cu_seqlens``
+metadata directly.
+"""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from functools import lru_cache
+from typing import TYPE_CHECKING, Literal
+
+import torch
+import torch.nn as nn
+
+from lerobot.utils.import_utils import _transformers_available
+
+if TYPE_CHECKING or _transformers_available:
+    from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
+        Qwen2_5_VLVisionAttention,
+        apply_rotary_pos_emb_vision,
+    )
+else:
+    Qwen2_5_VLVisionAttention = nn.Module
+    apply_rotary_pos_emb_vision = None
+
+try:
+    from torch.nn.attention.varlen import varlen_attn as _varlen_attn
+except ImportError:  # torch<2.10
+    _varlen_attn = None
+
+_VARLEN_USES_WINDOW_SIZE = (
+    _varlen_attn is not None and "window_size" in inspect.signature(_varlen_attn).parameters
+)
+
+
+VisionAttentionBackend = Literal["auto", "sdpa", "varlen"]
+
+logger = logging.getLogger(__name__)
+
+
+@lru_cache
+def _log_resolved_backend(requested: str, resolved: str) -> None:
+    logger.info("Wall-X vision attention backend: %s (requested: %s)", resolved, requested)
+
+
+def _varlen_unavailable_reason(
+    hidden_states: torch.Tensor,
+    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
+) -> str | None:
+    if _varlen_attn is None:
+        return "torch.nn.attention.varlen is unavailable (PyTorch 2.10 or newer is required)"
+    if position_embeddings is None:
+        return "precomputed vision position embeddings were not provided"
+    if hidden_states.device.type != "cuda" or torch.version.cuda is None:
+        return "packed varlen attention requires an NVIDIA CUDA device"
+    if hidden_states.dtype not in {torch.float16, torch.bfloat16}:
+        return f"packed varlen attention requires float16 or bfloat16 inputs, got {hidden_states.dtype}"
+    major, _minor = torch.cuda.get_device_capability(hidden_states.device)
+    if major < 8:
+        return "packed varlen attention requires an NVIDIA Ampere GPU or newer"
+    return None
+
+
+def _supports_varlen_attention(
+    hidden_states: torch.Tensor,
+    position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
+) -> bool:
+    return _varlen_unavailable_reason(hidden_states, position_embeddings) is None
+
+
+class WallXVisionAttention(Qwen2_5_VLVisionAttention):
+    """Qwen2.5-VL vision attention with packed varlen and native SDPA fallback."""
+
+    def __init__(self, config, backend: VisionAttentionBackend):
+        super().__init__(config)
+        self.wallx_backend = backend
+        self._resolved_backend_key = None
+        self._resolved_backend = None
+
+    def _resolve_backend(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> str:
+        key = (
+            hidden_states.device.type,
+            hidden_states.device.index,
+            hidden_states.dtype,
+            position_embeddings is not None,
+        )
+        if self._resolved_backend_key == key:
+            return self._resolved_backend
+
+        use_varlen = self.wallx_backend != "sdpa" and _supports_varlen_attention(
+            hidden_states, position_embeddings
+        )
+        if self.wallx_backend == "varlen" and not use_varlen:
+            reason = _varlen_unavailable_reason(hidden_states, position_embeddings)
+            raise RuntimeError(f"Wall-X vision_attn_implementation='varlen' cannot be used: {reason}")
+
+        resolved_backend = "varlen" if use_varlen else "sdpa"
+        self._resolved_backend_key = key
+        self._resolved_backend = resolved_backend
+        _log_resolved_backend(self.wallx_backend, resolved_backend)
+        return resolved_backend
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        rotary_pos_emb: torch.Tensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del rotary_pos_emb
+
+        if self._resolve_backend(hidden_states, position_embeddings) == "sdpa":
+            return super().forward(
+                hidden_states=hidden_states,
+                cu_seqlens=cu_seqlens,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        seq_length = hidden_states.shape[0]
+        query_states, key_states, value_states = (
+            self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+        )
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb_vision(
+            query_states,
+            key_states,
+            cos,
+            sin,
+        )
+
+        if cu_seqlens.dtype != torch.int32:
+            cu_seqlens = cu_seqlens.to(dtype=torch.int32)
+        max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max().item())
+        varlen_kwargs = {"scale": self.scaling}
+        if _VARLEN_USES_WINDOW_SIZE:
+            varlen_kwargs["window_size"] = (-1, -1)
+        else:  # Stable PyTorch 2.10 API; pre-release variants used window_size.
+            varlen_kwargs["is_causal"] = False
+        attn_output = _varlen_attn(
+            query_states,
+            key_states,
+            value_states,
+            cu_seqlens,
+            cu_seqlens,
+            max_seqlen,
+            max_seqlen,
+            **varlen_kwargs,
+        )
+        attn_output = attn_output.reshape(seq_length, -1).contiguous()
+        return self.proj(attn_output)
+
+
+def configure_wall_x_vision_attention(
+    vision_model: nn.Module,
+    backend: VisionAttentionBackend,
+) -> None:
+    """Install Wall-X's scoped packed attention without changing checkpoint keys."""
+    if backend == "sdpa":
+        _log_resolved_backend(backend, "sdpa")
+        return
+    if backend == "varlen" and _varlen_attn is None:
+        raise RuntimeError(
+            "Wall-X vision_attn_implementation='varlen' requires torch.nn.attention.varlen "
+            "from PyTorch 2.10 or newer"
+        )
+    if backend == "auto" and _varlen_attn is None:
+        _log_resolved_backend(backend, "sdpa")
+        return
+
+    for block in vision_model.blocks:
+        previous_attention = block.attn
+        replacement = WallXVisionAttention(previous_attention.config, backend=backend)
+        replacement.to(
+            device=previous_attention.qkv.weight.device,
+            dtype=previous_attention.qkv.weight.dtype,
+        )
+        replacement.load_state_dict(previous_attention.state_dict(), strict=True)
+        replacement.train(previous_attention.training)
+        block.attn = replacement

src/lerobot/policies/wall_x/utils.py

@@ -116,6 +116,7 @@ def preprocesser_call(
     images: list | Any | None = None,
     text: str | list[str] | None = None,
     videos: list | Any | None = None,
+    device: torch.device | str | None = None,
     padding: bool | str = False,
     truncation: bool | None = None,
     max_length: int | None = None,
@@ -134,6 +135,7 @@ def preprocesser_call(
         images: Input images (PIL, numpy arrays, or torch tensors)
         text: Text or list of texts to tokenize
         videos: Input videos (numpy arrays or torch tensors)
+        device: Device on which image/video preprocessing should run
         padding: Whether to pad sequences to same length
         truncation: Whether to truncate sequences longer than max_length
         max_length: Maximum length for truncation/padding
@@ -151,7 +153,11 @@ def preprocesser_call(
     """
     # Process image inputs
     if images is not None and len(images) > 0:
-        image_inputs = processor.image_processor(images=images, return_tensors=return_tensors)
+        image_inputs = processor.image_processor(
+            images=images,
+            return_tensors=return_tensors,
+            device=device,
+        )
         image_grid_thw = image_inputs["image_grid_thw"]
     else:
         image_inputs = {}
@@ -159,7 +165,11 @@ def preprocesser_call(
 
     # Process video inputs
     if videos is not None:
-        videos_inputs = processor.image_processor(videos=videos, return_tensors=return_tensors)
+        videos_inputs = processor.image_processor(
+            videos=videos,
+            return_tensors=return_tensors,
+            device=device,
+        )
         video_grid_thw = videos_inputs["video_grid_thw"]
     else:
         videos_inputs = {}
@@ -413,10 +423,7 @@ def get_task_instruction(
         }
     )
 
-    if priority_order is not None:
-        priority_order = OrderedDict(priority_order)
-    else:
-        priority_order = default_priority_order
+    priority_order = OrderedDict(priority_order) if priority_order is not None else default_priority_order
 
     got_instruction = False
     task_instruction = ""
@@ -424,8 +431,7 @@ def get_task_instruction(
     # Sample instruction components based on priority probabilities
     for key, prob in priority_order.items():
         if key in frame_instruction_info and frame_instruction_info[key] != "":
-            if got_instruction:
-                if random.random() >= prob:
+            if got_instruction and random.random() >= prob:
                 continue
 
             task_instruction += f"\n{frame_instruction_info[key]}"
@@ -538,10 +544,7 @@ def img_key_mapping(img_keys: list[str]) -> list[str]:
         if key in CAMERA_NAME_MAPPING:
             key = CAMERA_NAME_MAPPING[key]
         else:
-            if "view" in key:
-                key = key.replace("_", " ")
-            else:
-                key = key + " view"
+            key = key.replace("_", " ") if "view" in key else key + " view"
         processed_img_keys.append(key)
     return processed_img_keys
 

src/lerobot/policies/xvla/processor_xvla.py

@@ -22,19 +22,14 @@ import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
-    AddBatchDimensionProcessorStep,
-    DeviceProcessorStep,
-    NormalizerProcessorStep,
     ObservationProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
-    RenameObservationsProcessorStep,
     TokenizerProcessorStep,
-    UnnormalizerProcessorStep,
-    policy_action_to_transition,
-    transition_to_policy_action,
+    make_default_policy_processor_steps,
+    make_policy_processor_pipelines,
 )
 from lerobot.types import EnvTransition, TransitionKey
 from lerobot.utils.constants import (
@@ -42,8 +37,6 @@ from lerobot.utils.constants import (
     OBS_IMAGES,
     OBS_PREFIX,
     OBS_STATE,
-    POLICY_POSTPROCESSOR_DEFAULT_NAME,
-    POLICY_PREPROCESSOR_DEFAULT_NAME,
 )
 
 from .configuration_xvla import XVLAConfig
@@ -61,10 +54,11 @@ def make_xvla_pre_post_processors(
     Build the LeRobot processor pipelines for XVLA.
     """
 
-    features = {**config.input_features, **config.output_features}
+    steps = make_default_policy_processor_steps(config, dataset_stats)
+
     input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
+        steps.rename_observations,
+        steps.add_batch_dim,
         TokenizerProcessorStep(
             tokenizer_name=config.tokenizer_name,
             max_length=config.tokenizer_max_length,
@@ -74,32 +68,15 @@ def make_xvla_pre_post_processors(
         XVLAImageToFloatProcessorStep(),
         XVLAImageNetNormalizeProcessorStep(),
         XVLAAddDomainIdProcessorStep(),
-        DeviceProcessorStep(device=config.device),
-        NormalizerProcessorStep(
-            features=features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
+        steps.to_device,
+        steps.normalize,
     ]
     output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features,
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        DeviceProcessorStep(device="cpu"),
+        steps.unnormalize,
+        steps.to_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,
-        ),
-    )
+    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
 
 
 # Custom XVLA processor steps

src/lerobot/processor/__init__.py

@@ -42,10 +42,14 @@ from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorTo
 from .device_processor import DeviceProcessorStep
 from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
 from .factory import (
+    DefaultPolicyProcessorSteps,
+    make_default_policy_processor_steps,
+    make_default_pre_post_processors,
     make_default_processors,
     make_default_robot_action_processor,
     make_default_robot_observation_processor,
     make_default_teleop_action_processor,
+    make_policy_processor_pipelines,
 )
 from .gym_action_processor import (
     Numpy2TorchActionProcessorStep,
@@ -129,10 +133,14 @@ __all__ = [
     "ImageCropResizeProcessorStep",
     "InfoProcessorStep",
     "InterventionActionProcessorStep",
+    "DefaultPolicyProcessorSteps",
+    "make_default_policy_processor_steps",
+    "make_default_pre_post_processors",
     "make_default_processors",
     "make_default_teleop_action_processor",
     "make_default_robot_action_processor",
     "make_default_robot_observation_processor",
+    "make_policy_processor_pipelines",
     "AbsoluteActionsProcessorStep",
     "RelativeActionsProcessorStep",
     "MapDeltaActionToRobotActionStep",

src/lerobot/processor/batch_processor.py

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

src/lerobot/processor/factory.py

@@ -14,15 +14,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from lerobot.types import RobotAction, RobotObservation
+from dataclasses import dataclass
+from typing import Any
 
+import torch
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.types import PolicyAction, RobotAction, RobotObservation
+from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
+
+from .batch_processor import AddBatchDimensionProcessorStep
 from .converters import (
     observation_to_transition,
+    policy_action_to_transition,
     robot_action_observation_to_transition,
     transition_to_observation,
+    transition_to_policy_action,
     transition_to_robot_action,
 )
-from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
+from .device_processor import DeviceProcessorStep
+from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
+from .pipeline import (
+    IdentityProcessorStep,
+    PolicyProcessorPipeline,
+    ProcessorStep,
+    RobotProcessorPipeline,
+)
+from .rename_processor import RenameObservationsProcessorStep
 
 
 def make_default_teleop_action_processor() -> RobotProcessorPipeline[
@@ -61,3 +79,97 @@ def make_default_processors():
     robot_action_processor = make_default_robot_action_processor()
     robot_observation_processor = make_default_robot_observation_processor()
     return (teleop_action_processor, robot_action_processor, robot_observation_processor)
+
+
+@dataclass
+class DefaultPolicyProcessorSteps:
+    """The canonical processor steps shared by most policies' pre/post pipelines.
+
+    Policies compose these in their own order (step ORDER is a Hub-serialized contract
+    and intentionally stays explicit per policy) and interleave their custom steps.
+    """
+
+    rename_observations: RenameObservationsProcessorStep
+    add_batch_dim: AddBatchDimensionProcessorStep
+    to_device: DeviceProcessorStep
+    normalize: NormalizerProcessorStep
+    unnormalize: UnnormalizerProcessorStep
+    to_cpu: DeviceProcessorStep
+
+
+def make_default_policy_processor_steps(
+    config: PreTrainedConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    *,
+    normalizer_device: torch.device | str | None = None,
+) -> DefaultPolicyProcessorSteps:
+    """Construct the canonical policy processor steps from a policy config.
+
+    Args:
+        config: A `PreTrainedConfig` providing `device`, `input_features`,
+            `output_features` and `normalization_mapping`.
+        dataset_stats: Dataset statistics used for (un)normalization.
+        normalizer_device: Device passed to `NormalizerProcessorStep` (some policies pin
+            their normalization stats to the policy device; most leave it unset).
+    """
+    return DefaultPolicyProcessorSteps(
+        rename_observations=RenameObservationsProcessorStep(rename_map={}),
+        add_batch_dim=AddBatchDimensionProcessorStep(),
+        to_device=DeviceProcessorStep(device=config.device),
+        normalize=NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+            device=normalizer_device,
+        ),
+        unnormalize=UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
+        to_cpu=DeviceProcessorStep(device="cpu"),
+    )
+
+
+def make_policy_processor_pipelines(
+    input_steps: list[ProcessorStep],
+    output_steps: list[ProcessorStep],
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """Wrap pre/post step lists into the canonical policy pipeline pair.
+
+    Uses the standard pipeline names (which determine the serialized JSON filenames on
+    the Hub) and the standard policy-action converters on the postprocessor.
+    """
+    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 make_default_pre_post_processors(
+    config: PreTrainedConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    *,
+    normalizer_device: torch.device | str | None = None,
+) -> tuple[
+    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
+    PolicyProcessorPipeline[PolicyAction, PolicyAction],
+]:
+    """The pure-scaffold policy pipeline pair: Rename -> Batch -> Device -> Normalize,
+    and Unnormalize -> Device(cpu). Policies with custom steps or a different step order
+    compose `make_default_policy_processor_steps` themselves instead.
+    """
+    s = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=normalizer_device)
+    return make_policy_processor_pipelines(
+        input_steps=[s.rename_observations, s.add_batch_dim, s.to_device, s.normalize],
+        output_steps=[s.unnormalize, s.to_cpu],
+    )

src/lerobot/processor/pipeline.py

@@ -41,7 +41,7 @@ from pathlib import Path
 from typing import Any, TypedDict, TypeVar, cast
 
 import torch
-from huggingface_hub import hf_hub_download, snapshot_download
+from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file, save_file
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
@@ -205,10 +205,6 @@ class ProcessorStep(ABC):
         """
         return None
 
-    def save_artifacts(self, save_directory: Path) -> dict[str, str]:
-        """Save non-tensor assets and map constructor arguments to relative paths."""
-        return {}
-
     def reset(self) -> None:
         """Resets the internal state of the processor step, if any."""
         return None
@@ -553,22 +549,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         pipeline_config = self.get_config()
         pipeline_state_dict = self.state_dict()
 
-        for processor_step, step_entry in zip(self.steps, pipeline_config["steps"], strict=True):
-            artifacts = processor_step.save_artifacts(save_directory)
-            if artifacts:
-                for config_key, relative_path in artifacts.items():
-                    artifact_path = Path(relative_path)
-                    if artifact_path.is_absolute() or ".." in artifact_path.parts:
-                        raise ValueError(
-                            f"Processor artifact path must be relative to the checkpoint: {relative_path!r}"
-                        )
-                    if not (save_directory / artifact_path).exists():
-                        raise FileNotFoundError(
-                            f"Processor step did not save declared artifact '{relative_path}'"
-                        )
-                    step_entry["config"][config_key] = artifact_path.as_posix()
-                step_entry["artifacts"] = artifacts
-
         for state_key, step_state_dict in pipeline_state_dict.items():
             state_filename = f"{state_key}.safetensors"
             save_file(step_state_dict, save_directory / state_filename)
@@ -751,12 +731,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
 
         # 3. Build steps with overrides
         steps, validated_overrides = cls._build_steps_with_overrides(
-            loaded_config,
-            overrides or {},
-            model_id,
-            base_path,
-            config_filename,
-            hub_download_kwargs,
+            loaded_config, overrides or {}, model_id, base_path, hub_download_kwargs
         )
 
         # 4. Validate that all overrides were used
@@ -945,7 +920,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         overrides: dict[str, Any],
         model_id: str,
         base_path: Path | None,
-        config_filename: str,
         hub_download_kwargs: dict[str, Any],
     ) -> tuple[list[ProcessorStep], set[str]]:
         """Build all processor steps with overrides and state loading.
@@ -998,67 +972,13 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
             ImportError: If a step class cannot be imported or found in registry
             ValueError: If a step cannot be instantiated with its configuration
         """
-        loaded_config = deepcopy(loaded_config)
-        cls._resolve_artifact_paths(
-            loaded_config,
-            model_id,
-            base_path,
-            config_filename,
-            hub_download_kwargs,
-        )
         steps, remaining_override_keys = cls._build_steps_from_config(loaded_config, overrides)
 
         for step_instance, step_entry in zip(steps, loaded_config["steps"], strict=True):
-            cls._load_step_state(
-                step_instance,
-                step_entry,
-                model_id,
-                base_path,
-                config_filename,
-                hub_download_kwargs,
-            )
+            cls._load_step_state(step_instance, step_entry, model_id, base_path, hub_download_kwargs)
 
         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,
@@ -1218,7 +1138,6 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
         step_entry: dict[str, Any],
         model_id: str,
         base_path: Path | None,
-        config_filename: str,
         hub_download_kwargs: dict[str, Any],
     ) -> None:
         """Load state dictionary for a processor step if available.
@@ -1276,7 +1195,7 @@ class DataProcessorPipeline[TInput, TOutput](HubMixin):
             # Download from Hub
             state_path = hf_hub_download(
                 repo_id=model_id,
-                filename=(Path(config_filename).parent / state_filename).as_posix(),
+                filename=state_filename,
                 repo_type="model",
                 **hub_download_kwargs,
             )

src/lerobot/processor/render_messages_processor.py

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

src/lerobot/processor/tokenizer_processor.py

@@ -25,7 +25,6 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass, field
-from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -33,7 +32,6 @@ import torch
 from lerobot.configs import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.types import EnvTransition, RobotObservation, TransitionKey
 from lerobot.utils.constants import (
-    ACTION_CODE_TOKEN_MASK,
     ACTION_TOKEN_MASK,
     ACTION_TOKENS,
     OBS_LANGUAGE_ATTENTION_MASK,
@@ -351,7 +349,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
     max_action_tokens: int = 256
     fast_skip_tokens: int = 128
     paligemma_tokenizer_name: str = "google/paligemma-3b-pt-224"
-    allow_truncation: bool = True
     # Internal tokenizer instance (not part of the config)
     action_tokenizer: Any = field(default=None, init=False, repr=False)
     _paligemma_tokenizer: Any = field(default=None, init=False, repr=False)
@@ -415,15 +412,14 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             # During inference, no action is available, skip tokenization
             return new_transition
 
-        # Tokenize and get masks for the full formatted sequence and the discrete action codes.
-        tokens, mask, code_mask = self._tokenize_action(action)
+        # Tokenize and get both tokens and mask
+        tokens, mask = self._tokenize_action(action)
 
         # Store mask in complementary data
         complementary_data = new_transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
         if complementary_data is None:
             complementary_data = {}
         complementary_data[ACTION_TOKEN_MASK] = mask
-        complementary_data[ACTION_CODE_TOKEN_MASK] = code_mask
         complementary_data[ACTION_TOKENS] = tokens
         new_transition[TransitionKey.COMPLEMENTARY_DATA] = complementary_data
         return new_transition
@@ -434,7 +430,7 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
         """
         return self._paligemma_tokenizer.vocab_size - 1 - self.fast_skip_tokens - tokens
 
-    def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    def _tokenize_action(self, action: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Tokenizes the action tensor and creates a mask.
 
@@ -463,7 +459,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
         # The fast tokenizer expects action data and returns token IDs
         tokens_list = []
         masks_list = []
-        code_masks_list = []
 
         for i in range(batch_size):
             # Tokenize single action (move to CPU first as tokenizer uses scipy which requires numpy)
@@ -481,82 +476,65 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
             if tokens.dim() > 1:
                 tokens = tokens.flatten()
 
-            action_code_tokens = self._act_tokens_to_paligemma_tokens(tokens)
             bos_id = self._paligemma_tokenizer.bos_token_id
-            prompt_tokens = torch.tensor(
-                self._paligemma_tokenizer.encode("Action: ", add_special_tokens=False),
-                device=action.device,
-            )
-            end_tokens = torch.tensor(self._paligemma_tokenizer.encode("|"), device=action.device)
-
-            code_start = 1 + len(prompt_tokens)
-            code_end = code_start + len(action_code_tokens)
+            # add bos
             tokens = torch.cat(
                 [
                     torch.tensor([bos_id], device=action.device),
-                    prompt_tokens,
-                    action_code_tokens,
-                    end_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),
                 ]
             )
-            code_mask = torch.zeros(len(tokens), dtype=torch.bool, device=action.device)
-            code_mask[code_start:code_end] = True
 
             # Truncate or pad to max_action_tokens
             if len(tokens) > self.max_action_tokens:
-                if not self.allow_truncation:
-                    raise ValueError(
-                        f"FAST action sequence has {len(tokens)} tokens, exceeding "
-                        f"max_action_tokens={self.max_action_tokens}."
-                    )
                 logging.warning(
                     f"Token length ({len(tokens)}) exceeds max length ({self.max_action_tokens}), truncating. "
                     "Consider increasing the `max_action_tokens` in your model config if this happens frequently."
                 )
                 tokens = tokens[: self.max_action_tokens]
-                code_mask = code_mask[: self.max_action_tokens]
                 mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device)
             else:
-                pad_len = self.max_action_tokens - len(tokens)
                 mask = torch.cat(
                     [
                         torch.ones(len(tokens), dtype=torch.bool, device=action.device),
-                        torch.zeros(pad_len, dtype=torch.bool, device=action.device),
+                        torch.zeros(
+                            self.max_action_tokens - len(tokens), dtype=torch.bool, device=action.device
+                        ),
                     ]
                 )
-                code_mask = torch.nn.functional.pad(code_mask, (0, pad_len), value=False)
                 # Pad tokens with zeros
-                tokens = torch.nn.functional.pad(tokens, (0, pad_len), value=0)
+                tokens = torch.nn.functional.pad(tokens, (0, self.max_action_tokens - len(tokens)), value=0)
 
             tokens_list.append(tokens)
             masks_list.append(mask)
-            code_masks_list.append(code_mask)
 
         # Stack into batched tensors
         tokens_batch = torch.stack(tokens_list, dim=0)  # (B, max_action_tokens)
         masks_batch = torch.stack(masks_list, dim=0)  # (B, max_action_tokens)
-        code_masks_batch = torch.stack(code_masks_list, dim=0)  # (B, max_action_tokens)
 
         # Remove batch dimension if input was single sample
         if single_sample:
             tokens_batch = tokens_batch.squeeze(0)
             masks_batch = masks_batch.squeeze(0)
-            code_masks_batch = code_masks_batch.squeeze(0)
 
         # Move to the same device as the input
         if device is not None:
             tokens_batch = tokens_batch.to(device)
             masks_batch = masks_batch.to(device)
-            code_masks_batch = code_masks_batch.to(device)
 
-        return tokens_batch, masks_batch, code_masks_batch
+        return tokens_batch, masks_batch
 
     def action(self, action: torch.Tensor) -> torch.Tensor:
         """
         This method is not used since we override __call__.
         Required by ActionProcessorStep ABC.
         """
-        tokens, _, _ = self._tokenize_action(action)
+        tokens, _ = self._tokenize_action(action)
         return tokens
 
     def get_config(self) -> dict[str, Any]:
@@ -572,9 +550,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
         config = {
             "trust_remote_code": self.trust_remote_code,
             "max_action_tokens": self.max_action_tokens,
-            "fast_skip_tokens": self.fast_skip_tokens,
-            "paligemma_tokenizer_name": self.paligemma_tokenizer_name,
-            "allow_truncation": self.allow_truncation,
         }
 
         # Only save tokenizer_name if it was used to create the tokenizer
@@ -583,14 +558,6 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
 
         return config
 
-    def save_artifacts(self, save_directory: Path) -> dict[str, str]:
-        artifact_path = Path("action_tokenizer")
-        save_pretrained = getattr(self.action_tokenizer, "save_pretrained", None)
-        if save_pretrained is None:
-            raise TypeError("Action tokenizer must implement save_pretrained() to save a portable pipeline.")
-        save_pretrained(save_directory / artifact_path)
-        return {"action_tokenizer_name": artifact_path.as_posix()}
-
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:

src/lerobot/rewards/pretrained.py

@@ -21,8 +21,6 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, TypeVar
 
-import packaging
-import safetensors
 from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
@@ -30,6 +28,7 @@ from safetensors.torch import load_model as load_model_as_safetensor, save_model
 from torch import Tensor, nn
 
 from lerobot.configs.rewards import RewardModelConfig
+from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
 
 if TYPE_CHECKING:
@@ -129,29 +128,13 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        # Create base kwargs
-        kwargs = {"strict": strict}
-
-        # Add device parameter for newer versions that support it
-        if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
-            kwargs["device"] = map_location
-
-        # Load the model with appropriate kwargs
-        missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
+        missing_keys, unexpected_keys = load_model_as_safetensor(
+            model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
+        )
         if missing_keys:
             logging.warning(f"Missing key(s) when loading model: {missing_keys}")
         if unexpected_keys:
             logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
-
-        # For older versions, manually move to device if needed
-        if "device" not in kwargs and map_location != "cpu":
-            logging.warning(
-                "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
-                " This means that the model is loaded on 'cpu' first and then copied to the device."
-                " This leads to a slower loading time."
-                " Please update safetensors to version 0.4.3 or above for improved performance."
-            )
-            model.to(map_location)
         return model
 
     def get_optim_params(self):

src/lerobot/robots/utils.py

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

src/lerobot/runtime/__init__.py

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

src/lerobot/runtime/adapter.py

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

src/lerobot/runtime/language_runtime.py

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

src/lerobot/runtime/registry.py

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

src/lerobot/runtime/sim_robocasa.py

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

src/lerobot/scripts/lerobot_annotate.py

@@ -28,7 +28,12 @@ For distributed runs, see ``examples/annotations/run_hf_job.py``.
 """
 
 import logging
+from contextlib import suppress
 from pathlib import Path
+from typing import TYPE_CHECKING
+
+from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub.errors import RevisionNotFoundError
 
 from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
 from lerobot.annotations.steerable_pipeline.executor import Executor
@@ -42,6 +47,12 @@ from lerobot.annotations.steerable_pipeline.validator import StagingValidator
 from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
 from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
 from lerobot.configs import parser
+from lerobot.utils.import_utils import _datasets_available, require_package
+
+if TYPE_CHECKING or _datasets_available:
+    from lerobot.datasets.dataset_metadata import CODEBASE_VERSION
+    from lerobot.datasets.io_utils import load_info
+    from lerobot.datasets.utils import create_lerobot_dataset_card
 
 logger = logging.getLogger(__name__)
 
@@ -50,8 +61,6 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
     if cfg.root is not None:
         return Path(cfg.root)
     if cfg.repo_id is not None:
-        from huggingface_hub import snapshot_download
-
         return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset"))
     raise ValueError("Either --root or --repo_id must be provided.")
 
@@ -125,10 +134,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
 
     Pushes to ``cfg.new_repo_id`` when set, otherwise back to ``cfg.repo_id``.
     """
-    from huggingface_hub import HfApi  # noqa: PLC0415
-
-    from lerobot.datasets.io_utils import load_info  # noqa: PLC0415
-    from lerobot.datasets.utils import create_lerobot_dataset_card  # noqa: PLC0415
+    require_package("datasets", "dataset")
 
     repo_id = cfg.new_repo_id or cfg.repo_id
     commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
@@ -163,8 +169,6 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
     # ``RevisionNotFoundError``. Read the version straight from the
     # dataset's own ``meta/info.json`` so we tag whatever the writer
     # actually wrote (no accidental drift if the codebase floor moves).
-    from lerobot.datasets.dataset_metadata import CODEBASE_VERSION  # noqa: PLC0415
-
     version_tag = (
         dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
     )
@@ -178,10 +182,6 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         tag_kwargs["revision"] = revision
 
     try:
-        from contextlib import suppress  # noqa: PLC0415
-
-        from huggingface_hub.errors import RevisionNotFoundError  # noqa: PLC0415
-
         with suppress(RevisionNotFoundError):
             api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
         api.create_tag(**tag_kwargs)

src/lerobot/scripts/lerobot_eval.py

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

src/lerobot/scripts/lerobot_rollout.py

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

src/lerobot/scripts/lerobot_train.py

@@ -20,11 +20,9 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
-import os
 import sys
 import time
 from contextlib import nullcontext
-from datetime import timedelta
 from pprint import pformat
 from typing import TYPE_CHECKING, Any
 
@@ -83,7 +81,6 @@ def update_policy(
     lr_scheduler=None,
     lock=None,
     sample_weighter=None,
-    log_metrics: bool = True,
 ) -> tuple[MetricsTracker, dict | None]:
     """
     Performs a single training step to update the policy's weights.
@@ -101,7 +98,6 @@ def update_policy(
         lr_scheduler: An optional learning rate scheduler.
         lock: An optional lock for thread-safe optimizer updates.
         sample_weighter: Optional SampleWeighter instance for per-sample loss weighting.
-        log_metrics: Whether to synchronize and record GPU metrics this step.
 
     Returns:
         A tuple containing:
@@ -169,20 +165,15 @@ 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.
+    # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
     if output_dict:
-            output_dict = {
-                k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items()
-            }
+        train_metrics.update_metrics(output_dict)
     return train_metrics, output_dict
 
 
@@ -210,7 +201,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     require_package("accelerate", extra="training")
     from accelerate import Accelerator
-    from accelerate.utils import DistributedDataParallelKwargs, DistributedType, InitProcessGroupKwargs
+    from accelerate.utils import DistributedDataParallelKwargs, DistributedType
 
     cfg.validate()
 
@@ -219,16 +210,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes
     # We set find_unused_parameters=True to handle models with conditional computation
     if accelerator is None:
-        # Static graphs restore DDP overlap when conditional parameter usage is stable.
-        # Environment flags retain the existing defaults.
-        ddp_find_unused = os.environ.get("LEROBOT_DDP_FIND_UNUSED", "1") == "1"
-        ddp_static_graph = os.environ.get("LEROBOT_DDP_STATIC_GRAPH", "0") == "1"
-        ddp_kwargs = DistributedDataParallelKwargs(
-            find_unused_parameters=ddp_find_unused and not ddp_static_graph,
-            static_graph=ddp_static_graph,
-        )
-        # Allow rank 0 enough time to index large datasets before other ranks leave the barrier.
-        ipg_kwargs = InitProcessGroupKwargs(timeout=timedelta(hours=2))
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
         # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting.
         # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training).
         force_cpu = cfg.trainable_config.device == "cpu"
@@ -238,7 +220,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         accelerator = Accelerator(
             step_scheduler_with_optimizer=False,
             mixed_precision=mixed_precision,
-            kwargs_handlers=[ddp_kwargs, ipg_kwargs],
+            kwargs_handlers=[ddp_kwargs],
             cpu=force_cpu,
         )
 
@@ -334,14 +316,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
 
     active_cfg = cfg.trainable_config
     processor_pretrained_path = active_cfg.pretrained_path
-    # A weight checkpoint may contain PI05 or differently configured PI052 processors.
-    if cfg.policy.type == "pi052" and processor_pretrained_path is not None and not cfg.resume:
-        logging.warning(
-            "pi052 is loading pretrained weights from %s, but building processors from the current "
-            "pi052 config so recipe text labels and FAST action labels are generated.",
-            processor_pretrained_path,
-        )
-        processor_pretrained_path = None
 
     processor_kwargs = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
@@ -350,13 +324,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
-    if cfg.policy.type in {"pi0_fast", "pi052"}:
-        processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
-        processor_kwargs["dataset_revision"] = cfg.dataset.revision
-        processor_kwargs["dataset_episodes"] = cfg.dataset.episodes
-        processor_kwargs["dataset_exclude_episodes"] = cfg.dataset.exclude_episodes
-        processor_kwargs["dataset_root"] = cfg.dataset.root
-
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
         preprocessor_overrides = {
             "device_processor": {"device": device.type},
@@ -453,17 +420,13 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         # same permutation. accelerate then shards it disjointly across ranks via BatchSamplerShard
         # without needing a `generator` attribute to synchronize an RNG, and resume is sample-exact.
         shuffle = False
-        from_indices = dataset.meta.episodes["dataset_from_index"]
-        to_indices = dataset.meta.episodes["dataset_to_index"]
-        seed = cfg.seed if cfg.seed is not None else 0
-
         sampler = EpisodeAwareSampler(
-            from_indices,
-            to_indices,
+            dataset.meta.episodes["dataset_from_index"],
+            dataset.meta.episodes["dataset_to_index"],
             episode_indices_to_use=dataset.episodes,
             drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0),
             shuffle=True,
-            seed=seed,
+            seed=cfg.seed if cfg.seed is not None else 0,
             absolute_to_relative_idx=dataset.absolute_to_relative_idx,
         )
         if cfg.resume and step > 0:
@@ -501,8 +464,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # declares language columns; otherwise stay on PyTorch's default
     # collate so non-language training runs are unaffected.
     collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
-    # Allow spawn/forkserver workers where forking large rank processes exhausts memory.
-    mp_context = os.environ.get("LEROBOT_DATALOADER_MP_CONTEXT") or None
     dataloader = torch.utils.data.DataLoader(
         dataset,
         num_workers=cfg.num_workers,
@@ -514,7 +475,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         collate_fn=collate_fn,
         prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
         persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
-        multiprocessing_context=mp_context if cfg.num_workers > 0 else None,
     )
 
     # Build eval dataloader if a held-out split exists
@@ -542,7 +502,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             collate_fn=eval_collate_fn,
             prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
             persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
-            multiprocessing_context=mp_context if cfg.num_workers > 0 else None,
         )
 
     # Prepare everything with accelerator
@@ -616,10 +575,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        # Synchronize GPU metrics only for updates that will be logged.
-        log_metrics = cfg.log_freq > 0 and (step + 1) % cfg.log_freq == 0
-
-        train_tracker, output_dict = update_policy(
+        train_tracker, _ = update_policy(
             train_tracker,
             policy,
             batch,
@@ -628,7 +584,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             accelerator=accelerator,
             lr_scheduler=lr_scheduler,
             sample_weighter=sample_weighter,
-            log_metrics=log_metrics,
         )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
@@ -653,9 +608,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     train_tracker.samples_per_s = effective_batch_size / step_time
                 logging.info(train_tracker)
                 if wandb_logger:
+                    # Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
+                    # tracker by update_policy, so to_dict() already carries their windowed,
+                    # rank-reduced averages — no per-step output_dict passthrough needed.
                     wandb_log_dict = train_tracker.to_dict()
-                    if output_dict:
-                        wandb_log_dict.update(output_dict)
                     # Log sample weighting statistics if enabled
                     if sample_weighter is not None:
                         weighter_stats = sample_weighter.get_stats()
@@ -728,11 +684,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             if is_main_process:
                 step_id = get_step_identifier(step, cfg.steps)
                 logging.info(f"Eval policy at step {step}")
-                eval_target_policy = accelerator.unwrap_model(policy)
                 with torch.no_grad(), accelerator.autocast():
                     eval_info = eval_policy_all(
                         envs=eval_env,  # dict[suite][task_id] -> vec_env
-                        policy=eval_target_policy,
+                        policy=accelerator.unwrap_model(policy),
                         env_preprocessor=env_preprocessor,
                         env_postprocessor=env_postprocessor,
                         preprocessor=preprocessor,

src/lerobot/utils/collate.py

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

src/lerobot/utils/constants.py

@@ -26,7 +26,6 @@ OBS_IMAGES = OBS_IMAGE + "s"
 OBS_LANGUAGE = OBS_STR + ".language"
 OBS_LANGUAGE_TOKENS = OBS_LANGUAGE + ".tokens"
 OBS_LANGUAGE_ATTENTION_MASK = OBS_LANGUAGE + ".attention_mask"
-OBS_LANGUAGE_CAUSAL_MARKS = OBS_LANGUAGE + ".causal_marks"
 OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
 OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
 OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
@@ -35,7 +34,6 @@ ACTION = "action"
 ACTION_PREFIX = ACTION + "."
 ACTION_TOKENS = ACTION + ".tokens"
 ACTION_TOKEN_MASK = ACTION + ".token_mask"
-ACTION_CODE_TOKEN_MASK = ACTION + ".code_token_mask"
 REWARD = "next.reward"
 TRUNCATED = "next.truncated"
 DONE = "next.done"

src/lerobot/utils/device_utils.py

@@ -59,6 +59,20 @@ def get_safe_torch_device(try_device: str, log: bool = False) -> torch.device:
     return device
 
 
+def resolve_safetensors_device(map_location: str | torch.device) -> str:
+    """Resolve a device string for a safetensors load, working around a device-mapping quirk.
+
+    safetensors' load maps the bare string "cuda" to cuda:0 regardless of the current device
+    (unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under multi-GPU
+    accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it before sharding.
+    Resolve "cuda" to the concrete current-device index so each rank loads onto its own GPU.
+    """
+    map_location = str(map_location)
+    if map_location == "cuda" and torch.cuda.is_available():
+        return f"cuda:{torch.cuda.current_device()}"
+    return map_location
+
+
 def get_safe_dtype(dtype: torch.dtype, device: str | torch.device):
     """
     mps is currently not compatible with float64

src/lerobot/utils/io_utils.py

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

src/lerobot/utils/logging_utils.py

@@ -104,8 +104,7 @@ class MetricsTracker:
         "episodes",
         "epochs",
         "accelerator",
-        "_tensor_sums",
-        "_tensor_counts",
+        "_caller_metrics",
     ]
 
     def __init__(
@@ -131,8 +130,9 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
         self.accelerator = accelerator
-        self._tensor_sums: dict[str, torch.Tensor] = {}
-        self._tensor_counts: dict[str, int] = {}
+        # 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)
 
     def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
         if name in self.__dict__:
@@ -160,21 +160,20 @@ 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 update_metrics(self, values: dict[str, Any]) -> None:
+        """Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
 
-    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()
+        Non-numeric values and bools are ignored.
+        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)):
+                continue
+            if name in self._caller_metrics:
+                continue
+            if name not in self.metrics:
+                self.metrics[name] = AverageMeter(name, ":.3f", reduction="mean")
+            self.metrics[name].update(float(value))
 
     def reduce_across_ranks(self) -> None:
         """
@@ -196,21 +195,10 @@ class MetricsTracker:
         if not buckets:
             return
 
-        # NB: don't use ``accelerator.reduce(..., reduction="max")`` — accelerate only implements
-        # "sum"/"mean" (it always all-reduces with SUM and divides for "mean"), so "max" silently
-        # returns the SUM across ranks, inflating every "max" metric by ``num_processes`` (e.g. a
-        # 3.5s step reported as 28s on 8 GPUs). Gather per-rank values and reduce them explicitly.
         device = self.accelerator.device
-        num_processes = self.accelerator.num_processes
         for reduction, names in buckets.items():
-            local = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
-            gathered = self.accelerator.gather(local).view(num_processes, len(names))
-            if reduction == "max":
-                reduced = gathered.amax(dim=0)
-            elif reduction == "sum":
-                reduced = gathered.sum(dim=0)
-            else:  # "mean"
-                reduced = gathered.mean(dim=0)
+            tensor = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
+            reduced = self.accelerator.reduce(tensor, reduction=reduction)
             for name, value in zip(names, reduced.tolist(), strict=True):
                 meter = self.metrics[name]
                 # Preserve avg == sum / count so a later .update() on this meter accumulates
@@ -247,5 +235,3 @@ class MetricsTracker:
         """Resets average meters."""
         for m in self.metrics.values():
             m.reset()
-        self._tensor_sums.clear()
-        self._tensor_counts.clear()

src/lerobot/utils/rerun_visualization.py

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

src/lerobot/utils/video_annotation.py

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

tests/configs/test_recipe.py

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

tests/datasets/test_language_render.py

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

tests/datasets/test_sampler.py

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

tests/envs/test_robocasa_ordering.py

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

tests/policies/common/test_flow_matching.py

@@ -0,0 +1,245 @@
+#!/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.
+
+"""Behavior-pinning tests for the shared flow-matching sampling primitives.
+
+``euler_integrate`` is compared against a verbatim copy of the historical pi0/pi05/
+smolvla sampling loop (including its RTC hook semantics): any divergence from that
+reference is a behavior change for released checkpoints.
+"""
+
+import torch
+
+from lerobot.policies.common.flow_matching import (
+    _beta_distribution,
+    euler_integrate,
+    sample_beta,
+    sample_noise,
+    sample_time_beta,
+)
+
+
+def test_sample_beta_range_dtype_and_reproducibility():
+    torch.manual_seed(0)
+    s1 = sample_beta(1.5, 1.0, 4096, "cpu")
+    torch.manual_seed(0)
+    s2 = sample_beta(1.5, 1.0, 4096, "cpu")
+    assert torch.equal(s1, s2)
+    assert s1.shape == (4096,) and s1.dtype == torch.float32
+    assert s1.min() >= 0.0 and s1.max() <= 1.0
+    # Beta(1.5, 1.0) mean is 1.5/2.5 = 0.6.
+    assert abs(s1.mean().item() - 0.6) < 0.02
+
+
+def test_sample_time_beta_openpi_convention():
+    torch.manual_seed(1)
+    time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, offset=0.001)
+    assert time.dtype == torch.float32
+    assert time.min() >= 0.001 and time.max() <= 1.0
+    # Exact composition: Beta sample * scale + offset, same RNG stream.
+    torch.manual_seed(1)
+    expected = sample_beta(1.5, 1.0, 4096, "cpu") * 0.999 + 0.001
+    torch.testing.assert_close(time, expected, rtol=0, atol=0)
+
+
+def test_sample_noise_seeded():
+    torch.manual_seed(2)
+    n1 = sample_noise((2, 8, 4), "cpu")
+    torch.manual_seed(2)
+    n2 = sample_noise((2, 8, 4), "cpu")
+    assert torch.equal(n1, n2)
+    assert n1.dtype == torch.float32 and n1.shape == (2, 8, 4)
+
+
+def test_sample_noise_normal_is_default():
+    torch.manual_seed(2)
+    default = sample_noise((2, 8, 4), "cpu")
+    torch.manual_seed(2)
+    normal = sample_noise((2, 8, 4), "cpu", distribution="normal")
+    assert torch.equal(default, normal)
+
+
+def test_sample_noise_uniform_evo1():
+    torch.manual_seed(2)
+    n = sample_noise((4096,), "cpu", distribution="uniform")
+    assert n.dtype == torch.float32
+    assert n.min() >= -1.0 and n.max() < 1.0
+    # evo1's rand_like * 2 - 1 has mean ~0 over [-1, 1).
+    assert abs(n.mean().item()) < 0.05
+    # Exact match to the historical evo1 expression on the same RNG stream.
+    torch.manual_seed(2)
+    expected = torch.rand((4096,), dtype=torch.float32) * 2 - 1
+    torch.testing.assert_close(n, expected, rtol=0, atol=0)
+
+
+def test_sample_noise_invalid_distribution():
+    import pytest
+
+    with pytest.raises(ValueError, match="Unknown noise distribution"):
+        sample_noise((2, 2), "cpu", distribution="bogus")
+
+
+def test_sample_time_beta_forward_complement_convention():
+    # groot/wall_x forward convention: t = (1 - beta) * 0.999.
+    torch.manual_seed(9)
+    time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, complement=True)
+    torch.manual_seed(9)
+    expected = (1.0 - sample_beta(1.5, 1.0, 4096, "cpu")) * 0.999
+    torch.testing.assert_close(time, expected, rtol=0, atol=0)
+
+
+def test_sample_time_beta_evo1_clamp():
+    # evo1: Beta(2, 2) clamped to [0.02, 0.98].
+    torch.manual_seed(10)
+    time = sample_time_beta(4096, "cpu", alpha=2.0, beta=2.0, clamp_min=0.02, clamp_max=0.98)
+    assert time.min() >= 0.02 and time.max() <= 0.98
+
+
+def test_sample_beta_distribution_is_cached():
+    a = _beta_distribution(1.5, 1.0)
+    b = _beta_distribution(1.5, 1.0)
+    assert a is b
+    assert _beta_distribution(2.0, 2.0) is not a
+
+
+def test_euler_integrate_constant_velocity_is_exact():
+    # With v_t == c constant, x_0 = x_1 + sum(dt * c) = x_1 - c exactly (num_steps * dt = -1).
+    noise = torch.randn(3, 5, 2)
+    c = torch.randn(3, 5, 2)
+    out = euler_integrate(lambda x_t, time: c, noise, num_steps=10)
+    torch.testing.assert_close(out, noise - c, rtol=0, atol=1e-6)
+
+
+def _reference_pi0_loop(denoise_fn, noise, num_steps, rtc_enabled, rtc_processor, kw):
+    """Verbatim structure of the historical pi0/pi05/smolvla sample_actions loop."""
+    bsize = noise.shape[0]
+    device = noise.device
+    dt = -1.0 / num_steps
+    x_t = noise
+    for step in range(num_steps):
+        time = 1.0 + step * dt
+        time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+
+        def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
+            return denoise_fn(input_x_t, current_timestep)
+
+        if rtc_enabled:
+            v_t = rtc_processor.denoise_step(
+                x_t=x_t,
+                prev_chunk_left_over=kw.get("prev_chunk_left_over"),
+                inference_delay=kw.get("inference_delay"),
+                time=time,
+                original_denoise_step_partial=denoise_step_partial_call,
+                execution_horizon=kw.get("execution_horizon"),
+            )
+        else:
+            v_t = denoise_step_partial_call(x_t)
+        x_t = x_t + dt * v_t
+        if rtc_processor is not None and rtc_processor.is_debug_enabled():
+            rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
+    return x_t
+
+
+class _StubRTCProcessor:
+    def __init__(self, debug_enabled: bool):
+        self._debug = debug_enabled
+        self.tracked = []
+        self.guidance_calls = []
+
+    def is_debug_enabled(self):
+        return self._debug
+
+    def denoise_step(
+        self,
+        x_t,
+        prev_chunk_left_over,
+        inference_delay,
+        time,
+        original_denoise_step_partial,
+        execution_horizon,
+    ):
+        self.guidance_calls.append(
+            {
+                "time": time,
+                "inference_delay": inference_delay,
+                "execution_horizon": execution_horizon,
+                "x_t": x_t.clone(),
+            }
+        )
+        return original_denoise_step_partial(x_t) * 0.5
+
+    def track(self, time, x_t, v_t):
+        self.tracked.append({"time": time, "x_t": x_t.clone(), "v_t": v_t.clone()})
+
+
+def _make_denoise_fn():
+    weight = torch.randn(4, 4) * 0.1
+
+    def denoise_fn(x_t, time_tensor):
+        return x_t @ weight + time_tensor[:, None, None]
+
+    return denoise_fn
+
+
+def test_euler_integrate_matches_historical_loop():
+    torch.manual_seed(3)
+    denoise_fn = _make_denoise_fn()
+    noise = torch.randn(2, 6, 4)
+    ref = _reference_pi0_loop(denoise_fn, noise, 10, rtc_enabled=False, rtc_processor=None, kw={})
+    out = euler_integrate(denoise_fn, noise, 10)
+    assert torch.equal(out, ref)
+
+
+def test_euler_integrate_rtc_guidance_and_kwarg_forwarding():
+    torch.manual_seed(4)
+    denoise_fn = _make_denoise_fn()
+    noise = torch.randn(2, 6, 4)
+    leftover = torch.randn(2, 6, 4)
+    kw = {"inference_delay": 3, "prev_chunk_left_over": leftover, "execution_horizon": 25}
+
+    ref_proc, new_proc = _StubRTCProcessor(False), _StubRTCProcessor(False)
+    ref = _reference_pi0_loop(denoise_fn, noise, 6, rtc_enabled=True, rtc_processor=ref_proc, kw=kw)
+    out = euler_integrate(
+        denoise_fn,
+        noise,
+        6,
+        rtc_processor=new_proc,
+        rtc_enabled=True,
+        inference_delay=3,
+        prev_chunk_left_over=leftover,
+        execution_horizon=25,
+    )
+    assert torch.equal(out, ref)
+    assert len(new_proc.guidance_calls) == 6
+    for ref_call, new_call in zip(ref_proc.guidance_calls, new_proc.guidance_calls, strict=True):
+        assert ref_call["time"] == new_call["time"]
+        assert new_call["inference_delay"] == 3 and new_call["execution_horizon"] == 25
+        # Guidance sees the PRE-update x_t.
+        assert torch.equal(ref_call["x_t"], new_call["x_t"])
+
+
+def test_euler_integrate_debug_tracking_fires_even_when_rtc_disabled():
+    # Historical behavior: track() fires whenever the processor exists and has debugging
+    # enabled, independent of whether RTC guidance is active.
+    torch.manual_seed(5)
+    denoise_fn = _make_denoise_fn()
+    noise = torch.randn(2, 6, 4)
+    proc = _StubRTCProcessor(True)
+    out = euler_integrate(denoise_fn, noise, 4, rtc_processor=proc, rtc_enabled=False)
+    assert len(proc.guidance_calls) == 0
+    assert len(proc.tracked) == 4
+    # track() receives the POST-update x_t; the last one is the returned sample.
+    assert torch.equal(proc.tracked[-1]["x_t"], out)

tests/policies/common/test_vla_utils.py

@@ -0,0 +1,195 @@
+#!/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.
+
+"""Behavior-pinning tests for the shared VLA helpers.
+
+These helpers are the canonical versions of functions that used to be copy-pasted across
+the openpi-derived policies (pi0, pi05, pi0_fast, smolvla, eo1, xvla). The expected
+values below encode the historical per-policy behavior exactly; a failure here means a
+behavior change that would silently affect released checkpoints.
+"""
+
+import math
+
+import pytest
+import torch
+
+from lerobot.policies.common.vla_utils import (
+    create_sinusoidal_pos_embedding,
+    make_att_2d_masks,
+    pad_vector,
+    prepare_attention_masks_4d,
+    resize_with_pad,
+    resize_with_pad_torch,
+)
+from lerobot.utils.constants import OPENPI_ATTENTION_MASK_VALUE
+
+
+def test_create_sinusoidal_pos_embedding_matches_openpi_formula():
+    time = torch.tensor([0.0, 0.25, 1.0])
+    dim, min_period, max_period = 8, 4e-3, 4.0
+    emb = create_sinusoidal_pos_embedding(time, dim, min_period, max_period, device=torch.device("cpu"))
+
+    assert emb.shape == (3, dim)
+    # Independent recomputation of the openpi formula in float64.
+    fraction = torch.linspace(0.0, 1.0, dim // 2, dtype=torch.float64)
+    period = min_period * (max_period / min_period) ** fraction
+    scaling = 1.0 / period * 2 * math.pi
+    sin_input = scaling[None, :] * time.to(torch.float64)[:, None]
+    expected = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+    torch.testing.assert_close(emb, expected, rtol=1e-9, atol=1e-9)
+
+
+def test_create_sinusoidal_pos_embedding_validation():
+    with pytest.raises(ValueError, match="divisible by 2"):
+        create_sinusoidal_pos_embedding(torch.zeros(2), 7, 4e-3, 4.0, device=torch.device("cpu"))
+    with pytest.raises(ValueError, match="batch_size"):
+        create_sinusoidal_pos_embedding(torch.zeros(2, 2), 8, 4e-3, 4.0, device=torch.device("cpu"))
+
+
+def test_make_att_2d_masks_docstring_cases():
+    # Pure causal attention: [[1 1 1]]
+    pad = torch.ones(1, 3, dtype=torch.bool)
+    att = torch.tensor([[1, 1, 1]], dtype=torch.int32)
+    expected = torch.tensor([[[1, 0, 0], [1, 1, 0], [1, 1, 1]]], dtype=torch.bool)
+    assert torch.equal(make_att_2d_masks(pad, att), expected)
+
+    # Prefix-LM: [[0 0 1 1]] -> first two tokens attend bidirectionally, rest causal.
+    att = torch.tensor([[0, 0, 1, 1]], dtype=torch.int32)
+    pad = torch.ones(1, 4, dtype=torch.bool)
+    expected = torch.tensor([[[1, 1, 0, 0], [1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]]], dtype=torch.bool)
+    assert torch.equal(make_att_2d_masks(pad, att), expected)
+
+    # Padding removes rows and columns.
+    pad = torch.tensor([[True, True, False]])
+    att = torch.tensor([[0, 1, 1]], dtype=torch.int32)
+    out = make_att_2d_masks(pad, att)
+    assert not out[0, :, 2].any() and not out[0, 2, :].any()
+
+
+def test_make_att_2d_masks_validation():
+    with pytest.raises(ValueError):
+        make_att_2d_masks(torch.ones(3, dtype=torch.bool), torch.ones(1, 3, dtype=torch.int32))
+    with pytest.raises(ValueError):
+        make_att_2d_masks(torch.ones(1, 3, dtype=torch.bool), torch.ones(3, dtype=torch.int32))
+
+
+def test_prepare_attention_masks_4d():
+    masks = torch.tensor([[[True, False], [False, True]]])
+    out = prepare_attention_masks_4d(masks)
+    assert out.shape == (1, 1, 2, 2)
+    expected = torch.tensor([[[[0.0, OPENPI_ATTENTION_MASK_VALUE], [OPENPI_ATTENTION_MASK_VALUE, 0.0]]]])
+    assert torch.equal(out, expected)
+
+    out_bf16 = prepare_attention_masks_4d(masks, dtype=torch.bfloat16)
+    assert out_bf16.dtype == torch.bfloat16
+    assert torch.equal(out_bf16, expected.to(torch.bfloat16))
+
+
+def test_pad_vector_openpi_semantics():
+    v = torch.arange(6.0).reshape(2, 3)
+    padded = pad_vector(v, 5)
+    assert padded.shape == (2, 5)
+    assert torch.equal(padded[:, :3], v) and not padded[:, 3:].any()
+    # Already large enough (>=): returned unchanged, same object.
+    assert pad_vector(v, 3) is v
+    assert pad_vector(v, 2) is v
+    # 3D input.
+    v3 = torch.ones(2, 4, 3)
+    assert pad_vector(v3, 7).shape == (2, 4, 7)
+
+
+def test_pad_vector_truncate_semantics():
+    v = torch.arange(6.0).reshape(2, 3)
+    out = pad_vector(v, 2, truncate=True)
+    assert out.shape == (2, 2) and torch.equal(out, v[:, :2])
+    out = pad_vector(v, 5, truncate=True)
+    assert out.shape == (2, 5) and torch.equal(out[:, :3], v) and not out[:, 3:].any()
+    assert pad_vector(v, 0, truncate=True).shape == (2, 0)
+    assert pad_vector(v, 3, truncate=True) is v
+
+
+@pytest.mark.parametrize("channels_last", [True, False])
+def test_resize_with_pad_torch_centered(channels_last):
+    img = torch.rand(2, 3, 30, 60) if not channels_last else torch.rand(2, 30, 60, 3)
+    out = resize_with_pad_torch(img, 64, 64)
+    if channels_last:
+        assert out.shape == (2, 64, 64, 3)
+        # Aspect ratio preserved: 30x60 -> 32x64, padded 16 top and 16 bottom (centered).
+        assert not out[:, :16].any() and not out[:, -16:].any()
+        assert out[:, 16:48].abs().sum() > 0
+    else:
+        assert out.shape == (2, 3, 64, 64)
+        assert not out[:, :, :16].any() and not out[:, :, -16:].any()
+
+
+def test_resize_with_pad_torch_uint8_roundtrip():
+    img = (torch.rand(1, 3, 20, 20) * 255).to(torch.uint8)
+    out = resize_with_pad_torch(img, 40, 40)
+    assert out.dtype == torch.uint8 and out.shape == (1, 3, 40, 40)
+    with pytest.raises(ValueError, match="Unsupported image dtype"):
+        resize_with_pad_torch(torch.rand(1, 3, 8, 8, dtype=torch.float64), 16, 16)
+
+
+def test_resize_with_pad_top_left():
+    img = torch.rand(2, 3, 30, 60)
+    out = resize_with_pad(img, 64, 64, pad_value=-1.0)
+    assert out.shape == (2, 3, 64, 64)
+    # 30x60 -> 32x64; this variant pads on the TOP only (32 rows of pad_value).
+    assert torch.equal(out[:, :, :32], torch.full((2, 3, 32, 64), -1.0))
+    assert out[:, :, 32:].min() >= 0
+    # No-op fast path returns the same object.
+    assert resize_with_pad(img, 30, 60, pad_value=0.0) is img
+    with pytest.raises(ValueError, match="expected"):
+        resize_with_pad(torch.rand(3, 8, 8), 16, 16, pad_value=0.0)
+
+
+def test_clone_past_key_values():
+    pytest.importorskip("transformers")
+    from transformers import DynamicCache
+
+    from lerobot.policies.common.vla_utils import clone_past_key_values
+
+    cache = DynamicCache()
+    keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8)
+    cache.update(keys, values, 0)
+    cloned = clone_past_key_values(cache)
+    (ck, cv, _), (ok, ov, _) = next(iter(cloned)), next(iter(cache))
+    assert torch.equal(ck, ok) and torch.equal(cv, ov)
+    # Deep copy: mutating the clone must not touch the original.
+    ck.zero_()
+    assert not torch.equal(ck, ok)
+
+
+def test_clone_past_key_values_is_fullgraph_compilable():
+    pytest.importorskip("transformers")
+    from transformers import DynamicCache
+
+    from lerobot.policies.common.vla_utils import clone_past_key_values
+
+    cache = DynamicCache()
+    keys, values = torch.rand(1, 2, 4, 8), torch.rand(1, 2, 4, 8)
+    cache.update(keys, values, 0)
+
+    compiled_clone = torch.compile(clone_past_key_values, backend="eager", fullgraph=True)
+    cloned = compiled_clone(cache)
+
+    (cloned_keys, cloned_values, _), (original_keys, original_values, _) = (
+        next(iter(cloned)),
+        next(iter(cache)),
+    )
+    assert torch.equal(cloned_keys, original_keys)
+    assert torch.equal(cloned_values, original_values)

tests/policies/pi052/test_pi052_attention_masking.py

@@ -1,151 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""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

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

tests/policies/pi052/test_pi052_checkpoint_portability.py

@@ -1,152 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import 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

@@ -1,162 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""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

@@ -1,122 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import numpy as np
-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

@@ -1,65 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import 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