feat(unitree_g1): use captured neutral SONIC token instead of zeros

The all-zero token is off the encoder's learned FSQ manifold and decodes to a
slightly goofy stance. Replace it with a NEUTRAL_TOKEN captured from the encoder's
own idle output in sim (stored as integer FSQ codes, rescaled by the encoder's
1/16 quantization step to an exact on-grid token). token_mode now seeds this
neutral, and the onboard sender starts observation.state from it so the first
inference sees the token the decoder is actually holding.

docs/source/annotation_pipeline.mdx

@@ -81,6 +81,12 @@ merged. Both prompts also carry a causal **event-boundary** definition (a
 new event starts when an object becomes held / is released / reaches a new
 location / a lid changes state / contents move) to sharpen where cuts land.
 
+Optionally, a third **seeded-relabel** pass (`--plan.subtask_seeded_relabel`)
+revisits each span with its previous/current/next segment contact sheets and
+minimally corrects the label, using the first label as a prior — it keeps the
+boundaries fixed and only sharpens wording, at the cost of one extra call per
+subtask.
+
 The resulting spans are then stitched into a gap-free, full-episode
 cover, so **every frame has exactly one active subtask**. See
 [`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
@@ -158,7 +164,7 @@ Every module is on by default and can be toggled independently (set to
 ### The VLM (`--vlm.*`)
 
 | Flag                       | Default            | What it does                                                                         |
-| -------------------------- | ------------------ | ----------------------------------------------------------------------------------- |
+| -------------------------- | ------------------ | ------------------------------------------------------------------------------------ |
 | `--vlm.model_id`           | `Qwen/Qwen3.6-27B` | The model to serve and prompt.                                                       |
 | `--vlm.camera_key`         | first `images.*`   | Which camera every prompt is grounded on.                                            |
 | `--vlm.serve_command`      | auto               | The exact `vllm serve …` command (set TP size, GPU memory, `--max-model-len` here).  |
@@ -167,16 +173,19 @@ Every module is on by default and can be toggled independently (set to
 | `--vlm.client_concurrency` | `16`               | In-flight requests across all servers.                                               |
 | `--vlm.max_new_tokens`     | `512`              | Generation cap per call.                                                             |
 | `--vlm.temperature`        | `0.2`              | Sampling temperature.                                                                |
+| `--vlm.reasoning_effort`   | `null`             | Thinking-budget hint (`low`/`medium`/`high`) forwarded to OpenAI-compatible servers. |
 
 ### Subtasks / plan / memory (`--plan.*`)
 
 | Flag                            | Default    | What it does                                                                                                                 |
-| ------------------------------- | ---------- | ------------------------------------------------------------------------------------------------------------------------- |
+| ------------------------------- | ---------- | ---------------------------------------------------------------------------------------------------------------------------- |
 | `--plan.frames_per_second`      | `2.0`      | Frame sampling rate for the contact sheets (`2.0` = one frame every 0.5s).                                                   |
 | `--plan.max_frames_per_prompt`  | `60`       | Frame budget per VLM call. Episodes whose sampling exceeds this are auto-windowed at the same density, then stitched.        |
 | `--plan.contact_sheet_columns`  | `5`        | Columns per contact-sheet grid (`contact_sheet_frames_per_sheet` tiles, time row-major).                                     |
 | `--plan.plan_max_steps`         | `8`        | Upper bound on subtasks per episode.                                                                                         |
 | `--plan.subtask_describe_first` | `true`     | Run the describe→segment grounding pass (best subtask quality; +1 call/episode).                                             |
+| `--plan.subtask_seeded_relabel` | `false`    | Second pass: re-label each subtask from its prev/current/next contact sheets, seeded with the first label (+1 call/subtask). |
+| `--plan.subtask_relabel_frames` | `5`        | Frames sampled uniformly per segment sheet in the relabel pass (only used when `subtask_seeded_relabel=true`).               |
 | `--plan.emit_plan`              | `true`     | Emit the numbered `plan` rows (`false` = subtasks + memory only).                                                            |
 | `--plan.emit_memory`            | `true`     | Emit the `memory` rows (`false` = subtasks + plan only); symmetric to `emit_plan`.                                           |
 | `--plan.n_task_rephrasings`     | `10`       | How many `task_aug` rephrasings to emit (`0` disables).                                                                      |

docs/source/bring_your_own_policies.mdx

@@ -151,12 +151,12 @@ class MyPolicy(PreTrainedPolicy):
 The methods called by the train/eval loops:
 
 | Method                                                            | Used by           | What it does                                                                                                                                                                                                                                                                                 |
-| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | `reset() -> None`                                                 | `lerobot-eval`    | Clear per-episode state at the start of each episode.                                                                                                                                                                                                                                        |
 | `select_action(batch, **kwargs) -> Tensor`                        | `lerobot-eval`    | Return the next action `(B, action_dim)`. Called every step.                                                                                                                                                                                                                                 |
 | `predict_action_chunk(batch, **kwargs) -> Tensor`                 | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk.                                                                                                                                       |
 | `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train`   | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting.                                                                                                                                                                                         |
-| `get_optim_params() -> dict`                                      | the optimizer     | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). |
+| `get_optim_params() -> dict`                                      | the optimizer     | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
 | `update() -> None` _(optional)_                                   | `lerobot-train`   | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this).                                                                                                                                                                                   |
 
 Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
@@ -295,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.
 
@@ -332,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.
@@ -367,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.

examples/annotations/run_hf_job.py

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

pyproject.toml

@@ -374,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"]
@@ -413,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/annotations/steerable_pipeline/config.py

@@ -65,6 +65,14 @@ class PlanConfig:
     # invented from the task text (+1 VLM call/episode).
     subtask_describe_first: bool = True
 
+    # Seeded relabeling: after segmentation, re-label each span with a focused
+    # pass that sees the previous / current / next segment contact sheets and
+    # minimally corrects the seed label (macrodata's best end-to-end labeling
+    # step). Costs +1 VLM call per subtask; off by default.
+    subtask_seeded_relabel: bool = False
+    # Frames sampled uniformly per segment sheet in the relabel pass.
+    subtask_relabel_frames: int = 5
+
     # Emit ``style="plan"`` rows at each boundary; False = subtasks + memory only.
     emit_plan: bool = True
 
@@ -160,6 +168,11 @@ class VlmConfig:
     # Forwarded as extra_body.chat_template_kwargs (e.g. {"enable_thinking": false}).
     chat_template_kwargs: dict[str, Any] | None = None
 
+    # OpenAI-style thinking budget hint ("low"/"medium"/"high"); forwarded to
+    # the server when set. Used to cap a thinking model's reasoning so it
+    # leaves tokens for the actual JSON answer on OpenAI-compatible endpoints.
+    reasoning_effort: str | None = None
+
 
 @dataclass
 class ExecutorConfig:

src/lerobot/annotations/steerable_pipeline/frames.py

@@ -413,6 +413,15 @@ def _draw_timestamp_badge(image: PIL.Image.Image, timestamp: float) -> PIL.Image
 
     result = image.copy()
     draw = ImageDraw.Draw(result)
+    # Scale the timestamp to the tile so it stays legible after the model
+    # downsamples the full sheet into 768px tiles — a tiny bitmap font blurs
+    # at contact-sheet resolution and the VLM can no longer read the exact
+    # source time, which is what the boundary score depends on. ``size=`` is
+    # supported by Pillow's bitmap default since 10.1; fall back otherwise.
+    badge_px = max(14, round(image.height * 0.12))
+    try:
+        font = ImageFont.load_default(size=badge_px)
+    except TypeError:
         font = ImageFont.load_default()
     label = f"{timestamp:06.2f}s"
     left, top, right, bottom = draw.textbbox((0, 0), label, font=font)

src/lerobot/annotations/steerable_pipeline/modules/plan_subtasks_memory.py

@@ -116,6 +116,8 @@ class PlanSubtasksMemoryModule:
             rows.extend(self._task_aug_rows([effective_task, *variants], t0))
 
         subtask_spans = self._generate_subtasks(record, task=effective_task)
+        if self.config.subtask_seeded_relabel and subtask_spans:
+            subtask_spans = self._seeded_relabel(record, subtask_spans, effective_task)
 
         # subtask rows
         for span in subtask_spans:
@@ -509,6 +511,51 @@ class PlanSubtasksMemoryModule:
 
         return cleaned
 
+    def _seeded_relabel(
+        self, record: EpisodeRecord, spans: list[dict[str, Any]], task: str
+    ) -> list[dict[str, Any]]:
+        """Re-label each span using prev/current/next segment contact sheets.
+
+        Boundaries are kept fixed; only ``text`` is refined. The original
+        ("seed") label is passed as a strong prior so the model verifies and
+        minimally corrects it rather than re-describing from scratch — the
+        macrodata seeded-relabeling step. One VLM call per span.
+        """
+        n = len(spans)
+        out: list[dict[str, Any]] = []
+        for i, span in enumerate(spans):
+            content: list[dict[str, Any]] = []
+            if i > 0:
+                content += self._segment_sheet(record, spans[i - 1])
+            content += self._segment_sheet(record, span)
+            if i < n - 1:
+                content += self._segment_sheet(record, spans[i + 1])
+            prompt = load_prompt("plan_subtask_relabel").format(
+                episode_task=task,
+                seed_label=span["text"],
+                segment_index=i + 1,
+                segment_count=n,
+                start=float(span["start"]),
+                end=float(span["end"]),
+            )
+            content.append({"type": "text", "text": prompt})
+            label = self._vlm_field([{"role": "user", "content": content}], "label")
+            text = label.strip() if isinstance(label, str) and label.strip() else span["text"]
+            out.append({**span, "text": text})
+        return out
+
+    def _segment_sheet(self, record: EpisodeRecord, span: dict[str, Any]) -> list[dict[str, Any]]:
+        """Contact-sheet block(s) for one span: up to N frames sampled uniformly."""
+        s, e = float(span["start"]), float(span["end"])
+        n = max(1, int(self.config.subtask_relabel_frames))
+        if e <= s or n == 1:
+            timestamps = [s]
+        else:
+            step = (e - s) / (n - 1)
+            timestamps = [s + i * step for i in range(n)]
+        frames = self.frame_provider.frames_at(record, timestamps)
+        return self._contact_sheet_blocks(frames, timestamps[: len(frames)])
+
     def _generate_subtasks_windowed(
         self, record: EpisodeRecord, task: str, window_s: float
     ) -> list[dict[str, Any]]:

src/lerobot/annotations/steerable_pipeline/prompts/__init__.py

@@ -22,12 +22,23 @@ plain editors and roundtrip cleanly through ``ruff format``.
 
 from __future__ import annotations
 
+import os
 from pathlib import Path
 
 _DIR = Path(__file__).parent
 
 
 def load(name: str) -> str:
-    """Read prompt template ``name.txt`` from the ``prompts/`` directory."""
+    """Read prompt template ``name.txt`` from the ``prompts/`` directory.
+
+    A ``LEROBOT_PROMPT_OVERRIDE_<name>`` environment variable, when set to a
+    non-empty value, takes precedence over the packaged file. This lets prompt
+    search (e.g. GEPA) inject candidate templates into a remote job without
+    rebuilding the package; the override must keep the same ``{placeholder}``
+    fields the call site formats in.
+    """
+    override = os.environ.get(f"LEROBOT_PROMPT_OVERRIDE_{name}")
+    if override and override.strip():
+        return override
     path = _DIR / f"{name}.txt"
     return path.read_text(encoding="utf-8")

src/lerobot/annotations/steerable_pipeline/prompts/plan_subtask_relabel.txt

@@ -0,0 +1,35 @@
+Annotate one fixed segment from a longer robot demonstration.
+
+Return only JSON:
+  {{"label": "<short descriptive subtask label>"}}
+
+You are shown up to three timestamped contact sheets, in order:
+- The FIRST sheet is the PREVIOUS segment (context only); it may be absent.
+- The SECOND sheet is the CURRENT target segment.
+- The THIRD sheet is the NEXT segment (context only); it may be absent.
+Each tile has its timestamp (seconds, absolute video time) burned into its
+top-left corner.
+
+Episode instruction: "{episode_task}"
+Target segment: {segment_index} of {segment_count}
+Target time: {start:.2f}s to {end:.2f}s
+Original predicted label for this exact segment: "{seed_label}"
+
+Rules:
+- Label ONLY the current target segment (the second sheet). Use the
+  previous/next sheets only to disambiguate what changed.
+- Treat the original predicted label as a STRONG PRIOR, not ground truth:
+  verify it against the current segment and correct it minimally.
+- If it already names the right action and main object, keep it; only fix
+  grammar or add a clearly visible essential detail.
+- If it is vague but directionally correct, make it more specific.
+- If it describes the previous/next segment, the wrong action, wrong
+  object, wrong destination, or a wrong state change, replace it.
+- Do not describe the previous or next segment, and do not split, merge,
+  or move the fixed segment.
+- Do not introduce an action that is not clearly visible in the current
+  target segment.
+- Use one concise imperative phrase. Name the manipulated object and the
+  action / state change. Include source, destination, side, direction,
+  final placement, or opened/closed state when visible and central.
+- Do not mention timestamps, frame numbers, uncertainty, or intent.

src/lerobot/annotations/steerable_pipeline/prompts/plan_subtasks.txt

@@ -1,112 +1,68 @@
-You are labeling a teleoperated robot demonstration.
+You are annotating a teleoperated robot demonstration shown as
+timestamped contact sheets (each tile has its time in seconds burned
+into the top-left corner). The operator's goal was: "{episode_task}"
 
-The user originally asked: "{episode_task}"
+{observation_block}Reconstruct the sequence of COMPLETED manipulation events the robot
+performs, in chronological order. Output one segment per event with a
+[start, end] time in seconds and a short action label.
 
-You are shown the entire demonstration as a single video. Watch the
-whole clip, then segment it into a list of consecutive atomic subtasks
-the robot performs.
+GROUNDING — read first, it overrides everything below:
+- Label ONLY events you can SEE in the frames. The instruction is the
+  goal; the VIDEO is the ground truth for what actually happened.
+- Do NOT invent, anticipate, or pad steps that are not shown.
 
-{observation_block}GROUNDING — read this first, it overrides everything below:
-- Label ONLY what the robot actually does in the video. Every subtask
-  you emit must correspond to motion you can SEE in specific frames.
-- Do NOT invent, anticipate, or pad. If the robot only does one thing
-  (e.g. it just navigates to a location and the clip ends), emit
-  EXACTLY ONE subtask. Many demonstrations are a single atomic skill.
-- ``max_steps`` below is a hard CEILING, not a target. Emitting fewer
-  subtasks than the ceiling is not just allowed, it is expected for
-  short / atomic demonstrations. One correct subtask is far better
-  than several invented ones.
-- If the video does not clearly show the action implied by the task,
-  describe what you actually see — do NOT fabricate the task's steps
-  from the instruction text. The instruction tells you the goal; the
-  VIDEO is the ground truth for what happened.
+Granularity — segment by completed events, not by motion:
+- Start a NEW segment whenever the world state changes: an object is
+  grasped, lifted, transported, placed, or released; a held object
+  changes; a drawer/door/lid/container opens or closes; contents move
+  between containers (poured); a tool starts or stops acting on a
+  surface. Watch the gripper open/close transitions — they usually mark
+  boundaries.
+- Do NOT split approach, reach, grasp adjustment, small repositioning,
+  hesitation, or retreat into their own segments. Fold each into the
+  event it belongs to (the approach is part of the pick; the retreat is
+  part of the place).
+- Do NOT merge separate completed events. Each distinct pick, place,
+  open, close, pour, push, wipe, or insert is its own segment, even when
+  they repeat on different objects or locations.
+- Most segments last 2-10 seconds. Shorter segments are okay ONLY for
+  fast pick / place / open / close / release events. Never emit a
+  segment shorter than {min_subtask_seconds} seconds; merge a too-short
+  candidate into its neighbour instead.
+- Skip idle time, pure camera motion, and tiny hand jitter.
 
-Authoring rules — Hi Robot atom granularity, pi0.7-style short prompts:
+Labels — short imperative phrases:
+- One concise command naming the action and the manipulated object, e.g.
+  "pick up the red cup", "put the cup on the shelf", "open the top
+  drawer", "pour water into the glass", "insert the plug into the
+  socket".
+- Include source, destination, side, direction, or the final
+  open/closed state when it is visible and central to the event.
+- Prefer these verbs (extend only when none fits): pick up, put, place,
+  push, pull, turn, press, open, close, pour, insert, wipe, stack.
+  Disambiguate by what you SEE:
+    * STACK vs PUT: object placed ON TOP OF another object -> "stack".
+    * INSERT vs PUT: object pushed INTO a fitted slot/hole/socket -> "insert".
+    * PICK UP vs PUT (direction): gripper CLOSES and object moves WITH
+      the hand -> "pick up"; gripper OPENS and object stays -> "put".
+    * POUR vs PUT: source is tilted and contents flow -> "pour".
+- Use the exact object nouns implied by the task; stay consistent across
+  the episode (don't switch "cube" to "block").
+- Write imperative commands, never third person ("the robot ..."), and
+  drop articles/adverbs.
 
-- Each subtask = one COMPOSITE atomic skill the low-level policy can
-  execute end-to-end. A "skill" bundles its own approach motion with
-  its terminal action — do NOT split the approach off as its own
-  subtask. The whole-arm policy already learns to reach as part of
-  every manipulation primitive.
-- Write each subtask as an IMPERATIVE COMMAND, starting with one of
-  these verbs (extend only when none fits):
-    pick up <obj>           — approach + grasp + lift in one subtask
-    put <obj> on/in <loc>   — transport + release in one subtask
-    place <obj> on/in <loc> — synonym of "put"; pick one and stay consistent
-    push <obj>              — contact + linear shove
-    pull <obj>              — contact + linear retract
-    turn <knob/dial/handle> — rotary actuation
-    press <button>          — single-press contact
-    open <drawer/door/lid>  — full open motion
-    close <drawer/door/lid> — full close motion
-    pour <src> into <dst>   — tilt + flow
-    insert <obj> into <slot>— alignment + push-fit
-    go to <loc>             — ONLY when no grasp / actuation follows
-                             (e.g. a pure relocation between phases).
-                             If the next subtask grasps something at
-                             that location, drop "go to ..." and just
-                             write "pick up ..." instead.
-- Forbidden ultra-fine splits — the VLM is NOT allowed to emit these
-  as standalone subtasks; fold them into the parent composite:
-    "move to X"   → fold into "pick up X" (or whatever follows)
-    "reach for X" → fold into "pick up X"
-    "grasp X"     → fold into "pick up X"
-    "lift X"      → fold into "pick up X" (or "put X on Y" if it's
-                    the transport phase of a place)
-    "release X"   → fold into "put X on Y" (or "place X in Y")
-- Keep it SHORT — a verb phrase, not a sentence. Drop articles
-  ("the", "a") and adverbs ("carefully", "slowly"). Add a "how"
-  detail (which hand, which grasp point) ONLY when it is needed to
-  disambiguate. Every subtask must begin with one of the verbs
-  above (no leading nouns, no "then", no "first").
-- NEVER use third person. Never write "the robot", "the arm", "the
-  gripper moves", "it picks up" — the robot is implied. Command it,
-  do not describe it.
-- Use the exact object nouns from the task above. If the task says
-  "cube", every subtask says "cube" — never switch to "block". If it
-  says "box", never switch to "bin"/"container". Keep vocabulary
-  consistent across the whole episode.
-- Good: "pick up blue cube", "put blue cube in box", "open drawer",
-  "turn red knob", "press start button", "go to sink".
-- Bad: "move to blue cube" (approach as its own subtask — forbidden,
-  must be folded into "pick up blue cube"); "the robot arm moves
-  towards the blue cube" (third person, too long); "carefully pick
-  up the cube" (adverb, article); "release the yellow block"
-  ("block" when the task said "cube", and "release" must be folded
-  into a "put"/"place" subtask).
-- Subtasks are non-overlapping and cover the full episode in order.
-  Choose the cut points yourself based on what you see in the video
-  (gripper open/close events, contact, regrasps, transitions).
-- Each subtask spans at least {min_subtask_seconds} seconds. If a
-  candidate span would be shorter, merge it into its neighbour
-  rather than emitting it.
-- Do not exceed {max_steps} subtasks total. Fewer, larger composites
-  are preferred over many micro-steps.
-- Every subtask's [start_time, end_time] must lie within
-  [0.0, {episode_duration}] seconds.
-
-SPECIAL CASES — verb disambiguation (each rule is narrowly visual and
-fires ONLY on the spatial situation it names; it must not change how you
-label any other situation):
-- STACK vs PUT: if an object is placed ON TOP OF another specific object
-  (not on a flat table / shelf / counter), use "stack ... on ...", not
-  "put". "stack blue book on green book", NOT "put blue book on table".
-- INSERT vs PUT: if an object goes INTO a fitted slot / hole / socket /
-  receptacle (push-fit), use "insert ... into ...", not "put".
-- RETRIEVE/PICK-UP vs PUT (direction): watch the gripper. If it CLOSES
-  on the object and the object moves WITH the hand, it is "pick up" /
-  "retrieve" (object leaves its location). If the gripper OPENS and the
-  object stays where the hand left it, it is "put" / "place" (object
-  arrives at a location). Decide by which way the object moves, not by
-  where the hand ends up.
-- POUR vs PUT: only use "pour" when the source is tilted and contents
-  flow out; moving a full container without tilting is "put"/"place".
+Timing:
+- Use the burned-in timestamps to set start and end. Boundaries should
+  land on or near a printed time, and every [start, end] must lie within
+  [0.0, {episode_duration}] seconds, be non-overlapping, and cover the
+  episode in order.
+- Emit at most {max_steps} segments.
 
 Output strictly valid JSON of shape:
 
   {{
     "subtasks": [
-      {{"text": "<short imperative verb phrase>", "start": <float>, "end": <float>}},
+      {{"text": "<short imperative action label>", "start": <float>, "end": <float>}},
       ...
     ]
   }}

src/lerobot/annotations/steerable_pipeline/vlm_client.py

@@ -285,6 +285,8 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
             "max_tokens": max_tok,
             "temperature": temp,
         }
+        if config.reasoning_effort:
+            kwargs["reasoning_effort"] = config.reasoning_effort
         extra_body: dict[str, Any] = {}
         if send_mm_kwargs and mm_kwargs:
             extra_body["mm_processor_kwargs"] = {**mm_kwargs, "do_sample_frames": True}
@@ -296,7 +298,13 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
             chosen = clients[rr_counter["i"] % len(clients)]
             rr_counter["i"] += 1
         response = chosen.chat.completions.create(**kwargs)
-        return response.choices[0].message.content or ""
+        # Some OpenAI-compatible servers can return a choice with no message
+        # (safety filter, or a "thinking" model that spends the whole budget
+        # before emitting content). Treat that as an empty reply so the
+        # JSON-retry path handles it instead of crashing the run.
+        choice = response.choices[0] if response.choices else None
+        message = choice.message if choice is not None else None
+        return (message.content if message is not None else None) or ""
 
     def _gen(batch: Sequence[Sequence[dict[str, Any]]], max_tok: int, temp: float) -> list[str]:
         if len(batch) <= 1 or config.client_concurrency <= 1:

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/policies/__init__.py

@@ -32,6 +32,7 @@ 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
@@ -57,6 +58,7 @@ __all__ = [
     "PI05Config",
     "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/evo1/internvl3_embedder.py

@@ -27,9 +27,11 @@ from lerobot.utils.import_utils import _transformers_available, require_package
 
 if TYPE_CHECKING or _transformers_available:
     from transformers import AutoModel, AutoTokenizer
+    from transformers.utils import is_flash_attn_2_available
 else:
     AutoModel = None
     AutoTokenizer = None
+    is_flash_attn_2_available = None
 
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
@@ -135,9 +137,13 @@ class InternVL3Embedder(nn.Module):
                 raise ValueError(f"Unsupported EVO1 vlm_dtype '{model_dtype}'") from exc
         self.model_dtype = model_dtype
 
-        attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager"
+        attn_implementation = (
+            "flash_attention_2" if (use_flash_attn and is_flash_attn_2_available()) else "eager"
+        )
         if use_flash_attn and attn_implementation == "eager":
-            logger.warning("flash_attn is not installed. Falling back to eager attention.")
+            logger.warning(
+                "Flash Attention 2 is unavailable on this runtime. Falling back to eager attention."
+            )
 
         self.model = AutoModel.from_pretrained(
             model_name,
@@ -359,11 +365,3 @@ class InternVL3Embedder(nn.Module):
     @property
     def device(self) -> torch.device:
         return next(self.model.parameters()).device
-
-
-def _flash_attn_available() -> bool:
-    try:
-        import flash_attn  # noqa: F401
-    except ModuleNotFoundError:
-        return False
-    return True

src/lerobot/policies/evo1/processor_evo1.py

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

src/lerobot/policies/factory.py

@@ -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,100 +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 == "gaussian_actor":
-        from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
-
-        return GaussianActorPolicy
-    elif name == "smolvla":
-        from .smolvla.modeling_smolvla import SmolVLAPolicy
-
-        return SmolVLAPolicy
-    elif name == "groot":
-        from .groot.modeling_groot import GrootPolicy
-
-        return GrootPolicy
-    elif name == "xvla":
-        from .xvla.modeling_xvla import XVLAPolicy
-
-        return XVLAPolicy
-    elif name == "wall_x":
-        from .wall_x.modeling_wall_x import WallXPolicy
-
-        return WallXPolicy
-    elif name == "eo1":
-        from .eo1.modeling_eo1 import EO1Policy
-
-        return EO1Policy
-    elif name == "molmoact2":
-        from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
-
-        return MolmoAct2Policy
-    elif name == "vla_jepa":
-        from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
-
-        return VLAJEPAPolicy
-    elif name == "lingbot_va":
-        from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
-
-        return LingBotVAPolicy
-    elif name == "fastwam":
-        from .fastwam.modeling_fastwam import FastWAMPolicy
-
-        return FastWAMPolicy
-    elif name == "evo1":
-        from .evo1.modeling_evo1 import Evo1Policy
-
-        return Evo1Policy
-    else:
-        try:
     return _get_policy_cls_from_policy_name(name=name)
-        except Exception as e:
-            raise ValueError(f"Policy type '{name}' is not available.") from e
 
 
 def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
@@ -192,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", "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:
@@ -203,48 +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 == "gaussian_actor":
-        return GaussianActorConfig(**kwargs)
-    elif policy_type == "smolvla":
-        return SmolVLAConfig(**kwargs)
-    elif policy_type == "groot":
-        return GrootConfig(**kwargs)
-    elif policy_type == "xvla":
-        return XVLAConfig(**kwargs)
-    elif policy_type == "wall_x":
-        return WallXConfig(**kwargs)
-    elif policy_type == "eo1":
-        return EO1Config(**kwargs)
-    elif policy_type == "molmoact2":
-        return MolmoAct2Config(**kwargs)
-    elif policy_type == "vla_jepa":
-        return VLAJEPAConfig(**kwargs)
-    elif policy_type == "lingbot_va":
-        return LingBotVAConfig(**kwargs)
-    elif policy_type == "fastwam":
-        return FastWAMConfig(**kwargs)
-    elif policy_type == "evo1":
-        return Evo1Config(**kwargs)
-    else:
     try:
         config_cls = PreTrainedConfig.get_choice_class(policy_type)
-            return config_cls(**kwargs)
     except Exception as e:
         raise ValueError(f"Policy type '{policy_type}' is not available.") from e
+    return config_cls(**kwargs)
 
 
 class ProcessorConfigKwargs(TypedDict, total=False):
@@ -298,8 +168,7 @@ def make_pre_post_processors(
         A tuple containing the input (pre-processor) and output (post-processor) pipelines.
 
     Raises:
-        NotImplementedError: If a processor factory is not implemented for the given
-            policy configuration type.
+        ValueError: If no processor factory exists for the given policy configuration type.
     """
     if pretrained_path:
         if isinstance(policy_cfg, GrootConfig):
@@ -351,167 +220,14 @@ def make_pre_post_processors(
             )
         return preprocessor, postprocessor
 
-    # Create a new processor based on policy type
-    if isinstance(policy_cfg, TDMPCConfig):
-        from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
-
-        processors = make_tdmpc_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, DiffusionConfig):
-        from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
-
-        processors = make_diffusion_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, ACTConfig):
-        from .act.processor_act import make_act_pre_post_processors
-
-        processors = make_act_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, MultiTaskDiTConfig):
-        from .multi_task_dit.processor_multi_task_dit import (
-            make_multi_task_dit_pre_post_processors,
-        )
-
-        processors = make_multi_task_dit_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, VQBeTConfig):
-        from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
-
-        processors = make_vqbet_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, PI0Config):
-        from .pi0.processor_pi0 import make_pi0_pre_post_processors
-
-        processors = make_pi0_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, PI05Config):
-        from .pi05.processor_pi05 import make_pi05_pre_post_processors
-
-        processors = make_pi05_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, GaussianActorConfig):
-        from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
-
-        processors = make_gaussian_actor_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, SmolVLAConfig):
-        from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
-
-        processors = make_smolvla_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, GrootConfig):
-        from .groot.processor_groot import make_groot_pre_post_processors
-
-        processors = make_groot_pre_post_processors(
+    # Create new processors from the policy config, resolving the per-policy factory
+    # function by naming convention (lazy import keeps optional dependencies optional).
+    return _make_processors_from_policy_config(
         config=policy_cfg,
         dataset_stats=kwargs.get("dataset_stats"),
         dataset_meta=kwargs.get("dataset_meta"),
     )
 
-    elif isinstance(policy_cfg, XVLAConfig):
-        from .xvla.processor_xvla import (
-            make_xvla_pre_post_processors,
-        )
-
-        processors = make_xvla_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, WallXConfig):
-        from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
-
-        processors = make_wall_x_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, EO1Config):
-        from .eo1.processor_eo1 import make_eo1_pre_post_processors
-
-        processors = make_eo1_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-    elif isinstance(policy_cfg, Evo1Config):
-        from .evo1.processor_evo1 import make_evo1_pre_post_processors
-
-        processors = make_evo1_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, MolmoAct2Config):
-        from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
-
-        processors = make_molmoact2_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-            dataset_meta=kwargs.get("dataset_meta"),
-        )
-
-    elif isinstance(policy_cfg, VLAJEPAConfig):
-        from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
-
-        processors = make_vla_jepa_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, LingBotVAConfig):
-        from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
-
-        processors = make_lingbot_va_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    elif isinstance(policy_cfg, FastWAMConfig):
-        from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
-
-        processors = make_fastwam_pre_post_processors(
-            config=policy_cfg,
-            dataset_stats=kwargs.get("dataset_stats"),
-        )
-
-    else:
-        try:
-            processors = _make_processors_from_policy_config(
-                config=policy_cfg,
-                dataset_stats=kwargs.get("dataset_stats"),
-            )
-        except Exception as e:
-            raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
-
-    return processors
-
 
 def make_policy(
     cfg: PreTrainedConfig,
@@ -654,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.
@@ -683,22 +401,39 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
         "configuration_", "modeling_"
     )  # e.g., configuration_diffusion -> modeling_diffusion
 
+    try:
         module = importlib.import_module(module_path)
-    policy_cls = getattr(module, cls_name)
+    except ModuleNotFoundError as e:
+        if e.name == module_path:
+            # The modeling_* module itself does not exist for this policy type. A missing
+            # optional dependency inside an existing module propagates unchanged instead,
+            # so its actionable install hint stays visible.
+            raise ValueError(f"Policy class for '{name}' is not implemented.") from e
+        raise
+    policy_cls = getattr(module, cls_name, None)
+    if policy_cls is None:
+        raise ValueError(
+            f"Policy class '{cls_name}' not found in '{module_path}'. "
+            f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
+        )
     return policy_cls
 
 
 def _make_processors_from_policy_config(
     config: PreTrainedConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    dataset_meta: Any | None = None,
 ) -> tuple[Any, Any]:
     """Create pre- and post-processors from a policy configuration using dynamic imports.
 
-    This is used as a helper function to import processor factories from 3rd party lerobot plugins.
+    Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
+    by naming convention. Works for built-in policies and 3rd party lerobot plugins.
 
     Args:
         config: The policy configuration object.
         dataset_stats: Dataset statistics for normalization.
+        dataset_meta: Dataset metadata, forwarded only to factories that declare a
+            ``dataset_meta`` parameter (e.g. groot, molmoact2).
     Returns:
         A tuple containing the input (pre-processor) and output (post-processor) pipelines.
     """
@@ -711,6 +446,19 @@ def _make_processors_from_policy_config(
     logging.debug(
         f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
     )
+    try:
         module = importlib.import_module(module_path)
-    function = getattr(module, function_name)
-    return function(config, dataset_stats=dataset_stats)
+    except ModuleNotFoundError as e:
+        if e.name == module_path:
+            # The processor_* module itself does not exist for this policy type. A missing
+            # optional dependency inside an existing module propagates unchanged instead,
+            # so its actionable install hint stays visible.
+            raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
+        raise
+    function = getattr(module, function_name, None)
+    if function is None:
+        raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
+    call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
+    if "dataset_meta" in inspect.signature(function).parameters:
+        call_kwargs["dataset_meta"] = dataset_meta
+    return function(config, **call_kwargs)

src/lerobot/policies/fastwam/processor_fastwam.py

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

src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py

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

src/lerobot/policies/lingbot_va/processor_lingbot_va.py

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

src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py

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

src/lerobot/policies/pi0/processor_pi0.py

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

src/lerobot/policies/pi05/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/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
 
@@ -135,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,14 +137,10 @@ def make_pi0_fast_pre_post_processors(
     # before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
     # continues to receive normalized state in [-1, 1] as expected.
     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,26 +154,13 @@ def make_pi0_fast_pre_post_processors(
             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/pretrained.py

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

src/lerobot/policies/smolvla/processor_smolvla.py

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

src/lerobot/policies/tdmpc/processor_tdmpc.py

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

src/lerobot/policies/vla_jepa/processor_vla_jepa.py

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

src/lerobot/policies/vqbet/processor_vqbet.py

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

src/lerobot/policies/wall_x/configuration_wall_x.py

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

src/lerobot/policies/wall_x/modeling_wall_x.py

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

src/lerobot/policies/wall_x/processor_wall_x.py

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

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

@@ -0,0 +1,45 @@
+#!/usr/bin/env python
+
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .configuration_qwen2_5_vl import (
+    Qwen2_5_VLConfig,
+    Qwen2_5_VLTextConfig,
+    Qwen2_5_VLVisionConfig,
+)
+from .qwen2_5_vl_moe import (
+    BlockSparseMLP,
+    Qwen2_5_VLACausalLMOutputWithPast,
+    Qwen2_5_VLDecoderLayer_with_MoE,
+    Qwen2_5_VLMoEModel,
+    SparseMoeBlock,
+)
+from .vision_attention import (
+    WallXVisionAttention,
+    configure_wall_x_vision_attention,
+)
+
+__all__ = [
+    "BlockSparseMLP",
+    "Qwen2_5_VLACausalLMOutputWithPast",
+    "Qwen2_5_VLConfig",
+    "Qwen2_5_VLDecoderLayer_with_MoE",
+    "Qwen2_5_VLMoEModel",
+    "Qwen2_5_VLTextConfig",
+    "Qwen2_5_VLVisionConfig",
+    "SparseMoeBlock",
+    "WallXVisionAttention",
+    "configure_wall_x_vision_attention",
+]

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

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

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

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

src/lerobot/policies/wall_x/utils.py

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

src/lerobot/policies/xvla/processor_xvla.py

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

src/lerobot/processor/__init__.py

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

src/lerobot/processor/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/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,89 @@
+# Unitree G1 — SONIC encoder/decoder whole-body control
+
+This package runs NVIDIA's **SONIC** encoder/decoder on the Unitree G1, in MuJoCo
+simulation or on real hardware, driven by a dense **34-D whole-body command** (the
+OpenHLM / pi0.5 action layout). It is a pure-Python/ONNX reimplementation of the
+reference-tracking half of the SONIC deploy stack (no `gear_sonic`/torch dependency, and
+no motion planner): the encoder compresses a reference motion window into a latent token
+and the decoder maps that token + proprioception history into 50 Hz joint-position
+targets for the robot's PD controller.
+
+## Controllers
+
+Selected with `--robot.controller=<ClassName>`:
+
+| Controller                     | Purpose                                                      |
+| ------------------------------ | ------------------------------------------------------------ |
+| `SonicWholeBodyController`     | SONIC encoder/decoder driven by a 34-D OpenHLM/pi0.5 command |
+| `GrootLocomotionController`    | GR00T locomotion policy                                      |
+| `HolosomaLocomotionController` | Holosoma locomotion policy                                   |
+
+The rest of this document covers the SONIC whole-body path.
+
+Each tick the `SonicWholeBodyController` takes a 34-D command (`wb.0.pos … wb.33.pos`) in the OpenHLM
+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)]
+```
+
+The 29 joint targets become the SONIC encode-mode-0 reference (accumulated into a rolling
+50-frame trajectory with finite-difference velocities so the encoder's lookahead sees a
+real motion sequence), the root roll/pitch set the anchor orientation, and the two grip
+scalars can drive the Dex3 hands (see below). On startup the controller **interpolates**
+from the robot's measured pose into the policy's commanded target over ~3 s (no snap).
+
+## Requirements
+
+- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended). Verify with:
+  ```bash
+  python -c "import onnxruntime as ort; print(ort.get_available_providers())"
+  ```
+- `mujoco` for simulation (`is_simulation=True`).
+- The SONIC encoder/decoder ONNX models download automatically from the
+  `nvidia/GEAR-SONIC` Hub repo.
+
+## Running a rollout
+
+Drive the G1 with a 34-D VLA policy (OpenHLM / pi0.5) via `lerobot-rollout`:
+
+```bash
+lerobot-rollout \
+  --strategy.type=base \
+  --policy.path=<pi05_openhlm_dir> \
+  --robot.type=unitree_g1 \
+  --robot.controller=SonicWholeBodyController \
+  --robot.is_simulation=true \
+  --robot.publish_hands=true \
+  --task="<language instruction>" \
+  --duration=45 --device=cuda
+```
+
+### Cameras
+
+Image-conditioned policies need camera frames. Two options are available without live
+cameras:
+
+- **Black frames**: `--robot.empty_cameras='[base, left_wrist, right_wrist]'`.
+- **Replay a recorded episode** as the camera feed:
+  ```bash
+  --robot.replay_camera_parquet=<episode.parquet> \
+  --robot.replay_camera_map='{base: head_image_left, left_wrist: left_wrist_image, right_wrist: right_wrist_image}'
+  ```
+
+### Hands (Dex3)
+
+`--robot.publish_hands=true` publishes `rt/dex3/{left,right}/cmd` from the two grip
+scalars (`wb.7.pos` left, `wb.15.pos` right). The scalar is interpolated between
+`hand_open_grip_value` (default 1.0 = open) and `hand_closed_grip_value` (default 0.0 =
+closed) and scaled onto `hand_closed_pose` (7 joints:
+`thumb_0, thumb_1, thumb_2, middle_0, middle_1, index_0, index_1`). Flip the signs in
+`hand_closed_pose` if the fingers curl the wrong way, or raise `hand_kp` for a firmer
+grip.
+
+## Observation state
+
+When the whole-body controller is active the robot exposes a 34-D proprio state
+(`wb_state.0.pos … wb_state.33.pos`) in the same OpenHLM layout as the action, which the
+rollout aggregates into `observation.state` for the policy.

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -62,12 +62,76 @@ class UnitreeG1Config(RobotConfig):
     # Socket config for ZMQ bridge
     robot_ip: str = "192.168.123.164"  # default G1 IP
 
+    # Run the locomotion / whole-body controller ONBOARD the robot (policy on the G1
+    # itself, against local DDS at full rate) instead of on the laptop over the ZMQ
+    # socket bridge. In this mode the robot object uses the real Unitree SDK channels
+    # and expects high-level actions (arm targets + joystick axes, or 64-D SONIC
+    # tokens) fed via send_action -- e.g. by run_g1_onboard.py, which receives them
+    # from the laptop over ZMQ. Mutually exclusive with is_simulation.
+    onboard: bool = False
+    # DDS network interface for onboard mode (None = SDK default, matching
+    # run_g1_server.py's ChannelFactoryInitialize(0)).
+    dds_interface: str | None = None
+    # Onboard sub-flags. On a real G1 both are True: the built-in motion services
+    # must be released before we can write lowcmd, and locomotion axes are read from
+    # the physical wireless remote. Against a DDS sim neither applies (no
+    # MotionSwitcher, no physical remote), so set both False so the controller takes
+    # its locomotion axes purely from send_action (ZMQ) input.
+    release_motion_control: bool = True
+    physical_remote: bool = True
+
     # 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
+    # When False, connect() does not start the background controller thread, so a
+    # caller can drive the controller synchronously (one decode per fed action),
+    # reproducing the deploy's single 50Hz control clock for faithful replay.
+    run_controller_thread: bool = True
+    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
+
+    # Token-output VLA interface for the SONIC decoder. When True (and the controller
+    # is ``SonicWholeBodyController``), the robot advertises a 64-D latent-token action
+    # space (``motion_token.{i}.pos``) instead of the 34-D whole-body command, and
+    # exposes the last commanded token as a 64-D ``observation.state``
+    # (``motion_token_state.{i}.pos``). This lets ``lerobot-rollout`` drive a policy
+    # that was trained with 64-D SONIC motion tokens as both state and action
+    # (e.g. nepyope/sonic_walk): the decoder consumes the token directly, encoder
+    # bypassed. Ignored unless a SONIC whole-body controller is active.
+    sonic_token_action: bool = False
+
     # 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

@@ -0,0 +1,28 @@
+#!/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.
+
+"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
+
+from .gr00t_locomotion import GrootLocomotionController
+from .holosoma_locomotion import HolosomaLocomotionController
+from .sonic_whole_body import SonicRuntime, SonicWholeBodyController
+
+__all__ = [
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicRuntime",
+    "SonicWholeBodyController",
+]

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__)
 
 
@@ -68,9 +77,15 @@ def load_groot_policies(
         filename="GR00T-WholeBodyControl-Walk.onnx",
     )
 
-    # Load ONNX policies
-    policy_balance = ort.InferenceSession(balance_path)
-    policy_walk = ort.InferenceSession(walk_path)
+    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
+    # background thread alongside the (torch) upper-body policy, IK and sim; letting
+    # ORT grab every core starves those and makes the whole rollout stutter. These
+    # are small MLPs, so 1 thread is both enough and lowest-latency.
+    from .sonic_pipeline import make_ort_session_options
+
+    so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
+    policy_balance = ort.InferenceSession(balance_path, sess_options=so)
+    policy_walk = ort.InferenceSession(walk_path, sess_options=so)
 
     logger.info("GR00T policies loaded successfully")
 
@@ -83,6 +98,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()
 
@@ -196,6 +212,16 @@ class GrootLocomotionController:
         # Transform action back to target joint positions
         target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
 
+        # Waist override: an external upper-body IK can command the 3 waist joints
+        # (indices 12/13/14) via ``kWaist{Yaw,Roll,Pitch}.q`` in the action dict. When
+        # present, we substitute the balance policy's waist target so the torso tracks
+        # the IK while the policy keeps only the legs balanced. Single-publisher stays
+        # intact (this thread still owns joints 0-14).
+        for idx in (G1_29_JointIndex.kWaistYaw, G1_29_JointIndex.kWaistRoll, G1_29_JointIndex.kWaistPitch):
+            key = f"{idx.name}.q"
+            if key in action and action[key] is not None:
+                target_dof_pos_15[idx.value] = float(action[key])
+
         # Build action dict
         action_dict = {}
         for i in range(15):

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

@@ -0,0 +1,670 @@
+#!/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 encoder/decoder pipeline for the Unitree G1 whole-body controller.
+
+Pure-Python/ONNX re-implementation of the reference-tracking half of NVIDIA's SONIC
+deploy stack (mirrors ``g1_deploy_onnx_ref.cpp``). Given a reference motion buffer
+(joint targets + body orientation per frame) it produces 50 Hz joint-position targets
+for the robot's PD controller. The upstream *motion planner* is intentionally absent:
+here the reference is supplied directly by the caller (e.g. a 34-D OpenHLM / pi0.5 VLA
+command per tick, in ``sonic_whole_body.py``).
+
+Two cooperating ONNX models:
+  * **encoder** – compresses the reference window into a 64-D latent ``token``
+    (refreshed every ``ENCODER_UPDATE_EVERY`` ticks).
+  * **decoder** – every tick, maps the token + recent proprioception history to a
+    residual action that is scaled and added to ``DEFAULT_ANGLES``.
+
+Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order)
+and **MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` convert
+between them. Quaternions are scalar-first ``(w, x, y, z)``.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from lerobot.utils.import_utils import _onnxruntime_available
+
+from ..g1_utils import (
+    ISAACLAB_TO_MUJOCO,
+    MUJOCO_TO_ISAACLAB,
+    G1_29_JointIndex,
+    get_gravity_orientation,
+)
+
+if TYPE_CHECKING or _onnxruntime_available:
+    import onnxruntime as ort
+else:
+    ort = None
+
+logger = logging.getLogger(__name__)
+
+# ── Constants ────────────────────────────────────────────────────────────────
+# Robot/motor physical constants and the joint-order permutation tables. All
+# 29-vectors are in IsaacLab joint order unless the name says ``_MUJOCO``.
+
+# Nominal standing pose (rad), 29 joints in IsaacLab order. Actions are residuals
+# added on top of this; also used as the planner/encoder standing reference.
+DEFAULT_ANGLES = np.array(
+    [
+        -0.312,
+        0.0,
+        0.0,
+        0.669,
+        -0.363,
+        0.0,
+        -0.312,
+        0.0,
+        0.0,
+        0.669,
+        -0.363,
+        0.0,
+        0.0,
+        0.0,
+        0.0,
+        0.2,
+        0.2,
+        0.0,
+        0.6,
+        0.0,
+        0.0,
+        0.0,
+        0.2,
+        -0.2,
+        0.0,
+        0.6,
+        0.0,
+        0.0,
+        0.0,
+    ],
+    dtype=np.float32,
+)
+
+# Per-motor-type parameters used to derive action scaling and PD gains. Keys are
+# Unitree motor model names; ARMATURE = rotor inertia, EFFORT = torque limit (N·m).
+NATURAL_FREQ = 10.0 * 2.0 * np.pi  # target closed-loop stiffness bandwidth (rad/s)
+ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
+EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
+
+
+def _action_scale(k):
+    """Per-motor residual-action scale (maps policy output to joint-angle delta)."""
+    return 0.25 * EFFORT[k] / (ARMATURE[k] * NATURAL_FREQ**2)
+
+
+# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of
+# truth for both ACTION_SCALE and compute_kp_kd().
+MOTOR_MODELS = (
+    ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+    + ["7520_14", "5020", "5020"]
+    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
+)
+ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32)  # (29,) IsaacLab order
+
+CONTROL_DT = 0.02  # 50 Hz control period (s)
+DEFAULT_HEIGHT = 0.788740  # nominal pelvis height (m)
+TOKEN_DIM = 64  # encoder latent size
+ENCODER_UPDATE_EVERY = 5  # refresh the encoder token every N ticks (decoder runs every tick)
+DEBUG_PRINT_EVERY = 100  # ticks between debug prints
+
+
+def _to_mujoco(a):
+    """Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
+
+    NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
+    permutation direction throughout this module are a fixed convention validated
+    against the deployed SONIC ONNX policy (the encoder/decoder consume vectors in
+    this order). Do not "correct" the table or rename toward the opposite direction
+    without re-validating on hardware — the labels are historical, the ordering is
+    load-bearing.
+    """
+    return a[MUJOCO_TO_ISAACLAB]
+
+
+DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
+ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
+
+# Joint-index subsets (IsaacLab order) used to slice encoder observations.
+LOWER_BODY_IL = np.array([0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18], dtype=np.int32)  # 12 leg joints
+WRIST_IL = np.array([23, 24, 25, 26, 27, 28], dtype=np.int32)  # 6 wrist joints
+VR_TARGET_DEF = np.zeros(9, dtype=np.float32)  # 3-point VR position targets (mode 1)
+VR_ORN_DEF = np.array([1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0], dtype=np.float32)  # VR orn targets (mode 1)
+SMPL_DEF = np.zeros(720, dtype=np.float32)  # SMPL whole-body window default (mode 2)
+
+# ── PD gains ─────────────────────────────────────────────────────────────────
+
+
+def compute_kp_kd():
+    """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth.
+
+    Ankle and waist joints get a x2 factor for extra stiffness. Returns two
+    (29,) float32 arrays in IsaacLab joint order.
+    """
+
+    def s(k):
+        return ARMATURE[k] * NATURAL_FREQ**2
+
+    def d(k):
+        return 2.0 * 2.0 * ARMATURE[k] * NATURAL_FREQ
+
+    _double = {4, 5, 10, 11, 13, 14}  # ankle + waist indices with factor 2
+    kp = np.array([2 * s(k) if i in _double else s(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32)
+    kd = np.array([2 * d(k) if i in _double else d(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32)
+    return kp, kd
+
+
+_kp_kd = compute_kp_kd  # backward-compatible alias
+
+
+# ── Quaternion helpers ────────────────────────────────────────────────────────
+# All quaternions are scalar-first (w, x, y, z). "heading" = yaw-only quaternion.
+
+
+def quat_conj(q):
+    """Quaternion conjugate (inverse for unit quaternions)."""
+    return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float32)
+
+
+def quat_mul(q1, q2):
+    """Hamilton product ``q1 ⊗ q2``."""
+    w1, x1, y1, z1 = q1
+    w2, x2, y2, z2 = q2
+    return np.array(
+        [
+            w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
+            w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
+            w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
+            w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
+        ],
+        dtype=np.float32,
+    )
+
+
+def quat_to_6d(q):
+    """Quaternion → 6-D rotation representation (first two rotated basis rows)."""
+    w, x, y, z = q
+    return np.array(
+        [
+            1 - 2 * (y * y + z * z),
+            2 * (x * y - z * w),
+            2 * (x * y + z * w),
+            1 - 2 * (x * x + z * z),
+            2 * (x * z - y * w),
+            2 * (y * z + x * w),
+        ],
+        dtype=np.float32,
+    )
+
+
+def calc_heading(q):
+    """Extract the yaw (heading) angle in radians from a quaternion."""
+    w, x, y, z = q
+    return float(np.arctan2(2 * (x * y + w * z), 1 - 2 * (y * y + z * z)))
+
+
+def heading_quat(q, sign=1.0):
+    """Yaw-only quaternion for ``q``'s heading (``sign=-1`` gives its inverse)."""
+    a = sign * calc_heading(q) / 2.0
+    return np.array([np.cos(a), 0, 0, np.sin(a)], dtype=np.float64)
+
+
+def heading_quat_inv(q):
+    """Inverse yaw-only quaternion for ``q``'s heading."""
+    return heading_quat(q, -1.0)
+
+
+def quat_slerp(q0, q1, t):
+    """Spherical linear interpolation between two quaternions (scalar ``t``)."""
+    q0 = q0 / (np.linalg.norm(q0) + 1e-12)
+    q1 = q1 / (np.linalg.norm(q1) + 1e-12)
+    dot = float(np.dot(q0, q1))
+    if dot < 0:
+        q1, dot = -q1, -dot
+    dot = min(dot, 1.0)
+    if dot > 0.9995:
+        r = q0 + t * (q1 - q0)
+        return r / (np.linalg.norm(r) + 1e-12)
+    th = np.arccos(dot)
+    st = np.sin(th)
+    return (np.sin((1 - t) * th) / st) * q0 + (np.sin(t * th) / st) * q1
+
+
+def quat_slerp_batch(q0, q1, t):
+    """Vectorized slerp over arrays of quaternions with a per-row parameter ``t``."""
+    q0 = q0 / (np.linalg.norm(q0, axis=1, keepdims=True) + 1e-12)
+    q1 = q1 / (np.linalg.norm(q1, axis=1, keepdims=True) + 1e-12)
+    dot = np.sum(q0 * q1, axis=1)
+    neg = dot < 0
+    q1 = q1.copy()
+    q1[neg] = -q1[neg]
+    dot[neg] = -dot[neg]
+    dot = np.clip(dot, -1, 1)
+    lin = dot > 0.9995
+    th = np.arccos(dot)
+    st = np.where(np.sin(th) == 0, 1, np.sin(th))
+    c0 = np.sin((1 - t) * th) / st
+    c1 = np.sin(t * th) / st
+    c0[lin] = 1 - t[lin]
+    c1[lin] = t[lin]
+    r = c0[:, None] * q0 + c1[:, None] * q1
+    return r / (np.linalg.norm(r, axis=1, keepdims=True) + 1e-12)
+
+
+def ort_providers(force_cpu: bool = False) -> list[str]:
+    """Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed)."""
+    avail = ort.get_available_providers()
+    if not force_cpu and "CUDAExecutionProvider" in avail:
+        return ["CUDAExecutionProvider", "CPUExecutionProvider"]
+    return ["CPUExecutionProvider"]
+
+
+def make_ort_session_options(intra_op_num_threads: int | None = None,
+                             inter_op_num_threads: int | None = None):
+    """Build ONNX Runtime SessionOptions (quiet logging).
+
+    Pass thread counts to cap ORT's CPU pool. These tiny MLP policies are latency-
+    bound, not throughput-bound, so letting ORT grab every core just starves the
+    real-time control loop / torch policy / IK solver and causes stutter. 1 intra +
+    1 inter thread is plenty and lowest-latency for a per-step MLP inference.
+    """
+    so = ort.SessionOptions()
+    so.log_severity_level = 3
+    if intra_op_num_threads is not None:
+        so.intra_op_num_threads = intra_op_num_threads
+    if inter_op_num_threads is not None:
+        so.inter_op_num_threads = inter_op_num_threads
+    return so
+
+
+# ── Encoder / Decoder ─────────────────────────────────────────────────────────
+
+
+class StandingEncoderDecoder:
+    """Runs the encoder + decoder ONNX models and owns the proprioception history.
+
+    Each tick it appends the latest robot state to 10-frame history buffers, builds
+    the encoder observation (1762-D, layout depends on ``encode_mode``) to refresh
+    the 64-D ``token``, then builds the decoder observation (994-D) and maps
+    ``token + history`` to a residual action added onto ``DEFAULT_ANGLES``.
+
+    ``PlannerController`` subclasses this to source the reference from a live,
+    planner-generated motion buffer instead of a fixed standing pose.
+    """
+
+    def __init__(self, encoder, decoder):
+        self.encoder, self.decoder = encoder, decoder
+        self.encoder_input = encoder.get_inputs()[0].name
+        self.decoder_input = decoder.get_inputs()[0].name
+        enc_dim = int(encoder.get_inputs()[0].shape[1])
+        dec_dim = int(decoder.get_inputs()[0].shape[1])
+        if enc_dim != 1762 or dec_dim != 994:
+            raise RuntimeError(f"Unexpected dims encoder={enc_dim}, decoder={dec_dim}")
+        self.token = np.zeros(TOKEN_DIM, np.float32)
+        self.last_action_mj = np.zeros(29, np.float32)
+        self.h_q_mj = [np.zeros(29, np.float32)] * 10
+        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
+        self.h_ang = [np.zeros(3, np.float32)] * 10
+        self.h_act_mj = [np.zeros(29, np.float32)] * 10
+        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
+        self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
+        self._heading_init = False
+        self.encode_mode = 0
+        self.vr_3point_local_target = VR_TARGET_DEF.copy()
+        self.vr_3point_local_orn_target = VR_ORN_DEF.copy()
+        self.smpl_joints_10frame_step1 = SMPL_DEF.copy()
+        # Optional per-frame SMPL root orientation (wxyz) for the mode-2 anchor.
+        # When None, the anchor falls back to the planner reference body quat.
+        self.smpl_root_quat = None
+        self.set_zero_reference()
+
+    def reset(self):
+        """Clear the token, 10-frame proprioception history and heading init.
+
+        ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new
+        episode are not contaminated by the previous episode's state.
+        """
+        self.token = np.zeros(TOKEN_DIM, np.float32)
+        self.last_action_mj = np.zeros(29, np.float32)
+        self.h_q_mj = [np.zeros(29, np.float32)] * 10
+        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
+        self.h_ang = [np.zeros(3, np.float32)] * 10
+        self.h_act_mj = [np.zeros(29, np.float32)] * 10
+        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
+        self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
+        self._heading_init = False
+
+    def update_history(self, q, dq, ang, quat):
+        """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
+        quat = quat / (np.linalg.norm(quat) + 1e-8)
+        q_mj = _to_mujoco(q)
+        dq_mj = _to_mujoco(dq)
+        self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + self.h_q_mj[:-1]
+        self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
+        self.h_ang = [ang.copy()] + self.h_ang[:-1]
+        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
+        self.h_quat = [quat.copy()] + self.h_quat[:-1]
+        if not self._heading_init:
+            self.init_base_quat = quat.copy()
+            self._heading_init = True
+
+    def _heading_quat(self, q):
+        h = calc_heading(q) / 2.0
+        return np.array([np.cos(h), 0, 0, np.sin(h)], np.float32)
+
+    def _heading_quat_inv(self, q):
+        h = calc_heading(q) / 2.0
+        return np.array([np.cos(-h), 0, 0, np.sin(-h)], np.float32)
+
+    def _anchor_6d(self, base_quat, ref_quat=None):
+        """6-D orientation error between the robot base and the (heading-aligned) reference."""
+        if ref_quat is None:
+            ref_quat = self.init_ref_quat
+        delta = quat_mul(self._heading_quat(self.init_base_quat), self._heading_quat_inv(self.init_ref_quat))
+        new_ref = quat_mul(delta, ref_quat)
+        return quat_to_6d(quat_mul(quat_conj(base_quat), new_ref))
+
+    def set_zero_reference(self):
+        """Initialize the reference to a single standing frame (used before a plan exists)."""
+        self.motion_joint_positions = [ENCODER_STANDING_REF.copy()]
+        self.motion_joint_velocities = [np.zeros(29, np.float32)]
+        self.motion_body_quats = [np.array([1, 0, 0, 0], np.float32)]
+        self.motion_body_z = [DEFAULT_HEIGHT]
+        self.motion_timesteps = 1
+        self.freeze_ref_frame = 0
+        self.init_ref_quat = self.motion_body_quats[0].copy()
+
+    def build_encoder_obs(self):
+        """Assemble the 1762-D encoder input; slot layout depends on ``encode_mode``.
+
+        mode 0 = locomotion (ref joint pos + anchor), 1 = 3-point VR teleop
+        (lower-body ref + VR targets), 2 = SMPL whole-body window + anchor/wrist.
+        """
+        obs = np.zeros(1762, np.float32)
+        obs[0] = float(self.encode_mode)
+        rf = min(self.freeze_ref_frame, self.motion_timesteps - 1)
+        ref_pos, ref_quat = self.motion_joint_positions[rf], self.motion_body_quats[rf]
+        if self.encode_mode == 0:
+            for f in range(10):
+                obs[4 + 29 * f : 4 + 29 * (f + 1)] = ref_pos
+                obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], ref_quat)
+        elif self.encode_mode == 1:
+            ref_lower = ref_pos[LOWER_BODY_IL]
+            for f in range(10):
+                obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
+            obs[901:910] = self.vr_3point_local_target
+            obs[910:922] = self.vr_3point_local_orn_target
+            obs[595:601] = self._anchor_6d(self.h_quat[0], ref_quat)
+        elif self.encode_mode == 2:
+            # Prefer the SMPL clip/stream root orientation for the anchor; fall
+            # back to the planner reference body quat when no root is provided.
+            anchor_ref = self.smpl_root_quat if self.smpl_root_quat is not None else ref_quat
+            obs[922:1642] = self.smpl_joints_10frame_step1
+            for f in range(10):
+                obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = self._anchor_6d(self.h_quat[0], anchor_ref)
+                obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = ref_pos[WRIST_IL]
+        else:
+            raise RuntimeError(f"Unsupported encoder mode: {self.encode_mode}")
+        return obs
+
+    def build_decoder_obs(self):
+        """Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
+        obs = np.zeros(994, np.float32)
+        off = 0
+        obs[off : off + 64] = self.token
+        off += 64
+        for h, sz in [
+            (list(reversed(self.h_ang)), 3),
+            (list(reversed(self.h_q_mj)), 29),
+            (list(reversed(self.h_dq_mj)), 29),
+            (list(reversed(self.h_act_mj)), 29),
+        ]:
+            for f in range(10):
+                obs[off : off + sz] = h[f]
+                off += sz
+        for q in reversed(self.h_quat):
+            obs[off : off + 3] = get_gravity_orientation(q)
+            off += 3
+        assert off == 994, f"Decoder obs mismatch: {off}"
+        return obs
+
+    def run_encoder(self):
+        """Run the encoder ONNX model and return the fresh 64-D token."""
+        return (
+            self.encoder.run(None, {self.encoder_input: self.build_encoder_obs().reshape(1, -1)})[0]
+            .squeeze()
+            .astype(np.float32)
+        )
+
+    def step(self, robot_obs, update_encoder, debug=False):
+        """One control tick: read robot obs, (optionally) re-encode, decode → joint targets.
+
+        Args:
+            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
+            update_encoder: refresh the token this tick (else reuse the cached one).
+            debug: print action/delta norms.
+
+        Returns:
+            dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
+        """
+        jnames = [m.name for m in G1_29_JointIndex]
+        q = np.array(
+            [
+                robot_obs.get(f"{n}.q", DEFAULT_ANGLES[m.value])
+                for m, n in zip(G1_29_JointIndex, jnames, strict=False)
+            ],
+            np.float32,
+        )
+        dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
+        quat = np.array(
+            [
+                robot_obs.get("imu.quat.w", 1),
+                robot_obs.get("imu.quat.x", 0),
+                robot_obs.get("imu.quat.y", 0),
+                robot_obs.get("imu.quat.z", 0),
+            ],
+            np.float32,
+        )
+        ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
+        self.update_history(q, dq, ang, quat)
+        if update_encoder:
+            self.token = self.run_encoder()
+        action_mj = (
+            self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
+            .squeeze()
+            .astype(np.float32)
+        )
+        self.last_action_mj = action_mj.copy()
+        target = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * ACTION_SCALE
+        if debug:
+            delta = target - q
+            logger.debug(
+                "token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
+                np.linalg.norm(self.token),
+                np.linalg.norm(action_mj),
+                np.max(np.abs(delta)),
+                np.sqrt(np.mean(delta**2)),
+            )
+        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
+
+
+class PlannerController(StandingEncoderDecoder):
+    """Encoder/decoder driven by a caller-supplied, rolling motion buffer.
+
+    Extends ``StandingEncoderDecoder`` so the reference comes from a motion buffer
+    (a lookahead window with per-frame velocities) instead of a single fixed pose,
+    and handles heading re-initialization on the first frame / after a reset.
+    ``motion_lock`` guards the buffer, which the whole-body controller rewrites each
+    tick from the incoming command. The class name is retained for continuity with
+    the SONIC reference; no motion planner is involved.
+    """
+
+    def __init__(self, encoder, decoder):
+        super().__init__(encoder, decoder)
+        self.ref_cursor = 0
+        self.motion_timesteps = 0
+        self.motion_joint_positions = np.zeros((1500, 29), np.float64)
+        self.motion_joint_velocities = np.zeros((1500, 29), np.float64)
+        self.motion_body_quats = np.zeros((1500, 4), np.float64)
+        self.motion_body_quats[:, 0] = 1.0
+        self.motion_body_pos = np.zeros((1500, 3), np.float64)
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
+        self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
+        self.delta_heading = 0.0
+        self.reinit_heading = False
+        self.playing = self.first_motion = False
+        self.motion_lock = threading.Lock()
+
+    def reset(self):
+        """Full reset: clear enc/dec state (super) plus the motion buffer and heading.
+
+        Forces a heading re-init on the next ``step`` so the reference frame is
+        re-latched to the post-reset robot orientation.
+        """
+        super().reset()
+        with self.motion_lock:
+            self.ref_cursor = 0
+            self.motion_timesteps = 0
+            self.motion_joint_positions[:] = 0.0
+            self.motion_joint_velocities[:] = 0.0
+            self.motion_body_quats[:] = 0.0
+            self.motion_body_quats[:, 0] = 1.0
+            self.motion_body_pos[:] = 0.0
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
+        self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
+        self.delta_heading = 0.0
+        self.first_motion = False
+        self.playing = False
+        self.reinit_heading = True
+
+    def _heading_apply_delta(self):
+        """Heading correction quaternion (init base-vs-ref heading + operator ``delta_heading``)."""
+        delta = quat_mul(
+            heading_quat(self.heading_init_base_quat).astype(np.float32),
+            heading_quat_inv(self.init_ref_quat).astype(np.float32),
+        )
+        if self.delta_heading:
+            h = self.delta_heading / 2.0
+            delta = quat_mul(np.array([np.cos(h), 0, 0, np.sin(h)], np.float32), delta)
+        return delta
+
+    def _anchor_6d(self, base_quat, ref_quat=None):
+        """6-D base-vs-reference orientation error, including the operator heading delta."""
+        if ref_quat is None:
+            ref_quat = self.init_ref_quat
+        new_ref = quat_mul(self._heading_apply_delta(), ref_quat.astype(np.float32))
+        return quat_to_6d(quat_mul(quat_conj(base_quat.astype(np.float32)), new_ref))
+
+    def build_encoder_obs(self):
+        """Encoder input sourced from the live motion buffer (mode 0/2), lock-protected."""
+        obs = np.zeros(1762, np.float32)
+        obs[0] = float(self.encode_mode)
+        with self.motion_lock:
+            if self.encode_mode == 2:
+                # SMPL whole-body imitation: the 720-dim SMPL window carries the
+                # target pose; the planner reference frame supplies anchor + wrist.
+                rf = min(self.ref_cursor, self.motion_timesteps - 1)
+                ref_pos = self.motion_joint_positions[rf].astype(np.float32)
+                ref_quat = self.motion_body_quats[rf].astype(np.float32)
+                # Prefer the SMPL clip/stream root orientation (if provided) so the
+                # anchor tracks the operator's/clip's heading; else planner ref.
+                if self.smpl_root_quat is not None:
+                    ref_quat = np.asarray(self.smpl_root_quat, np.float32)
+                anchor = self._anchor_6d(self.h_quat[0], ref_quat)
+                wrist = ref_pos[WRIST_IL]
+                obs[922:1642] = self.smpl_joints_10frame_step1
+                for f in range(10):
+                    obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = anchor
+                    obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = wrist
+                return obs
+            if self.encode_mode == 1:
+                # 3-point VR teleop: the upper body tracks the VR wrist/neck targets
+                # while the planner reference supplies the lower body + anchor. Lower
+                # body is per-frame (step 5) like mode 0; the VR targets are current.
+                rf = min(self.ref_cursor, self.motion_timesteps - 1)
+                obs[595:601] = self._anchor_6d(self.h_quat[0], self.motion_body_quats[rf].astype(np.float32))
+                for f in range(10):
+                    tf = min(
+                        self.ref_cursor + f * 5 if self.playing else self.ref_cursor,
+                        self.motion_timesteps - 1,
+                    )
+                    ref_lower = self.motion_joint_positions[tf].astype(np.float32)[LOWER_BODY_IL]
+                    obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
+                obs[901:910] = self.vr_3point_local_target
+                obs[910:922] = self.vr_3point_local_orn_target
+                return obs
+            for f in range(10):
+                tf = min(
+                    self.ref_cursor + f * 5 if self.playing else self.ref_cursor, self.motion_timesteps - 1
+                )
+                obs[4 + 29 * f : 4 + 29 * (f + 1)] = self.motion_joint_positions[tf].astype(np.float32)
+                if self.playing:
+                    obs[294 + 29 * f : 294 + 29 * (f + 1)] = self.motion_joint_velocities[tf].astype(
+                        np.float32
+                    )
+                obs[601 + 6 * f : 601 + 6 * (f + 1)] = self._anchor_6d(
+                    self.h_quat[0], self.motion_body_quats[tf].astype(np.float32)
+                )
+        return obs
+
+    def step(self, robot_obs, update_encoder, debug=False):
+        """Re-init the heading reference on first frame / after a reset, then run the base step."""
+        if robot_obs and (self.first_motion or self.reinit_heading):
+            q = None
+            if "imu.quat.w" in robot_obs:
+                q = np.array(
+                    [
+                        robot_obs["imu.quat.w"],
+                        robot_obs["imu.quat.x"],
+                        robot_obs["imu.quat.y"],
+                        robot_obs["imu.quat.z"],
+                    ],
+                    np.float64,
+                )
+            else:
+                q = robot_obs.get("imu.quaternion")
+                if q is not None:
+                    q = np.array(q, np.float64)
+            if q is not None:
+                self.heading_init_base_quat = np.array(q, np.float64)
+                with self.motion_lock:
+                    rf = min(self.ref_cursor, self.motion_timesteps - 1)
+                    if self.encode_mode == 2 and self.smpl_root_quat is not None:
+                        # Anchor the heading delta to the SMPL root at init so the
+                        # robot turns *relative* to the clip/operator start heading.
+                        self.init_ref_quat = np.asarray(self.smpl_root_quat, np.float64)
+                    else:
+                        self.init_ref_quat = self.motion_body_quats[rf].copy()
+                self.delta_heading = 0.0
+                self.first_motion = False
+                self.reinit_heading = False
+                logger.debug("[Heading] init quat: %s", self.heading_init_base_quat)
+        return super().step(robot_obs, update_encoder=update_encoder, debug=debug)
+
+    def advance_cursor(self):
+        """Advance the reference cursor one frame per 50 Hz tick (no wall-clock catch-up)."""
+        if not self.playing:
+            return
+        with self.motion_lock:
+            if self.motion_timesteps > 0:
+                self.ref_cursor = min(self.ref_cursor + 1, self.motion_timesteps - 1)

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

@@ -0,0 +1,415 @@
+#!/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
+
+from collections import deque
+import logging
+from typing import TYPE_CHECKING
+
+from huggingface_hub import hf_hub_download
+import numpy as np
+
+from lerobot.utils.import_utils import _onnxruntime_available, require_package
+
+from ..g1_utils import (
+    MUJOCO_TO_ISAACLAB,
+    WB_ACTION_DIM,
+    G1_29_JointIndex,
+    lowstate_to_obs,
+    wb_action_key,
+)
+from .sonic_pipeline import (
+    CONTROL_DT,
+    DEFAULT_ANGLES,
+    ENCODER_UPDATE_EVERY,
+    TOKEN_DIM,
+    PlannerController,
+    compute_kp_kd,
+    make_ort_session_options,
+    ort_providers,
+)
+
+# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
+TOKEN_ACTION_PREFIX = "motion_token"
+# Proprio-state prefix for the token interface: the robot echoes the last commanded
+# token here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
+TOKEN_STATE_PREFIX = "motion_token_state"
+
+
+def token_action_key(i: int) -> str:
+    """Action-dict key for the i-th component of the 64-D SONIC latent token.
+
+    The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``,
+    which only routes ``.pos`` scalar features onto the policy action vector.
+    """
+    return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
+
+
+def token_state_key(i: int) -> str:
+    """Observation key for the i-th component of the 64-D SONIC latent token state."""
+    return f"{TOKEN_STATE_PREFIX}.{i}.pos"
+
+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
+
+# Neutral ("zero pose") SONIC token, held by token_mode until the first real token
+# arrives. Captured from the encoder's own output while the robot stood idle in sim
+# (capture_neutral_token.py): the encoder is an FSQ bottleneck (~5 bit/dim, 15.5 half-
+# width, Div(16)), so its tokens live on the 1/16 grid. We store the integer FSQ codes
+# and rescale by the same 1/16 step, giving an exact on-grid token -- unlike the literal
+# all-zero token, which is off the encoder's learned manifold and decodes to a slightly
+# goofy stance. This one decodes to a stable, natural standing pose.
+_NEUTRAL_TOKEN_CODES = np.array(
+    [-1, 3, 1, -1, 1, -3, 6, 1, 1, 1, -2, -4, -2, 0, -3, -1,
+     2, -1, -3, -5, 3, 1, 1, -4, -1, -1, 1, -7, 0, 1, 2, -2,
+     5, -2, -2, -4, 0, -1, 3, -1, 0, -5, -1, 0, -4, 0, 0, -1,
+     -1, 2, -2, 1, 3, 3, 1, 0, 0, 6, 0, -7, 3, 0, 2, -2],
+    dtype=np.float32,
+)
+NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0  # FSQ Div(16): integer codes -> on-grid token
+
+
+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:
+        return None
+    keys = [wb_action_key(i) for i in range(WB_ACTION_DIM)]
+    # Require the full dense command: a partial action (e.g. only ``wb.0.pos``)
+    # must not be silently zero-filled, which would drive most joints toward 0.
+    if any(key not in action for key in keys):
+        return None
+    return np.fromiter(
+        (float(action[key]) for key in keys),
+        dtype=np.float32,
+        count=WB_ACTION_DIM,
+    )
+
+
+def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
+
+    This is the token-only replay interface: instead of a joint reference driving the
+    encoder, the caller supplies the 64-D encoder latent directly (e.g. a recorded
+    ``action.motion_token`` column), which the decoder consumes with the encoder
+    bypassed. Requires the full dense token; a partial one is ignored (returns None).
+    """
+    if not action:
+        return None
+    keys = [token_action_key(i) for i in range(TOKEN_DIM)]
+    if any(key not in action for key in keys):
+        return None
+    return np.fromiter(
+        (float(action[key]) for key in keys),
+        dtype=np.float32,
+        count=TOKEN_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) are not part of the 29-DoF SONIC
+    reference, and yaw-rate (33) is integrated into the heading by the caller (it
+    cannot be represented in this static per-tick anchor).
+    """
+    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
+
+
+class SonicRuntime:
+    """Loads the SONIC encoder/decoder ONNX models and owns the controller.
+
+    No motion planner: the reference motion buffer is written directly each tick by
+    :class:`SonicWholeBodyController` from the incoming 34-D whole-body command.
+    """
+
+    def __init__(self, force_cpu: bool = False):
+        require_package("onnxruntime", extra="unitree_g1")
+        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)
+        so = make_ort_session_options()
+
+        encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
+        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
+
+        # Report the provider actually bound, not the one requested: ORT silently falls
+        # back to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a
+        # CPU decoder drifts the closed-loop heading. Warn loudly so it can't hide.
+        self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider"
+        if not force_cpu and not self.use_gpu:
+            print(
+                "[SONIC] WARNING: decoder bound to CPUExecutionProvider (CUDA unavailable). "
+                "Closed-loop replay/control will drift. Ensure libcudnn is on LD_LIBRARY_PATH "
+                "(site-packages/nvidia/*/lib).",
+                flush=True,
+            )
+
+        self.kp, self.kd = compute_kp_kd()
+        self.controller = PlannerController(encoder_sess, decoder_sess)
+
+    @property
+    def pipeline(self):
+        return self.controller
+
+    def reset(self):
+        # Full pipeline reset: clears the encoder token, proprioception history and
+        # heading, and rewinds the motion buffer. reinit_heading is set so the next
+        # step re-latches the reference frame to the current robot orientation.
+        self.controller.reset()
+
+    def shutdown(self):
+        pass
+
+
+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):
+        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
+
+        # 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)
+        # Integrated heading (rad) from the whole-body command's yaw-rate (index 33),
+        # forwarded to the pipeline as ``delta_heading`` so turn commands take effect.
+        self._heading = 0.0
+
+        # Token-interface state. ``token_mode`` is set True by the robot when the deploy
+        # is token-driven (``UnitreeG1Config.sonic_token_action``): the controller then
+        # holds a stable *neutral* (all-zero) token until the first real token arrives,
+        # and afterwards holds the *last* token received between ticks (the async
+        # controller runs ~50 Hz while a token VLA streams ~30 Hz). This lives here (not
+        # in the entry-point script) so it applies uniformly to run_g1_onboard,
+        # lerobot-rollout and the sim replays. ``token_mode`` stays False for the dense
+        # 34-D whole-body / OpenHLM path, which keeps its own "hold last target" idle.
+        self.token_mode = False
+        self._last_token: np.ndarray | None = None
+
+        logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)")
+
+    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
+        # Index 33 is a yaw-rate (rad/s): integrate it into a heading offset and hand
+        # it to the pipeline as ``delta_heading`` so commanded turns are tracked rather
+        # than silently dropped (the anchor from _wb34_to_reference only carries r/p).
+        self._heading += float(wb[33]) * CONTROL_DT
+        c.delta_heading = self._heading
+        # 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 _run_token(self, obs: dict, token: np.ndarray) -> dict:
+        """Decode a supplied 64-D latent token directly (encoder bypassed).
+
+        Token-only replay: set the pipeline's cached token to the supplied one and run
+        a decode-only step (``update_encoder=False``). The decoder still closes the loop
+        on live proprioception (history is refreshed inside ``step`` from ``obs``); only
+        the encoder — which would recompute the token from a motion reference — is
+        skipped. Returns the ``<joint>.q`` target dict.
+        """
+        c = self.controller
+        c.token = np.asarray(token, np.float32)
+        self._wb_step += 1
+        return c.step(obs, update_encoder=False, debug=False)
+
+    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)
+
+        # Token-only interface (latent replay / token-output VLA): a dense 64-D
+        # ``motion_token.{i}`` command is decoded directly, bypassing the encoder.
+        # Checked before the joint path so a token action takes precedence.
+        token = _extract_token_from_action(action)
+        if token is not None:
+            self._last_token = token
+        elif self._last_token is None and self.token_mode:
+            # Token-driven deploy, but no token has arrived yet: hold the captured
+            # neutral token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural
+            # standing pose (the encoder's own idle output; see NEUTRAL_TOKEN).
+            self._last_token = NEUTRAL_TOKEN.copy()
+        if self._last_token is not None:
+            # Either a fresh token this tick or the last one received (held between the
+            # ~30 Hz token stream and the ~50 Hz control loop).
+            return self._startup_blend(obs, self._run_token(obs, self._last_token))
+
+        # Dense 34-D whole-body command (OpenHLM / pi0.5 joint interface): a single
+        # vector per tick drives the mode-0 encoder reference directly. Until the
+        # policy produces one, hold (no command) so the robot keeps its last target.
+        wb = _extract_wb34_from_action(action)
+        if wb is None:
+            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],
+                )
+            return {}
+        return self._startup_blend(obs, self._run_wholebody34(obs, wb))
+
+    def reset(self):
+        self._runtime.reset()
+        self._init_step = 0  # re-run the startup blend after a reset
+        self._start_pose = {}
+        self._wb_step = 0
+        self._wb_traj.clear()
+        self._wb_quat_traj.clear()
+        self._heading = 0.0
+        # Drop the held token so token_mode re-seeds the neutral token after a reset.
+        self._last_token = None
+
+    def shutdown(self):
+        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/infer_sonic_g1_onboard.py

@@ -0,0 +1,192 @@
+#!/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.
+
+"""Laptop-side sender for the SONIC whole-body walk policy, onboard deployment.
+
+This is the counterpart to ``run_g1_onboard.py`` (which runs the SONIC decoder on the
+robot). The heavy VLA (``nepyope/sonic_walk``, a pi0.5 token policy) runs here on the
+laptop GPU; only the resulting 64-D latent token is shipped to the robot over ZMQ:
+
+    laptop:  camera frame (ZMQ from robot :5555) + previous token
+             -> pi0.5 -> next 64-D token
+             -> PUSH JSON {motion_token.i.pos: ...} to robot :6004
+    robot:   run_g1_onboard receives the token, SonicWholeBodyController decodes it
+             into whole-body joint commands against local DDS at full rate.
+
+The policy's ``observation.state`` is the token currently being executed, so we close
+the loop by feeding back the *last token we sent* (the decoder holds it until a new one
+arrives). This mirrors what ``lerobot-rollout`` does via the robot's token echo, but
+without a controller / DDS on the laptop.
+
+The policy is pi0.5 with chunk_size=50, so a full diffusion inference runs only about
+once every 50 ticks; ``select_action`` pops one queued token per tick in between.
+
+Run ``run_g1_onboard.py --controller SonicWholeBodyController --sonic-token-action
+--cameras ...`` on the robot first, then this on the laptop:
+
+    python -m lerobot.robots.unitree_g1.infer_sonic_g1_onboard \
+        --policy-path nepyope/sonic_walk --robot-ip 192.168.123.164 \
+        --task "walk back and forth"
+"""
+
+import argparse
+import contextlib
+import json
+import logging
+import signal
+import time
+
+import numpy as np
+import torch
+
+from lerobot.cameras.zmq import ZMQCamera, ZMQCameraConfig
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.policies.factory import get_policy_class, make_pre_post_processors
+from lerobot.policies.utils import prepare_observation_for_inference
+from lerobot.robots.unitree_g1.controllers.sonic_whole_body import (
+    NEUTRAL_TOKEN,
+    TOKEN_DIM,
+    token_action_key,
+)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
+logger = logging.getLogger("sonic_sender")
+
+ACTION_PORT = 6004  # matches run_g1_onboard.py --action-port
+IMAGE_KEY = "observation.images.ego_view"  # pi05 sonic_walk VISUAL input
+STATE_KEY = "observation.state"
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--policy-path", default="nepyope/sonic_walk", help="Policy repo id or local path")
+    p.add_argument("--robot-ip", default="192.168.123.164", help="Robot IP (camera + action ports)")
+    p.add_argument("--action-port", type=int, default=ACTION_PORT, help="Onboard ZMQ PULL port for actions")
+    p.add_argument("--camera-port", type=int, default=5555, help="Onboard ZMQ camera PUB port")
+    p.add_argument("--camera-name", default="head_camera", help="Camera name served by run_g1_onboard")
+    p.add_argument("--camera-width", type=int, default=640, help="Camera width")
+    p.add_argument("--camera-height", type=int, default=480, help="Camera height")
+    p.add_argument("--task", default="walk back and forth", help="Language prompt for the VLA")
+    p.add_argument("--fps", type=float, default=30.0, help="Token send rate (matches training inference)")
+    p.add_argument("--device", default="cuda", help="Torch device")
+    p.add_argument("--max-ticks", type=int, default=0, help="Stop after N ticks (0 = run forever)")
+    p.add_argument("--dry-run", action="store_true", help="Run inference but do not PUSH tokens to the robot")
+    args = p.parse_args()
+
+    device = torch.device(args.device)
+
+    # --- Policy + processors (normalization stats baked into the checkpoint) ---
+    logger.info("Loading policy from '%s'...", args.policy_path)
+    policy_cfg = PreTrainedConfig.from_pretrained(args.policy_path)
+    policy_cfg.pretrained_path = args.policy_path
+    policy = get_policy_class(policy_cfg.type).from_pretrained(args.policy_path, config=policy_cfg)
+    policy = policy.to(device)
+    policy.eval()
+    policy.reset()
+
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=policy_cfg,
+        pretrained_path=args.policy_path,
+        preprocessor_overrides={"device_processor": {"device": str(device)}},
+    )
+    logger.info("Policy loaded (type=%s, device=%s, chunk=%s)", policy_cfg.type, device,
+                getattr(policy_cfg, "chunk_size", "?"))
+
+    # --- Camera (ZMQ from the robot's onboard image server) ---
+    cam = ZMQCamera(
+        ZMQCameraConfig(
+            server_address=args.robot_ip,
+            port=args.camera_port,
+            camera_name=args.camera_name,
+            width=args.camera_width,
+            height=args.camera_height,
+            fps=int(args.fps),
+        )
+    )
+    logger.info("Connecting camera %s@%s:%d ...", args.camera_name, args.robot_ip, args.camera_port)
+    cam.connect()
+
+    # --- Action PUSH socket to the onboard controller ---
+    import zmq
+
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.PUSH)
+    sock.setsockopt(zmq.SNDHWM, 2)
+    sock.setsockopt(zmq.LINGER, 0)
+    sock.connect(f"tcp://{args.robot_ip}:{args.action_port}")
+    logger.info("Sending tokens to tcp://%s:%d (dry_run=%s)", args.robot_ip, args.action_port, args.dry_run)
+
+    stop = {"flag": False}
+    signal.signal(signal.SIGINT, lambda *_: stop.__setitem__("flag", True))
+    signal.signal(signal.SIGTERM, lambda *_: stop.__setitem__("flag", True))
+
+    # observation.state = the token currently executing on the robot (last one we sent);
+    # start at the neutral token the decoder holds before the first send, so the very
+    # first inference sees the true executing token (not zeros).
+    prev_token = NEUTRAL_TOKEN.copy()
+    period = 1.0 / args.fps
+    n = 0
+    t_infer_total = 0.0
+    logger.info("Streaming tokens at %.0f Hz. Ctrl-C to stop.", args.fps)
+    try:
+        while not stop["flag"]:
+            t0 = time.time()
+            try:
+                frame = cam.read()  # HxWxC uint8 RGB
+            except Exception as e:  # noqa: BLE001
+                logger.warning("Camera read failed: %s", e)
+                time.sleep(period)
+                continue
+
+            raw_obs = {
+                IMAGE_KEY: np.ascontiguousarray(frame),
+                STATE_KEY: prev_token.copy(),
+            }
+            with torch.inference_mode():
+                obs = prepare_observation_for_inference(raw_obs, device, args.task, "unitree_g1")
+                obs = preprocessor(obs)
+                action = policy.select_action(obs)
+                action = postprocessor(action)
+            token = action.squeeze(0).to("cpu").numpy().astype(np.float32)
+            prev_token = token
+
+            if not args.dry_run:
+                msg = {token_action_key(i): float(token[i]) for i in range(TOKEN_DIM)}
+                with contextlib.suppress(zmq.Again):
+                    sock.send_string(json.dumps(msg), zmq.NOBLOCK)
+
+            n += 1
+            t_infer_total += time.time() - t0
+            if n % 30 == 0:
+                logger.info(
+                    "tick %d | avg %.1f ms/tick | token[:3]=%s",
+                    n, 1000.0 * t_infer_total / 30.0, np.round(token[:3], 3).tolist(),
+                )
+                t_infer_total = 0.0
+
+            if args.max_ticks and n >= args.max_ticks:
+                break
+            time.sleep(max(0.0, period - (time.time() - t0)))
+    finally:
+        logger.info("Stopping sender after %d ticks.", n)
+        with contextlib.suppress(Exception):
+            cam.disconnect()
+        with contextlib.suppress(Exception):
+            sock.close(linger=0)
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/robots/unitree_g1/run_g1_onboard.py

@@ -0,0 +1,254 @@
+#!/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.
+
+"""Run the G1 locomotion / whole-body controller ONBOARD, driven by high-level actions
+from a laptop.
+
+The controller (GR00T / Holosoma / SONIC whole-body) runs on the robot itself against
+local DDS, at full control rate. The laptop ships only the resulting high-level action
+(arm joint targets + joystick axes + gripper flags, or a 64-D SONIC motion token) as
+JSON over ZMQ. This process applies each action via ``UnitreeG1.send_action`` while the
+onboard controller thread keeps the legs balanced / decodes the token.
+
+This is the real-deploy counterpart to running ``lerobot-rollout`` on the laptop with
+``--robot.is_simulation=false`` (the ZMQ *socket bridge*): there the 50 Hz lowcmd
+crosses the network; here only compact high-level actions do, and the control loop stays
+local to the robot. Pair with a laptop client that produces actions (exo teleop, or a
+policy such as ``nepyope/sonic_walk`` emitting ``motion_token.{i}.pos``).
+
+Besides receiving actions, this process publishes ``observation.state`` (29 joint ``.q``)
+on a ZMQ PUB port so a laptop policy client has proprioception.
+
+Safety: type ``e`` then Enter in this terminal to stop immediately (zero-torque + exit).
+Ctrl-C does the normal graceful shutdown (kp ramp).
+
+Examples (on the robot):
+
+    # GR00T locomotion, arm targets from the laptop:
+    python -m lerobot.robots.unitree_g1.run_g1_onboard --controller GrootLocomotionController
+
+    # SONIC whole-body walk policy: laptop ships 64-D tokens, decoder runs here:
+    python -m lerobot.robots.unitree_g1.run_g1_onboard \
+        --controller SonicWholeBodyController --sonic-token-action \
+        --cameras "head_camera:/dev/v4l/by-path/platform-3610000.usb-usb-0:2.1:1.3-video-index0:640x480"
+"""
+
+import argparse
+import contextlib
+import json
+import logging
+import os
+import signal
+import sys
+import threading
+import time
+
+import numpy as np
+import zmq
+
+from lerobot.cameras.zmq.image_server import ImageServer
+from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
+from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
+from lerobot.robots.unitree_g1.run_g1_server import Gripper, build_gripper, parse_camera_specs
+from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
+logger = logging.getLogger("g1_onboard")
+
+ACTION_PORT = 6004
+STATE_PORT = 6005
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--controller", default="GrootLocomotionController", help="Controller class name")
+    p.add_argument("--dds-interface", default=None, help="DDS network interface (default: SDK default)")
+    p.add_argument(
+        "--sim",
+        action="store_true",
+        help="Attach to a DDS MuJoCo sim: skip MotionSwitcher + physical remote, default dds-interface 'lo'.",
+    )
+    p.add_argument(
+        "--sonic-token-action",
+        action="store_true",
+        help="SONIC token interface: actions carry a 64-D motion_token.{i}.pos that the decoder consumes.",
+    )
+    p.add_argument("--action-port", type=int, default=ACTION_PORT, help="ZMQ PULL port for laptop actions")
+    p.add_argument("--state-port", type=int, default=STATE_PORT, help="ZMQ PUB port for observation.state")
+    p.add_argument("--state-fps", type=float, default=30.0, help="observation.state publish rate; <=0 disables")
+    p.add_argument("--gravity-compensation", action="store_true", help="Enable arm gravity compensation")
+    # Gripper control (Damiao over CAN).
+    p.add_argument("--grippers", action="store_true", help="Drive Damiao grippers from action L3/R3 flags")
+    p.add_argument("--gripper-port-left", default="can1", help="CAN interface for LEFT gripper")
+    p.add_argument("--gripper-port-right", default="can0", help="CAN interface for RIGHT gripper")
+    p.add_argument("--gripper-send-id", type=lambda x: int(x, 0), default=0x08, help="Motor send CAN id")
+    p.add_argument("--gripper-recv-id", type=lambda x: int(x, 0), default=0x18, help="Motor recv CAN id")
+    p.add_argument("--gripper-motor-type", default="dm4310", help="Damiao motor type")
+    p.add_argument("--gripper-open-deg", type=float, default=-65.0, help="Gripper OPEN position (deg)")
+    p.add_argument("--gripper-close-deg", type=float, default=0.0, help="Gripper CLOSE position (deg)")
+    p.add_argument("--gripper-kp", type=float, default=15.0, help="MIT position gain (stiffness)")
+    p.add_argument("--gripper-kd", type=float, default=0.5, help="MIT damping gain")
+    p.add_argument("--gripper-no-fd", dest="gripper_fd", action="store_false", help="Classic CAN (non-FD)")
+    p.set_defaults(gripper_fd=True)
+    # Optional camera streaming (ZMQ) so the laptop policy client / viewer can connect.
+    p.add_argument("--cameras", default=None, help="Camera spec 'name:device[:WxH[:FOURCC]]', comma-sep")
+    p.add_argument("--camera-fps", type=int, default=30, help="Camera FPS")
+    p.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port")
+    p.add_argument("--camera-width", type=int, default=640, help="Default camera width")
+    p.add_argument("--camera-height", type=int, default=480, help="Default camera height")
+    args = p.parse_args()
+
+    dds_interface = args.dds_interface
+    if args.sim and dds_interface is None:
+        dds_interface = "lo"
+
+    cfg = UnitreeG1Config(
+        is_simulation=False,
+        onboard=True,
+        controller=args.controller,
+        dds_interface=dds_interface,
+        gravity_compensation=args.gravity_compensation,
+        release_motion_control=not args.sim,
+        physical_remote=not args.sim,
+        sonic_token_action=args.sonic_token_action,
+        cameras={},
+    )
+
+    # Optional camera server (background thread; independent of DDS/CAN).
+    camera_server = None
+    if args.cameras:
+        cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
+        camera_server = ImageServer({"fps": args.camera_fps, "cameras": cameras}, port=args.camera_port)
+        threading.Thread(target=camera_server.run, daemon=True).start()
+        cam_summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
+        logger.info("Camera server started on :%d: %s", args.camera_port, cam_summary)
+
+    robot = UnitreeG1(cfg)
+    logger.info("Connecting onboard robot (controller=%s, token=%s)...", args.controller, args.sonic_token_action)
+    robot.connect()
+    # Note: with --sonic-token-action the SonicWholeBodyController holds a neutral
+    # (all-zero) token until the first laptop token arrives, then holds the last token
+    # between ticks -- see SonicWholeBodyController.token_mode (set from config).
+
+    grippers: dict[str, Gripper] = {}
+    if args.grippers:
+        for side, port in (("L", args.gripper_port_left), ("R", args.gripper_port_right)):
+            grippers[side] = build_gripper(
+                side, port, args.gripper_send_id, args.gripper_recv_id, args.gripper_motor_type,
+                args.gripper_fd, args.gripper_open_deg, args.gripper_close_deg, args.gripper_kp, args.gripper_kd,
+            )
+        logger.info("Grippers enabled: L3 -> left, R3 -> right")
+
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.PULL)
+    sock.setsockopt(zmq.CONFLATE, 1)  # only ever act on the freshest command
+    sock.setsockopt(zmq.RCVTIMEO, 200)  # keeps the loop responsive to the stop event
+    sock.bind(f"tcp://0.0.0.0:{args.action_port}")
+    logger.info("Onboard controller live. Waiting for laptop actions on :%d ...", args.action_port)
+    logger.info("Type 'e' then Enter to STOP immediately (or Ctrl-C for graceful shutdown).")
+
+    stop = threading.Event()
+    signal.signal(signal.SIGINT, lambda *_: stop.set())
+    signal.signal(signal.SIGTERM, lambda *_: stop.set())
+
+    def estop_listener() -> None:
+        for line in sys.stdin:
+            if line.strip().lower() == "e":
+                logger.warning("E-STOP ('e'): going passive NOW.")
+                try:
+                    robot._shutdown_event.set()  # stop the controller loop publishing
+                    time.sleep(0.05)
+                    robot._send_zero_torque()  # motors limp; nothing overwrites it now
+                except Exception as e:  # noqa: BLE001
+                    logger.warning("E-stop zero-torque failed: %s", e)
+                os._exit(0)  # immediate hard exit, no slow cleanup
+
+    threading.Thread(target=estop_listener, daemon=True).start()
+
+    # Proprioception feedback: publish observation.state (29 joint .q) so a laptop
+    # inference client can feed it to a policy. DDS stays local; only compact JSON
+    # state crosses the network. (For a token policy the laptop closes the loop on the
+    # token instead, but publishing joint state is harmless and useful for logging.)
+    state_sock = None
+    if args.state_fps > 0:
+        state_sock = ctx.socket(zmq.PUB)
+        state_sock.setsockopt(zmq.SNDHWM, 2)
+        state_sock.setsockopt(zmq.LINGER, 0)
+        state_sock.bind(f"tcp://0.0.0.0:{args.state_port}")
+        logger.info("Publishing observation.state on :%d at %.0f Hz", args.state_port, args.state_fps)
+
+        def publish_state() -> None:
+            period = 1.0 / args.state_fps
+            joint_names = [j.name for j in G1_29_JointIndex]
+            while not stop.is_set():
+                t0 = time.time()
+                obs = robot.get_observation()
+                if obs:
+                    state = {f"{name}.q": float(obs.get(f"{name}.q", 0.0)) for name in joint_names}
+                    with contextlib.suppress(zmq.Again):
+                        state_sock.send_json(state, zmq.NOBLOCK)
+                time.sleep(max(0.0, period - (time.time() - t0)))
+
+        threading.Thread(target=publish_state, daemon=True).start()
+    else:
+        logger.info("observation.state PUB disabled (--state-fps<=0)")
+
+    n = 0
+    try:
+        while not stop.is_set():
+            try:
+                payload = sock.recv()
+            except zmq.Again:
+                continue
+            except zmq.ContextTerminated:
+                break
+
+            try:
+                action = json.loads(payload.decode("utf-8"))
+            except (ValueError, UnicodeDecodeError) as e:
+                logger.warning("Dropping malformed action: %s", e)
+                continue
+
+            robot.send_action(action)
+
+            if grippers:
+                # L3 = remote.button.4 -> left, R3 = remote.button.0 -> right.
+                if "L" in grippers and "remote.button.4" in action:
+                    grippers["L"].apply(bool(action["remote.button.4"]))
+                if "R" in grippers and "remote.button.0" in action:
+                    grippers["R"].apply(bool(action["remote.button.0"]))
+
+            n += 1
+            if n % 60 == 0:
+                axes = {k: round(float(action.get(k, 0.0)), 3) for k in ("remote.lx", "remote.ly", "remote.rx", "remote.ry")}
+                logger.info("Applied %d actions | axes=%s", n, axes)
+    finally:
+        logger.info("Shutting down onboard controller...")
+        stop.set()
+        if state_sock is not None:
+            with contextlib.suppress(Exception):
+                state_sock.close(linger=0)
+        if camera_server is not None:
+            with contextlib.suppress(Exception):
+                camera_server.stop()
+        for g in grippers.values():
+            with contextlib.suppress(Exception):
+                g.bus.disconnect()
+        robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/robots/unitree_g1/run_g1_server.py

@@ -28,9 +28,11 @@ import argparse
 import base64
 import contextlib
 import json
+import re
 import threading
 import time
-from typing import Any
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
 
 import zmq
 from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
@@ -41,6 +43,9 @@ from unitree_sdk2py.utils.crc import CRC
 
 from lerobot.cameras.zmq.image_server import ImageServer
 
+if TYPE_CHECKING:
+    from lerobot.motors.damiao.damiao import DamiaoMotorsBus
+
 # DDS topic names follow Unitree SDK naming conventions
 # ruff: noqa: N816
 kTopicLowCommand_Debug = "rt/lowcmd"  # action to robot
@@ -51,6 +56,105 @@ LOWSTATE_PORT = 6001
 NUM_MOTORS = 35
 
 
+@dataclass
+class Gripper:
+    """A single Damiao gripper that only writes to CAN when the open/close state changes."""
+
+    name: str
+    bus: "DamiaoMotorsBus"
+    open_deg: float
+    close_deg: float
+    _last_cmd: str | None = None  # "open" | "close"
+
+    def apply(self, want_close: bool) -> None:
+        want = "close" if want_close else "open"
+        if want == self._last_cmd:
+            return
+        target = self.close_deg if want_close else self.open_deg
+        self.bus.write("Goal_Position", "gripper", target)
+        self._last_cmd = want
+        print(f"[gripper] {self.name} -> {want.upper()} ({target:.1f} deg)")
+
+
+def build_gripper(
+    name: str,
+    port: str,
+    send_id: int,
+    recv_id: int,
+    motor_type: str,
+    use_can_fd: bool,
+    open_deg: float,
+    close_deg: float,
+    kp: float,
+    kd: float,
+) -> Gripper:
+    from lerobot.motors.damiao.damiao import DamiaoMotorsBus
+    from lerobot.motors.motors_bus import Motor, MotorNormMode
+
+    motors = {
+        "gripper": Motor(
+            id=send_id,
+            model=motor_type,
+            norm_mode=MotorNormMode.DEGREES,
+            motor_type_str=motor_type,
+            recv_id=recv_id,
+        )
+    }
+    bus = DamiaoMotorsBus(port=port, motors=motors, use_can_fd=use_can_fd)
+    print(f"Connecting {name} gripper on {port} (fd={use_can_fd})...")
+    bus.connect(handshake=True)
+    bus.write("Kp", "gripper", kp)
+    bus.write("Kd", "gripper", kd)
+    bus.write("Goal_Position", "gripper", open_deg)  # start open
+    print(f"  {name}: connected, torque enabled, opened.")
+    return Gripper(name, bus, open_deg, close_deg, _last_cmd="open")
+
+
+def parse_camera_specs(spec: str, default_width: int, default_height: int) -> dict[str, dict]:
+    """Parse a multi-camera spec string into an ImageServer ``cameras`` dict.
+
+    Format: comma-separated ``name:device[:WxH[:FOURCC]]`` entries, e.g.
+    ``head_camera:6,left_wrist:0``. ``device`` may be an integer index or an explicit
+    device path (e.g. ``/dev/video6``), including stable ``by-path`` names like
+    ``/dev/v4l/by-path/platform-...:2.1:1.3-video-index0`` which survive USB
+    re-enumeration (unlike bare ``/dev/videoN`` indices). Because a by-path name
+    itself contains colons, the optional ``WxH`` and ``FOURCC`` are parsed from the
+    *right* so the device-path colons are preserved.
+    """
+    wh_re = re.compile(r"\d+x\d+", re.IGNORECASE)
+    fourcc_re = re.compile(r"[A-Za-z0-9]{4}")
+
+    cameras: dict[str, dict] = {}
+    for entry in spec.split(","):
+        entry = entry.strip()
+        if not entry:
+            continue
+        if ":" not in entry:
+            raise ValueError(f"Invalid camera spec '{entry}', expected 'name:device[:WxH[:FOURCC]]'")
+        name, rest = entry.split(":", 1)
+        name = name.strip()
+        tokens = [t.strip() for t in rest.split(":")]
+
+        fourcc = None
+        if len(tokens) >= 3 and wh_re.fullmatch(tokens[-2]) and fourcc_re.fullmatch(tokens[-1]):
+            fourcc = tokens.pop().upper()
+        width, height = default_width, default_height
+        if len(tokens) >= 2 and wh_re.fullmatch(tokens[-1]):
+            w, h = tokens.pop().lower().split("x")
+            width, height = int(w), int(h)
+
+        raw_id = ":".join(tokens).strip()
+        if not raw_id:
+            raise ValueError(f"Invalid camera spec '{entry}', missing device")
+        device_id: int | str = int(raw_id) if raw_id.lstrip("-").isdigit() else raw_id
+        if name in cameras:
+            raise ValueError(f"Duplicate camera name '{name}' in --cameras")
+        cameras[name] = {"device_id": device_id, "shape": [height, width], "fourcc": fourcc}
+    if not cameras:
+        raise ValueError("No cameras parsed from --cameras spec")
+    return cameras
+
+
 def lowstate_to_dict(msg: hg_LowState) -> dict[str, Any]:
     """Convert LowState SDK message to a JSON-serializable dictionary."""
     motor_states = []
@@ -155,7 +259,11 @@ def main() -> None:
     """Main entry point for the robot server bridge."""
     parser = argparse.ArgumentParser(description="DDS-to-ZMQ bridge server for Unitree G1")
     parser.add_argument("--camera", action="store_true", help="Also launch camera server")
-    parser.add_argument("--camera-device", type=int, default=4, help="Camera device ID (default: 4)")
+    parser.add_argument("--camera-device", default="4",
+                        help="Camera device: index or /dev/video path or by-path name (default: 4)")
+    parser.add_argument("--cameras", default=None,
+                        help="Multi-camera spec 'name:device[:WxH[:FOURCC]]', comma-separated. Overrides "
+                             "--camera-device; device may be a by-path name to survive USB re-enumeration.")
     parser.add_argument("--camera-fps", type=int, default=30, help="Camera FPS (default: 30)")
     parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
     parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
@@ -164,20 +272,20 @@ def main() -> None:
 
     # Optionally start camera server in background thread
     camera_thread = None
-    if args.camera:
-        camera_config = {
-            "fps": args.camera_fps,
-            "cameras": {
-                "head_camera": {
-                    "device_id": args.camera_device,
-                    "shape": [args.camera_height, args.camera_width],
-                }
-            },
-        }
+    if args.camera or args.cameras:
+        if args.cameras:
+            cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
+        else:
+            # Single camera; accept an int index or a device/by-path string.
+            dev = args.camera_device
+            dev = int(dev) if str(dev).lstrip("-").isdigit() else dev
+            cameras = {"head_camera": {"device_id": dev, "shape": [args.camera_height, args.camera_width]}}
+        camera_config = {"fps": args.camera_fps, "cameras": cameras}
         camera_server = ImageServer(camera_config, port=args.camera_port)
         camera_thread = threading.Thread(target=camera_server.run, daemon=True)
         camera_thread.start()
-        print(f"Camera server started on port {args.camera_port} (device {args.camera_device})")
+        cam_summary = ", ".join(f"{n}(dev {c['device_id']})" for n, c in cameras.items())
+        print(f"Camera server started on port {args.camera_port}: {cam_summary}")
 
     # initialize DDS
     ChannelFactoryInitialize(0)

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
@@ -79,6 +87,14 @@ class LocomotionController(Protocol):
 kTopicLowCommand_Debug = "rt/lowcmd"
 kTopicLowState = "rt/lowstate"
 
+# Wireless-remote button byte layout, mapped to the positional button indices the
+# locomotion controllers expect. Used in onboard mode to read the physical Unitree
+# remote from lowstate (mirrors the exo teleoperator's RemoteController).
+_REMOTE_BUTTON_MAP: list[str] = [
+    "RB", "LB", "start", "back", "RT", "LT", "", "",
+    "A", "B", "X", "Y", "up", "right", "down", "left",
+]
+
 
 @dataclass
 class MotorState:
@@ -122,8 +138,10 @@ class UnitreeG1(Robot):
         # Initialize cameras config (ZMQ-based) - actual connection in connect()
         self._cameras = make_cameras_from_configs(config.cameras)
 
-        # Import channel classes based on mode
-        if config.is_simulation:
+        # Import channel classes based on mode. Simulation and onboard both talk to a
+        # real (local) DDS via the Unitree SDK; only the laptop-side bridge client uses
+        # the ZMQ socket shim.
+        if config.is_simulation or config.onboard:
             self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
             self._ChannelPublisher = _SDKChannelPublisher
             self._ChannelSubscriber = _SDKChannelSubscriber
@@ -151,19 +169,100 @@ class UnitreeG1(Robot):
         # Lower-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
 
+        # Token-driven deploy: let a SONIC controller hold a neutral token until the
+        # first real one arrives, then hold the last token between control ticks.
+        if config.sonic_token_action and hasattr(self.controller, "token_mode"):
+            self.controller.token_mode = True
+
         # Controller thread state
         self._controller_thread = None
+        # When set, the controller loop stops publishing low commands so reset() can
+        # drive the joints directly without two publishers fighting (single-publisher).
+        self._controller_paused = threading.Event()
         self._controller_action_lock = threading.Lock()
         self.controller_input = default_remote_input()
         self.controller_output = {}
 
+        # Onboard-only: parser for the physical Unitree wireless remote (read straight
+        # from local lowstate so joystick locomotion works without a laptop round-trip).
+        self._joystick = None
+
+        # Replay-camera state: keep the encoded (raw) cells per camera and decode
+        # frames lazily as the play cursor advances, with a small frame cache, so we
+        # don't materialize gigabytes of decoded RGB at construction time.
+        self._replay_raw: dict[str, list] = {}
+        self._replay_cache: dict[tuple[str, int], np.ndarray] = {}
+        self._replay_cache_cap = 8
+        self._replay_len = 0
+        self._replay_idx = 0
+        if config.replay_camera_parquet and config.replay_camera_map:
+            self._load_replay_frames()
+
+        # Token-mode state: last 64-D SONIC latent token commanded by the policy,
+        # echoed back as ``observation.state`` so a token-output VLA closes the loop
+        # on its own previous token (see ``sonic_token_action``). Seeded to zeros;
+        # the controller's startup blend eases joints in regardless.
+        self._last_token: np.ndarray | None = None
+        if config.sonic_token_action:
+            from .controllers.sonic_whole_body import TOKEN_DIM
+
+            self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
+
+    def _load_replay_frames(self) -> None:
+        """Load only the mapped parquet columns (encoded frames); decode on demand."""
+        import pyarrow.parquet as pq
+
+        cols_needed = list(dict.fromkeys(self.config.replay_camera_map.values()))
+        table = pq.read_table(self.config.replay_camera_parquet, columns=cols_needed)
+        self._replay_len = table.num_rows
+        self._replay_raw = {
+            cam_name: table.column(column).to_pylist()
+            for cam_name, column in self.config.replay_camera_map.items()
+        }
+        logger.info(
+            "Loaded %d replay frames (lazy-decode) for cameras %s from %s",
+            self._replay_len,
+            list(self.config.replay_camera_map),
+            self.config.replay_camera_parquet,
+        )
+
+    def _decode_replay_cell(self, cell) -> np.ndarray:
+        import io
+
+        from PIL import Image
+
+        data = cell["bytes"] if isinstance(cell, dict) else cell
+        return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
+
+    def _replay_frame(self, cam_name: str, idx: int) -> np.ndarray:
+        """Decode (and briefly cache) a single replay frame for a camera."""
+        key = (cam_name, idx)
+        cached = self._replay_cache.get(key)
+        if cached is not None:
+            return cached
+        frame = self._decode_replay_cell(self._replay_raw[cam_name][idx])
+        if len(self._replay_cache) >= self._replay_cache_cap:
+            self._replay_cache.pop(next(iter(self._replay_cache)))
+        self._replay_cache[key] = frame
+        return frame
+
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
             # Step simulation if in simulation mode
             if self.config.is_simulation and self.sim_env is not None:
+                try:
                     self.sim_env.step()
+                except ValueError as e:
+                    # Startup race: the sim thread can step once before reset() has
+                    # written a valid base pose, giving a zero-norm pelvis quaternion
+                    # (scipy>=1.11 raises instead of normalizing). Skip and retry so
+                    # the thread survives instead of dying and freezing the sim.
+                    if "zero norm" not in str(e).lower():
+                        raise
+                    time.sleep(self.control_dt)
+                    continue
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -231,15 +330,80 @@ 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 self.config.sonic_token_action:
+            return {}
+        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 _token_state_ft(self) -> dict[str, type]:
+        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
+
+        Exposed only in ``sonic_token_action`` mode; aggregated by the rollout into a
+        64-D ``observation.state`` (the last token the policy commanded).
+        """
+        if not self.config.sonic_token_action:
+            return {}
+        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
+
+        return {token_state_key(i): float for i in range(TOKEN_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 dict.fromkeys(self.config.empty_cameras, (h, w, 3))
+
+    @property
+    def _replay_cameras_ft(self) -> dict[str, tuple]:
+        """Replay cameras, shaped from their first (lazily decoded) frame."""
+        if not self._replay_len:
+            return {}
+        return {name: self._replay_frame(name, 0).shape for name in self._replay_raw}
+
     @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._token_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}
 
+        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
+        # (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
+        # straight onto the decoder, bypassing the encoder.
+        if self.config.sonic_token_action:
+            from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
+
+            return {token_action_key(i): float for i in range(TOKEN_DIM)}
+
+        # 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}
@@ -255,6 +419,11 @@ class UnitreeG1(Robot):
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
+            # Paused during reset() so the reset routine is the sole low-cmd publisher.
+            if self._controller_paused.is_set():
+                time.sleep(control_dt)
+                continue
+
             with self._lowstate_lock:
                 lowstate = self._lowstate
 
@@ -271,6 +440,13 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
+                # Onboard: the physical Unitree remote (in local lowstate) takes
+                # priority for locomotion when active; otherwise laptop/ZMQ axes stand.
+                if self.config.onboard:
+                    wl = self._wireless_remote_input(lowstate)
+                    if wl is not None:
+                        controller_input.update(wl)
+
                 # Run controller step
                 controller_action = self.controller.run_step(controller_input, lowstate)
 
@@ -293,15 +469,105 @@ class UnitreeG1(Robot):
     def configure(self) -> None:
         pass
 
+    def _wireless_remote_input(self, lowstate) -> dict | None:
+        """Parse the physical Unitree remote from lowstate into controller inputs.
+
+        Onboard only. Returns None when the remote is idle so the laptop-provided
+        (ZMQ) axes keep control; otherwise the physical remote takes priority.
+        """
+        js = self._joystick
+        if js is None:
+            return None
+        wr = getattr(lowstate, "wireless_remote", None)
+        if not wr or len(wr) < 24:
+            return None
+        try:
+            js.extract(wr)
+        except Exception:  # noqa: BLE001
+            return None
+
+        axes = {
+            "remote.lx": float(js.lx.data),
+            "remote.ly": float(js.ly.data),
+            "remote.rx": float(js.rx.data),
+            "remote.ry": float(js.ry.data),
+        }
+        active = any(abs(v) > 1e-2 for v in axes.values())
+        out = dict(axes)
+        for i, name in enumerate(_REMOTE_BUTTON_MAP):
+            if name:
+                val = float(getattr(js, name).data)
+                out[f"remote.button.{i}"] = val
+                if val:
+                    active = True
+        return out if active else None
+
+    def _release_motion_control(self) -> None:
+        """Release the robot's built-in motion services so we can send raw lowcmd.
+
+        Onboard-only. Mirrors run_g1_server.py: on the real robot the factory
+        locomotion/hand services must relinquish control before our controller can
+        write to ``rt/lowcmd``, otherwise commands are ignored or fought.
+        """
+        from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
+
+        msc = MotionSwitcherClient()
+        msc.SetTimeout(5.0)
+        msc.Init()
+        _, result = msc.CheckMode()
+        while result is not None and "name" in result and result["name"]:
+            logger.info("[UnitreeG1] Releasing built-in mode '%s'...", result["name"])
+            msc.ReleaseMode()
+            _, result = msc.CheckMode()
+            time.sleep(1.0)
+
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
         # Initialize DDS channel and simulation environment
         if self.config.is_simulation:
-            from lerobot.envs import make_env
+            from lerobot.envs.utils import (
+                _download_hub_file,
+                _import_hub_module,
+                _normalize_hub_result,
+            )
 
             self._ChannelFactoryInitialize(0, "lo")
-            self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
+            # Call the hub env's make_env directly so we can disable the offscreen
+            # head_camera renderer. We drive image-conditioned policies from recorded
+            # frames (see replay_camera_parquet / external obs), never the sim's own
+            # camera, so building a MuJoCo offscreen GL context is pure liability: it
+            # crashes with "Failed to make the EGL context current" when GLFW/SDL
+            # already own a context, killing the sim thread and hanging on
+            # "Waiting for robot state...". publish_images=False -> no renderer.
+            repo_id, _, local_file, _ = _download_hub_file(
+                "lerobot/unitree-g1-mujoco", True, None
+            )
+            hub_mod = _import_hub_module(local_file, repo_id)
+            raw = hub_mod.make_env(n_envs=1, use_async_envs=False, publish_images=False, cameras=[])
+            self._env_wrapper = _normalize_hub_result(raw)
             # Extract the actual gym env from the dict structure
             self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
+        elif self.config.onboard:
+            # Real robot, controller running onboard against local DDS. Initialize the
+            # real SDK channel factory on the robot's DDS interface and take low-level
+            # control from the built-in services before we start writing lowcmd.
+            if self.config.dds_interface:
+                self._ChannelFactoryInitialize(0, self.config.dds_interface)
+            else:
+                self._ChannelFactoryInitialize(0)
+            # Real robot: hand low-level control over from the built-in services.
+            # A DDS sim has no MotionSwitcher, so this is skipped there.
+            if self.config.release_motion_control:
+                self._release_motion_control()
+            # Real robot: read the physical wireless remote from lowstate for
+            # locomotion. A sim has no physical remote, so leave _joystick=None and
+            # let send_action (ZMQ) drive the locomotion axes instead.
+            if self.config.physical_remote:
+                from unitree_sdk2py.utils.joystick import Joystick
+
+                self._joystick = Joystick()
+                for axis in (self._joystick.lx, self._joystick.ly, self._joystick.rx, self._joystick.ry):
+                    axis.smooth = 1.0
+                    axis.deadzone = 0.0
         else:
             self._ChannelFactoryInitialize(0, config=self.config)
 
@@ -311,6 +577,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 +620,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
@@ -350,12 +630,16 @@ class UnitreeG1(Robot):
             self.msg.motor_cmd[joint].kd = self.kd[joint.value]
             self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
 
-        # Start controller thread if enabled
-        if self.controller is not None:
+        # Start controller thread if enabled. Skipped when run_controller_thread is
+        # False so a caller can step the controller synchronously (faithful replay).
+        if self.controller is not None and self.config.run_controller_thread:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
             self._controller_thread.start()
             fps = int(1.0 / self.controller.control_dt)
             logger.info(f"Controller thread started ({fps}Hz)")
+        elif self.controller is not None:
+            logger.info("Controller thread disabled (run_controller_thread=False); "
+                        "caller must drive controller.run_step synchronously.")
 
     def _send_zero_torque(self) -> None:
         """Send a zero-gain command to make joints passive before shutting down."""
@@ -371,13 +655,59 @@ 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:
+    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.
+
+        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()
 
-        # Signal thread to stop and unblock any waits
+    def disconnect(self):
+        # Stop the controller loop first so it isn't fighting the shutdown ramp.
         self._shutdown_event.set()
+        controller_stopped = True
+        if self._controller_thread is not None:
+            # Wait long enough for any in-flight inference tick to finish and the loop
+            # to observe the shutdown flag, so no stray low command is published while
+            # the ramp runs (the shutdown routine must be the single publisher).
+            self._controller_thread.join(timeout=5.0)
+            if self._controller_thread.is_alive():
+                controller_stopped = False
+                logger.error(
+                    "Controller thread did not stop; skipping graceful ramp to avoid "
+                    "concurrent low commands (fail-safe: joints keep last command until exit)"
+                )
+
+        # Soft, damped settle instead of an instant limp (real robot only; the
+        # subscribe thread is still alive here to supply the current pose). Only ramp
+        # once the controller thread has definitely exited.
+        if not self.config.is_simulation and controller_stopped:
+            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 +715,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 +746,41 @@ 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 self.config.sonic_token_action:
+            # Token mode: echo the last commanded latent token as observation.state
+            # so a token-output VLA closes the loop on its own previous token.
+            from .controllers.sonic_whole_body import token_state_key
 
-        # 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]
+            token = self._last_token if self._last_token is not None else []
+            for i, v in enumerate(token):
+                obs[token_state_key(i)] = float(v)
+        elif 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 - 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]
+        # 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 - 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 in self._replay_raw:
+                obs[name] = self._replay_frame(name, idx)
+            self._replay_idx += 1
 
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
@@ -473,9 +794,19 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
+            if self.config.sonic_token_action:
+                from .controllers.sonic_whole_body import _extract_token_from_action
+
+                token = _extract_token_from_action(action)
+                if token is not None:
+                    self._last_token = token
+            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 +834,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:
@@ -537,6 +924,18 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
+        # Full-body controllers (SONIC / OpenHLM) own the whole 29-DoF command and
+        # ignore ``<joint>.q`` in send_action(), so reset() must publish the default
+        # pose directly. Pause the background controller first so the two aren't both
+        # writing low commands while the robot moves to the default pose.
+        full_body = getattr(self.controller, "full_body", False)
+        paused = False
+        if full_body and self._controller_thread is not None:
+            self._controller_paused.set()
+            paused = True
+            time.sleep(control_dt)  # let any in-flight controller tick settle
+
+        try:
             if self.config.is_simulation and self.sim_env is not None:
                 self.sim_env.reset()
                 self.publish_lowcmd(
@@ -565,6 +964,11 @@ class UnitreeG1(Robot):
                         interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
                         action_dict[f"{motor.name}.q"] = float(interp_pos)
 
+                    # Full-body controllers no-op in send_action(); publish the pose
+                    # directly (arm-only controllers keep the send_action() path).
+                    if full_body:
+                        self.publish_lowcmd(action_dict)
+                    else:
                         self.send_action(action_dict)
 
                     # Maintain constant control rate
@@ -572,8 +976,12 @@ class UnitreeG1(Robot):
                     sleep_time = max(0, control_dt - elapsed)
                     time.sleep(sleep_time)
 
-        # Reset controller internal state (gait phase, obs history, etc.)
+            # Reset controller internal state (gait phase, obs history, etc.) before
+            # resuming so its buffers reflect the post-reset pose.
             if self.controller is not None and hasattr(self.controller, "reset"):
                 self.controller.reset()
+        finally:
+            if paused:
+                self._controller_paused.clear()
 
         logger.info("Reached default position")

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,7 +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
+    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)"
@@ -143,32 +152,25 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         repo_id=repo_id,
         repo_type="dataset",
         commit_message=commit_message,
-        ignore_patterns=[".annotate_staging/**", "**/.DS_Store"],
+        # README.md is excluded because when pushing to ``new_repo_id`` the
+        # source card's links (e.g. the visualize badge) would keep pointing
+        # at the source dataset; a fresh card is generated below instead.
+        ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
     )
     print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
 
+    dataset_info = load_info(root)
+    card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
+    card.push_to_hub(repo_id=repo_id, repo_type="dataset")
+
     # Tag the upload with the codebase version. ``LeRobotDatasetMetadata``
     # resolves the dataset revision via ``get_safe_version`` which scans
     # for tags like ``v3.0``; without a tag it raises
     # ``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
-
-    info_path = root / "meta" / "info.json"
-    version_tag = CODEBASE_VERSION
-    if info_path.exists():
-        try:
-            from lerobot.utils.io_utils import load_json  # noqa: PLC0415
-
-            info = load_json(info_path)
-            ds_version = info.get("codebase_version")
-            if isinstance(ds_version, str) and ds_version.startswith("v"):
-                version_tag = ds_version
-        except Exception as exc:  # noqa: BLE001
-            print(
-                f"[lerobot-annotate] could not read codebase_version from info.json ({exc}); falling back to {version_tag}",
-                flush=True,
+    version_tag = (
+        dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
     )
     revision = getattr(commit_info, "oid", None)
     tag_kwargs = {
@@ -180,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_train.py

@@ -171,6 +171,9 @@ def update_policy(
     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)
+    # 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
 
 
@@ -572,7 +575,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        train_tracker, output_dict = update_policy(
+        train_tracker, _ = update_policy(
             train_tracker,
             policy,
             batch,
@@ -605,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()

src/lerobot/utils/device_utils.py

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

src/lerobot/utils/import_utils.py

@@ -60,7 +60,17 @@ def is_package_available(
                     # If the package can't be imported, it's not available
                     package_exists = False
             else:
-                # For packages other than "torch", don't attempt the fallback and set as not available
+                # The distribution may be published under a name that differs from the
+                # import name (e.g. ``onnxruntime`` imports from ``onnxruntime-gpu`` /
+                # ``onnxruntime-silicon``). Resolve the import name to its actual
+                # distribution(s) and read the version from there before giving up.
+                try:
+                    dists = importlib.metadata.packages_distributions().get(import_name, [])
+                    if dists:
+                        package_version = importlib.metadata.version(dists[0])
+                    else:
+                        package_exists = False
+                except importlib.metadata.PackageNotFoundError:
                     package_exists = False
         logging.debug(f"Detected {pkg_name} version: {package_version}")
     if return_version:
@@ -123,6 +133,8 @@ _pyrealsense2_available = is_package_available("pyrealsense2") or is_package_ava
     "pyrealsense2-macosx", import_name="pyrealsense2"
 )
 _zmq_available = is_package_available("pyzmq", import_name="zmq")
+_onnxruntime_available = is_package_available("onnxruntime")
+_onnx_available = is_package_available("onnx")
 _hebi_available = is_package_available("hebi-py", import_name="hebi")
 _teleop_available = is_package_available("teleop")
 _placo_available = is_package_available("placo")

src/lerobot/utils/logging_utils.py

@@ -104,6 +104,7 @@ class MetricsTracker:
         "episodes",
         "epochs",
         "accelerator",
+        "_caller_metrics",
     ]
 
     def __init__(
@@ -129,6 +130,9 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
         self.accelerator = accelerator
+        # Meter names the caller registered up front. update_metrics() leaves these untouched, so a
+        # policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter.
+        self._caller_metrics: set[str] = set(self.metrics)
 
     def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
         if name in self.__dict__:
@@ -156,6 +160,21 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
 
+    def update_metrics(self, values: dict[str, Any]) -> None:
+        """Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
+
+        Non-numeric values and bools are ignored.
+        Caller-registered metrics (those passed to the constructor) are never overridden.
+        """
+        for name, value in values.items():
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                continue
+            if name in self._caller_metrics:
+                continue
+            if name not in self.metrics:
+                self.metrics[name] = AverageMeter(name, ":.3f", reduction="mean")
+            self.metrics[name].update(float(value))
+
     def reduce_across_ranks(self) -> None:
         """
         Synchronises the running averages of every metric whose ``reduction`` is not ``"none"``

tests/annotations/test_modules.py

@@ -85,7 +85,7 @@ def _spy_responder(captured: list[list[dict[str, Any]]], reply: Any):
 def test_module1_plan_memory_subtask_smoke(fixture_dataset_root: Path, tmp_path: Path) -> None:
     vlm = make_canned_responder(
         {
-            "atomic subtasks": {
+            "COMPLETED manipulation events": {
                 "subtasks": [
                     {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
                     {"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
@@ -126,7 +126,7 @@ def test_module1_emit_memory_false_skips_memory_keeps_subtasks_and_plan(
     leaving subtask + plan generation intact — symmetric to ``emit_plan``."""
     vlm = make_canned_responder(
         {
-            "atomic subtasks": {
+            "COMPLETED manipulation events": {
                 "subtasks": [
                     {"text": "grasp the handle of the sponge", "start": 0.0, "end": 0.4},
                     {"text": "wipe the counter from left to right", "start": 0.4, "end": 0.8},
@@ -318,7 +318,7 @@ def test_module1_attaches_contact_sheets_to_subtask_prompt(
                     return block.get("text", "")
         return ""
 
-    subtask_calls = [m for m in captured if "atomic subtasks" in _prompt_text(m)]
+    subtask_calls = [m for m in captured if "COMPLETED manipulation events" in _prompt_text(m)]
     assert len(subtask_calls) == 1, "expected exactly one subtask-prompt VLM call"
     content = subtask_calls[0][0]["content"]
     video_blocks = [b for b in content if isinstance(b, dict) and b.get("type") == "video"]

tests/policies/common/test_flow_matching.py

@@ -0,0 +1,193 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Behavior-pinning tests for the shared flow-matching sampling primitives.
+
+``euler_integrate`` is compared against a verbatim copy of the historical pi0/pi05/
+smolvla sampling loop (including its RTC hook semantics): any divergence from that
+reference is a behavior change for released checkpoints.
+"""
+
+import torch
+
+from lerobot.policies.common.flow_matching import (
+    euler_integrate,
+    sample_beta,
+    sample_noise,
+    sample_time_beta,
+)
+
+
+def test_sample_beta_range_dtype_and_reproducibility():
+    torch.manual_seed(0)
+    s1 = sample_beta(1.5, 1.0, 4096, "cpu")
+    torch.manual_seed(0)
+    s2 = sample_beta(1.5, 1.0, 4096, "cpu")
+    assert torch.equal(s1, s2)
+    assert s1.shape == (4096,) and s1.dtype == torch.float32
+    assert s1.min() >= 0.0 and s1.max() <= 1.0
+    # Beta(1.5, 1.0) mean is 1.5/2.5 = 0.6.
+    assert abs(s1.mean().item() - 0.6) < 0.02
+
+
+def test_sample_time_beta_openpi_convention():
+    torch.manual_seed(1)
+    time = sample_time_beta(4096, "cpu", alpha=1.5, beta=1.0, scale=0.999, offset=0.001)
+    assert time.dtype == torch.float32
+    assert time.min() >= 0.001 and time.max() <= 1.0
+    # Exact composition: Beta sample * scale + offset, same RNG stream.
+    torch.manual_seed(1)
+    expected = sample_beta(1.5, 1.0, 4096, "cpu") * 0.999 + 0.001
+    torch.testing.assert_close(time, expected, rtol=0, atol=0)
+
+
+def test_sample_noise_seeded():
+    torch.manual_seed(2)
+    n1 = sample_noise((2, 8, 4), "cpu")
+    torch.manual_seed(2)
+    n2 = sample_noise((2, 8, 4), "cpu")
+    assert torch.equal(n1, n2)
+    assert n1.dtype == torch.float32 and n1.shape == (2, 8, 4)
+
+
+def test_euler_integrate_constant_velocity_is_exact():
+    # With v_t == c constant, x_0 = x_1 + sum(dt * c) = x_1 - c exactly (num_steps * dt = -1).
+    noise = torch.randn(3, 5, 2)
+    c = torch.randn(3, 5, 2)
+    out = euler_integrate(lambda x_t, time: c, noise, num_steps=10)
+    torch.testing.assert_close(out, noise - c, rtol=0, atol=1e-6)
+
+
+def _reference_pi0_loop(denoise_fn, noise, num_steps, rtc_enabled, rtc_processor, kw):
+    """Verbatim structure of the historical pi0/pi05/smolvla sample_actions loop."""
+    bsize = noise.shape[0]
+    device = noise.device
+    dt = -1.0 / num_steps
+    x_t = noise
+    for step in range(num_steps):
+        time = 1.0 + step * dt
+        time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
+
+        def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
+            return denoise_fn(input_x_t, current_timestep)
+
+        if rtc_enabled:
+            v_t = rtc_processor.denoise_step(
+                x_t=x_t,
+                prev_chunk_left_over=kw.get("prev_chunk_left_over"),
+                inference_delay=kw.get("inference_delay"),
+                time=time,
+                original_denoise_step_partial=denoise_step_partial_call,
+                execution_horizon=kw.get("execution_horizon"),
+            )
+        else:
+            v_t = denoise_step_partial_call(x_t)
+        x_t = x_t + dt * v_t
+        if rtc_processor is not None and rtc_processor.is_debug_enabled():
+            rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
+    return x_t
+
+
+class _StubRTCProcessor:
+    def __init__(self, debug_enabled: bool):
+        self._debug = debug_enabled
+        self.tracked = []
+        self.guidance_calls = []
+
+    def is_debug_enabled(self):
+        return self._debug
+
+    def denoise_step(
+        self,
+        x_t,
+        prev_chunk_left_over,
+        inference_delay,
+        time,
+        original_denoise_step_partial,
+        execution_horizon,
+    ):
+        self.guidance_calls.append(
+            {
+                "time": time,
+                "inference_delay": inference_delay,
+                "execution_horizon": execution_horizon,
+                "x_t": x_t.clone(),
+            }
+        )
+        return original_denoise_step_partial(x_t) * 0.5
+
+    def track(self, time, x_t, v_t):
+        self.tracked.append({"time": time, "x_t": x_t.clone(), "v_t": v_t.clone()})
+
+
+def _make_denoise_fn():
+    weight = torch.randn(4, 4) * 0.1
+
+    def denoise_fn(x_t, time_tensor):
+        return x_t @ weight + time_tensor[:, None, None]
+
+    return denoise_fn
+
+
+def test_euler_integrate_matches_historical_loop():
+    torch.manual_seed(3)
+    denoise_fn = _make_denoise_fn()
+    noise = torch.randn(2, 6, 4)
+    ref = _reference_pi0_loop(denoise_fn, noise, 10, rtc_enabled=False, rtc_processor=None, kw={})
+    out = euler_integrate(denoise_fn, noise, 10)
+    assert torch.equal(out, ref)
+
+
+def test_euler_integrate_rtc_guidance_and_kwarg_forwarding():
+    torch.manual_seed(4)
+    denoise_fn = _make_denoise_fn()
+    noise = torch.randn(2, 6, 4)
+    leftover = torch.randn(2, 6, 4)
+    kw = {"inference_delay": 3, "prev_chunk_left_over": leftover, "execution_horizon": 25}
+
+    ref_proc, new_proc = _StubRTCProcessor(False), _StubRTCProcessor(False)
+    ref = _reference_pi0_loop(denoise_fn, noise, 6, rtc_enabled=True, rtc_processor=ref_proc, kw=kw)
+    out = euler_integrate(
+        denoise_fn,
+        noise,
+        6,
+        rtc_processor=new_proc,
+        rtc_enabled=True,
+        inference_delay=3,
+        prev_chunk_left_over=leftover,
+        execution_horizon=25,
+    )
+    assert torch.equal(out, ref)
+    assert len(new_proc.guidance_calls) == 6
+    for ref_call, new_call in zip(ref_proc.guidance_calls, new_proc.guidance_calls, strict=True):
+        assert ref_call["time"] == new_call["time"]
+        assert new_call["inference_delay"] == 3 and new_call["execution_horizon"] == 25
+        # Guidance sees the PRE-update x_t.
+        assert torch.equal(ref_call["x_t"], new_call["x_t"])
+
+
+def test_euler_integrate_debug_tracking_fires_even_when_rtc_disabled():
+    # Historical behavior: track() fires whenever the processor exists and has debugging
+    # enabled, independent of whether RTC guidance is active.
+    torch.manual_seed(5)
+    denoise_fn = _make_denoise_fn()
+    noise = torch.randn(2, 6, 4)
+    proc = _StubRTCProcessor(True)
+    out = euler_integrate(denoise_fn, noise, 4, rtc_processor=proc, rtc_enabled=False)
+    assert len(proc.guidance_calls) == 0
+    assert len(proc.tracked) == 4
+    # track() receives the POST-update x_t; the last one is the returned sample.
+    assert torch.equal(proc.tracked[-1]["x_t"], out)

tests/policies/common/test_vla_utils.py

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

tests/policies/evo1/test_evo1.py

@@ -496,60 +496,6 @@ def test_evo1_processor_save_load_round_trip_applies_config_overrides(tmp_path):
     assert "embodiment_id" in processed
 
 
-def test_reconcile_evo1_processors_repads_overridden_stats(tmp_path):
-    """Loading a checkpoint and injecting raw (unpadded) dataset stats must be re-padded.
-
-    Regression test: lerobot-train passes the raw dataset stats as normalizer/unnormalizer
-    overrides when resuming from a checkpoint (e.g. stage2 from a stage1 checkpoint). Those stats
-    are at the dataset dims (e.g. LIBERO state=8/action=7), but EVO1 pads state/action to
-    max_state_dim/max_action_dim before normalization, so reconcile_evo1_processors must re-pad the
-    stats or normalization crashes with a shape mismatch.
-    """
-    config = make_config()
-    preprocessor, postprocessor = make_evo1_pre_post_processors(config, dataset_stats=make_stats())
-    preprocessor.save_pretrained(tmp_path)
-    postprocessor.save_pretrained(tmp_path)
-
-    # Reload with the generic override path injecting raw, unpadded dataset stats.
-    raw_stats = make_stats()
-    loaded_pre = PolicyProcessorPipeline.from_pretrained(
-        tmp_path,
-        config_filename=f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
-        overrides={"normalizer_processor": {"stats": raw_stats}},
-        to_transition=batch_to_transition,
-        to_output=transition_to_batch,
-    )
-    loaded_post = PolicyProcessorPipeline.from_pretrained(
-        tmp_path,
-        config_filename=f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json",
-        overrides={"unnormalizer_processor": {"stats": raw_stats}},
-        to_transition=policy_action_to_transition,
-        to_output=transition_to_policy_action,
-    )
-
-    # Sanity: the override really injected unpadded stats before reconciliation.
-    normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
-    assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (STATE_DIM,)
-
-    loaded_pre, loaded_post = reconcile_evo1_processors(config, loaded_pre, loaded_post)
-
-    normalizer = next(step for step in loaded_pre.steps if isinstance(step, NormalizerProcessorStep))
-    unnormalizer = next(step for step in loaded_post.steps if isinstance(step, UnnormalizerProcessorStep))
-    assert normalizer._tensor_stats[OBS_STATE]["min"].shape == (MAX_STATE_DIM,)
-    assert normalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
-    assert unnormalizer._tensor_stats[ACTION]["min"].shape == (MAX_ACTION_DIM,)
-
-    # Normalizing a padded state must not raise (this is the exact runtime path that crashed).
-    processed = loaded_pre(
-        {
-            "task": "pick the block",
-            OBS_STATE: torch.zeros(STATE_DIM),
-            f"{OBS_IMAGES}.front": torch.rand(3, 16, 16),
-        }
-    )
-    assert processed[OBS_STATE].shape == (1, MAX_STATE_DIM)
-
-
 def test_evo1_policy_forward_and_inference_use_batched_embedding(monkeypatch):
     monkeypatch.setattr(modeling_evo1, "Evo1Model", DummyEvo1Model)
     policy = modeling_evo1.Evo1Policy(make_config())

tests/policies/wall_x/test_wallx.py

@@ -25,13 +25,57 @@ pytest.importorskip("transformers")
 pytest.importorskip("torchdiffeq")
 
 from lerobot.policies.factory import make_policy_config  # noqa: E402
-from lerobot.policies.wall_x import WallXConfig  # noqa: E402
+from lerobot.policies.wall_x import (
+    WallXConfig,  # noqa: E402
+)
 from lerobot.policies.wall_x.modeling_wall_x import WallXPolicy  # noqa: E402
 from lerobot.policies.wall_x.processor_wall_x import make_wall_x_pre_post_processors  # noqa: E402
+from lerobot.policies.wall_x.qwen_model import Qwen2_5_VLMoEModel, Qwen2_5_VLTextConfig  # noqa: E402
 from lerobot.utils.random_utils import set_seed  # noqa: E402
 from tests.utils import require_cuda, require_hf_token  # noqa: E402
 
 
+def test_moe_model_captures_requested_hidden_states_and_attentions():
+    hidden_size = 16
+    expert_config = {
+        "hidden_size": hidden_size,
+        "intermediate_size": 32,
+        "hidden_act": "silu",
+    }
+    config = Qwen2_5_VLTextConfig(
+        vocab_size=32,
+        hidden_size=hidden_size,
+        intermediate_size=32,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        max_position_embeddings=32,
+        layer_types=["full_attention", "full_attention"],
+        rope_parameters={
+            "rope_type": "default",
+            "rope_theta": 1_000_000.0,
+            "mrope_section": [1, 1, 0],
+        },
+        num_experts=2,
+        experts=[expert_config, expert_config],
+        dim_inputs=(hidden_size, hidden_size),
+        mlp_moe=True,
+    )
+    config._attn_implementation = "eager"
+    model = Qwen2_5_VLMoEModel(config)
+    input_ids = torch.tensor([[1, 2, 3]])
+
+    output = model(
+        input_ids=input_ids,
+        moe_token_types=torch.zeros_like(input_ids),
+        output_hidden_states=True,
+        output_attentions=True,
+    )
+
+    assert len(output.hidden_states) == config.num_hidden_layers + 1
+    assert len(output.attentions) == config.num_hidden_layers
+
+
 @require_cuda
 @require_hf_token
 def test_policy_instantiation():

tests/scripts/test_lerobot_annotate.py

@@ -18,6 +18,8 @@ import json
 from types import SimpleNamespace
 
 import pytest
+import requests
+from huggingface_hub.errors import RevisionNotFoundError
 
 # ``lerobot.scripts.lerobot_annotate`` (and the ``_push_to_hub`` path it
 # exercises) imports ``lerobot.datasets``, which only ships under the
@@ -26,11 +28,13 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
 
 
 def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
-    from lerobot.scripts.lerobot_annotate import _push_to_hub
+    from lerobot.scripts import lerobot_annotate
 
     root = tmp_path / "dataset"
     (root / "meta").mkdir(parents=True)
-    (root / "meta" / "info.json").write_text(json.dumps({"codebase_version": "v3.0"}))
+    (root / "meta" / "info.json").write_text(
+        json.dumps({"codebase_version": "v3.0", "fps": 30, "features": {}})
+    )
 
     calls = {}
 
@@ -43,9 +47,6 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
             return SimpleNamespace(oid="abc123")
 
         def delete_tag(self, repo_id, **kwargs):
-            import requests
-            from huggingface_hub.errors import RevisionNotFoundError
-
             calls["delete_tag"] = {"repo_id": repo_id, **kwargs}
             # Simulate the common case: no stale tag to delete.
             raise RevisionNotFoundError("no such tag", response=requests.Response())
@@ -53,7 +54,12 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
         def create_tag(self, **kwargs):
             calls["create_tag"] = kwargs
 
-    monkeypatch.setattr("huggingface_hub.HfApi", FakeHfApi)
+    monkeypatch.setattr(lerobot_annotate, "HfApi", FakeHfApi)
+
+    def fake_card_push(self, **kwargs):
+        calls["card_push"] = {"content": str(self), **kwargs}
+
+    monkeypatch.setattr("huggingface_hub.DatasetCard.push_to_hub", fake_card_push)
 
     cfg = SimpleNamespace(
         repo_id="source/dataset",
@@ -62,7 +68,7 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
         push_commit_message=None,
     )
 
-    _push_to_hub(root, cfg)
+    lerobot_annotate._push_to_hub(root, cfg)
 
     assert calls["create_repo"] == {
         "repo_id": "annotated/dataset",
@@ -71,6 +77,13 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
         "exist_ok": True,
     }
     assert calls["upload_folder"]["repo_id"] == "annotated/dataset"
+    # The source README must not be copied over: its links (e.g. the
+    # visualize badge) point at the source dataset. A card regenerated for
+    # the target repo is pushed instead.
+    assert "README.md" in calls["upload_folder"]["ignore_patterns"]
+    assert calls["card_push"]["repo_id"] == "annotated/dataset"
+    assert "visualize_dataset?path=annotated/dataset" in calls["card_push"]["content"]
+    assert "source/dataset" not in calls["card_push"]["content"]
     # A stale tag (e.g. from a previous annotation run) is deleted first so
     # the new tag always points at the upload we just made.
     assert calls["delete_tag"] == {

tests/utils/test_logging_utils.py

@@ -233,3 +233,37 @@ def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
     # accumulate against the cluster view rather than the stale per-rank sum.
     meter = tracker.update_s
     assert meter.sum / meter.count == pytest.approx(meter.avg)
+
+
+def test_metrics_tracker_update_metrics_registers_and_averages():
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
+    tracker.update_metrics({"latent_loss": 0.2, "action_loss": 0.4})
+    tracker.update_metrics({"latent_loss": 0.4, "action_loss": 0.6})
+
+    # New keys are auto-registered as mean-reduced meters and averaged over the window.
+    assert tracker.metrics["latent_loss"].reduction == "mean"
+    assert tracker.metrics["latent_loss"].avg == pytest.approx(0.3)
+    assert tracker.metrics["action_loss"].avg == pytest.approx(0.5)
+    assert tracker.to_dict()["latent_loss"] == pytest.approx(0.3)
+
+
+def test_metrics_tracker_update_metrics_skips_non_numeric():
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
+    tracker.update_metrics({"loss": 0.5, "head_mode": "sparse", "enabled": True})
+
+    # strings and bools ignored
+    assert "loss" in tracker.metrics
+    assert "head_mode" not in tracker.metrics
+    assert "enabled" not in tracker.metrics
+
+
+def test_metrics_tracker_update_metrics_does_not_override_caller_meter():
+    # A policy that echoes "loss" in its output dict must not overwrite the caller-owned,
+    # already-aggregated loss meter.
+    metrics = {"loss": AverageMeter("loss", ":.3f", reduction="mean")}
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
+    tracker.loss = 1.0  # caller-set optimized loss
+    tracker.update_metrics({"loss": 99.0, "latent_loss": 0.2})
+
+    assert tracker.metrics["loss"].avg == pytest.approx(1.0)  # snapshot ignored
+    assert tracker.metrics["latent_loss"].avg == pytest.approx(0.2)