feat(unitree_g1): drive SONIC whole-body from a 34-D OpenHLM/pi0.5 VLA

Add a dense 34-D whole-body command path so lerobot-rollout can drive the
G1 directly with an OpenHLM / pi0.5 policy through the SONIC encoder/decoder:

- SonicWholeBodyController: wb.{i}.pos action interface, mode-0 reference with
  a rolling 50-frame trajectory (finite-diff velocities) and first-tick anchor
  init; correct MuJoCo->IsaacLab joint reordering.
- unitree_g1: expose 34-D wb_state.{i}.pos proprio; empty/replay camera feeds
  for image-conditioned policies; Dex3 hand publishing from the grip scalars.
- g1_utils: obs_to_wb34_state + WB action constants.

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

README.md

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

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

@@ -150,14 +150,14 @@ 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). |
-| `update() -> None` _(optional)_                                   | `lerobot-train`   | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this).                                                                                                                                           |
+| 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 (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,11 +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.
-
 ## Graceful absence
 
 If both language columns are missing, `None`, or empty, `RenderMessagesStep` is a no-op.

docs/source/multi_gpu_training.mdx

@@ -113,19 +113,6 @@ accelerate launch --num_processes=2 $(which lerobot-train) \
   --policy.type=act
 ```
 
-When the desired global batch is larger than the per-GPU batch that fits in memory, use gradient
-accumulation. `steps` continues to count optimizer updates:
-
-```bash
-# 8 samples/GPU × 4 GPUs × 2 microbatches = effective global batch 64.
-accelerate launch --num_processes=4 $(which lerobot-train) \
-  --batch_size=8 \
-  --gradient_accumulation_steps=2 \
-  --steps=80000 \
-  --dataset.repo_id=lerobot/pusht \
-  --policy.type=act
-```
-
 ## Training Large Models with FSDP
 
 DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under

docs/source/pi052.mdx

@@ -1,255 +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_mem.yaml`            | `subtask`, `memory`                                                     | Subtasks, actions, and memory updates              |
-| `recipes/subtask_mem_vqa_speech.yaml` | `subtask`, `memory`, `vqa`; interjection/speech rows for those branches | Subtasks, actions, memory, VQA, and spoken replies |
-
-Use `lerobot-annotate` to generate these columns. The repository includes a
-Hugging Face Jobs launcher that you can edit for your source and destination
-datasets. For a local annotation run, first install
-`pip install -e ".[annotations]"`:
-
-```bash
-HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
-```
-
-Before a long training run, inspect several episodes and verify that subtasks
-are temporally correct and cover the full demonstration. See
-[Annotation Pipeline](./annotation_pipeline) for generation and validation, and
-[Language Columns and Recipes](./language_and_recipes) for the schema and
-recipe resolver.
-
-<Tip>
-  If a dataset has no language columns, recipe rendering becomes a no-op and
-  Pi052 falls back to the plain Pi05 prompt path. This is useful for
-  compatibility but does not train the language planner.
-</Tip>
-
-## Train Pi052
-
-This example initializes Pi052 from the public Pi05 base checkpoint and trains
-the default subtask-and-memory recipe:
-
-```bash
-lerobot-train \
-    --dataset.repo_id=${HF_USER}/my_language_annotated_dataset \
-    --policy.type=pi052 \
-    --policy.pretrained_path=lerobot/pi05_base \
-    --policy.recipe_path=recipes/subtask_mem.yaml \
-    --policy.dtype=bfloat16 \
-    --policy.device=cuda \
-    --policy.freeze_vision_encoder=false \
-    --policy.gradient_checkpointing=true \
-    --batch_size=8 \
-    --steps=30000 \
-    --output_dir=outputs/pi052 \
-    --job_name=pi052 \
-    --wandb.enable=true
-```
-
-For subtask-only data, change the recipe to `recipes/subtask.yaml` and disable
-memory during rollout. Start with a small run and confirm that W&B examples show
-the expected prompt, text target, and action endpoints before scaling up.
-
-### Main training controls
-
-| Option                           |                    Default | Purpose                                                        |
-| -------------------------------- | -------------------------: | -------------------------------------------------------------- |
-| `policy.recipe_path`             | `recipes/subtask_mem.yaml` | Selects the language/action objective mixture                  |
-| `policy.text_loss_weight`        |                      `1.0` | Language-head cross-entropy weight; `0` disables text training |
-| `policy.flow_loss_weight`        |                     `10.0` | Continuous action flow-loss weight                             |
-| `policy.enable_fast_action_loss` |                     `true` | Adds discrete FAST action-token supervision                    |
-| `policy.fast_action_loss_weight` |                      `1.0` | FAST cross-entropy weight                                      |
-| `policy.knowledge_insulation`    |                     `true` | Blocks action-loss gradients through the VLM K/V path          |
-| `policy.flow_num_repeats`        |                        `5` | Reuses one VLM prefix for independent denoising targets        |
-| `policy.lm_head_lr_scale`        |                      `1.0` | Scales language-head learning rate; `1.0` uses the base rate   |
-
-The loss weights are starting points, not dataset-independent constants. Track
-flow loss and text/FAST losses separately, and inspect generated subtasks rather
-than selecting a checkpoint from total loss alone.
-
-### Dataset-specific FAST tokenizer
-
-The universal FAST tokenizer works out of the box. For a large or
-embodiment-specific dataset, Pi052 can fit and cache a tokenizer on normalized
-actions before training:
-
-```bash
-lerobot-train \
-    ... \
-    --policy.auto_fit_fast_tokenizer=true \
-    --policy.fast_tokenizer_fit_samples=4096
-```
-
-The fit runs once per dataset/tokenizer configuration. Keep
-`auto_fit_fast_tokenizer=false` when you do not want the extra preprocessing
-pass.
-
-## Training performance
-
-Pi052 uses optimized training paths by default:
-
-- batches repeated flow targets and suffix projections instead of replaying
-  small operations in Python;
-- caches constant action masks and computes RoPE positions once per forward;
-- selects the text/FAST cross-entropy implementation from target shape and
-  sparsity;
-- skips the mathematically dead VLM/vision backward on knowledge-insulated,
-  flow-only batches;
-- uses native non-reentrant SigLIP layer checkpointing when gradient
-  checkpointing is enabled; and
-- retains the Liger RoPE/GeGLU kernels while avoiding the slower LayerNorm
-  patch at SigLIP shapes.
-
-Optional training backends are disabled by default:
-
-| Option                                 | When to try it                                                                                    |
-| -------------------------------------- | ------------------------------------------------------------------------------------------------- |
-| `policy.use_flashrt_adarms=true`       | Fused adaptive RMSNorm and gated residuals on supported CUDA GPUs                                 |
-| `policy.use_compiled_text_ce=true`     | Compiled materialized-logit CE buckets                                                            |
-| `policy.use_compiled_vision=true`      | Compiled vision only when the vision pass has no gradients                                        |
-| `policy.use_flex_attention=true`       | Profiled CUDA setups with knowledge insulation and `flow_num_repeats > 1`; otherwise SDPA is used |
-| `policy.use_manual_attention=true`     | Explicitly profiled shapes where materialized attention is faster                                 |
-| `policy.manual_attention_scope=action` | Restricts manual attention to action queries                                                      |
-
-Do not enable every backend blindly. Flex and manual attention are mutually
-exclusive, and attention/AdaRMS alternatives require knowledge insulation.
-The benchmark-best configuration used compiled text CE and FlashRT AdaRMS,
-with Flex/manual attention and compiled vision disabled.
-
-### Reported training benchmarks
-
-These benchmarks measure complete optimizer steps with three real camera
-inputs, BF16 transformer/action execution, FP32 vision, fused AdamW, and no
-video decoding or network I/O. Results vary with GPU, batch shape, annotation
-mixture, and checkpointing:
-
-| Workload                   |     RTX PRO 6000 Blackwell |                 A100 80 GB |
-| -------------------------- | -------------------------: | -------------------------: |
-| Full flow + text, batch 1  | 4.75× vs checkpointing off | 3.33× vs checkpointing off |
-| Full flow + text, batch 8  | 2.16× vs checkpointing off | 1.66× vs checkpointing off |
-| Full flow + text, batch 64 |  1.24× vs checkpointing on |  1.15× vs checkpointing on |
-| Flow-only, batch 1         | 3.70× vs checkpointing off | 3.58× vs checkpointing off |
-| Flow-only, batch 64        |  3.76× vs checkpointing on |  3.61× vs checkpointing on |
-
-On those 80 GB GPUs, full training was fastest without gradient checkpointing
-through batch 8, then required checkpointing at batch 16 and above. Treat that
-as a tuning rule to test on your hardware, not a universal threshold. Flow-only
-means both text and FAST supervision are disabled; it is useful for action-only
-ablation or post-training but does not learn the language runtime.
-
-## Inference performance
-
-Pi052 has two inference loops, and both avoid repeatedly encoding the expensive
-multimodal prefix:
-
-1. **Action denoising** encodes the image/language prefix once, reuses its KV
-   cache across flow steps, precomputes the timestep schedule on-device, and
-   crops temporary suffix K/V instead of cloning the prefix cache.
-2. **Language decoding** uses autoregressive KV caching, so each new token only
-   processes the sampled token against cached image/language keys instead of
-   rerunning the full prefix.
-
-The runtime also runs language and actions at different rates. Increase
-`--subtask_chunks_per_gen` when a subtask remains valid across several action
-chunks, lower `--high_level_hz`, or use `--direct_subtask` to bypass language
-generation entirely. These settings reduce compute but also slow replanning.
-
-`--fp8` enables the optional FlashRT inference MLP swap on supported CUDA GPUs.
-It calibrates on the first observation and falls back to BF16 when unavailable;
-because FP8 can change outputs slightly, validate task success before using it
-for production rollouts.
-
-## Run a checkpoint
-
-RoboCasa:
-
-```bash
-MUJOCO_GL=egl lerobot-rollout \
-    --policy.path=lerobot/pi052_robocasa \
-    --sim --sim.task=CloseFridge --sim.split=pretrain \
-    --task="close the fridge" \
-    --disable_memory \
-    --sim.render_size=384 \
-    --sim.views=robot0_agentview_left,robot0_eye_in_hand,robot0_agentview_right \
-    --mode=action --ctrl_hz=20
-```
-
-Open `http://localhost:8010` for the live view. Without
-`--sim.direct_subtask`, Pi052 generates the low-level subtask; with it, each
-prompt becomes the action policy's subtask directly.
-
-The same runtime supports real robots. See [Interactive language
-control](./inference#interactive-language-control) for the real-arm command,
-safety behavior, and runtime controls.
-
-## Troubleshooting
-
-- **No text loss or generated subtasks:** confirm the selected recipe can bind
-  the annotations on sampled frames and that `policy.text_loss_weight > 0`.
-- **Subtasks look plausible but actions fail:** verify subtask boundaries,
-  normalized state/action statistics, and that low-level recipe samples are
-  present.
-- **Text collapses to repeated or location tokens:** inspect text-target
-  coverage, language-head learning rate, and the balance between flow, FAST,
-  and text losses.
-- **Out of memory:** reduce batch size first, then enable gradient
-  checkpointing. Do not enable compiled or alternative attention backends
-  without profiling their memory on your camera count.
-- **Slow rollout:** separate action latency from language latency, then tune
-  `--subtask_chunks_per_gen`, `--high_level_hz`, and the number of flow
-  inference steps.

docs/source/pi0fast.mdx

@@ -109,21 +109,15 @@ lerobot-train \
 
 ### Key Training Parameters
 
-| Parameter                               | Description                                        | Default                         |
-| --------------------------------------- | -------------------------------------------------- | ------------------------------- |
-| `--policy.gradient_checkpointing=true`  | Reduces memory usage significantly during training | `false`                         |
-| `--policy.dtype=bfloat16`               | Use mixed precision training for efficiency        | `float32`                       |
-| `--policy.chunk_size`                   | Number of action steps to predict (action horizon) | `50`                            |
-| `--policy.n_action_steps`               | Number of action steps to execute                  | `50`                            |
-| `--policy.max_action_tokens`            | Maximum number of FAST tokens per action chunk     | `256`                           |
-| `--policy.action_tokenizer_name`        | FAST tokenizer to use                              | `lerobot/fast-action-tokenizer` |
-| `--policy.auto_fit_fast_tokenizer=true` | Fit and cache a tokenizer for the training dataset | `false`                         |
-| `--policy.compile_model=true`           | Enable torch.compile for faster training           | `false`                         |
-
-Set `--policy.auto_fit_fast_tokenizer=true` to sample action chunks from the
-training dataset and cache a fitted tokenizer under
-`~/.cache/lerobot/fast_tokenizers`. This also works when fine-tuning with
-`--policy.path`; leave it disabled to retain the checkpoint's tokenizer.
+| Parameter                              | Description                                        | Default                         |
+| -------------------------------------- | -------------------------------------------------- | ------------------------------- |
+| `--policy.gradient_checkpointing=true` | Reduces memory usage significantly during training | `false`                         |
+| `--policy.dtype=bfloat16`              | Use mixed precision training for efficiency        | `float32`                       |
+| `--policy.chunk_size`                  | Number of action steps to predict (action horizon) | `50`                            |
+| `--policy.n_action_steps`              | Number of action steps to execute                  | `50`                            |
+| `--policy.max_action_tokens`           | Maximum number of FAST tokens per action chunk     | `256`                           |
+| `--policy.action_tokenizer_name`       | FAST tokenizer to use                              | `lerobot/fast-action-tokenizer` |
+| `--policy.compile_model=true`          | Enable torch.compile for faster training           | `false`                         |
 
 ## Inference
 

docs/source/robomme.mdx

@@ -44,18 +44,6 @@ docker build -f docker/Dockerfile.benchmark.robomme -t lerobot-robomme .
 
 The `docker/Dockerfile.benchmark.robomme` image overrides `gymnasium==0.29.1` and `numpy==1.26.4` after lerobot's install. Both versions are runtime-safe for lerobot's actual API usage.
 
-When evaluating a checkpoint saved on the host, run the container with the host UID. Checkpoint weights are intentionally private by default (`0600`), so the image's built-in user cannot otherwise read a bind-mounted `model.safetensors` file. Mount the Hugging Face cache at an accessible path as well:
-
-```bash
-docker run --gpus all --rm --ipc=host \
-  --user "$(id -u):$(id -g)" \
-  -e HF_HOME=/tmp/hf-cache \
-  -v "$HOME/.cache/huggingface:/tmp/hf-cache:ro" \
-  -v "$PWD/outputs:/results" \
-  lerobot-robomme \
-  lerobot-eval --policy.path=/results/<checkpoint>/pretrained_model # ...
-```
-
 ## Running Evaluation
 
 ### Default (single task, single episode)
@@ -73,7 +61,7 @@ lerobot-eval \
 
 ### Multi-task evaluation
 
-Evaluate multiple tasks in one run by comma-separating task names. Use `task_ids` to select the dataset-backed episodes evaluated for every task. For the standard 50-episode test protocol, select IDs 0 through 49 and run each selected episode once.
+Evaluate multiple tasks in one run by comma-separating task names. Use `task_ids` to control which episodes are evaluated per task. Recommended: 50 episodes per task for the test split.
 
 ```bash
 lerobot-eval \
@@ -81,22 +69,20 @@ lerobot-eval \
     --env.type=robomme \
     --env.task=PickXtimes,BinFill,StopCube,MoveCube,InsertPeg \
     --env.dataset_split=test \
-    --env.task_ids=[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49] \
+    --env.task_ids=[0,1,2,3,4,5,6,7,8,9] \
     --eval.batch_size=1 \
-    --eval.n_episodes=1
+    --eval.n_episodes=50
 ```
 
 ### Key CLI options for `env.type=robomme`
 
-| Option               | Default       | Description                                      |
-| -------------------- | ------------- | ------------------------------------------------ |
-| `env.task`           | `PickXtimes`  | Any of the 16 task names above (comma-separated) |
-| `env.dataset_split`  | `test`        | `train`, `val`, or `test`                        |
-| `env.action_space`   | `joint_angle` | `joint_angle` (8-D) or `ee_pose` (7-D)           |
-| `env.episode_length` | `300`         | Max steps per episode                            |
-| `env.task_ids`       | `null`        | Dataset episode indices (null = `[0]`)           |
-
-`eval.n_episodes` repeats each selected `task_id`; it does not advance to the next dataset episode. Keep it at `1` when enumerating the 50 distinct test episodes with `env.task_ids`.
+| Option               | Default       | Description                                        |
+| -------------------- | ------------- | -------------------------------------------------- |
+| `env.task`           | `PickXtimes`  | Any of the 16 task names above (comma-separated)   |
+| `env.dataset_split`  | `test`        | `train`, `val`, or `test`                          |
+| `env.action_space`   | `joint_angle` | `joint_angle` (8-D) or `ee_pose` (7-D)             |
+| `env.episode_length` | `300`         | Max steps per episode                              |
+| `env.task_ids`       | `null`        | List of episode indices to evaluate (null = `[0]`) |
 
 ## Dataset
 
@@ -110,68 +96,22 @@ dataset = LeRobotDataset("lerobot/robomme")
 
 ### Dataset features
 
-| Feature            | Shape         | Description                      |
-| ------------------ | ------------- | -------------------------------- |
-| `image`            | (256, 256, 3) | Front camera RGB                 |
-| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                 |
-| `actions`          | (8,)          | Joint angles + gripper           |
-| `state`            | (8,)          | Joint positions + gripper state  |
-| `simple_subgoal`   | str           | High-level language annotation   |
-| `grounded_subgoal` | str           | Grounded language annotation     |
-| `exec_start_idx`   | scalar        | First action-execution frame     |
-| `is_demo`          | bool          | Video-demonstration prefix frame |
-| `episode_index`    | int           | Episode ID                       |
-| `frame_index`      | int           | Frame within episode             |
+| Feature            | Shape         | Description                     |
+| ------------------ | ------------- | ------------------------------- |
+| `image`            | (256, 256, 3) | Front camera RGB                |
+| `wrist_image`      | (256, 256, 3) | Wrist camera RGB                |
+| `actions`          | (8,)          | Joint angles + gripper          |
+| `state`            | (8,)          | Joint positions + gripper state |
+| `simple_subgoal`   | str           | High-level language annotation  |
+| `grounded_subgoal` | str           | Grounded language annotation    |
+| `episode_index`    | int           | Episode ID                      |
+| `frame_index`      | int           | Frame within episode            |
 
 ### Feature key alignment (training)
 
 The env wrapper exposes `pixels/image` and `pixels/wrist_image` as observation keys. The `features_map` in `RoboMMEEnv` maps these to `observation.images.image` and `observation.images.wrist_image` for the policy. State is exposed as `agent_pos` and maps to `observation.state`.
 
-The published dataset uses the raw keys shown above. Policies that expect canonical LeRobot keys should map them at training time. For example, the pretrained SmolVLA checkpoint expects three cameras, while RoboMME provides two:
-
-```bash
-uv run lerobot-train \
-  --policy.path=lerobot/smolvla_base \
-  --policy.empty_cameras=1 \
-  --dataset.repo_id=lerobot/robomme \
-  --dataset.training_target_start_feature=exec_start_idx \
-  '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
-```
-
-RoboMME episodes begin with video-demonstration/context frames that should remain available to a
-memory policy but should not be sampled as action-training targets. Setting
-`dataset.training_target_start_feature=exec_start_idx` starts target sampling at each episode's
-execution boundary while preserving earlier frames for temporal observation deltas. In the published
-training split this keeps 476,857 execution targets out of 768,897 total frames. One execution-target
-epoch therefore takes `ceil(476857 / effective_batch_size)` optimizer steps.
-
-For a sample-matched SmolVLA visual-memory ablation, use
-`examples/robomme/smolvla_visual_memory_ablation.sh`. `TARGET_SAMPLES` counts examples across all
-GPUs, and `NUM_PROCESSES` is included when the script converts that target into optimizer steps. For
-example, the following reproduces 5.12 million example exposures (the exposure of RoboMME's
-80,000-step, global-batch-64 memory-policy recipe) with four GPUs, two accumulated microbatches,
-and an effective global batch of 64:
-
-```bash
-TARGET_SAMPLES=5120000 \
-NUM_PROCESSES=4 \
-BATCH_SIZE=8 \
-GRADIENT_ACCUMULATION_STEPS=2 \
-VARIANT=visual-memory \
-RUN_EVAL=false \
-bash examples/robomme/smolvla_visual_memory_ablation.sh
-```
-
-This becomes 80,000 optimizer steps, or about 10.74 execution-target epochs. Run the baseline with
-the same `TARGET_SAMPLES`, effective batch size, seed, and scheduler settings to isolate the visual
-memory change. RoboMME's released baseline uses a different global batch (128), so its nominal
-80,000-step recipe is not sample-matched to its global-batch-64 memory recipe.
-
-At evaluation time the environment wrapper has already converted observations to canonical keys. Only the two camera suffixes need to be aligned with the checkpoint:
-
-```bash
-'--rename_map={"observation.images.image":"observation.images.camera1","observation.images.wrist_image":"observation.images.camera2"}'
-```
+The dataset's `image` and `wrist_image` columns already align with the policy input keys, so no renaming is needed when fine-tuning.
 
 ## Action Spaces
 

docs/source/smolvla.mdx

@@ -76,37 +76,6 @@ Fine-tuning is an art. For a complete overview of the options for finetuning, ru
 lerobot-train --help
 ```
 
-### Experimental MEM visual memory
-
-SmolVLA can optionally fuse a short history of camera frames using the
-space-time separable vision encoder from [MEM](https://arxiv.org/abs/2603.03596).
-Every fourth SigLIP layer adds causal attention across time for matching image
-patches. Historical tokens are discarded inside the vision tower, so the VLM
-receives the same number of image tokens as the baseline policy.
-
-The option is disabled by default. The following example uses six observations
-spaced one second apart for a 10 fps dataset:
-
-```bash
-lerobot-train \
-  --policy.path=lerobot/smolvla_base \
-  --policy.use_visual_memory=true \
-  --policy.visual_memory_frames=6 \
-  --policy.visual_memory_stride=10 \
-  --policy.visual_memory_temporal_attention_every=4 \
-  --policy.freeze_vision_encoder=false \
-  --policy.train_expert_only=false \
-  --dataset.repo_id=${HF_USER}/mydataset \
-  --output_dir=outputs/train/my_smolvla_mem
-```
-
-`visual_memory_stride` is measured in dataset or environment steps. Keep all
-other settings and the random seed fixed when comparing against a baseline.
-Because the public SmolVLA checkpoint was not pretrained with this temporal
-attention pattern, this is a post-training-only MEM ablation; the MEM paper
-reports that memory-aware pretraining performs better than introducing visual
-memory only during task-specific fine-tuning.
-
 <p align="center">
   <img
     src="https://cdn-uploads.huggingface.co/production/uploads/640e21ef3c82bd463ee5a76d/S-3vvVCulChREwHDkquoc.gif"

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

examples/robomme/smolvla_visual_memory_ablation.sh

@@ -1,131 +0,0 @@
-#!/usr/bin/env bash
-# Matched SmolVLA ablation for the 10 fps lerobot/robomme dataset.
-# Run from the LeRobot repository root on a CUDA machine.
-
-set -euo pipefail
-
-BATCH_SIZE="${BATCH_SIZE:-4}"
-NUM_PROCESSES="${NUM_PROCESSES:-1}"
-GRADIENT_ACCUMULATION_STEPS="${GRADIENT_ACCUMULATION_STEPS:-1}"
-# The published training split has 476,857 execution frames. By default, train on one
-# execution-frame epoch; set STEPS explicitly to use a different optimizer-step budget.
-TARGET_SAMPLES="${TARGET_SAMPLES:-476857}"
-EFFECTIVE_BATCH_SIZE=$((BATCH_SIZE * NUM_PROCESSES * GRADIENT_ACCUMULATION_STEPS))
-STEPS="${STEPS:-$(((TARGET_SAMPLES + EFFECTIVE_BATCH_SIZE - 1) / EFFECTIVE_BATCH_SIZE))}"
-SCHEDULER_WARMUP_STEPS="${SCHEDULER_WARMUP_STEPS:-$(((STEPS + 29) / 30))}"
-SCHEDULER_DECAY_STEPS="${SCHEDULER_DECAY_STEPS:-${STEPS}}"
-SEED="${SEED:-1000}"
-OUTPUT_ROOT="${OUTPUT_ROOT:-outputs/robomme-smolvla-mem-ablation}"
-WANDB_ENABLE="${WANDB_ENABLE:-false}"
-RUN_TRAIN="${RUN_TRAIN:-true}"
-RUN_EVAL="${RUN_EVAL:-true}"
-VARIANT="${VARIANT:-both}"
-TASKS="BinFill,PickXtimes,SwingXtimes,StopCube,VideoUnmask,VideoUnmaskSwap,ButtonUnmask,ButtonUnmaskSwap,PickHighlight,VideoRepick,VideoPlaceButton,VideoPlaceOrder,MoveCube,InsertPeg,PatternLock,RouteStick"
-TASK_IDS="[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49]"
-
-case "${VARIANT}" in
-  baseline) VARIANTS=(baseline) ;;
-  visual-memory) VARIANTS=(visual-memory) ;;
-  both) VARIANTS=(baseline visual-memory) ;;
-  *) echo "VARIANT must be baseline, visual-memory, or both" >&2; exit 2 ;;
-esac
-
-TRAIN_COMMAND=()
-if [[ "${RUN_TRAIN}" == "true" ]]; then
-  if ((NUM_PROCESSES > 1)); then
-    TRAIN_ENTRYPOINT="$(uv run which lerobot-train)"
-    TRAIN_COMMAND=(
-      uv run accelerate launch
-      --multi_gpu
-      --num_processes="${NUM_PROCESSES}"
-      --mixed_precision=bf16
-      "${TRAIN_ENTRYPOINT}"
-    )
-  else
-    TRAIN_COMMAND=(uv run lerobot-train)
-  fi
-fi
-
-echo "Training ${VARIANT}: ${TARGET_SAMPLES} target samples, global batch ${EFFECTIVE_BATCH_SIZE}, ${STEPS} optimizer steps"
-
-COMMON_TRAIN_ARGS=(
-  --policy.path=lerobot/smolvla_base
-  --policy.device=cuda
-  --policy.push_to_hub=false
-  --policy.empty_cameras=1
-  --policy.freeze_vision_encoder=false
-  --policy.train_expert_only=false
-  --dataset.repo_id=lerobot/robomme
-  --dataset.training_target_start_feature=exec_start_idx
-  '--rename_map={"image":"observation.images.camera1","wrist_image":"observation.images.camera2","state":"observation.state","actions":"action"}'
-  --batch_size="${BATCH_SIZE}"
-  --gradient_accumulation_steps="${GRADIENT_ACCUMULATION_STEPS}"
-  --steps="${STEPS}"
-  --policy.scheduler_warmup_steps="${SCHEDULER_WARMUP_STEPS}"
-  --policy.scheduler_decay_steps="${SCHEDULER_DECAY_STEPS}"
-  --seed="${SEED}"
-  --env_eval_freq=0
-  --save_freq=5000
-  --wandb.enable="${WANDB_ENABLE}"
-)
-
-if [[ "${RUN_TRAIN}" == "true" ]]; then
-  for variant in "${VARIANTS[@]}"; do
-    MEMORY_ARGS=(--policy.use_visual_memory=false)
-    if [[ "${variant}" == "visual-memory" ]]; then
-      MEMORY_ARGS=(
-        --policy.use_visual_memory=true
-        --policy.visual_memory_frames=6
-        --policy.visual_memory_stride=10
-        --policy.visual_memory_temporal_attention_every=4
-      )
-    fi
-    "${TRAIN_COMMAND[@]}" \
-      "${COMMON_TRAIN_ARGS[@]}" \
-      "${MEMORY_ARGS[@]}" \
-      --output_dir="${OUTPUT_ROOT}/${variant}" \
-      --job_name="robomme-smolvla-${variant}"
-  done
-fi
-
-if [[ "${RUN_EVAL}" == "true" ]]; then
-  for variant in "${VARIANTS[@]}"; do
-    uv run lerobot-eval \
-      --policy.path="${OUTPUT_ROOT}/${variant}/checkpoints/last/pretrained_model" \
-      --env.type=robomme \
-      --env.task="${TASKS}" \
-      --env.dataset_split=test \
-      --env.task_ids="${TASK_IDS}" \
-      '--rename_map={"observation.images.image":"observation.images.camera1","observation.images.wrist_image":"observation.images.camera2"}' \
-      --eval.batch_size=1 \
-      --eval.n_episodes=1 \
-      --seed="${SEED}" \
-      --output_dir="${OUTPUT_ROOT}/eval-${variant}"
-  done
-fi
-
-if [[ "${RUN_EVAL}" == "true" ]]; then
-  uv run python - "${OUTPUT_ROOT}" <<'PY'
-import json
-import pathlib
-import sys
-
-root = pathlib.Path(sys.argv[1])
-for variant in ("baseline", "visual-memory"):
-    result_path = root / f"eval-{variant}" / "eval_info.json"
-    if not result_path.exists():
-        continue
-    with result_path.open() as handle:
-        info = json.load(handle)
-    overall = info["overall"]
-    print(
-        f"{variant}: success={overall['pc_success']:.2f}% "
-        f"avg_reward={overall['avg_sum_reward']:.4f}"
-    )
-    for task, metrics in info["per_group"].items():
-        print(
-            f"  {task}: success={metrics['pc_success']:.2f}% "
-            f"avg_reward={metrics['avg_sum_reward']:.4f}"
-        )
-PY
-fi

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]"]
@@ -375,7 +374,11 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 
 [tool.setuptools.package-data]
-lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
+lerobot = [
+    "envs/*.json",
+    "annotations/steerable_pipeline/prompts/*.txt",
+    "teleoperators/pico_headset/assets/*.npz",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]
@@ -414,8 +417,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"]

src/lerobot/configs/default.py

@@ -33,10 +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
-    # Optional per-frame feature containing the episode-relative index of the first frame that may be
-    # sampled as a training target. Earlier frames remain in the dataset and can still be loaded through
-    # temporal observation deltas. This is useful for datasets with demonstration/context prefixes.
-    training_target_start_feature: str | None = None
     image_transforms: ImageTransformsConfig = field(default_factory=ImageTransformsConfig)
     revision: str | None = None
     use_imagenet_stats: bool = True

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

@@ -99,9 +99,6 @@ class TrainPipelineConfig(HubMixin):
     # Number of workers for the dataloader.
     num_workers: int = 4
     batch_size: int = 8
-    # Number of microbatches accumulated before each optimizer update. The effective global batch is
-    # batch_size * accelerator.num_processes * gradient_accumulation_steps.
-    gradient_accumulation_steps: int = 1
     prefetch_factor: int = 4
     persistent_workers: bool = True
     steps: int = 100_000
@@ -224,8 +221,6 @@ class TrainPipelineConfig(HubMixin):
             )
 
         active_cfg = self.trainable_config
-        if self.gradient_accumulation_steps < 1:
-            raise ValueError("gradient_accumulation_steps must be at least 1.")
         if self.rename_map and active_cfg.pretrained_path is None:
             raise ValueError(
                 "`rename_map` requires a pretrained policy checkpoint. "

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

@@ -32,9 +32,7 @@ from .streaming_dataset import StreamingLeRobotDataset
 
 
 def resolve_delta_timestamps(
-    cfg: PreTrainedConfig | RewardModelConfig,
-    ds_meta: LeRobotDatasetMetadata,
-    rename_map: dict[str, str] | None = None,
+    cfg: PreTrainedConfig | RewardModelConfig, ds_meta: LeRobotDatasetMetadata
 ) -> dict[str, list] | None:
     """Resolves delta_timestamps by reading from the 'delta_indices' properties of the config.
 
@@ -55,12 +53,11 @@ def resolve_delta_timestamps(
     """
     delta_timestamps = {}
     for key in ds_meta.features:
-        policy_key = (rename_map or {}).get(key, key)
-        if policy_key == REWARD and cfg.reward_delta_indices is not None:
+        if key == REWARD and cfg.reward_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.reward_delta_indices]
-        if policy_key == ACTION and cfg.action_delta_indices is not None:
+        if key == ACTION and cfg.action_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.action_delta_indices]
-        if policy_key.startswith(OBS_PREFIX) and cfg.observation_delta_indices is not None:
+        if key.startswith(OBS_PREFIX) and cfg.observation_delta_indices is not None:
             delta_timestamps[key] = [i / ds_meta.fps for i in cfg.observation_delta_indices]
 
     if len(delta_timestamps) == 0:
@@ -89,7 +86,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         ds_meta = LeRobotDatasetMetadata(
             cfg.dataset.repo_id, root=cfg.dataset.root, revision=cfg.dataset.revision
         )
-        delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta, cfg.rename_map)
+        delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
         if not cfg.dataset.streaming:
             dataset = LeRobotDataset(
                 cfg.dataset.repo_id,
@@ -133,9 +130,6 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
         for key in dataset.meta.camera_keys:
             if key in dataset.meta.depth_keys:
                 continue  # Exclude depth keys from ImageNet stats
-            # Some video-only datasets omit camera statistics entirely. Visual
-            # normalization can still use the requested ImageNet defaults.
-            dataset.meta.stats.setdefault(key, {})
             for stats_type, stats in IMAGENET_STATS.items():
                 dataset.meta.stats[key][stats_type] = torch.tensor(stats, dtype=torch.float32)
 
@@ -181,7 +175,7 @@ def make_train_eval_datasets(
         f"(eval_split={cfg.dataset.eval_split}, {len(task_to_episodes)} tasks)"
     )
 
-    delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, full_dataset.meta, cfg.rename_map)
+    delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, full_dataset.meta)
 
     train_image_transforms = (
         ImageTransforms(cfg.dataset.image_transforms) if cfg.dataset.image_transforms.enable else None

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/datasets/sampler.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 import logging
 import math
-from collections.abc import Iterator, Sequence
+from collections.abc import Iterator
 
 import numpy as np
 import torch
@@ -49,7 +49,7 @@ class EpisodeAwareSampler:
         dataset_from_indices: list[int],
         dataset_to_indices: list[int],
         episode_indices_to_use: list | None = None,
-        drop_n_first_frames: int | Sequence[int] = 0,
+        drop_n_first_frames: int = 0,
         drop_n_last_frames: int = 0,
         shuffle: bool = False,
         seed: int = 0,
@@ -60,12 +60,13 @@ class EpisodeAwareSampler:
             dataset_from_indices: Start index of each episode in the dataset.
             dataset_to_indices: End index of each episode in the dataset.
             episode_indices_to_use: Episode indices to use; None means all.
-            drop_n_first_frames: Frames to drop from the start of each episode. An integer applies the
-                same offset to every episode; a sequence supplies one offset per episode.
+            drop_n_first_frames: Frames to drop from the start of each episode.
             drop_n_last_frames: Frames to drop from the end of each episode.
             shuffle: Whether to shuffle the indices.
             seed: Seed the permutation is derived from (together with the epoch).
         """
+        if drop_n_first_frames < 0:
+            raise ValueError(f"drop_n_first_frames must be >= 0, got {drop_n_first_frames}")
         if drop_n_last_frames < 0:
             raise ValueError(f"drop_n_last_frames must be >= 0, got {drop_n_last_frames}")
 
@@ -77,27 +78,12 @@ class EpisodeAwareSampler:
                 f"got {len(from_indices)} and {len(to_indices)}"
             )
 
-        if isinstance(drop_n_first_frames, int):
-            first_frame_offsets = np.full(len(from_indices), drop_n_first_frames, dtype=np.int64)
-        else:
-            first_frame_offsets = np.asarray(drop_n_first_frames, dtype=np.int64)
-            if first_frame_offsets.shape != from_indices.shape:
-                raise ValueError(
-                    "drop_n_first_frames must be an integer or have one value per episode; "
-                    f"got {len(first_frame_offsets)} values for {len(from_indices)} episodes"
-                )
-        if np.any(first_frame_offsets < 0):
-            raise ValueError(
-                "drop_n_first_frames must be >= 0, got "
-                f"{first_frame_offsets[first_frame_offsets < 0].tolist()}"
-            )
-
         used = np.ones(len(from_indices), dtype=bool)
         if episode_indices_to_use is not None:
             used = np.zeros(len(from_indices), dtype=bool)
             used[np.asarray(episode_indices_to_use, dtype=np.int64)] = True
 
-        starts = from_indices + first_frame_offsets
+        starts = from_indices + drop_n_first_frames
         lengths = to_indices - drop_n_last_frames - starts
         for episode_idx in np.flatnonzero(used & (lengths <= 0)):
             logger.warning(
@@ -105,7 +91,7 @@ class EpisodeAwareSampler:
                 "drop_n_last_frames=%d removes all frames. Skipping.",
                 episode_idx,
                 to_indices[episode_idx] - from_indices[episode_idx],
-                first_frame_offsets[episode_idx],
+                drop_n_first_frames,
                 drop_n_last_frames,
             )
         used &= lengths > 0

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/envs/robomme.py

@@ -63,8 +63,6 @@ class RoboMMEGymEnv(gym.Env):
         from robomme.env_record_wrapper import BenchmarkEnvBuilder
 
         self._task = task
-        self.task = task
-        self.task_description = task
         self._action_space_type = action_space_type
         self._dataset = dataset
         self._episode_idx = episode_idx
@@ -107,12 +105,6 @@ class RoboMMEGymEnv(gym.Env):
         )
         obs, info = self._env.reset()
         self._last_raw_obs = obs
-        task_goal = info.get("task_goal")
-        # RoboMME returns [simple_subgoal, grounded_subgoal]. The published
-        # LeRobot training dataset uses the simple subgoal as its `task` text.
-        if isinstance(task_goal, (list, tuple)):
-            task_goal = task_goal[0] if task_goal else ""
-        self.task_description = str(task_goal or self._task)
         return self._convert_obs(obs), self._convert_info(info)
 
     def step(self, action):

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,122 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""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 typing import TYPE_CHECKING
+
+import torch
+from torch import Tensor
+
+if TYPE_CHECKING:
+    from lerobot.policies.rtc.modeling_rtc import RTCProcessor
+
+
+def sample_beta(alpha: float, beta: float, bsize: int, device) -> Tensor:  # see openpi (exact copy)
+    # Beta sampling uses _sample_dirichlet which isn't implemented for MPS, so sample on CPU
+    alpha_t = torch.tensor(alpha, dtype=torch.float32)
+    beta_t = torch.tensor(beta, dtype=torch.float32)
+    dist = torch.distributions.Beta(alpha_t, beta_t)
+    return dist.sample((bsize,)).to(device)
+
+
+def sample_noise(shape, device) -> Tensor:
+    """Standard-normal float32 noise, the flow-matching x_1 sample."""
+    return torch.normal(
+        mean=0.0,
+        std=1.0,
+        size=shape,
+        dtype=torch.float32,
+        device=device,
+    )
+
+
+def sample_time_beta(bsize: int, device, *, alpha: float, beta: float, scale: float, offset: float) -> Tensor:
+    """Beta-distributed flow-matching timesteps: ``Beta(alpha, beta) * scale + offset`` (openpi convention)."""
+    time_beta = sample_beta(alpha, beta, bsize, device)
+    time = time_beta * scale + offset
+    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
+    return _get_policy_cls_from_policy_name(name=name)
 
 
 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
+    try:
+        config_cls = PreTrainedConfig.get_choice_class(policy_type)
+    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,8 +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_meta: Any | None
 
 
@@ -308,25 +168,8 @@ 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
-        and getattr(policy_cfg, "type", None) in {"pi0_fast", "pi052"}
-        and getattr(policy_cfg, "auto_fit_fast_tokenizer", False)
-        and kwargs.get("dataset_repo_id") is not None
-    ):
-        from .pi052.fit_fast_tokenizer import resolve_fast_tokenizer
-
-        overrides = dict(kwargs.get("preprocessor_overrides") or {})
-        action_tokenizer_override = {
-            **overrides.get("action_tokenizer_processor", {}),
-            "action_tokenizer_name": resolve_fast_tokenizer(policy_cfg, kwargs.get("dataset_repo_id")),
-        }
-        overrides["action_tokenizer_processor"] = action_tokenizer_override
-        kwargs["preprocessor_overrides"] = overrides
-
     if pretrained_path:
         if isinstance(policy_cfg, GrootConfig):
             from .groot.processor_groot import make_groot_pre_post_processors_from_pretrained
@@ -377,186 +220,13 @@ 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"),
-        )
-
-    elif policy_cfg.type == "pi052":
-        # PI052 must precede PI05 because its config subclasses PI05Config.
-        from .pi052.processor_pi052 import make_pi052_pre_post_processors
-
-        processors = make_pi052_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            # Without a dataset repo, FAST auto-fit falls back to the universal tokenizer.
-            dataset_repo_id=kwargs.get("dataset_repo_id"),
-        )
-
-    elif isinstance(policy_cfg, PI05Config):
-        from .pi05.processor_pi05 import make_pi05_pre_post_processors
-
-        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(
-            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
+    # 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"),
+    )
 
 
 def make_policy(
@@ -623,9 +293,6 @@ def make_policy(
             raise ValueError("env_cfg cannot be None when ds_meta is not provided")
         features = env_to_policy_features(env_cfg)
 
-    if rename_map:
-        features = {rename_map.get(key, key): feature for key, feature in features.items()}
-
     cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION}
     if not cfg.input_features:
         cfg.input_features = {key: ft for key, ft in features.items() if key not in cfg.output_features}
@@ -703,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.
@@ -732,22 +401,39 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
         "configuration_", "modeling_"
     )  # e.g., configuration_diffusion -> modeling_diffusion
 
-    module = importlib.import_module(module_path)
-    policy_cls = getattr(module, cls_name)
+    try:
+        module = importlib.import_module(module_path)
+    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.
     """
@@ -760,6 +446,19 @@ def _make_processors_from_policy_config(
     logging.debug(
         f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
     )
-    module = importlib.import_module(module_path)
-    function = getattr(module, function_name)
-    return function(config, dataset_stats=dataset_stats)
+    try:
+        module = importlib.import_module(module_path)
+    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,14 +734,8 @@ 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))
+        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
+        att_masks = att_masks[None, :].expand(bsize, len(att_masks))
 
         return embs, pad_masks, att_masks, adarms_cond
 
@@ -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)
+            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 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,16 +956,16 @@ class PI05Policy(PreTrainedPolicy):
         strict: bool = True,
         **kwargs,
     ) -> T:
-        """Load PI05-compatible single-file or sharded safetensors checkpoints."""
-        if cls.show_openpi_disclaimer:
-            print(
-                "The PI05 model is a direct port of the OpenPI implementation. \n"
-                "This implementation follows the original OpenPI structure for compatibility. \n"
-                "Original implementation: https://github.com/Physical-Intelligence/openpi"
-            )
+        """Override the from_pretrained method to handle key remapping and display important disclaimer."""
+        print(
+            "The PI05 model is a direct port of the OpenPI implementation. \n"
+            "This implementation follows the original OpenPI structure for compatibility. \n"
+            "Original implementation: https://github.com/Physical-Intelligence/openpi"
+        )
         if pretrained_name_or_path is None:
             raise ValueError("pretrained_name_or_path is required")
 
+        # Use provided config if available, otherwise create default config
         if config is None:
             config = PreTrainedConfig.from_pretrained(
                 pretrained_name_or_path=pretrained_name_or_path,
@@ -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(
-            pretrained_name_or_path,
-            force_download=force_download,
-            resume_download=resume_download,
-            proxies=proxies,
-            token=token,
-            cache_dir=cache_dir,
-            local_files_only=local_files_only,
-            revision=revision,
-        )
-        fixed_state_dict = model._fix_pytorch_state_dict_keys(_load_weight_files(files), model.config)
-        remapped_state_dict = {
-            key if key.startswith("model.") else f"model.{key}": value
-            for key, value in fixed_state_dict.items()
-        }
-        remapped_state_dict = model._prepare_pretrained_state_dict(remapped_state_dict)
-        missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
-        if missing_keys:
-            logging.warning("Missing %s checkpoint keys: %s", cls.name, missing_keys)
-        if unexpected_keys:
-            logging.warning("Unexpected %s checkpoint keys: %s", cls.name, unexpected_keys)
-        if model.eval_after_pretrained_load:
-            model.eval()
-        return model
 
-    def _prepare_pretrained_state_dict(self, state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
-        return state_dict
+        # Load state dict (expects keys with "model." prefix)
+        try:
+            print(f"Loading model from: {pretrained_name_or_path}")
+            try:
+                from transformers.utils import cached_file
+
+                resolved_file = cached_file(
+                    pretrained_name_or_path,
+                    "model.safetensors",
+                    cache_dir=kwargs.get("cache_dir"),
+                    force_download=kwargs.get("force_download", False),
+                    resume_download=kwargs.get("resume_download"),
+                    proxies=kwargs.get("proxies"),
+                    token=kwargs.get("token"),
+                    revision=kwargs.get("revision"),
+                    local_files_only=kwargs.get("local_files_only", False),
+                )
+                from safetensors.torch import load_file
+
+                original_state_dict = load_file(resolved_file)
+                print("✓ Loaded state dict from model.safetensors")
+            except Exception as e:
+                print(f"Could not load state dict from remote files: {e}")
+                print("Returning model without loading pretrained weights")
+                return model
+
+            # First, fix any key differences (see openpi model.py, _fix_pytorch_state_dict_keys)
+            fixed_state_dict = model._fix_pytorch_state_dict_keys(original_state_dict, model.config)
+
+            # Then add "model." prefix for all keys that don't already have it
+            remapped_state_dict = {}
+            remap_count = 0
+
+            for key, value in fixed_state_dict.items():
+                if not key.startswith("model."):
+                    new_key = f"model.{key}"
+                    remapped_state_dict[new_key] = value
+                    remap_count += 1
+                else:
+                    remapped_state_dict[key] = value
+
+            if remap_count > 0:
+                print(f"Remapped {remap_count} state dict keys")
+
+            # Load the remapped state dict into the model
+            missing_keys, unexpected_keys = model.load_state_dict(remapped_state_dict, strict=strict)
+
+            if missing_keys:
+                print(f"Missing keys when loading state dict: {len(missing_keys)} keys")
+                if len(missing_keys) <= 5:
+                    for key in missing_keys:
+                        print(f"  - {key}")
+                else:
+                    for key in missing_keys[:5]:
+                        print(f"  - {key}")
+                    print(f"  ... and {len(missing_keys) - 5} more")
+
+            if unexpected_keys:
+                print(f"Unexpected keys when loading state dict: {len(unexpected_keys)} keys")
+                if len(unexpected_keys) <= 5:
+                    for key in unexpected_keys:
+                        print(f"  - {key}")
+                else:
+                    for key in unexpected_keys[:5]:
+                        print(f"  - {key}")
+                    print(f"  ... and {len(unexpected_keys) - 5} more")
+
+            if not missing_keys and not unexpected_keys:
+                print("All keys loaded successfully!")
+
+        except Exception as e:
+            print(f"Warning: Could not load state dict: {e}")
+
+        return model
 
     def _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/configuration_pi052.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.
-
-"""PI0.5 with hierarchical text generation and flow-matched actions."""
-
-from dataclasses import dataclass
-
-from lerobot.configs import PreTrainedConfig
-from lerobot.optim.optimizers import AdamWConfig
-
-from ..pi05.configuration_pi05 import PI05Config
-
-
-@PreTrainedConfig.register_subclass("pi052")
-@dataclass
-class PI052Config(PI05Config):
-    """PI0.5 configuration for recipe-driven text and action supervision."""
-
-    # Recipe / language stack ---------------------------------------------
-    recipe_path: str | None = "recipes/subtask_mem.yaml"
-    """Recipe path relative to ``src/lerobot/configs/``, or ``None`` for the plain PI0.5 prompt."""
-
-    apply_chat_template: bool = False
-    """Whether to apply a tokenizer chat template.
-
-    PaliGemma defaults to plain recipe-rendered prefixes because it is not chat-pretrained.
-    """
-
-    # Balance frequent recipe text supervision against the paper's α=10 flow weight.
-    text_loss_weight: float = 1.0
-    """LM-head cross-entropy weight; ``0`` disables text training."""
-
-    flow_loss_weight: float = 10.0
-    """Weight on action-expert flow matching relative to text supervision."""
-
-    # Backbone training ---------------------------------------------------
-    unfreeze_lm_head: bool = True
-    """Keep PaliGemma's language head trainable for hierarchical inference."""
-
-    # Optional context dropout improves tolerance to missing or stale language state.
-    plan_dropout_prob: float = 0.0
-    memory_dropout_prob: float = 0.0
-    subtask_dropout_prob: float = 0.0
-
-    # FAST adds discrete-action CE to the text and flow objectives from paper §III.B-C.
-    enable_fast_action_loss: bool = True
-    """Add FAST-tokenized action cross-entropy to text CE and flow matching."""
-
-    action_tokenizer_name: str = "physical-intelligence/fast"
-    """HF identifier for the FAST action tokenizer."""
-
-    max_action_tokens: int = 256
-    """Maximum number of FAST tokens per action chunk."""
-
-    fast_skip_tokens: int = 128
-    """Number of low-vocab tokens the FAST tokenizer skips to avoid
-    collisions with PaliGemma's text vocabulary."""
-
-    fast_action_loss_weight: float = 1.0
-    """Weight on FAST action-token CE relative to continuous-flow supervision."""
-
-    subtask_replan_steps: int = 0
-    """Environment steps between subtask generations during evaluation.
-
-    Non-positive values regenerate each action chunk while still refreshing the action prompt every chunk.
-    """
-
-    auto_fit_fast_tokenizer: bool = False
-    """Fit and cache a dataset-specific FAST tokenizer before training.
-
-    Disabled by default to avoid the extra dataset pass and use the universal tokenizer.
-    """
-
-    fast_tokenizer_cache_dir: str = "~/.cache/lerobot/fast_tokenizers"
-    """Where fitted FAST tokenizers are stored. ``~`` expands."""
-
-    fast_tokenizer_fit_samples: int = 1024
-    """Number of action chunks sampled when fitting FAST."""
-
-    # Knowledge insulation detaches VLM K/V from action-loss gradients (paper §III.B).
-    knowledge_insulation: bool = True
-    """Block action-loss gradients through VLM keys and values."""
-
-    # Optional training backends. Defaults preserve the eager/SDPA path.
-    use_flashrt_adarms: bool = False
-    """Use FlashRT adaptive RMSNorm kernels when available."""
-
-    use_compiled_text_ce: bool = False
-    """Compile the materialized-logits text and FAST CE path."""
-
-    use_compiled_vision: bool = False
-    """Compile the SigLIP tower for no-grad flow and inference passes."""
-
-    use_flex_attention: bool = False
-    """Use FlexAttention for amortized KI, with SDPA fallback where unsupported."""
-
-    use_manual_attention: bool = False
-    """Use materialized-logits attention for explicitly profiled KI shapes."""
-
-    manual_attention_scope: str = "all"
-    """Apply manual attention to all KI queries or only action queries."""
-
-    # Scale language-head updates relative to the base optimizer schedule.
-    lm_head_lr_scale: float = 1.0
-
-    # Scale backbone and action-expert optimizer groups independently.
-    backbone_lr_scale: float = 1.0
-    action_expert_lr_scale: float = 1.0
-
-    # Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow.
-    flow_num_repeats: int = 5
-
-    # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
-    text_ce_z_loss_weight: float = 1e-4
-
-    use_flashrt_fp8_mlp: bool = False
-    """Enable calibrated FlashRT FP8 kernels for Gemma and SigLIP MLPs.
-
-    Apply after loading with ``PI052Policy.apply_flashrt_fp8_mlp``; unavailable kernels keep BF16.
-    """
-
-    # Keep serialized PI052 AdamW options local because PI05Config lacks them.
-    optimizer_foreach: bool | None = False
-    optimizer_fused: bool | None = True
-
-    def get_optimizer_preset(self) -> AdamWConfig:
-        return AdamWConfig(
-            lr=self.optimizer_lr,
-            betas=self.optimizer_betas,
-            eps=self.optimizer_eps,
-            weight_decay=self.optimizer_weight_decay,
-            grad_clip_norm=self.optimizer_grad_clip_norm,
-            foreach=self.optimizer_foreach,
-            fused=self.optimizer_fused,
-        )
-
-    def __post_init__(self) -> None:
-        super().__post_init__()
-        if self.text_loss_weight > 0 and self.unfreeze_lm_head:
-            self.train_expert_only = False
-        if self.flow_num_repeats < 1:
-            raise ValueError(f"flow_num_repeats must be >= 1, got {self.flow_num_repeats}")
-        if self.manual_attention_scope not in {"all", "action"}:
-            raise ValueError(
-                f"manual_attention_scope must be 'all' or 'action', got {self.manual_attention_scope!r}"
-            )
-        if self.use_flex_attention and self.use_manual_attention:
-            raise ValueError("use_flex_attention and use_manual_attention are mutually exclusive")
-        if self.use_flex_attention and self.flow_num_repeats == 1:
-            raise ValueError("use_flex_attention requires flow_num_repeats > 1")
-        if not self.knowledge_insulation and (
-            self.use_flex_attention or self.use_manual_attention or self.use_flashrt_adarms
-        ):
-            raise ValueError("KI attention and AdaRMS optimizations require knowledge_insulation=True")

src/lerobot/policies/pi052/fit_fast_tokenizer.py

@@ -1,256 +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 logging
-import os
-import time
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-
-logger = logging.getLogger(__name__)
-
-# ``ProcessorMixin.save_pretrained`` writes this shared cache sentinel.
-_CACHE_SENTINEL = "processor_config.json"
-
-
-def _is_local_leader() -> bool:
-    return int(os.environ.get("LOCAL_RANK", "0")) == 0
-
-
-def _dataset_signature(
-    dataset_repo_id: str,
-    base_tokenizer_name: str,
-    n_samples: int,
-    chunk_size: int,
-) -> str:
-    """Deterministic short hash for naming the cache directory.
-
-    Keys on (dataset, base tokenizer, sample count, chunk size) so any
-    of those changing re-runs the fit. ``chunk_size`` matters because
-    the tokenizer is fit on chunks of that length.
-    """
-    h = hashlib.sha256()
-    h.update(dataset_repo_id.encode("utf-8"))
-    h.update(b"\0")
-    h.update(base_tokenizer_name.encode("utf-8"))
-    h.update(b"\0")
-    h.update(str(n_samples).encode("utf-8"))
-    h.update(b"\0")
-    h.update(str(chunk_size).encode("utf-8"))
-    return h.hexdigest()[:16]
-
-
-def fit_fast_tokenizer(
-    *,
-    dataset_repo_id: str,
-    cache_dir: str | Path,
-    base_tokenizer_name: str = "physical-intelligence/fast",
-    n_samples: int = 1024,
-    chunk_size: int = 50,
-    seed: int = 42,
-) -> str:
-    """Fit a FAST tokenizer on a LeRobot dataset's action distribution.
-
-    Args:
-        dataset_repo_id: HF Hub repo id of the LeRobotDataset to fit on.
-        cache_dir: Directory under which to save (and look up) fitted
-            tokenizers. The actual save path is
-            ``{cache_dir}/{signature}``.
-        base_tokenizer_name: HF identifier for the base FAST tokenizer
-            to finetune from. ``physical-intelligence/fast`` is the
-            universal one.
-        n_samples: Number of action chunks to sample for the fit. The
-            FAST paper uses a few thousand; ``1024`` is a good default
-            for medium datasets.
-        chunk_size: Length of each action chunk (matches
-            ``policy.chunk_size``). The FAST tokenizer is fit on
-            sequences of this length.
-        seed: RNG seed for sample selection.
-
-    Returns:
-        The local path to the fitted tokenizer. Passed directly to
-        ``--policy.action_tokenizer_name`` for the training run.
-
-    Raises:
-        ImportError: If the ``transformers`` library doesn't expose
-            ``AutoProcessor`` or the FAST tokenizer doesn't have a
-            ``.fit()`` method (then you're on an older FAST snapshot —
-            update to the current published model).
-        FileNotFoundError: If the dataset can't be loaded.
-    """
-    cache_dir = Path(cache_dir)
-    sig = _dataset_signature(dataset_repo_id, base_tokenizer_name, n_samples, chunk_size)
-    out_dir = cache_dir / sig
-
-    if out_dir.exists() and (out_dir / _CACHE_SENTINEL).exists():
-        logger.info(
-            "FAST tokenizer cache hit: %s — re-using fitted tokenizer for dataset=%s base=%s n_samples=%d",
-            out_dir,
-            dataset_repo_id,
-            base_tokenizer_name,
-            n_samples,
-        )
-        return str(out_dir)
-
-    # Each node fits its node-local cache once; its other local ranks wait.
-    is_leader = _is_local_leader()
-    if not is_leader:
-        timeout_s = 1800.0  # 30 min — covers ~1024-sample fits on cold caches
-        start = time.monotonic()
-        while not (out_dir / _CACHE_SENTINEL).exists():
-            if time.monotonic() - start > timeout_s:
-                raise RuntimeError(
-                    f"FAST tokenizer fit: non-leader rank timed out after "
-                    f"{timeout_s:.0f}s waiting for {out_dir / _CACHE_SENTINEL}. "
-                    "Leader rank likely crashed during the fit."
-                )
-            time.sleep(2.0)
-        logger.info("FAST tokenizer ready (leader populated cache): %s", out_dir)
-        return str(out_dir)
-
-    logger.info(
-        "FAST tokenizer cache miss — fitting on dataset=%s base=%s n_samples=%d chunk_size=%d → %s",
-        dataset_repo_id,
-        base_tokenizer_name,
-        n_samples,
-        chunk_size,
-        out_dir,
-    )
-
-    from transformers import AutoProcessor  # noqa: PLC0415
-
-    # Read action columns directly to avoid video decoding and bound memory to sampled episodes.
-    rng = np.random.default_rng(seed)
-    actions_buf: list[np.ndarray] = []
-
-    # Read v3 parquet shards directly to avoid split lookup failures and repeated metadata parsing.
-    import pyarrow as _pa  # noqa: PLC0415
-    import pyarrow.parquet as _pq  # noqa: PLC0415
-    from huggingface_hub import snapshot_download  # noqa: PLC0415
-
-    snap = Path(snapshot_download(repo_id=dataset_repo_id, repo_type="dataset"))
-    data_files = sorted((snap / "data").glob("chunk-*/file-*.parquet"))
-    if not data_files:
-        raise RuntimeError(f"FAST fit: no ``data/chunk-*/file-*.parquet`` shards found under {snap!s}.")
-
-    # Load only episode indices and fixed-width actions across all shards.
-    tables = [_pq.read_table(f, columns=["episode_index", "action"]) for f in data_files]
-    table = _pa.concat_tables(tables)
-    eps = table["episode_index"].to_numpy()
-    acts_col = table["action"]
-    # Normalize Arrow action representations into an (N, D) array.
-    try:
-        acts = np.stack(acts_col.to_numpy(zero_copy_only=False)).astype(np.float32)
-    except Exception:  # noqa: BLE001
-        # Fallback path for nested-list types: flatten via to_pylist().
-        acts = np.asarray(acts_col.to_pylist(), dtype=np.float32)
-    if acts.ndim != 2:
-        raise RuntimeError(f"FAST fit: expected ``action`` rows to be 1-D vectors; got shape {acts.shape}.")
-
-    # Sort once because episode order is only guaranteed within each shard.
-    order = np.argsort(eps, kind="stable")
-    eps_sorted = eps[order]
-    boundaries = np.searchsorted(eps_sorted, np.arange(int(eps_sorted.max()) + 2))
-    ep_to_slice: dict[int, tuple[int, int]] = {
-        int(ep): (int(boundaries[ep]), int(boundaries[ep + 1]))
-        for ep in range(len(boundaries) - 1)
-        if boundaries[ep] < boundaries[ep + 1]
-    }
-    num_episodes = len(ep_to_slice)
-    # ``acts`` is in original (un-sorted-by-episode) row order; reorder
-    # so per-episode slices are contiguous.
-    acts = acts[order]
-
-    samples_per_episode = max(1, n_samples // max(num_episodes, 1))
-    collected = 0
-    eps_visited = 0
-    short_episodes = 0
-    ep_indices = list(ep_to_slice.keys())
-    for ep_idx in rng.permutation(ep_indices):
-        if collected >= n_samples:
-            break
-        start, stop = ep_to_slice[int(ep_idx)]
-        ep_actions = acts[start:stop]
-        if ep_actions.shape[0] < chunk_size:
-            short_episodes += 1
-            continue
-        starts = rng.integers(0, ep_actions.shape[0] - chunk_size + 1, size=samples_per_episode)
-        for s in starts:
-            actions_buf.append(ep_actions[int(s) : int(s) + chunk_size])
-            collected += 1
-            if collected >= n_samples:
-                break
-        eps_visited += 1
-
-    if not actions_buf:
-        raise RuntimeError(
-            f"FAST fit collected zero action chunks from {dataset_repo_id!r}: "
-            f"all {num_episodes} episodes were shorter than chunk_size="
-            f"{chunk_size} ({short_episodes} too short) or had an unreadable "
-            "``action`` column. Lower ``chunk_size`` to match your episode "
-            "lengths."
-        )
-
-    actions = np.stack(actions_buf, axis=0).astype(np.float32)  # (N, H, D)
-    logger.info(
-        "FAST fit: collected %d chunks of shape %s from %d episodes",
-        actions.shape[0],
-        actions.shape[1:],
-        eps_visited,
-    )
-
-    # Match training-time quantile normalization so FAST sees the same bounded action space.
-    flat = actions.reshape(-1, actions.shape[-1])
-    q01 = np.quantile(flat, 0.01, axis=0)
-    q99 = np.quantile(flat, 0.99, axis=0)
-    span = np.where((q99 - q01) > 1e-6, q99 - q01, 1.0)
-    actions = np.clip((actions - q01) / span * 2.0 - 1.0, -1.0, 1.0).astype(np.float32)
-
-    base = AutoProcessor.from_pretrained(base_tokenizer_name, trust_remote_code=True)
-    if not hasattr(base, "fit"):
-        raise ImportError(
-            f"Base FAST tokenizer {base_tokenizer_name!r} has no ``.fit()`` "
-            "method — your transformers / model snapshot is too old. Update "
-            "to the current ``physical-intelligence/fast`` revision."
-        )
-
-    fitted = base.fit(actions)
-    out_dir.mkdir(parents=True, exist_ok=True)
-    fitted.save_pretrained(str(out_dir))
-    logger.info("FAST fit: saved fitted tokenizer to %s", out_dir)
-    return str(out_dir)
-
-
-def resolve_fast_tokenizer(config: Any, dataset_repo_id: str | None) -> str:
-    """Return the configured tokenizer, fitting a cached dataset-specific one when requested."""
-    if not getattr(config, "auto_fit_fast_tokenizer", False) or dataset_repo_id is None:
-        return config.action_tokenizer_name
-
-    return fit_fast_tokenizer(
-        dataset_repo_id=dataset_repo_id,
-        cache_dir=Path(config.fast_tokenizer_cache_dir).expanduser(),
-        base_tokenizer_name=config.action_tokenizer_name,
-        n_samples=config.fast_tokenizer_fit_samples,
-        chunk_size=config.chunk_size,
-    )

src/lerobot/policies/pi052/flashrt_fp8.py

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

@@ -1,205 +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.
-        content = subtask
-        obs_state = observation.get(OBS_STATE)
-        if isinstance(obs_state, torch.Tensor) and obs_state.numel() > 0:
-            from lerobot.policies.pi052.text_processor_pi052 import discretize_state_str  # noqa: PLC0415
-
-            state_row = obs_state[0] if obs_state.ndim > 1 else obs_state
-            content = f"{subtask}, State: {discretize_state_str(state_row)};"
-
-        text_batch = _build_text_batch(
-            self.policy,
-            [{"role": "user", "content": content}],
-            add_generation_prompt=False,
-        )
-        batch = dict(observation)
-        batch[OBS_LANGUAGE_TOKENS] = text_batch["lang_tokens"]
-        batch[OBS_LANGUAGE_ATTENTION_MASK] = text_batch["lang_masks"]
-        return self.policy.predict_action_chunk(batch)
-
-    def generate_text(
-        self,
-        kind: str,
-        observation: dict[str, Any] | None,
-        state: RuntimeState,
-        user_text: str | None = None,
-    ) -> str:
-        messages = self.build_messages(kind, state, user_text=user_text)
-        return _generate_with_policy(
-            self.policy,
-            messages,
-            observation=observation,
-            state=state,
-            label=f"{kind} gen",
-            min_new_tokens=self.gen.min_new_tokens,
-            temperature=self.gen.temperature,
-            top_p=self.gen.top_p,
-            suppress_loc_tokens=True,  # all runtime text is prose; never emit <loc>
-        )
-
-    def build_messages(
-        self,
-        kind: str,
-        state: RuntimeState,
-        *,
-        user_text: str | None = None,
-    ) -> list[dict[str, Any]]:
-        if kind in ("subtask", "plan"):
-            return [{"role": "user", "content": state.task or ""}]
-        if kind == "memory":
-            messages = [{"role": "user", "content": state.task or ""}]
-            if state.language_context.get("memory"):
-                messages.append(
-                    {"role": "assistant", "content": f"Previous memory: {state.language_context['memory']}"}
-                )
-            if state.extra.get("prior_subtask"):
-                messages.append(
-                    {"role": "user", "content": f"Completed subtask: {state.extra['prior_subtask']}"}
-                )
-            return messages
-        if kind == "interjection":
-            messages = [{"role": "user", "content": state.task or ""}]
-            if state.language_context.get("plan"):
-                messages.append(
-                    {"role": "assistant", "content": f"Previous plan:\n{state.language_context['plan']}"}
-                )
-            if user_text:
-                messages.append({"role": "user", "content": user_text})
-            return messages
-        raise ValueError(f"Unknown PI052 text kind: {kind}")
-
-
-def _get_loc_tokenizer(tok_name: str, auto_tokenizer_cls: Any, register_loc_fn: Any) -> Any:
-    tokenizer = _LOC_TOKENIZER_CACHE.get(tok_name)
-    if tokenizer is None:
-        tokenizer = register_loc_fn(auto_tokenizer_cls.from_pretrained(tok_name))
-        _LOC_TOKENIZER_CACHE[tok_name] = tokenizer
-    return tokenizer
-
-
-def _build_text_batch(
-    policy: Any,
-    prompt_messages: list[dict[str, Any]],
-    *,
-    add_generation_prompt: bool = True,
-) -> dict[str, Any]:
-    import torch  # noqa: PLC0415
-    from transformers import AutoTokenizer  # noqa: PLC0415
-
-    from lerobot.policies.pi052.text_processor_pi052 import (  # noqa: PLC0415
-        _flatten_say_tool_calls,
-        _format_messages,
-        _strip_blocks,
-        register_paligemma_loc_tokens,
-    )
-
-    tok_name = getattr(policy.config, "tokenizer_name", None) or "google/paligemma-3b-pt-224"
-    tokenizer = _get_loc_tokenizer(tok_name, AutoTokenizer, register_paligemma_loc_tokens)
-
-    messages = [_strip_blocks(_flatten_say_tool_calls(m)) for m in prompt_messages]
-    prompt, _spans = _format_messages(messages)
-    if add_generation_prompt:
-        prompt = prompt + "Assistant: "
-
-    encoded = tokenizer(prompt, return_tensors="pt")
-    ids = encoded["input_ids"]
-    attn = encoded.get("attention_mask")
-    if attn is None and tokenizer.pad_token_id is not None:
-        attn = ids != tokenizer.pad_token_id
-    if attn is not None and hasattr(attn, "dtype") and attn.dtype != torch.bool:
-        attn = attn.bool()
-
-    device = getattr(getattr(policy, "config", None), "device", None)
-    if device is not None:
-        try:
-            ids = ids.to(device)
-            if attn is not None and hasattr(attn, "to"):
-                attn = attn.to(device)
-        except Exception as exc:  # noqa: BLE001
-            logger.debug("could not move pi052 lang tokens to %s: %s", device, exc)
-    return {"lang_tokens": ids, "lang_masks": attn, "tokenizer": tokenizer}
-
-
-def _generate_with_policy(
-    policy: Any,
-    messages: list[dict[str, Any]],
-    *,
-    observation: dict[str, Any] | None = None,
-    state: RuntimeState | None = None,
-    label: str = "select_message",
-    min_new_tokens: int = 0,
-    temperature: float = 0.0,
-    top_p: float = 1.0,
-    suppress_loc_tokens: bool = False,
-) -> str:
-    if not hasattr(policy, "select_message"):
-        if state is not None:
-            state.log(f"  [warn] policy has no select_message — skipping {label}")
-        return ""
-    text_batch = _build_text_batch(policy, messages)
-    try:
-        from lerobot.utils.constants import OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS  # noqa: PLC0415
-
-        batch: dict[str, Any] = {
-            OBS_LANGUAGE_TOKENS: text_batch["lang_tokens"],
-            OBS_LANGUAGE_ATTENTION_MASK: text_batch["lang_masks"],
-        }
-        if observation:
-            for k, v in observation.items():
-                if isinstance(k, str) and k.startswith("observation.") and k not in batch:
-                    batch[k] = v
-        return policy.select_message(
-            batch,
-            tokenizer=text_batch["tokenizer"],
-            min_new_tokens=min_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            suppress_loc_tokens=suppress_loc_tokens,
-        )
-    except Exception as exc:  # noqa: BLE001
-        logger.warning("%s failed: %s", label, exc, exc_info=logger.isEnabledFor(logging.DEBUG))
-        if state is not None:
-            state.log(f"  [warn] {label} failed: {type(exc).__name__}: {exc}")
-        return ""

src/lerobot/policies/pi052/processor_pi052.py

@@ -1,149 +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,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """Build PI0.5-v2's pre/post-processor pipelines.
-
-    Falls through to π0.5's stock pipeline when ``recipe_path`` is unset.
-    """
-    if not config.recipe_path:
-        return make_pi05_pre_post_processors(config, dataset_stats=dataset_stats)
-
-    recipe = _load_recipe(config.recipe_path)
-
-    relative_step = RelativeActionsProcessorStep(
-        enabled=config.use_relative_actions,
-        exclude_joints=getattr(config, "relative_exclude_joints", []),
-        action_names=getattr(config, "action_feature_names", None),
-    )
-
-    input_steps = [
-        RenameObservationsProcessorStep(rename_map={}),
-        AddBatchDimensionProcessorStep(),
-        relative_step,
-        NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        RenderMessagesStep(recipe=recipe),
-        PI052TextTokenizerStep(
-            tokenizer_name="google/paligemma-3b-pt-224",
-            max_length=config.tokenizer_max_length,
-            plan_dropout_prob=getattr(config, "plan_dropout_prob", 0.0),
-            memory_dropout_prob=getattr(config, "memory_dropout_prob", 0.0),
-            subtask_dropout_prob=getattr(config, "subtask_dropout_prob", 0.0),
-        ),
-    ]
-
-    # Add FAST action-token supervision only when explicitly enabled.
-    if getattr(config, "enable_fast_action_loss", False):
-        from .fit_fast_tokenizer import resolve_fast_tokenizer  # noqa: PLC0415
-
-        input_steps.append(
-            ActionTokenizerProcessorStep(
-                action_tokenizer_name=resolve_fast_tokenizer(config, dataset_repo_id),
-                max_action_tokens=config.max_action_tokens,
-                fast_skip_tokens=config.fast_skip_tokens,
-                paligemma_tokenizer_name="google/paligemma-3b-pt-224",
-            )
-        )
-
-    input_steps.append(DeviceProcessorStep(device=config.device))
-
-    output_steps = [
-        UnnormalizerProcessorStep(
-            features=config.output_features,
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-        ),
-        AbsoluteActionsProcessorStep(
-            enabled=config.use_relative_actions,
-            relative_step=relative_step,
-        ),
-        DeviceProcessorStep(device="cpu"),
-    ]
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )
-
-
-def _load_recipe(path_str: str) -> TrainingRecipe:
-    """Resolve ``path_str`` to a ``TrainingRecipe``.
-
-    Accepts an absolute path or a path relative to
-    ``src/lerobot/configs/``.
-    """
-    p = Path(path_str)
-    if not p.is_absolute() and not p.exists():
-        from lerobot.configs import recipe as _recipe_module  # noqa: PLC0415
-
-        configs_dir = Path(_recipe_module.__file__).resolve().parent
-        candidate = configs_dir / path_str
-        if candidate.exists():
-            p = candidate
-    return TrainingRecipe.from_yaml(p)

src/lerobot/policies/pi052/text_processor_pi052.py

@@ -1,483 +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
-
-
-@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,7 +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,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -136,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,
@@ -144,19 +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)
-
     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,
@@ -165,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)
+        missing_keys, unexpected_keys = load_model_as_safetensor(
+            model, model_file, strict=strict, device=resolve_safetensors_device(map_location)
+        )
         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."
-            )
-            model.to(map_location)
         return model
 
     @abc.abstractmethod

src/lerobot/policies/smolvla/configuration_smolvla.py

@@ -44,14 +44,6 @@ class SmolVLAConfig(PreTrainedConfig):
     # Image preprocessing
     resize_imgs_with_padding: tuple[int, int] = (512, 512)
 
-    # MEM short-horizon visual memory (https://arxiv.org/abs/2603.03596).
-    # Past frames are fused inside the vision tower and dropped before the VLM,
-    # keeping the number of prefix tokens identical to single-frame SmolVLA.
-    use_visual_memory: bool = False
-    visual_memory_frames: int = 6
-    visual_memory_stride: int = 10
-    visual_memory_temporal_attention_every: int = 4
-
     # Add empty images. Used by smolvla_aloha_sim which adds the empty
     # left and right wrist cameras in addition to the top camera.
     empty_cameras: int = 0
@@ -127,12 +119,6 @@ class SmolVLAConfig(PreTrainedConfig):
             raise NotImplementedError(
                 "`use_delta_joint_actions_aloha` is used by smolvla for aloha real models. It is not ported yet in LeRobot."
             )
-        if self.visual_memory_frames < 1:
-            raise ValueError("visual_memory_frames must be at least 1")
-        if self.visual_memory_stride < 1:
-            raise ValueError("visual_memory_stride must be at least 1")
-        if self.visual_memory_temporal_attention_every < 1:
-            raise ValueError("visual_memory_temporal_attention_every must be at least 1")
 
     def validate_features(self) -> None:
         for i in range(self.empty_cameras):
@@ -162,10 +148,7 @@ class SmolVLAConfig(PreTrainedConfig):
 
     @property
     def observation_delta_indices(self) -> list:
-        if not self.use_visual_memory:
-            return [0]
-        horizon = (self.visual_memory_frames - 1) * self.visual_memory_stride
-        return list(range(-horizon, 1, self.visual_memory_stride))
+        return [0]
 
     @property
     def action_delta_indices(self) -> list:

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -71,7 +71,6 @@ from ..utils import (
 )
 from .configuration_smolvla import SmolVLAConfig
 from .smolvlm_with_expert import SmolVLMWithExpertModel
-from .visual_memory import sample_visual_history
 
 
 class ActionSelectKwargs(TypedDict, total=False):
@@ -254,10 +253,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
         self._queues = {
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
-        if self.config.use_visual_memory:
-            history_length = (self.config.visual_memory_frames - 1) * self.config.visual_memory_stride + 1
-            self._queues.update({key: deque(maxlen=history_length) for key in self.config.image_features})
-            self._visual_memory_steps_seen = 0
 
     def init_rtc_processor(self):
         """Initialize RTC processor if RTC is enabled in config."""
@@ -286,15 +281,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
         # In the case of offline inference, we have the action in the batch
         # that why without the k != ACTION check, it will raise an error because we are trying to stack
         # on an empty container.
-        for k in list(batch):
+        for k in batch:
             if k in self._queues and k != ACTION:
-                history = list(self._queues[k])
-                batch[k], batch[f"{k}_is_pad"] = sample_visual_history(
-                    history,
-                    num_frames=self.config.visual_memory_frames,
-                    stride=self.config.visual_memory_stride,
-                    steps_seen=self._visual_memory_steps_seen,
-                )
+                batch[k] = torch.stack(list(self._queues[k]), dim=1)
 
         images, img_masks = self.prepare_images(batch)
         state = self.prepare_state(batch)
@@ -320,11 +309,6 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         return batch
 
-    def _populate_observation_queues(self, batch: dict[str, Tensor]) -> None:
-        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
-        if self.config.use_visual_memory:
-            self._visual_memory_steps_seen += 1
-
     @torch.no_grad()
     def predict_action_chunk(
         self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs: Unpack[ActionSelectKwargs]
@@ -332,7 +316,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
         self.eval()
 
         batch = self._prepare_batch(batch)
-        self._populate_observation_queues(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
         actions = self._get_action_chunk(batch, noise, **kwargs)
         return actions
@@ -354,7 +338,7 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         self.eval()
         batch = self._prepare_batch(batch)
-        self._populate_observation_queues(batch)
+        self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
         if self._check_get_actions_condition():
             actions = self._get_action_chunk(batch, noise)
@@ -443,30 +427,19 @@ class SmolVLAPolicy(PreTrainedPolicy):
             )
         # Preprocess image features present in the batch
         for key in present_img_keys:
-            img = batch[key]
-            if img.ndim == 5 and not self.config.use_visual_memory:
-                img = img[:, -1, :, :, :]
+            img = batch[key][:, -1, :, :, :] if batch[key].ndim == 5 else batch[key]
             if self.config.resize_imgs_with_padding is not None:
-                if img.ndim == 5:
-                    batch_size, num_frames = img.shape[:2]
-                    img = resize_with_pad(
-                        img.flatten(0, 1), *self.config.resize_imgs_with_padding, pad_value=0
-                    ).unflatten(0, (batch_size, num_frames))
-                else:
-                    img = resize_with_pad(img, *self.config.resize_imgs_with_padding, pad_value=0)
+                img = resize_with_pad(img, *self.config.resize_imgs_with_padding, pad_value=0)
 
             # Normalize from range [0,1] to [-1,1] as expacted by siglip
             img = img * 2.0 - 1.0
 
             bsize = img.shape[0]
             device = img.device
-            if f"{key}_is_pad" in batch:
-                mask = ~batch[f"{key}_is_pad"].bool()
-            elif f"{key}_padding_mask" in batch:
+            if f"{key}_padding_mask" in batch:
                 mask = batch[f"{key}_padding_mask"].bool()
             else:
-                mask_shape = img.shape[:2] if img.ndim == 5 else (bsize,)
-                mask = torch.ones(mask_shape, dtype=torch.bool, device=device)
+                mask = torch.ones(bsize, dtype=torch.bool, device=device)
             images.append(img)
             img_masks.append(mask)
 
@@ -689,11 +662,7 @@ class VLAFlowMatching(nn.Module):
                 embs.append(image_start_token)
                 pad_masks.append(image_start_mask)
 
-            img_emb = self.vlm_with_expert.embed_image(
-                img,
-                frame_mask=img_mask if img.ndim == 5 else None,
-                temporal_attention_every=self.config.visual_memory_temporal_attention_every,
-            )
+            img_emb = self.vlm_with_expert.embed_image(img)
             img_emb = img_emb
 
             # Normalize image embeddings
@@ -701,8 +670,6 @@ class VLAFlowMatching(nn.Module):
             img_emb = img_emb * torch.tensor(img_emb_dim**0.5, dtype=img_emb.dtype, device=img_emb.device)
 
             bsize, num_img_embs = img_emb.shape[:2]
-            if img_mask.ndim == 2:
-                img_mask = img_mask[:, -1]
             img_mask = img_mask[:, None].expand(bsize, num_img_embs)
 
             embs.append(img_emb)

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/smolvla/smolvlm_with_expert.py

@@ -20,8 +20,6 @@ from torch import nn
 
 from lerobot.utils.import_utils import _transformers_available, require_package
 
-from .visual_memory import encode_video_with_mem
-
 if TYPE_CHECKING or _transformers_available:
     from transformers import (
         AutoConfig,
@@ -190,31 +188,17 @@ class SmolVLMWithExpertModel(nn.Module):
         if self.train_expert_only:
             self.vlm.eval()
 
-    def embed_image(
-        self,
-        image: torch.Tensor,
-        *,
-        frame_mask: torch.Tensor | None = None,
-        temporal_attention_every: int = 4,
-    ):
+    def embed_image(self, image: torch.Tensor):
         patch_attention_mask = None
         # Get sequence from the vision encoder
-        vision_model = self.get_vlm_model().vision_model
-        image = image.to(dtype=vision_model.dtype)
-        if image.ndim == 5:
-            if frame_mask is None:
-                frame_mask = torch.ones(image.shape[:2], dtype=torch.bool, device=image.device)
-            image_hidden_states = encode_video_with_mem(
-                vision_model,
-                image,
-                frame_mask,
-                temporal_attention_every=temporal_attention_every,
-            )
-        else:
-            image_hidden_states = vision_model(
-                pixel_values=image,
+        image_hidden_states = (
+            self.get_vlm_model()
+            .vision_model(
+                pixel_values=image.to(dtype=self.get_vlm_model().vision_model.dtype),
                 patch_attention_mask=patch_attention_mask,
-            ).last_hidden_state
+            )
+            .last_hidden_state
+        )
         # Modality projection & resampling
         image_hidden_states = self.get_vlm_model().connector(image_hidden_states)
         return image_hidden_states

src/lerobot/policies/smolvla/visual_memory.py

@@ -1,161 +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.
-
-"""Short-horizon visual memory from MEM (arXiv:2603.03596).
-
-The encoder keeps SigLIP's pretrained parameters and alternates its ordinary
-per-frame spatial attention with causal attention across time for matching
-spatial patches. Only the current frame's patch tokens leave the vision tower,
-so enabling memory does not increase the VLM prefix length.
-"""
-
-import math
-
-import torch
-from torch import Tensor
-
-
-def sample_visual_history(
-    history: list[Tensor], *, num_frames: int, stride: int, steps_seen: int
-) -> tuple[Tensor, Tensor]:
-    """Subsample an inference queue and mark pre-episode padding frames."""
-    sampled = history[::stride][-num_frames:]
-    if len(sampled) != num_frames:
-        raise ValueError(f"visual history has {len(sampled)} samples, expected {num_frames}")
-    video = torch.stack(sampled, dim=1)
-    required_ages = range((num_frames - 1) * stride, -1, -stride)
-    valid = torch.tensor([steps_seen > age for age in required_ages], dtype=torch.bool, device=video.device)
-    padding_mask = (~valid)[None, :].expand(video.shape[0], -1)
-    return video, padding_mask
-
-
-def temporal_sinusoidal_embedding(
-    num_frames: int, hidden_size: int, *, device: torch.device, dtype: torch.dtype
-) -> Tensor:
-    """Return fixed temporal embeddings with an exactly-zero current position."""
-    if hidden_size % 2:
-        raise ValueError(f"hidden_size must be even, got {hidden_size}")
-
-    # History is ordered oldest -> current, matching t in [-K, 0] in MEM.
-    positions = torch.arange(1 - num_frames, 1, device=device, dtype=torch.float32)[:, None]
-    frequencies = torch.exp(
-        torch.arange(0, hidden_size, 2, device=device, dtype=torch.float32)
-        * (-math.log(10_000.0) / hidden_size)
-    )[None, :]
-    angles = positions * frequencies
-    embedding = torch.zeros(num_frames, hidden_size, device=device, dtype=torch.float32)
-    embedding[:, 0::2] = torch.sin(angles)
-    embedding[:, 1::2] = torch.cos(angles) - 1.0
-    return embedding.to(dtype=dtype)
-
-
-def causal_temporal_mask(frame_mask: Tensor, *, dtype: torch.dtype, num_patches: int) -> Tensor:
-    """Build an additive causal/key-padding mask for per-patch temporal attention."""
-    if frame_mask.ndim != 2:
-        raise ValueError(f"frame_mask must have shape (batch, frames), got {tuple(frame_mask.shape)}")
-
-    batch_size, num_frames = frame_mask.shape
-    allowed = torch.ones(num_frames, num_frames, dtype=torch.bool, device=frame_mask.device).tril()
-    allowed = allowed[None, :, :] & frame_mask[:, None, :].bool()
-    # The current frame is always present in normal use. Keeping the diagonal
-    # valid also prevents NaNs for fully padded historical query rows.
-    diagonal = torch.eye(num_frames, dtype=torch.bool, device=frame_mask.device)[None, :, :]
-    allowed = allowed | diagonal
-    mask = torch.zeros(batch_size, 1, num_frames, num_frames, dtype=dtype, device=frame_mask.device)
-    mask.masked_fill_(~allowed[:, None, :, :], torch.finfo(dtype).min)
-    return mask.repeat_interleave(num_patches, dim=0)
-
-
-def encode_video_with_mem(
-    vision_model,
-    pixel_values: Tensor,
-    frame_mask: Tensor,
-    *,
-    temporal_attention_every: int,
-) -> Tensor:
-    """Encode ``(B,T,C,H,W)`` video with MEM's space-time separable attention.
-
-    No modules or learnable parameters are added. At temporal layers, the same
-    layer-normalization and Q/K/V/out projections are reused for a second,
-    causal attention operation along time. For a one-frame input the temporal
-    branch is skipped, making this exactly equivalent to the original image
-    encoder.
-    """
-    if pixel_values.ndim != 5:
-        raise ValueError(f"pixel_values must have shape (B,T,C,H,W), got {tuple(pixel_values.shape)}")
-    if temporal_attention_every < 1:
-        raise ValueError("temporal_attention_every must be at least 1")
-
-    batch_size, num_frames, channels, height, width = pixel_values.shape
-    if frame_mask.shape != (batch_size, num_frames):
-        raise ValueError(
-            f"frame_mask must have shape {(batch_size, num_frames)}, got {tuple(frame_mask.shape)}"
-        )
-
-    required = ("embeddings", "encoder", "post_layernorm")
-    if any(not hasattr(vision_model, name) for name in required):
-        raise TypeError("MEM visual memory currently requires a SigLIP-compatible vision tower")
-
-    flat_pixels = pixel_values.reshape(batch_size * num_frames, channels, height, width)
-    if hasattr(vision_model, "patch_size"):
-        patch_size = vision_model.patch_size
-        patch_attention_mask = torch.ones(
-            batch_size * num_frames,
-            height // patch_size,
-            width // patch_size,
-            dtype=torch.bool,
-            device=pixel_values.device,
-        )
-        hidden_states = vision_model.embeddings(
-            pixel_values=flat_pixels,
-            patch_attention_mask=patch_attention_mask,
-        )
-    else:
-        hidden_states = vision_model.embeddings(flat_pixels)
-    num_patches, hidden_size = hidden_states.shape[1:]
-    hidden_states = hidden_states.reshape(batch_size, num_frames, num_patches, hidden_size)
-
-    temporal_positions = temporal_sinusoidal_embedding(
-        num_frames, hidden_size, device=hidden_states.device, dtype=hidden_states.dtype
-    )[None, :, None, :]
-    temporal_mask = causal_temporal_mask(frame_mask, dtype=hidden_states.dtype, num_patches=num_patches)
-
-    for layer_index, layer in enumerate(vision_model.encoder.layers):
-        spatial_input = hidden_states.reshape(batch_size * num_frames, num_patches, hidden_size)
-        if num_frames == 1 or (layer_index + 1) % temporal_attention_every:
-            hidden_states = layer(spatial_input, attention_mask=None).reshape(
-                batch_size, num_frames, num_patches, hidden_size
-            )
-            continue
-
-        residual = hidden_states
-        spatial_norm = layer.layer_norm1(spatial_input)
-        spatial_output, _ = layer.self_attn(hidden_states=spatial_norm, attention_mask=None)
-        spatial_output = spatial_output.reshape(batch_size, num_frames, num_patches, hidden_size)
-
-        temporal_input = (hidden_states + temporal_positions).permute(0, 2, 1, 3)
-        temporal_input = temporal_input.reshape(batch_size * num_patches, num_frames, hidden_size)
-        temporal_norm = layer.layer_norm1(temporal_input)
-        temporal_output, _ = layer.self_attn(
-            hidden_states=temporal_norm,
-            attention_mask=temporal_mask,
-        )
-        temporal_output = temporal_output.reshape(batch_size, num_patches, num_frames, hidden_size)
-        temporal_output = temporal_output.permute(0, 2, 1, 3)
-
-        hidden_states = residual + spatial_output + temporal_output
-        hidden_states = hidden_states + layer.mlp(layer.layer_norm2(hidden_states))
-
-    # MEM drops all historical tokens before the VLA backbone.
-    return vision_model.post_layernorm(hidden_states[:, -1])

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,10 +337,7 @@ 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 = torch.cat([proprioception, dof_mask], dim=-1)
 
         proprioception = proprioception.to(device=self.propri_proj.weight.device).to(
             dtype=self.propri_proj.weight.dtype
@@ -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
-            )
-            matches = (input_ids[0, :-1] == generation_prompt_ids[0]) & (
-                input_ids[0, 1:] == generation_prompt_ids[1]
+                self.processor.tokenizer.encode(generation_prompt, add_special_tokens=False),
+                device=input_ids.device,
+                dtype=input_ids.dtype,
             )
+            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.
+
+"""Wall-X configuration extensions for the native Transformers Qwen2.5-VL config."""
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from huggingface_hub.dataclasses import strict
+
+from lerobot.utils.import_utils import _transformers_available
+
+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:
+
+    @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
+
+# 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",
+}
 
 
-class Qwen2_5_VLVisionConfig(PretrainedConfig):
-    model_type = "qwen2_5_vl"
-    base_config_key = "vision_config"
+@strict
+class Qwen2_5_VLTextConfig(TransformersQwen2_5_VLTextConfig):  # noqa: N801
+    """Native Qwen2.5-VL text config plus Wall-X's hard-routed MoE settings."""
 
-    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)
+    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.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
+    def __post_init__(self, **kwargs):
+        self.dim_inputs = tuple(self.dim_inputs)
+        super().__post_init__(**kwargs)
 
 
-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).
+@strict
+class Qwen2_5_VLConfig(TransformersQwen2_5_VLConfig):  # noqa: N801
+    """Native composite Qwen2.5-VL config with a Wall-X text sub-config.
 
-    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
-    documentation from [`PretrainedConfig`] for more information.
+    The native composite loader supports both current nested configs and the
+    flat layout used by existing ``wall-oss-flow`` checkpoints.
+    """
 
-
-    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"]),
+    sub_configs = {
+        "vision_config": Qwen2_5_VLVisionConfig,
+        "text_config": Qwen2_5_VLTextConfig,
     }
 
-    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"]()
+    def __getattr__(self, name):
+        """Keep legacy direct access to fields now owned by ``text_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
-
-        # for backward compatibility
-        if num_key_value_heads is None:
-            num_key_value_heads = num_attention_heads
-
-        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
-
-
-__all__ = ["Qwen2_5_VLConfig"]
+        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,9 +431,8 @@ 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:
-                    continue
+            if got_instruction and random.random() >= prob:
+                continue
 
             task_instruction += f"\n{frame_instruction_info[key]}"
             got_instruction = True
@@ -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/rename_processor.py

@@ -21,23 +21,6 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
 
 from .pipeline import ObservationProcessorStep, ProcessorStepRegistry
 
-_AUXILIARY_KEY_SUFFIXES = ("_is_pad", "_padding_mask")
-
-
-def rename_transition_key(key: str, rename_map: dict[str, str]) -> str:
-    """Rename a feature key and any sampling metadata derived from it."""
-    if key in rename_map:
-        return rename_map[key]
-    for suffix in _AUXILIARY_KEY_SUFFIXES:
-        if key.endswith(suffix) and key[: -len(suffix)] in rename_map:
-            return f"{rename_map[key[: -len(suffix)]]}{suffix}"
-    return key
-
-
-def rename_transition_keys(data: dict[str, Any], rename_map: dict[str, str]) -> dict[str, Any]:
-    """Rename all transition keys, including delta-sampling padding masks."""
-    return {rename_transition_key(key, rename_map): value for key, value in data.items()}
-
 
 @dataclass
 @ProcessorStepRegistry.register(name="rename_observations_processor")
@@ -58,7 +41,14 @@ class RenameObservationsProcessorStep(ObservationProcessorStep):
     rename_map: dict[str, str] = field(default_factory=dict)
 
     def observation(self, observation):
-        return rename_transition_keys(observation, self.rename_map)
+        processed_obs = {}
+        for key, value in observation.items():
+            if key in self.rename_map:
+                processed_obs[self.rename_map[key]] = value
+            else:
+                processed_obs[key] = value
+
+        return processed_obs
 
     def get_config(self) -> dict[str, Any]:
         return {"rename_map": self.rename_map}

src/lerobot/processor/render_messages_processor.py

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

src/lerobot/processor/tokenizer_processor.py

@@ -32,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,
@@ -413,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
@@ -432,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.
 
@@ -461,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)
@@ -479,26 +476,19 @@ 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:
@@ -507,49 +497,44 @@ class ActionTokenizerProcessorStep(ActionProcessorStep):
                     "Consider increasing the `max_action_tokens` in your model config if this happens frequently."
                 )
                 tokens = tokens[: self.max_action_tokens]
-                code_mask = code_mask[: self.max_action_tokens]
                 mask = torch.ones(self.max_action_tokens, dtype=torch.bool, device=action.device)
             else:
-                pad_len = self.max_action_tokens - len(tokens)
                 mask = torch.cat(
                     [
                         torch.ones(len(tokens), dtype=torch.bool, device=action.device),
-                        torch.zeros(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]:
@@ -565,8 +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,
         }
 
         # Only save tokenizer_name if it was used to create the tokenizer

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/unitree_g1/README.md

@@ -0,0 +1,162 @@
+# Unitree G1 — SONIC whole-body control
+
+This package runs NVIDIA's **SONIC** whole-body controller (and the GR00T/Holosoma
+locomotion controllers) on the Unitree G1, in MuJoCo simulation or on real hardware.
+SONIC turns a high-level movement intent — or a streamed **SMPL** whole-body pose — into
+50 Hz joint-position targets. It is a pure-Python/ONNX reimplementation of the SONIC
+deploy stack (no `gear_sonic`/torch dependency).
+
+## Controllers
+
+Selected with `--robot.controller=<ClassName>`:
+
+| Controller                     | Purpose                                                                                    |
+| ------------------------------ | ------------------------------------------------------------------------------------------ |
+| `SonicWholeBodyController`     | SONIC whole-body: locomotion (mode 0), 3-point VR teleop (mode 1), SMPL imitation (mode 2) |
+| `GrootLocomotionController`    | GR00T locomotion policy                                                                    |
+| `HolosomaLocomotionController` | Holosoma locomotion policy                                                                 |
+
+On startup the controller **interpolates** from the robot's measured pose into the
+policy's commanded target over ~3 s (no snap), and on disconnect (Ctrl-C) it performs a
+**graceful damped settle** — holding pose while ramping stiffness to zero over
+`--robot.graceful_stop_s` (default 1.5 s) instead of going instantly limp. Both apply in
+every mode.
+
+## Requirements
+
+- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended — SONIC runs three ONNX
+  sessions and is much smoother on GPU). Install the CUDA build that matches your
+  driver (e.g. `onnxruntime-gpu==1.26.0` for a CUDA-12.x driver). Verify with:
+  ```bash
+  python -c "import onnxruntime as ort; print(ort.get_available_providers())"
+  # expect CUDAExecutionProvider in the list for GPU
+  ```
+- `mujoco` for simulation (`is_simulation=True`).
+- `pyzmq` only if you use the live SMPL stream (pico headset).
+- The SONIC ONNX models are downloaded automatically from the `nvidia/GEAR-SONIC` Hub repo.
+
+## Running
+
+**Replay an SMPL dataset (motion imitation):**
+
+```bash
+lerobot-replay \
+  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
+  --dataset.repo_id=<user>/<smpl_dataset> --dataset.episode=0
+```
+
+**Keyboard teleop** (drives locomotion via the native keyboard teleoperator):
+
+```bash
+lerobot-teleoperate \
+  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
+  --teleop.type=keyboard
+```
+
+Controls: `WASD` move · `Q`/`E` turn · `1`–`8` mode · `9`/`0` speed · `-`/`=` height ·
+`R` replan · `Space` emergency-stop.
+
+**PICO headset teleop — SMPL whole-body** (mode 2, needs PICO Motion Trackers):
+
+```bash
+# 1) publisher (streams rt/smpl from full-body tracking)
+python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50
+# 2) controller
+lerobot-teleoperate \
+  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
+  --teleop.type=pico_headset
+```
+
+**PICO headset teleop — 3-point VR** (mode 1, head + 2 controllers only, **no trackers**):
+
+```bash
+# 1) publisher (head + controllers -> 3-point targets + stick locomotion)
+python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --headset-source devices
+# 2) controller
+lerobot-teleoperate \
+  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
+  --teleop.type=pico_headset --teleop.mode=vr3
+```
+
+3-point controls: left stick move · right stick X turn · right stick Y height ·
+`A`+`B` / `X`+`Y` cycle locomotion mode (walk/run/squat/kneel/…) · hands+head track the
+upper body. **Calibration**: stand in a neutral rest pose and press `A`+`B`+`X`+`Y` — the
+publisher status line flips from `UNCALIBRATED` to `calibrated`. This maps your rest pose
+onto the G1's neutral stance and is required before the hands track well; the SMPL
+(mode 2) path is self-calibrating and needs no such step.
+
+Both require the XRoboToolkit stack — see below.
+
+## PICO headset / XRoboToolkit install
+
+Live full-body teleop needs the **XRoboToolkit** system (a PC Service on your
+workstation + a PICO app on the headset) and its Python binding, `xrobotoolkit_sdk`.
+The full hardware + software walkthrough lives in the SONIC repo:
+[`docs/source/getting_started/vr_teleop_setup.md`](https://nvlabs.github.io/GR00T-WholeBodyControl/getting_started/vr_teleop_setup.html).
+
+Summary:
+
+1. **PC Service** (workstation) — install and run it before connecting the headset.
+   - Ubuntu 22.04 / 24.04 (x86_64): prebuilt `.deb` from the
+     [XRoboToolkit-PC-Service releases](https://github.com/XR-Robotics/XRoboToolkit-PC-Service/releases).
+   - Jetson (aarch64): the arm64 `.deb`.
+   - Windows (x64): the Windows PC Service build.
+2. **PICO app** — install `XRoboToolkit-PICO-*.apk` on the headset (see the guide),
+   enable Developer Mode. For **SMPL whole-body** (mode 2) you also need the PICO Motion
+   Trackers paired/calibrated and "Full body" enabled; for **3-point** (mode 1,
+   `--headset-source devices`) only Head + Controller + Send are required — no trackers.
+3. **`xrobotoolkit_sdk`** — a pybind11/CMake build (not a pip package), from
+   [`XRoboToolkit-PC-Service-Pybind`](https://github.com/XR-Robotics/XRoboToolkit-PC-Service-Pybind):
+   - Linux x86_64: `pip install pybind11 cmake` then `bash setup_ubuntu.sh` (or the
+     SONIC repo's `install_scripts/install_pico.sh`, which builds everything into a
+     `.venv_teleop`).
+   - Jetson aarch64: `bash setup_orin.sh` (builds `libPXREARobotSDK.so` from source).
+   - Windows x64: `pip install pybind11` then `setup_windows.bat` (needs git + an
+     MSVC/CMake toolchain; uses the prebuilt `PXREARobotSDK.dll`/`.lib`).
+4. Connect PICO and workstation to the **same Wi-Fi**, open the XRoboToolkit app, enter
+   the PC IP, and enable Head/Controller/Send (plus Full-body for SMPL mode 2).
+
+### Platform support
+
+| Platform                    | Live headset teleop | Notes                                       |
+| --------------------------- | ------------------- | ------------------------------------------- |
+| Linux x86_64                | ✅                  | Guided `install_pico.sh` (SONIC repo)       |
+| Linux aarch64 (Jetson Orin) | ✅                  | `setup_orin.sh` builds the native lib       |
+| Windows x64                 | ✅ (manual)         | `setup_windows.bat`; no one-shot env script |
+| macOS                       | ❌                  | No PC Service / SDK build for Darwin        |
+
+### No hardware required (any platform, incl. macOS/Windows)
+
+The SMPL pipeline can be exercised without a headset or the SDK — the publisher emits
+`rt/smpl` frames that the controller consumes exactly as it would from the headset:
+
+```bash
+# synthetic motion
+python -m lerobot.teleoperators.pico_headset.pico_publisher --fake
+
+# replay a canned SMPL clip
+python -m lerobot.teleoperators.pico_headset.pico_publisher --motion-file <clip>.npz
+```
+
+## Notes
+
+- SMPL **root motion** into the mode-2 anchor is opt-in (`SonicWholeBodyController(enable_smpl_root=True)`);
+  it stays off by default (untested on hardware). When enabled, the per-frame root quat is
+  spherically smoothed (`root_smoothing_alpha`, default 0.15) before it reaches the anchor,
+  which removes the base-acceleration spikes the raw 30 Hz→50 Hz trajectory used to cause.
+- Direct `rt/smpl` subscription without the pico teleoperator is available via
+  `SonicWholeBodyController(enable_smpl_stream=True, smpl_host=..., smpl_port=...)`.
+- 3-point (mode 1) uses the **headset-yaw frame** as its reference and the `A`+`B`+`X`+`Y`
+  calibration to align to the G1 neutral stance. Calibration maps the operator's rest pose
+  onto the G1's **standing** (`default_angles`) wrist/neck key-frame poses (position **and**
+  orientation) computed by FK — the `default_angles` stand-in for gear_sonic's live
+  measured-q recalibration, since the robot holds `default_angles` at calibration time.
+  Re-aligning the arms only (preserving the neck level) is available via the calibrator's
+  `recalibrate_wrists()`.
+- 3-point **locomotion** from the PICO sticks follows gear_sonic's `PlannerLoop` exactly:
+  a yaw accumulator on the right stick and **mode-dependent speed curves** on the left
+  (slow `0.1+0.5·mag`, run `1.5+3·mag`, walk = planner default). Stick signs replicate
+  gear_sonic's `get_controller_axes` usage (forward `+ly`, strafe `-lx`, turn `-rx`); since
+  the publisher forwards the same raw SDK axes, this is the correct convention by construction.
+- Startup interpolation and the graceful-stop settle are mode-agnostic; set
+  `--robot.graceful_stop_s=0` to restore the old instant zero-torque on disconnect.

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -65,9 +65,43 @@ class UnitreeG1Config(RobotConfig):
     # Cameras (ZMQ-based remote cameras)
     cameras: dict[str, CameraConfig] = field(default_factory=dict)
 
+    # Synthetic zero-image cameras exposed as ``observation.images.{name}`` (H×W×3
+    # black frames). Lets image-conditioned policies (e.g. pi0.5 / OpenHLM) run in
+    # sim before real cameras are wired. Empty = disabled.
+    empty_cameras: list[str] = field(default_factory=list)
+    empty_camera_hw: tuple[int, int] = (224, 224)
+
+    # Publish Dex3 hand commands (``rt/dex3/{left,right}/cmd``) driven by the OpenHLM
+    # gripper scalars (``wb.7.pos`` left, ``wb.15.pos`` right). Lets the 43-DoF sim
+    # (or a real Dex3-equipped G1) show grasping. The scalar in [0, 1] is remapped to
+    # a curl amount (``hand_open_grip_value`` -> open) and scaled onto
+    # ``hand_closed_pose`` (7 joints: thumb_0/1/2, middle_0/1, index_0/1). Flip signs
+    # in ``hand_closed_pose`` if fingers curl the wrong way.
+    publish_hands: bool = False
+    hand_open_grip_value: float = 1.0
+    hand_closed_grip_value: float = 0.0
+    hand_closed_pose: list[float] = field(
+        default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3]
+    )
+    hand_kp: float = 1.5
+    hand_kd: float = 0.1
+
+    # Replay recorded camera frames from a LeRobot parquet episode as the camera
+    # feed (e.g. OpenHLM-data episode). Maps a robot camera name to a parquet image
+    # column; frames advance one per observation and loop. Lets a VLA see the real
+    # task video in sim without live cameras. Empty map = disabled.
+    replay_camera_parquet: str | None = None
+    replay_camera_map: dict[str, str] = field(default_factory=dict)
+    replay_camera_loop: bool = True
+
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Lower-body controller class name, e.g. "GrootLocomotionController" or
-    # "HolosomaLocomotionController". None disables it.
+    # Locomotion controller class name, e.g. "GrootLocomotionController",
+    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
     controller: str | None = None
+
+    # On disconnect (e.g. Ctrl-C), seconds to hold the current pose while ramping joint
+    # stiffness (kp) to zero — a soft, damped settle instead of an instant limp /
+    # free-fall. 0 disables it (immediate zero-torque). Real robot only.
+    graceful_stop_s: float = 1.5

src/lerobot/robots/unitree_g1/controllers/__init__.py

@@ -1,4 +1,6 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,8 +14,11 @@
 # 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."""
+"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
 
-from .configuration_pi052 import PI052Config
-
-__all__ = ["PI052Config"]
+__all__ = [
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicWholeBodyController",
+    "SonicRuntime",
+]

src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py

@@ -14,20 +14,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import logging
 from collections import deque
+from typing import TYPE_CHECKING
 
 import numpy as np
-import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from lerobot.utils.import_utils import _onnxruntime_available, require_package
+
+from ..g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,
     get_gravity_orientation,
 )
 
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -83,6 +92,7 @@ class GrootLocomotionController:
     control_dt = CONTROL_DT  # Expose for unitree_g1.py
 
     def __init__(self):
+        require_package("onnxruntime", extra="unitree_g1")
         # Load policies
         self.policy_balance, self.policy_walk = load_groot_policies()
 

src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py

@@ -14,21 +14,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import json
 import logging
+from typing import TYPE_CHECKING
 
 import numpy as np
-import onnx
-import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from lerobot.utils.import_utils import _onnx_available, _onnxruntime_available, require_package
+
+from ..g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     get_gravity_orientation,
 )
 
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
+if TYPE_CHECKING or _onnx_available:
+    import onnx
+else:
+    onnx = None
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_ANGLES = np.zeros(29, dtype=np.float32)
@@ -101,6 +114,8 @@ class HolosomaLocomotionController:
     control_dt = CONTROL_DT  # Expose for unitree_g1.py
 
     def __init__(self):
+        require_package("onnxruntime", extra="unitree_g1")
+        require_package("onnx", extra="unitree_g1")
         # Load policy and gains
         self.policy, self.kp, self.kd = load_policy()
 

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -0,0 +1,718 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SONIC full-body controller for Unitree G1."""
+
+from __future__ import annotations
+
+import logging
+import math
+from collections import deque
+from typing import TYPE_CHECKING
+
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+from lerobot.teleoperators.pico_headset.smpl_constants import (
+    LOCO_AXES_PREFIX,
+    LOCO_BTN_PREFIX,
+    LOCO_N_AXES,
+    LOCO_N_BTN,
+    ROOT_ACTION_DIM,
+    ROOT_ACTION_PREFIX,
+    SMPL_ACTION_PREFIX,
+    SMPL_OBS_DIM as SMPL_ACTION_DIM,
+    VR3_ORN_DIM,
+    VR3_ORN_PREFIX,
+    VR3_POS_DIM,
+    VR3_POS_PREFIX,
+    WB_ACTION_DIM,
+    wb_action_key,
+)
+from lerobot.utils.import_utils import _onnxruntime_available, require_package
+
+from ..g1_utils import MUJOCO_TO_ISAACLAB, KEYBOARD_KEYS_FIELD, G1_29_JointIndex, lowstate_to_obs
+from .sonic_pipeline import (
+    CONTROL_DT,
+    DEBUG_PRINT_EVERY,
+    DEFAULT_ANGLES,
+    DEFAULT_HEIGHT,
+    ENCODER_UPDATE_EVERY,
+    LM,
+    MOTION_SETS,
+    MovementState,
+    PlannerController,
+    SonicPlanner,
+    apply_pico_loco_axes,
+    clamp_mode_params,
+    compute_kp_kd,
+    make_ort_session_options,
+    ort_providers,
+    process_joystick,
+    should_replan_request,
+    snapshot_ms,
+)
+
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
+logger = logging.getLogger(__name__)
+
+# Startup blend duration: over the first control ticks, linearly interpolate every joint
+# from the robot's initial measured pose into the policy's commanded target, so control
+# eases in without a snap on the first command.
+INIT_RAMP_S = 3.0
+
+
+def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a (720,) SMPL window from ``smpl.{i}`` action keys, or None.
+
+    The pico_headset teleoperator emits the whole-body reference as flat floats so
+    it flows unchanged through the standard lerobot action pipeline.
+    """
+    # The keys are smpl.0 .. smpl.719; presence of the first element (smpl.0) is the
+    # sentinel that a full SMPL window was sent this tick. If it's absent, there's no
+    # whole-body reference, so bail out.
+    if not action or f"{SMPL_ACTION_PREFIX}0" not in action:
+        return None
+    arr = np.fromiter(
+        (float(action.get(f"{SMPL_ACTION_PREFIX}{i}", 0.0)) for i in range(SMPL_ACTION_DIM)),
+        dtype=np.float32,
+        count=SMPL_ACTION_DIM,
+    )
+    return arr
+
+
+def _extract_root_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a (4,) SMPL root quaternion (wxyz) from ``root.{i}`` keys, or None."""
+    if not action or f"{ROOT_ACTION_PREFIX}0" not in action:
+        return None
+    q = np.fromiter(
+        (float(action.get(f"{ROOT_ACTION_PREFIX}{i}", 0.0)) for i in range(ROOT_ACTION_DIM)),
+        dtype=np.float32,
+        count=ROOT_ACTION_DIM,
+    )
+    n = float(np.linalg.norm(q))
+    if n < 1e-6:
+        return None
+    return q / n
+
+
+def _extract_vr3_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None:
+    """Reassemble the 3-point VR targets from ``vr3_pos.{i}`` / ``vr3_orn.{i}`` keys.
+
+    Returns ``(pos (9,), orn (12,))`` for the [l-wrist, r-wrist, neck] keypoints, or
+    None when no VR3 reference was sent this tick. Presence of ``vr3_pos.0`` is the
+    sentinel that a full 3-point frame is available (mirrors the SMPL sentinel).
+    """
+    if not action or f"{VR3_POS_PREFIX}0" not in action:
+        return None
+    pos = np.fromiter(
+        (float(action.get(f"{VR3_POS_PREFIX}{i}", 0.0)) for i in range(VR3_POS_DIM)),
+        dtype=np.float32,
+        count=VR3_POS_DIM,
+    )
+    orn = np.fromiter(
+        (float(action.get(f"{VR3_ORN_PREFIX}{i}", 0.0)) for i in range(VR3_ORN_DIM)),
+        dtype=np.float32,
+        count=VR3_ORN_DIM,
+    )
+    return pos, orn
+
+
+def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a dense (34,) whole-body command from ``wb.{i}.pos`` keys, or None.
+
+    This is the OpenHLM / pi0.5 joint-based interface: one 34-D vector per tick
+    (sentinel: presence of ``wb.0.pos``) carrying absolute joint targets in real
+    units. The ``.pos`` suffix lets these flow through ``lerobot-rollout`` as normal
+    joint-position action features.
+    """
+    if not action or wb_action_key(0) not in action:
+        return None
+    return np.fromiter(
+        (float(action.get(wb_action_key(i), 0.0)) for i in range(WB_ACTION_DIM)),
+        dtype=np.float32,
+        count=WB_ACTION_DIM,
+    )
+
+
+def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+    """Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
+
+    Returns ``(ref29, anchor_quat)`` where ``ref29`` is the 29 joint targets in
+    IsaacLab order (what SONIC's ``motion_joint_positions`` expects) and
+    ``anchor_quat`` (wxyz) encodes the root roll/pitch (yaw=0).
+
+    OpenHLM layout : [L-arm 0:7, L-grip 7, R-arm 8:15, R-grip 15,
+                      L-leg 16:22, R-leg 22:28, waist 28:31, root rp+yaw 31:34]
+    The 29 joints are first assembled in MuJoCo / Unitree-SDK order
+    ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the
+    ``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via
+    ``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) and yaw-rate (33) are not part of the
+    29-DoF SONIC reference.
+    """
+    ref_mj = np.zeros(29, np.float32)  # MuJoCo / Unitree-SDK grouped order
+    ref_mj[0:6] = wb[16:22]   # left leg
+    ref_mj[6:12] = wb[22:28]  # right leg
+    ref_mj[12:15] = wb[28:31]  # waist
+    ref_mj[15:22] = wb[0:7]   # left arm
+    ref_mj[22:29] = wb[8:15]  # right arm
+    ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32)  # -> IsaacLab order for SONIC
+    roll, pitch = float(wb[31]), float(wb[32])
+    cr, sr, cp, sp = np.cos(roll / 2), np.sin(roll / 2), np.cos(pitch / 2), np.sin(pitch / 2)
+    anchor = np.array([cr * cp, sr * cp, cr * sp, sr * sp], np.float32)  # Rx(roll)·Ry(pitch)
+    return ref, anchor
+
+
+def _extract_loco_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None:
+    """Reassemble controller-stick locomotion from ``loco_axes.{i}`` / ``loco_btn.{i}``.
+
+    Returns ``(axes (4,) = [lx, ly, rx, ry], buttons (4,) = [A, B, X, Y])`` or None
+    when no locomotion state was sent this tick (sentinel: ``loco_axes.0``).
+    """
+    if not action or f"{LOCO_AXES_PREFIX}0" not in action:
+        return None
+    axes = np.fromiter(
+        (float(action.get(f"{LOCO_AXES_PREFIX}{i}", 0.0)) for i in range(LOCO_N_AXES)),
+        dtype=np.float32,
+        count=LOCO_N_AXES,
+    )
+    buttons = np.fromiter(
+        (float(action.get(f"{LOCO_BTN_PREFIX}{i}", 0.0)) for i in range(LOCO_N_BTN)),
+        dtype=np.float32,
+        count=LOCO_N_BTN,
+    )
+    return axes, buttons
+
+
+class SonicRuntime:
+    """Shared SONIC control loop state (standalone demo + locomotion controller)."""
+
+    def __init__(self, force_cpu: bool = False):
+        require_package("onnxruntime", extra="unitree_g1")
+        planner_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="planner_sonic.onnx")
+        encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx")
+        decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
+
+        providers = ort_providers(force_cpu=force_cpu)
+        self.use_gpu = providers[0] == "CUDAExecutionProvider"
+        so = make_ort_session_options()
+
+        planner_sess = ort.InferenceSession(planner_path, sess_options=so, providers=providers)
+        encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
+        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
+
+        self.kp, self.kd = compute_kp_kd()
+        self.ms = MovementState()
+        self.planner = SonicPlanner(planner_sess, planner_path)
+        self.controller = PlannerController(self.planner, encoder_sess, decoder_sess)
+
+        motion = self.planner.initialize(DEFAULT_ANGLES, self.ms)
+        self.controller.load_initial_motion(motion)
+        self.planner.start_subprocess(self.controller, use_gpu=self.use_gpu)
+
+        self.step = 0
+        self.replan_timer = 0.0
+        self.last_ms = snapshot_ms(self.ms)
+
+    @property
+    def pipeline(self):
+        return self.controller
+
+    def tick(self, obs: dict, *, debug: bool | None = None, use_joystick: bool = True) -> dict:
+        if not obs:
+            self.step += 1
+            return {}
+
+        if use_joystick:
+            process_joystick(obs, self.ms, self.controller)
+        clamp_mode_params(self.ms)
+
+        if self.step > 0:
+            self.replan_timer += CONTROL_DT
+        if should_replan_request(self.ms, self.last_ms, self.replan_timer, self.step):
+            self.planner.request_replan(self.controller.ref_cursor, self.ms)
+            self.replan_timer = 0.0
+            self.ms.needs_replan = False
+            self.last_ms = snapshot_ms(self.ms)
+
+        do_enc = self.step % ENCODER_UPDATE_EVERY == 0
+        if debug is None:
+            debug = self.step % DEBUG_PRINT_EVERY == 0
+        action = self.controller.step(obs, update_encoder=do_enc, debug=debug)
+
+        result = self.planner.try_get_new_motion()
+        if result:
+            self.controller.blend_new_motion(*result)
+
+        self.controller.advance_cursor()
+        self.step += 1
+        return action
+
+    def reset(self):
+        self.ms = MovementState()
+        self.controller.reinit_heading = True
+        self.controller.playing = True
+        self.step = 0
+        self.replan_timer = 0.0
+        self.last_ms = snapshot_ms(self.ms)
+
+    def shutdown(self):
+        self.planner.stop_subprocess()
+
+
+class SonicWholeBodyController:
+    """Full-body SONIC controller for UnitreeG1's background controller thread."""
+
+    control_dt = CONTROL_DT
+    full_body = True
+    # Advertise a dense 34-D whole-body action space (OpenHLM / pi0.5) so the robot
+    # exposes ``wb.{i}.pos`` action features and ``lerobot-rollout`` can drive it
+    # directly with a 34-D VLA policy.
+    wb_action = True
+
+    def __init__(
+        self,
+        force_cpu: bool = False,
+        *,
+        enable_smpl_root: bool = False,
+        root_smoothing_alpha: float = 0.15,
+        enable_smpl_stream: bool = False,
+        smpl_host: str | None = None,
+        smpl_port: int | None = None,
+    ):
+        logger.info("Loading SONIC whole-body controller...")
+        self._runtime = SonicRuntime(force_cpu=force_cpu)
+        self.kp = self._runtime.kp
+        self.kd = self._runtime.kd
+        self.controller = self._runtime.controller
+        self.ms = self._runtime.ms
+
+        # When True, the per-frame SMPL root quaternion steers the mode-2 anchor.
+        # Off by default: even with smoothing this changes the anchor/heading and is
+        # untested on hardware, so it stays opt-in. When enabled, the raw per-frame
+        # root quat (from a 30 Hz dataset resampled to a 50 Hz loop) is spherically
+        # smoothed by :meth:`_smooth_root_quat` before it reaches the anchor, which
+        # removes the root-acceleration spikes (NaN QACC at DOF 0) the unsmoothed
+        # trajectory caused. ``root_smoothing_alpha`` in (0, 1] is the per-tick blend
+        # toward the incoming quat (smaller = smoother/laggier, 1 = no smoothing).
+        self.enable_smpl_root = enable_smpl_root
+        self._root_smoothing_alpha = float(np.clip(root_smoothing_alpha, 1e-3, 1.0))
+        self._smoothed_root_quat: np.ndarray | None = None
+
+        # Tracks the previous keyboard held-key set so discrete controls (mode,
+        # motion set, replan, e-stop, WASD direction) fire once per physical press
+        # instead of every 50 Hz tick while the key is held.
+        self._prev_keys: set[str] = set()
+        # Edge state for the PICO A+B / X+Y locomotion-mode cycle (3-point teleop).
+        self._prev_loco_mode_pair: tuple[bool, bool] = (False, False)
+
+        # Startup blend: ease from the robot's initial pose into the first commanded
+        # policy targets over INIT_RAMP_S (captured on the first control tick).
+        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
+        self._init_step = 0
+        self._start_pose: dict[str, float] = {}
+
+        # Tick counter for the dense whole-body (OpenHLM, mode-0) path's encoder cadence.
+        self._wb_step = 0
+        # Rolling 50-frame reference trajectory (ref29 + anchor quat) built from the
+        # stream of per-tick whole-body commands, fed to the encoder as a batch.
+        self._wb_traj: deque[np.ndarray] = deque(maxlen=50)
+        self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50)
+
+        # Optional: subscribe directly to the rt/smpl headset stream so full-body
+        # teleop works with ANY teleoperator (e.g. --teleop.type=unitree_g1 for the
+        # estop/joystick) before the dedicated pico_headset teleop exists.
+        self._smpl_host = smpl_host
+        self._smpl_port = smpl_port
+        self._smpl_stream = None
+        if enable_smpl_stream:
+            self._init_smpl_stream()
+
+        logger.info(
+            "SONIC ready: %s (default mode: %s, smpl_stream=%s)",
+            MOTION_SETS[0][0],
+            LM(self.ms.mode).name,
+            self._smpl_stream is not None,
+        )
+
+    def _init_smpl_stream(self) -> None:
+        # Lazy import so the zmq dependency is only required when streaming is on.
+        from lerobot.teleoperators.pico_headset.smpl_stream import (
+            DEFAULT_SMPL_HOST,
+            DEFAULT_SMPL_PORT,
+            SmplStream,
+        )
+
+        host = self._smpl_host or DEFAULT_SMPL_HOST
+        port = self._smpl_port or DEFAULT_SMPL_PORT
+        self._smpl_stream = SmplStream(host=host, port=port)
+        logger.info("SONIC subscribed to rt/smpl @ tcp://%s:%d", host, port)
+
+    def _enter_wholebody(self) -> None:
+        """Switch into SMPL whole-body tracking (encode_mode 2)."""
+        self.controller.encode_mode = 2
+        self.controller.reinit_heading = True
+        logger.info("SONIC: SMPL stream active -> whole-body tracking (mode 2)")
+
+    def _enter_3point(self) -> None:
+        """Switch into 3-point VR upper-body teleop (encode_mode 1).
+
+        The upper body tracks the VR wrist/neck targets while the lower body /
+        locomotion keeps running off the planner (joystick/keyboard-driven).
+        """
+        self.controller.encode_mode = 1
+        self.controller.playing = True
+        self.controller.reinit_heading = True
+        self.ms.needs_replan = True
+        logger.info("SONIC: 3-point VR active -> upper-body tracking + planner locomotion (mode 1)")
+
+    def _exit_wholebody(self) -> None:
+        """Revert to locomotion/standing (encode_mode 0) after a teleop reference is lost.
+
+        Mirrors the 'M' toggle in sonic.py so the handoff is clean: the robot holds
+        a standing reference and (if a joystick teleop is attached) can be driven.
+        """
+        self.controller.encode_mode = 0
+        self.controller.playing = True
+        self.controller.reinit_heading = True
+        self.ms.needs_replan = True
+        logger.warning("SONIC: teleop reference lost/stale -> reverting to locomotion (standing)")
+
+    def _process_keyboard(self, action: dict | None) -> None:
+        """Translate a native KeyboardTeleop's held-key set into MovementState.
+
+        Mirrors the standalone SONIC demo's keyboard mapping so locomotion (mode 0/1)
+        can be driven with ``--teleop.type=keyboard`` instead of the PICO SMPL stream.
+        Discrete controls act on newly-pressed keys (edge-detected against the previous
+        tick); inherently-continuous controls (facing turn, height, speed) integrate a
+        small per-tick delta while the key is held so they feel smooth at 50 Hz.
+
+        Controls: WASD move, Q/E turn, 1-8 select mode, 9/0 speed down/up,
+        -/= height down/up, R replan, Space emergency-stop -> IDLE.
+        """
+        if action is None:
+            return
+        keys = action.get(KEYBOARD_KEYS_FIELD)
+        if keys is None:
+            return  # No KeyboardTeleop attached; leave joystick/SMPL paths untouched.
+
+        ms, controller = self.ms, self.controller
+        held = {k.lower() if isinstance(k, str) and len(k) == 1 else k for k in keys}
+        prev = self._prev_keys
+        pressed = held - prev  # newly-pressed this tick (edge)
+        self._prev_keys = held
+
+        # ── Discrete: fire once per press ────────────────────────────────────
+        if "space" in pressed:
+            ms.mode = LM.IDLE
+            ms.speed = ms.height = -1.0
+            ms.has_movement = False
+            ms.needs_replan = True
+            controller.playing = False
+            controller.reinit_heading = True
+            logger.info("SONIC keyboard: EMERGENCY STOP -> IDLE")
+        if "r" in pressed:
+            ms.needs_replan = True
+        if "n" in pressed or "p" in pressed:
+            step = 1 if "n" in pressed else -1
+            ms.motion_set_idx = (ms.motion_set_idx + step) % len(MOTION_SETS)
+            logger.info("SONIC keyboard: motion set -> %s", MOTION_SETS[ms.motion_set_idx][0])
+        for digit in ("1", "2", "3", "4", "5", "6", "7", "8"):
+            if digit in pressed:
+                idx = int(digit) - 1
+                modes = MOTION_SETS[ms.motion_set_idx][1]
+                if 0 <= idx < len(modes):
+                    ms.mode = modes[idx]
+                    ms.has_movement = False
+                    ms.needs_replan = True
+                    controller.playing = True
+                    controller.reinit_heading = True
+                    logger.info("SONIC keyboard: mode -> %s", LM(ms.mode).name)
+        # WASD sets the movement direction relative to current facing (press to set,
+        # Space to stop) to match the standalone demo.
+        if "w" in pressed:
+            ms.movement_angle = ms.facing_angle
+        elif "s" in pressed:
+            ms.movement_angle = ms.facing_angle + math.pi
+        elif "a" in pressed:
+            ms.movement_angle = ms.facing_angle + math.pi / 2
+        elif "d" in pressed:
+            ms.movement_angle = ms.facing_angle - math.pi / 2
+        if pressed & {"w", "a", "s", "d"}:
+            ms.has_movement = True
+            ms.needs_replan = True
+
+        # ── Continuous: integrate a small delta while held ───────────────────
+        if "q" in held:
+            ms.facing_angle += 0.02
+            controller.delta_heading += 0.02
+        if "e" in held:
+            ms.facing_angle -= 0.02
+            controller.delta_heading -= 0.02
+        if "0" in held:
+            ms.speed = min(5.0, (ms.speed if ms.speed >= 0 else 1.0) + 0.02)
+        if "9" in held:
+            ms.speed = max(0.0, (ms.speed if ms.speed >= 0 else 1.0) - 0.02)
+        if "=" in held:
+            ms.height = min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + 0.005)
+        if "-" in held:
+            ms.height = max(0.1, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) - 0.005)
+
+    def _process_pico_loco(self, axes: np.ndarray, buttons: np.ndarray) -> None:
+        """Drive locomotion from the PICO controller sticks/buttons (encode_mode 1).
+
+        Mirrors gear_sonic's ``PlannerLoop`` VR-3PT tick: left/right sticks steer
+        movement/facing/speed via :func:`apply_pico_loco_axes` (the faithful gear_sonic
+        yaw-accumulator + mode-dependent speed curves, not the keyboard-parity map), and
+        A+B / X+Y edge-cycle the locomotion mode within the current motion set.
+        """
+        lx, ly, rx, ry = (float(v) for v in axes)
+        apply_pico_loco_axes(lx, ly, rx, ry, self.ms)
+
+        # Mode cycling: step linearly through the LocomotionMode enum (A+B = next,
+        # X+Y = previous), exactly like gear_sonic's PlannerLoop — so the operator can
+        # reach squat/kneel/crawl, not just the modes in one UI motion set.
+        a, b, x, y = (v > 0.5 for v in buttons)
+        ab_now, xy_now = (a and b), (x and y)
+        ab_prev, xy_prev = self._prev_loco_mode_pair
+        mode = int(self.ms.mode)
+        if ab_now and not ab_prev:
+            mode = min(int(LM.INJURED_WALK), mode + 1)
+        elif xy_now and not xy_prev:
+            mode = max(int(LM.IDLE), mode - 1)
+        if mode != int(self.ms.mode):
+            self.ms.mode = LM(mode)
+            self.ms.needs_replan = True
+            self.controller.playing = True
+            logger.info("SONIC 3-point: locomotion mode -> %s", LM(self.ms.mode).name)
+        self._prev_loco_mode_pair = (ab_now, xy_now)
+
+    def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict:
+        """Feed a dense 34-D OpenHLM whole-body command as the mode-0 encoder reference.
+
+        The 29 joint targets are held across the encoder lookahead window (zero
+        velocity) and the root roll/pitch set the anchor orientation, then the
+        encoder/decoder run directly (planner bypassed). One command per tick, so the
+        VLA's commanded pose is what SONIC tracks.
+        """
+        ref, anchor = _wb34_to_reference(wb)
+        c = self.controller
+        if c.encode_mode != 0:
+            c.encode_mode = 0
+            c.reinit_heading = True
+        # Capture the heading/anchor reference on the first whole-body tick. The
+        # controller only latches ``init_ref_quat`` (and the base heading) inside
+        # ``step()`` when ``first_motion or reinit_heading`` — but it already boots in
+        # mode 0, so the mode-switch guard above misses the very first command and the
+        # anchor would stay identity. This mirrors the GEAR reference, which seeds
+        # ``init_ref_quat`` from the first anchor. Must run before the buffers below so
+        # ``step()`` latches ``motion_body_quats[0]`` = this tick's anchor.
+        if self._wb_step == 0:
+            c.reinit_heading = True
+
+        # Accumulate the per-tick commands into a rolling 50-frame reference
+        # trajectory so the encoder's 10-frame, step-5 lookahead sees an actual
+        # motion sequence (with velocities) instead of one repeated pose. 50 frames
+        # == chunk horizon == 10 lookahead frames × step 5.
+        self._wb_traj.append(ref)
+        self._wb_quat_traj.append(anchor)
+        traj = np.asarray(self._wb_traj, np.float32)  # (L, 29), oldest -> newest
+        quats = np.asarray(self._wb_quat_traj, np.float32)  # (L, 4)
+        n = len(traj)
+        # Per-frame velocities from finite differences (rad/s at the control rate).
+        vel = np.zeros_like(traj)
+        if n > 1:
+            vel[1:] = (traj[1:] - traj[:-1]) / CONTROL_DT
+            vel[0] = vel[1]
+        with c.motion_lock:
+            c.motion_joint_positions[:n] = traj
+            c.motion_joint_velocities[:n] = vel
+            c.motion_body_quats[:n] = quats
+            c.motion_body_pos[:n] = 0.0
+            c.motion_timesteps = n
+            c.ref_cursor = 0
+        c.playing = True
+        do_enc = self._wb_step % ENCODER_UPDATE_EVERY == 0
+        out = c.step(obs, update_encoder=do_enc, debug=False)
+        if self._wb_step % 25 == 0:
+            tgt = np.array([out[f"{m.name}.q"] for m in G1_29_JointIndex], np.float32)
+            logger.info(
+                "[WB34] step=%d |ref|mean=%.3f |target|mean=%.3f target_std=%.3f init_ref_quat=%s",
+                self._wb_step,
+                float(np.abs(ref).mean()),
+                float(np.abs(tgt).mean()),
+                float(tgt.std()),
+                np.round(c.init_ref_quat, 3).tolist(),
+            )
+        self._wb_step += 1
+        return out
+
+    def _smooth_root_quat(self, root_quat: np.ndarray | None) -> np.ndarray | None:
+        """Spherically smooth the per-frame SMPL root quaternion (mode-2 anchor).
+
+        The reference root trajectory is authored at ~30 Hz and consumed at 50 Hz, so
+        the raw per-tick quat steps unevenly and injects root-acceleration spikes into
+        the anchor. This keeps a persistent estimate and shortest-path nlerp-slerps it
+        toward each incoming (unit) quat by ``root_smoothing_alpha``, yielding a
+        continuous, rate-matched heading. Quaternions are scalar-first (w, x, y, z).
+        Returns ``None`` (leaving the anchor self-driven) for an invalid/zero input.
+        """
+        if root_quat is None:
+            self._smoothed_root_quat = None
+            return None
+        q = np.asarray(root_quat, np.float64)
+        n = np.linalg.norm(q)
+        if n < 1e-8:
+            return self._smoothed_root_quat
+        q = q / n
+        if self._smoothed_root_quat is None:
+            self._smoothed_root_quat = q
+        else:
+            prev = self._smoothed_root_quat
+            if np.dot(prev, q) < 0.0:  # shortest-path: quats double-cover SO(3)
+                q = -q
+            blended = prev + self._root_smoothing_alpha * (q - prev)
+            self._smoothed_root_quat = blended / (np.linalg.norm(blended) + 1e-12)
+        return self._smoothed_root_quat.astype(np.float32)
+
+    def _startup_blend(self, obs: dict, out: dict) -> dict:
+        """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
+        interpolate between the robot's pose captured on the first tick and the policy's
+        live commanded target, so the handoff has no snap.
+
+        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
+        climbs 0->1 over the ramp, after which the raw policy target passes through.
+        """
+        if self._init_step >= self._init_ramp_steps or not out:
+            return out
+        if self._init_step == 0:
+            # Capture the robot's actual pose as the interpolation start point.
+            self._start_pose = {
+                f"{m.name}.q": float(obs.get(f"{m.name}.q", DEFAULT_ANGLES[m.value]))
+                for m in G1_29_JointIndex
+            }
+        self._init_step += 1
+        ratio = min(1.0, self._init_step / self._init_ramp_steps)
+        blended = {
+            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
+            for k, tgt in out.items()
+        }
+        if self._init_step >= self._init_ramp_steps:
+            logger.info("SONIC startup blend complete -> full policy control")
+        return blended
+
+    def run_step(self, action: dict, lowstate) -> dict:
+        if lowstate is None:
+            return {}
+        obs = lowstate_to_obs(lowstate)
+
+        # Keyboard teleop (native KeyboardTeleop) drives the same locomotion intent
+        # the joystick does; applied before the SMPL check so whole-body tracking
+        # still takes priority when a headset stream is present.
+        self._process_keyboard(action)
+
+        # Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a
+        # direct rt/smpl subscription when enabled (enable_smpl_stream). A stale
+        # stream (headset silent past its timeout) is treated as "no SMPL" so the
+        # robot doesn't stay frozen tracking the last pose.
+        # Dense whole-body command (OpenHLM / pi0.5 joint interface) takes priority:
+        # a single 34-D vector drives the mode-0 joint reference directly.
+        wb = _extract_wb34_from_action(action)
+        if wb is not None:
+            return self._startup_blend(obs, self._run_wholebody34(obs, wb))
+        self._wb_miss = getattr(self, "_wb_miss", 0) + 1
+        if self._wb_miss % 50 == 1:
+            akeys = [k for k in action if isinstance(k, str)]
+            logger.info(
+                "[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s",
+                self._wb_miss,
+                akeys[:8],
+            )
+
+        smpl = _extract_smpl_from_action(action)
+        root_quat = _extract_root_from_action(action)
+        vr3 = _extract_vr3_from_action(action)
+        loco = _extract_loco_from_action(action)
+        if smpl is None and vr3 is None and self._smpl_stream is not None:
+            window = self._smpl_stream.step()
+            if self._smpl_stream.has_data and not self._smpl_stream.is_stale:
+                smpl = window
+                root_quat = np.asarray(self._smpl_stream.root_quat, np.float32)
+            # VR3 is independent of the SMPL window: the controller-state source
+            # (head + controllers only) sends 3-point targets with no SMPL frame.
+            elif self._smpl_stream.has_fresh_vr3:
+                vr3 = (self._smpl_stream.vr3_pos, self._smpl_stream.vr3_orn)
+                if self._smpl_stream.has_fresh_loco:
+                    loco = (self._smpl_stream.loco_axes, self._smpl_stream.loco_buttons)
+
+        if smpl is not None:
+            # Full-body whole-body tracking: SMPL drives the reference, not joystick.
+            if self.controller.encode_mode != 2:
+                self._enter_wholebody()
+            self.controller.smpl_joints_10frame_step1 = smpl
+            # Root orientation steers the mode-2 anchor/heading, but only when
+            # explicitly enabled (see enable_smpl_root); the raw per-frame quat is
+            # spherically smoothed first so the 30->50 Hz resample doesn't spike the
+            # anchor. Disabled -> anchor stays self-driven.
+            self.controller.smpl_root_quat = (
+                self._smooth_root_quat(root_quat) if self.enable_smpl_root else None
+            )
+            out = self._runtime.tick(obs, debug=False, use_joystick=False)
+        elif vr3 is not None:
+            # 3-point VR teleop: upper body tracks the wrist/neck targets; the lower
+            # body / locomotion keeps running off the planner, so the joystick (and
+            # keyboard) still steer walking/turning underneath.
+            if self.controller.encode_mode != 1:
+                self._enter_3point()
+            self.controller.vr_3point_local_target = vr3[0]
+            self.controller.vr_3point_local_orn_target = vr3[1]
+            # Replicate the original encode_mode-1 handling: when the PICO controller
+            # sticks are forwarded, drive locomotion from them directly (and skip the
+            # wireless-remote joystick read). Otherwise leave the remote/keyboard path.
+            if loco is not None:
+                self._process_pico_loco(loco[0], loco[1])
+                out = self._runtime.tick(obs, debug=False, use_joystick=False)
+            else:
+                out = self._runtime.tick(obs, debug=False, use_joystick=True)
+        else:
+            # No (or stale) teleop reference: fall back to locomotion so the robot stays balanced.
+            if self.controller.encode_mode != 0:
+                self.controller.smpl_root_quat = None
+                self._smoothed_root_quat = None
+                self._exit_wholebody()
+            out = self._runtime.tick(obs, debug=False)
+
+        # Startup interpolation: blend from the robot's initial pose into the policy's
+        # commanded target over INIT_RAMP_S, regardless of mode.
+        return self._startup_blend(obs, out)
+
+    def reset(self):
+        self._runtime.reset()
+        self._init_step = 0  # re-run the startup blend after a reset
+        self._start_pose = {}
+        self._smoothed_root_quat = None
+        self._wb_step = 0
+        self._wb_traj.clear()
+        self._wb_quat_traj.clear()
+
+    def shutdown(self):
+        if self._smpl_stream is not None:
+            self._smpl_stream.close()
+        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,10 +23,102 @@ import numpy as np
 
 NUM_MOTORS = 29
 
+# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
+# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
+# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
+ISAACLAB_TO_MUJOCO = np.array(
+    [
+        0,
+        3,
+        6,
+        9,
+        13,
+        17,
+        1,
+        4,
+        7,
+        10,
+        14,
+        18,
+        2,
+        5,
+        8,
+        11,
+        15,
+        19,
+        21,
+        23,
+        25,
+        27,
+        12,
+        16,
+        20,
+        22,
+        24,
+        26,
+        28,
+    ],
+    dtype=np.int32,
+)
+MUJOCO_TO_ISAACLAB = np.array(
+    [
+        0,
+        6,
+        12,
+        1,
+        7,
+        13,
+        2,
+        8,
+        14,
+        3,
+        9,
+        15,
+        22,
+        4,
+        10,
+        16,
+        23,
+        5,
+        11,
+        17,
+        24,
+        18,
+        25,
+        19,
+        26,
+        20,
+        27,
+        21,
+        28,
+    ],
+    dtype=np.int32,
+)
+
 REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
 REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
 REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
 
+# Reserved action-dict field used to forward the set of currently-pressed keyboard
+# keys from a KeyboardTeleop through the standard action pipeline to the SONIC
+# whole-body controller (see SonicWholeBodyController._process_keyboard).
+KEYBOARD_KEYS_FIELD = "keyboard.keys"
+
+# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM / pi0.5) ──────
+# A single 34-D whole-body command per tick, in the OpenHLM action layout:
+#   [L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
+#    root roll/pitch + yaw-rate(3)]
+# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix makes these
+# behave like ordinary joint-position action features so ``lerobot-rollout`` routes
+# them straight from a 34-D VLA (OpenHLM / pi0.5) onto the robot.
+WB_ACTION_PREFIX = "wb."
+WB_ACTION_DIM = 34
+
+
+def wb_action_key(i: int) -> str:
+    """Action-dict key for the ``i``-th whole-body command scalar (``wb.{i}.pos``)."""
+    return f"{WB_ACTION_PREFIX}{i}.pos"
+
 
 def default_remote_input() -> dict[str, float]:
     """Return a zeroed-out remote input dict (axes + buttons)."""
@@ -63,13 +155,92 @@ class G1_29_JointArmIndex(IntEnum):
     kRightWristYaw = 28
 
 
+def lowstate_to_obs(lowstate) -> dict:
+    """Build a robot observation dict from a Unitree lowstate.
+
+    Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
+    lowstate -> obs mapping lives in exactly one place. Keys match the
+    ``<joint>.q``/``imu.*`` schema consumed across the controllers.
+    """
+    obs: dict = {}
+
+    for motor in G1_29_JointIndex:
+        idx = motor.value
+        obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
+        obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
+        obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
+
+    imu = lowstate.imu_state
+    if imu.gyroscope:
+        obs["imu.gyro.x"] = imu.gyroscope[0]
+        obs["imu.gyro.y"] = imu.gyroscope[1]
+        obs["imu.gyro.z"] = imu.gyroscope[2]
+    if imu.accelerometer:
+        obs["imu.accel.x"] = imu.accelerometer[0]
+        obs["imu.accel.y"] = imu.accelerometer[1]
+        obs["imu.accel.z"] = imu.accelerometer[2]
+    if imu.quaternion:
+        obs["imu.quat.w"] = imu.quaternion[0]
+        obs["imu.quat.x"] = imu.quaternion[1]
+        obs["imu.quat.y"] = imu.quaternion[2]
+        obs["imu.quat.z"] = imu.quaternion[3]
+    if imu.rpy:
+        obs["imu.rpy.roll"] = imu.rpy[0]
+        obs["imu.rpy.pitch"] = imu.rpy[1]
+        obs["imu.rpy.yaw"] = imu.rpy[2]
+
+    wr = getattr(lowstate, "wireless_remote", None)
+    if wr:
+        obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
+
+    return obs
+
+
+def obs_to_wb34_state(obs: dict) -> np.ndarray:
+    """Build the 34-D OpenHLM / pi0.5 proprio state from a G1 observation dict.
+
+    Mirrors the whole-body *action* layout so the policy sees state and action in
+    the same coordinates::
+
+        [L-arm(7), L-grip(1), R-arm(7), R-grip(1),
+         L-leg(6), R-leg(6), waist(3), root roll/pitch + yaw-rate(3)]
+
+    Joint positions come from the ``<joint>.q`` obs keys, which are already in
+    MuJoCo / Unitree-SDK order — the same body-part grouping OpenHLM uses
+    ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29]) — so they are
+    regrouped directly (no IsaacLab permutation). The G1 has no grippers in its
+    29-DoF body, so both gripper slots are 0. Root roll/pitch are the IMU RPY and
+    the last slot is the IMU yaw rate (gyro z).
+    """
+    q_mj = np.array(
+        [float(obs.get(f"{m.name}.q", 0.0)) for m in G1_29_JointIndex],
+        dtype=np.float32,
+    )
+    lleg, rleg, waist = q_mj[0:6], q_mj[6:12], q_mj[12:15]
+    larm, rarm = q_mj[15:22], q_mj[22:29]
+
+    state = np.zeros(34, dtype=np.float32)
+    state[0:7] = larm
+    # state[7] left gripper — none on 29-DoF G1
+    state[8:15] = rarm
+    # state[15] right gripper — none on 29-DoF G1
+    state[16:22] = lleg
+    state[22:28] = rleg
+    state[28:31] = waist
+    state[31] = float(obs.get("imu.rpy.roll", 0.0))
+    state[32] = float(obs.get("imu.rpy.pitch", 0.0))
+    state[33] = float(obs.get("imu.gyro.z", 0.0))
+    return state
+
+
 def make_locomotion_controller(name: str | None):
     """Instantiate a locomotion controller by class name. Returns None if name is None."""
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
+        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -33,12 +33,14 @@ from ..robot import Robot
 from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
+    KEYBOARD_KEYS_FIELD,
     REMOTE_AXES,
-    REMOTE_KEYS,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
+    lowstate_to_obs,
     make_locomotion_controller,
+    obs_to_wb34_state,
 )
 
 if TYPE_CHECKING or _unitree_sdk_available:
@@ -47,8 +49,12 @@ if TYPE_CHECKING or _unitree_sdk_available:
         ChannelPublisher as _SDKChannelPublisher,
         ChannelSubscriber as _SDKChannelSubscriber,
     )
-    from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
+    from unitree_sdk2py.idl.default import (
+        unitree_hg_msg_dds__HandCmd_ as hg_HandCmd_default,
+        unitree_hg_msg_dds__LowCmd_,
+    )
     from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
+        HandCmd_ as hg_HandCmd,
         LowCmd_ as hg_LowCmd,
         LowState_ as hg_LowState,
     )
@@ -58,6 +64,8 @@ else:
     _SDKChannelPublisher = None
     _SDKChannelSubscriber = None
     unitree_hg_msg_dds__LowCmd_ = None
+    hg_HandCmd_default = None
+    hg_HandCmd = None
     hg_LowCmd = None
     hg_LowState = None
     CRC = None
@@ -157,6 +165,37 @@ class UnitreeG1(Robot):
         self.controller_input = default_remote_input()
         self.controller_output = {}
 
+        # Replay-camera state (decoded frames per robot camera name + play cursor).
+        self._replay_frames: dict[str, list[np.ndarray]] = {}
+        self._replay_len = 0
+        self._replay_idx = 0
+        if config.replay_camera_parquet and config.replay_camera_map:
+            self._load_replay_frames()
+
+    def _load_replay_frames(self) -> None:
+        """Decode recorded episode frames from a parquet into per-camera image lists."""
+        import io
+
+        import pyarrow.parquet as pq
+        from PIL import Image
+
+        table = pq.read_table(self.config.replay_camera_parquet)
+        cols = {col: table.column(col).to_pylist() for col in self.config.replay_camera_map.values()}
+        self._replay_len = table.num_rows
+
+        def decode(cell) -> np.ndarray:
+            data = cell["bytes"] if isinstance(cell, dict) else cell
+            return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
+
+        for cam_name, column in self.config.replay_camera_map.items():
+            self._replay_frames[cam_name] = [decode(c) for c in cols[column]]
+        logger.info(
+            "Loaded %d replay frames for cameras %s from %s",
+            self._replay_len,
+            list(self.config.replay_camera_map),
+            self.config.replay_camera_parquet,
+        )
+
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
@@ -231,15 +270,54 @@ class UnitreeG1(Robot):
                 features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
         return features
 
+    @property
+    def _wb_state_ft(self) -> dict[str, type]:
+        """34-D whole-body proprio state (``wb_state.{i}.pos``) for dense controllers.
+
+        Exposed only when the controller consumes a dense whole-body command
+        (OpenHLM / pi0.5). These ``.pos`` scalars are aggregated by the rollout
+        pipeline into a single 34-D ``observation.state`` for the policy.
+        """
+        if not getattr(self.controller, "wb_action", False):
+            return {}
+        from .g1_utils import WB_ACTION_DIM
+
+        return {f"wb_state.{i}.pos": float for i in range(WB_ACTION_DIM)}
+
+    @property
+    def _empty_cameras_ft(self) -> dict[str, tuple]:
+        """Synthetic zero-image cameras (see ``UnitreeG1Config.empty_cameras``)."""
+        h, w = self.config.empty_camera_hw
+        return {name: (h, w, 3) for name in self.config.empty_cameras}
+
+    @property
+    def _replay_cameras_ft(self) -> dict[str, tuple]:
+        """Replay cameras, shaped from their first decoded frame."""
+        return {name: frames[0].shape for name, frames in self._replay_frames.items() if frames}
+
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._cameras_ft}
+        return {
+            **self._motors_ft,
+            **self._wb_state_ft,
+            **self._empty_cameras_ft,
+            **self._replay_cameras_ft,
+            **self._cameras_ft,
+        }
 
     @cached_property
     def action_features(self) -> dict[str, type]:
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
+        # Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
+        # 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
+        # so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
+        if getattr(self.controller, "wb_action", False):
+            from .g1_utils import WB_ACTION_DIM, wb_action_key
+
+            return {wb_action_key(i): float for i in range(WB_ACTION_DIM)}
+
         arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
         remote_features = dict.fromkeys(REMOTE_AXES, float)
         return {**arm_features, **remote_features}
@@ -311,6 +389,17 @@ class UnitreeG1(Robot):
         self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState)
         self.lowstate_subscriber.Init()
 
+        # Dex3 hand command publishers (grasping). Driven by the OpenHLM grip scalars.
+        self._hand_publishers = {}
+        if self.config.publish_hands:
+            self._left_hand_cmd = hg_HandCmd_default()
+            self._right_hand_cmd = hg_HandCmd_default()
+            self._hand_publishers["left"] = self._ChannelPublisher("rt/dex3/left/cmd", hg_HandCmd)
+            self._hand_publishers["right"] = self._ChannelPublisher("rt/dex3/right/cmd", hg_HandCmd)
+            for pub in self._hand_publishers.values():
+                pub.Init()
+            logger.info("Dex3 hand command publishers initialized (rt/dex3/{left,right}/cmd)")
+
         # Start subscribe thread to read robot state
         self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
         self.subscribe_thread.start()
@@ -343,6 +432,9 @@ class UnitreeG1(Robot):
 
         self.kp = np.array(self.config.kp, dtype=np.float32)
         self.kd = np.array(self.config.kd, dtype=np.float32)
+        if self.controller is not None and hasattr(self.controller, "kp"):
+            self.kp = np.array(self.controller.kp, dtype=np.float32)
+            self.kd = np.array(self.controller.kd, dtype=np.float32)
 
         for joint in G1_29_JointIndex:
             self.msg.motor_cmd[joint].mode = 1
@@ -371,13 +463,50 @@ class UnitreeG1(Robot):
         except Exception as e:
             logger.warning(f"Failed to send zero-torque on disconnect: {e}")
 
-    def disconnect(self):
-        # Put robot in passive mode before stopping threads
-        if not self.config.is_simulation:
-            self._send_zero_torque()
+    def _graceful_stop(self) -> None:
+        """Soft shutdown: hold the current pose and ramp joint stiffness (kp) to zero
+        over ``graceful_stop_s`` while keeping damping (kd), then go passive.
 
-        # Signal thread to stop and unblock any waits
+        Prevents the robot from collapsing the instant control ends (a bare
+        zero-torque command is kp=kd=0 ≈ free-fall). Must run after the controller
+        loop has stopped so the two aren't publishing at once.
+        """
+        if self.config.graceful_stop_s <= 0:
+            self._send_zero_torque()
+            return
+        with self._lowstate_lock:
+            lowstate = self._lowstate
+        if lowstate is None:
+            self._send_zero_torque()
+            return
+        q_hold = {f"{motor.name}.q": lowstate.motor_state[motor.value].q for motor in G1_29_JointIndex}
+        kp = np.array(self.kp, dtype=np.float32)
+        kd = np.array(self.kd, dtype=np.float32)
+        zeros = np.zeros(29, dtype=np.float32)
+        dt = self.controller.control_dt if self.controller is not None else self.config.control_dt
+        steps = max(1, int(self.config.graceful_stop_s / dt))
+        logger.info("Graceful stop: damping down over %.1fs", self.config.graceful_stop_s)
+        for i in range(steps):
+            ratio = (i + 1) / steps
+            self.publish_lowcmd(q_hold, kp=kp * (1.0 - ratio), kd=kd, tau=zeros)
+            time.sleep(dt)
+        self._send_zero_torque()
+
+    def disconnect(self):
+        # Stop the controller loop first so it isn't fighting the shutdown ramp.
         self._shutdown_event.set()
+        if self._controller_thread is not None:
+            self._controller_thread.join(timeout=2.0)
+            if self._controller_thread.is_alive():
+                logger.warning("Controller thread did not stop cleanly")
+
+        # Soft, damped settle instead of an instant limp (real robot only; the
+        # subscribe thread is still alive here to supply the current pose).
+        if not self.config.is_simulation:
+            self._graceful_stop()
+
+        if self.controller is not None and hasattr(self.controller, "shutdown"):
+            self.controller.shutdown()
 
         # Wait for subscribe thread to finish
         if self.subscribe_thread is not None:
@@ -385,12 +514,6 @@ class UnitreeG1(Robot):
             if self.subscribe_thread.is_alive():
                 logger.warning("Subscribe thread did not stop cleanly")
 
-        # Wait for controller thread to finish
-        if self._controller_thread is not None:
-            self._controller_thread.join(timeout=2.0)
-            if self._controller_thread.is_alive():
-                logger.warning("Controller thread did not stop cleanly")
-
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:
@@ -422,44 +545,33 @@ class UnitreeG1(Robot):
         if lowstate is None:
             return {}
 
-        obs = {}
+        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
+        obs = lowstate_to_obs(lowstate)
 
-        # Motors - q, dq, tau for all joints
-        for motor in G1_29_JointIndex:
-            name = motor.name
-            idx = motor.value
-            obs[f"{name}.q"] = lowstate.motor_state[idx].q
-            obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
-            obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
+        # Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
+        # state as ``wb_state.{i}.pos`` so the rollout aggregates it into
+        # ``observation.state`` for the policy.
+        if getattr(self.controller, "wb_action", False):
+            wb_state = obs_to_wb34_state(obs)
+            for i, v in enumerate(wb_state):
+                obs[f"wb_state.{i}.pos"] = float(v)
 
-        # IMU - gyroscope
-        if lowstate.imu_state.gyroscope:
-            obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
-            obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
-            obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
+        # Synthetic empty cameras: black frames so image-conditioned policies run
+        # before real cameras are wired.
+        if self.config.empty_cameras:
+            h, w = self.config.empty_camera_hw
+            black = np.zeros((h, w, 3), dtype=np.uint8)
+            for name in self.config.empty_cameras:
+                obs[name] = black
 
-        # IMU - accelerometer
-        if lowstate.imu_state.accelerometer:
-            obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
-            obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
-            obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
-
-        # IMU - quaternion
-        if lowstate.imu_state.quaternion:
-            obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
-            obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
-            obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
-            obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
-
-        # IMU - rpy
-        if lowstate.imu_state.rpy:
-            obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
-            obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
-            obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
-
-        # Wireless remote (raw bytes for teleoperator)
-        if lowstate.wireless_remote:
-            obs["wireless_remote"] = lowstate.wireless_remote
+        # Replay cameras: serve the current recorded frame per camera, then advance.
+        if self._replay_len:
+            idx = self._replay_idx
+            if idx >= self._replay_len:
+                idx = self._replay_len - 1 if not self.config.replay_camera_loop else idx % self._replay_len
+            for name, frames in self._replay_frames.items():
+                obs[name] = frames[idx]
+            self._replay_idx += 1
 
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
@@ -473,9 +585,13 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
+            self._update_controller_action(action)
+            if self.config.publish_hands and getattr(self.controller, "wb_action", False):
+                self._publish_hand_cmds(action)
+            if getattr(self.controller, "full_body", False):
+                return action
             # Controller thread owns legs/waist. Here we only update joystick inputs
             # and publish arm targets from the teleoperator.
-            self._update_controller_action(action)
             arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
             action_to_publish = {
                 key: value
@@ -503,11 +619,67 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from incoming teleop action."""
+        """Update controller input state from an incoming teleop action.
+
+        Controller-agnostic: every value-carrying key is forwarded verbatim into
+        ``controller_input`` (whole-body ``wb.{i}.pos`` from a 34-D VLA, or whatever a
+        future controller expects), and each controller extracts only the keys it
+        understands. The robot deliberately does not enumerate any controller's key
+        schema here.
+
+        KeyboardTeleop is the one special case: it emits the currently-pressed keys as
+        bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so
+        those are collected into a single held-key set under ``KEYBOARD_KEYS_FIELD``,
+        rebuilt each tick so releases clear. Special keys arrive as pynput objects and
+        are normalised to their name ("space", ...).
+        """
         with self._controller_action_lock:
-            for key in REMOTE_KEYS:
-                if key in action:
-                    self.controller_input[key] = action[key]
+            self.controller_input[KEYBOARD_KEYS_FIELD] = {
+                (k if isinstance(k, str) else getattr(k, "name", str(k)))
+                for k, value in action.items()
+                if value is None
+            }
+            for key, value in action.items():
+                if isinstance(key, str) and value is not None:
+                    self.controller_input[key] = value
+
+    def _publish_hand_cmds(self, action: RobotAction) -> None:
+        """Drive the Dex3 hands from the OpenHLM grip scalars in a 34-D wb action.
+
+        ``wb.7.pos`` is the left grip and ``wb.15.pos`` the right grip. Each scalar in
+        [0, 1] (``hand_open_grip_value`` == fully open) is turned into a curl amount and
+        scaled onto ``hand_closed_pose`` (7 joints), then published as a PD target on
+        ``rt/dex3/{left,right}/cmd`` so the fingers close when the policy grips.
+        """
+        if not self._hand_publishers:
+            return
+        from .g1_utils import wb_action_key
+
+        open_val = float(self.config.hand_open_grip_value)
+        closed_val = float(self.config.hand_closed_grip_value)
+        closed_pose = self.config.hand_closed_pose
+        kp, kd = float(self.config.hand_kp), float(self.config.hand_kd)
+        span = (closed_val - open_val) or 1.0
+
+        def curl_amount(grip: float) -> float:
+            # Fraction of the way from the open scalar to the closed scalar, in [0, 1].
+            return float(min(max((grip - open_val) / span, 0.0), 1.0))
+
+        for side, grip_idx, cmd in (
+            ("left", 7, self._left_hand_cmd),
+            ("right", 15, self._right_hand_cmd),
+        ):
+            grip = action.get(wb_action_key(grip_idx))
+            if grip is None:
+                continue
+            amount = curl_amount(float(grip))
+            for i, closed_q in enumerate(closed_pose):
+                cmd.motor_cmd[i].q = float(closed_q) * amount
+                cmd.motor_cmd[i].dq = 0.0
+                cmd.motor_cmd[i].kp = kp
+                cmd.motor_cmd[i].kd = kd
+                cmd.motor_cmd[i].tau = 0.0
+            self._hand_publishers[side].Write(cmd)
 
     @property
     def is_calibrated(self) -> bool:

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

@@ -116,6 +116,7 @@ from lerobot.teleoperators import (  # noqa: F401
     omx_leader,
     openarm_leader,
     openarm_mini,
+    pico_headset,
     reachy2_teleoperator,
     rebot_102_leader,
     so_leader,

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
 
@@ -57,7 +55,6 @@ from lerobot.envs import close_envs, make_env, make_env_pre_post_processors
 from lerobot.jobs import submit_to_hf
 from lerobot.optim.factory import make_optimizer_and_scheduler
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
-from lerobot.processor.rename_processor import rename_transition_keys
 from lerobot.rewards import make_reward_pre_post_processors
 from lerobot.utils.collate import lerobot_collate_fn
 from lerobot.utils.import_utils import register_third_party_plugins
@@ -84,8 +81,6 @@ def update_policy(
     lr_scheduler=None,
     lock=None,
     sample_weighter=None,
-    log_metrics: bool = True,
-    track_update_time: bool = True,
 ) -> tuple[MetricsTracker, dict | None]:
     """
     Performs a single training step to update the policy's weights.
@@ -103,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:
@@ -149,16 +143,13 @@ def update_policy(
     # Use accelerator's backward method
     accelerator.backward(loss)
 
-    # Accelerate suppresses gradient synchronization and optimizer updates on intermediate
-    # microbatches. Clip and report the norm only when the accumulated update is complete.
-    grad_norm = None
-    if accelerator.sync_gradients:
-        if grad_clip_norm > 0:
-            grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
-        else:
-            grad_norm = torch.nn.utils.clip_grad_norm_(
-                policy.parameters(), float("inf"), error_if_nonfinite=False
-            )
+    # Clip gradients if specified
+    if grad_clip_norm > 0:
+        grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
+    else:
+        grad_norm = torch.nn.utils.clip_grad_norm_(
+            policy.parameters(), float("inf"), error_if_nonfinite=False
+        )
 
     # Optimizer step
     with lock if lock is not None else nullcontext():
@@ -167,32 +158,22 @@ def update_policy(
     optimizer.zero_grad()
 
     # Step through pytorch scheduler at every batch instead of epoch
-    if lr_scheduler is not None and accelerator.sync_gradients:
+    if lr_scheduler is not None:
         lr_scheduler.step()
 
     # Update internal buffers if policy has update method
-    if accelerator.sync_gradients and has_method(
-        accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"
-    ):
+    if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"):
         accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update()
 
-    if accelerator.sync_gradients:
-        train_metrics.lr = optimizer.param_groups[0]["lr"]
-    if torch.cuda.is_available() and accelerator.sync_gradients:
+    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)
-    if grad_norm is not None:
-        train_metrics.accumulate_tensor("grad_norm", grad_norm)
-    if track_update_time:
-        train_metrics.update_s = time.perf_counter() - start_time
-    # Synchronize accumulated GPU metrics only when logging.
-    if log_metrics:
-        train_metrics.materialize_tensors()
-        # Materialize detached loss components during the same logging synchronization.
-        if output_dict:
-            output_dict = {
-                k: (v.item() if isinstance(v, torch.Tensor) else v) for k, v in output_dict.items()
-            }
+    # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
+    if output_dict:
+        train_metrics.update_metrics(output_dict)
     return train_metrics, output_dict
 
 
@@ -220,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()
 
@@ -229,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"
@@ -247,9 +219,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         mixed_precision = {"bfloat16": "bf16", "float16": "fp16", "float32": "no"}.get(policy_dtype)
         accelerator = Accelerator(
             step_scheduler_with_optimizer=False,
-            gradient_accumulation_steps=cfg.gradient_accumulation_steps,
             mixed_precision=mixed_precision,
-            kwargs_handlers=[ddp_kwargs, ipg_kwargs],
+            kwargs_handlers=[ddp_kwargs],
             cpu=force_cpu,
         )
 
@@ -345,31 +316,19 @@ 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
 
-    dataset_stats = {cfg.rename_map.get(key, key): value for key, value in dataset.meta.stats.items()}
     processor_kwargs = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
-        processor_kwargs["dataset_stats"] = dataset_stats
+        processor_kwargs["dataset_stats"] = dataset.meta.stats
 
     if cfg.is_reward_model_training:
         processor_kwargs["dataset_meta"] = dataset.meta
 
-    if cfg.policy.type in {"pi0_fast", "pi052"}:
-        processor_kwargs["dataset_repo_id"] = cfg.dataset.repo_id
-
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
         preprocessor_overrides = {
             "device_processor": {"device": device.type},
             "normalizer_processor": {
-                "stats": dataset_stats,
+                "stats": dataset.meta.stats,
                 "features": {**policy.config.input_features, **policy.config.output_features},
                 "norm_map": policy.config.normalization_mapping,
             },
@@ -377,7 +336,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         }
         postprocessor_overrides = {
             "unnormalizer_processor": {
-                "stats": dataset_stats,
+                "stats": dataset.meta.stats,
                 "features": policy.config.output_features,
                 "norm_map": policy.config.normalization_mapping,
             },
@@ -449,11 +408,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})")
         logging.info(f"{dataset.num_episodes=}")
         num_processes = accelerator.num_processes
-        effective_bs = cfg.batch_size * num_processes * cfg.gradient_accumulation_steps
-        logging.info(
-            "Effective batch size: "
-            f"{cfg.batch_size} x {num_processes} x {cfg.gradient_accumulation_steps} = {effective_bs}"
-        )
+        effective_bs = cfg.batch_size * num_processes
+        logging.info(f"Effective batch size: {cfg.batch_size} x {num_processes} = {effective_bs}")
         logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
         logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
@@ -464,53 +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
-
-        drop_n_first_frames: int | list[int] = 0
-        target_start_feature = cfg.dataset.training_target_start_feature
-        if target_start_feature is not None:
-            if not hasattr(dataset, "hf_dataset"):
-                raise ValueError(
-                    "dataset.training_target_start_feature is only supported for a single map-style "
-                    "LeRobotDataset."
-                )
-            if target_start_feature not in dataset.hf_dataset.column_names:
-                raise ValueError(
-                    f"Training target start feature {target_start_feature!r} is not present in the dataset. "
-                    f"Available features: {dataset.hf_dataset.column_names}"
-                )
-            start_column = dataset.hf_dataset.data.column(target_start_feature)
-            absolute_to_relative_idx = dataset.absolute_to_relative_idx
-            episode_indices = dataset.episodes if dataset.episodes is not None else range(len(from_indices))
-            drop_n_first_frames = [0] * len(from_indices)
-            for episode_index in episode_indices:
-                absolute_index = int(from_indices[episode_index])
-                relative_index = (
-                    absolute_to_relative_idx[absolute_index]
-                    if absolute_to_relative_idx is not None
-                    else absolute_index
-                )
-                drop_n_first_frames[episode_index] = int(start_column[relative_index].as_py())
-            logging.info(
-                "Restricting training targets with %s: keeping %d of %d frames",
-                target_start_feature,
-                sum(
-                    int(to_indices[index]) - int(from_indices[index]) - drop_n_first_frames[index]
-                    for index in episode_indices
-                ),
-                sum(int(to_indices[index]) - int(from_indices[index]) for index in episode_indices),
-            )
-
         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_first_frames=drop_n_first_frames,
             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:
@@ -533,12 +449,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     f"batch_size={saved_batch_size}. The data order resumes at the right epoch/offset, "
                     "but per-rank sample-exactness requires the same batch size."
                 )
-            sampler_state = compute_sampler_state(
-                step,
-                len(sampler),
-                ckpt_batch_size * cfg.gradient_accumulation_steps,
-                ckpt_num_processes,
-            )
+            sampler_state = compute_sampler_state(step, len(sampler), ckpt_batch_size, ckpt_num_processes)
             sampler.load_state_dict(sampler_state)
             if is_main_process:
                 logging.info(
@@ -553,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,
@@ -566,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
@@ -635,9 +543,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         train_metrics["gpu_mem_gb"] = AverageMeter("mem_gb", ":.2f", reduction="max")
 
     # Keep global batch size for logging; MetricsTracker handles world size internally.
-    effective_batch_size = cfg.batch_size * accelerator.num_processes * cfg.gradient_accumulation_steps
+    effective_batch_size = cfg.batch_size * accelerator.num_processes
     train_tracker = MetricsTracker(
-        cfg.batch_size * cfg.gradient_accumulation_steps,
+        cfg.batch_size,
         dataset.num_frames,
         dataset.num_episodes,
         train_metrics,
@@ -659,42 +567,24 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         )
 
     for _ in range(step, cfg.steps):
-        update_start_time = time.perf_counter()
-        dataloading_s = 0.0
-        output_dict = None
-        for microbatch_idx in range(cfg.gradient_accumulation_steps):
-            start_time = time.perf_counter()
-            batch = next(dl_iter)
-            for cam_key in dataset.meta.camera_keys:
-                if cam_key in batch and batch[cam_key].dtype == torch.uint8:
-                    batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
-            if cfg.rename_map:
-                batch = rename_transition_keys(batch, cfg.rename_map)
-            batch = preprocessor(batch)
-            dataloading_s += time.perf_counter() - start_time
+        start_time = time.perf_counter()
+        batch = next(dl_iter)
+        for cam_key in dataset.meta.camera_keys:
+            if cam_key in batch and batch[cam_key].dtype == torch.uint8:
+                batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0
+        batch = preprocessor(batch)
+        train_tracker.dataloading_s = time.perf_counter() - start_time
 
-            # Synchronize GPU metrics only on the final microbatch of logged optimizer updates.
-            log_metrics = (
-                cfg.log_freq > 0
-                and (step + 1) % cfg.log_freq == 0
-                and microbatch_idx == cfg.gradient_accumulation_steps - 1
-            )
-
-            with accelerator.accumulate(policy):
-                train_tracker, output_dict = update_policy(
-                    train_tracker,
-                    policy,
-                    batch,
-                    optimizer,
-                    cfg.optimizer.grad_clip_norm,
-                    accelerator=accelerator,
-                    lr_scheduler=lr_scheduler,
-                    sample_weighter=sample_weighter,
-                    log_metrics=log_metrics,
-                    track_update_time=False,
-                )
-        train_tracker.dataloading_s = dataloading_s
-        train_tracker.update_s = time.perf_counter() - update_start_time - dataloading_s
+        train_tracker, _ = update_policy(
+            train_tracker,
+            policy,
+            batch,
+            optimizer,
+            cfg.optimizer.grad_clip_norm,
+            accelerator=accelerator,
+            lr_scheduler=lr_scheduler,
+            sample_weighter=sample_weighter,
+        )
 
         # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we
         # increment `step` here.
@@ -718,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()
@@ -793,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/teleoperators/pico_headset/__init__.py

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 # Copyright 2026 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,8 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""PI052 adapter for the policy-agnostic language runtime."""
+from .config_pico_headset import PicoHeadsetConfig
+from .pico_headset import PicoHeadset
 
-from .pi052_adapter import PI052PolicyAdapter
-
-__all__ = ["PI052PolicyAdapter"]
+__all__ = ["PicoHeadset", "PicoHeadsetConfig"]

src/lerobot/teleoperators/pico_headset/config_pico_headset.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+from ..config import TeleoperatorConfig
+
+
+@TeleoperatorConfig.register_subclass("pico_headset")
+@dataclass
+class PicoHeadsetConfig(TeleoperatorConfig):
+    """PICO full-body headset teleop: live SMPL over the rt/smpl ZMQ stream.
+
+    Consumes the ``rt/smpl`` channel published by the GEAR PICO manager
+    (``gear_sonic/scripts/pico_manager_thread_server.py``) and emits the whole-body
+    SONIC reference window (``encode_mode == 2``) for SonicWholeBodyController.
+    """
+
+    smpl_host: str = "127.0.0.1"
+    """Host of the pico_manager rt/smpl publisher (the laptop bridging the PICO)."""
+    smpl_port: int = 5560
+    """Port of the rt/smpl publisher."""
+    stale_after_s: float = 0.5
+    """Warn if no fresh headset frame arrives within this many seconds."""
+    mode: str = "smpl"
+    """Teleop reference to emit: ``"smpl"`` for whole-body imitation (SONIC
+    encode_mode 2) or ``"vr3"`` for sparse 3-point upper-body teleop (encode_mode 1,
+    lower body driven by the joystick/keyboard planner)."""