Merge branch 'main' into fix/zero-shaped-features

Per review, log a warning when compute_episode_stats skips a feature with a
zero-width shape, so users know stats were intentionally not computed for it.

.github/workflows/documentation.yml

@@ -55,7 +55,7 @@ jobs:
       github.repository == 'huggingface/lerobot'
     permissions:
       contents: read
-    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@e60a538eea9817ab312196d0d233604b01697265  # main
+    uses: huggingface/doc-builder/.github/workflows/build_main_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391  # main
     with:
       commit_sha: ${{ github.sha }}
       package: lerobot
@@ -78,7 +78,7 @@ jobs:
     permissions:
       contents: read
       pull-requests: write
-    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@e60a538eea9817ab312196d0d233604b01697265  # main
+    uses: huggingface/doc-builder/.github/workflows/build_pr_documentation.yml@2430c1ec91d04667414e2fa31ecfc36c153ea391  # main
     with:
       commit_sha: ${{ github.event.pull_request.head.sha }}
       pr_number: ${{ github.event.number }}

docs/source/annotation_pipeline.mdx

@@ -81,12 +81,6 @@ 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)
@@ -163,33 +157,30 @@ 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).  |
-| `--vlm.parallel_servers`   | `1`                | Independent servers for round-robin routing (one per GPU).                           |
-| `--vlm.num_gpus`           | `0`                | GPUs per server (`0` = one each).                                                    |
-| `--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. |
+| 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). |
+| `--vlm.parallel_servers`   | `1`                | Independent servers for round-robin routing (one per GPU).                          |
+| `--vlm.num_gpus`           | `0`                | GPUs per server (`0` = one each).                                                   |
+| `--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.                                                               |
 
 ### 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).                                                                      |
-| `--plan.derive_task_from_video` | `if_short` | Use the dataset task as-is (`off`), only when it's missing/short (`if_short`), or always re-derive from video (`always`).    |
+| 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.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).                                                                   |
+| `--plan.derive_task_from_video` | `if_short` | Use the dataset task as-is (`off`), only when it's missing/short (`if_short`), or always re-derive from video (`always`). |
 
 ### Interjections + VQA
 

docs/source/groot.mdx

@@ -162,11 +162,11 @@ Preliminary LeRobot integration results (GR00T-LeRobot, `eval.n_episodes >= 50`
 
 | Suite            | Success rate | Checkpoint                                                                                                    |
 | ---------------- | -----------: | ------------------------------------------------------------------------------------------------------------- |
-| LIBERO Spatial   |          95% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
-| LIBERO Object    |         100% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640)   |
-| LIBERO Goal      |          98% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640)       |
-| LIBERO 10 (Long) |          93% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640)           |
-| **Average**      |    **96.5%** |                                                                                                               |
+| LIBERO Spatial   |          91% | [nvidia/gr00t17-lerobot-libero_spatial-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_spatial-640) |
+| LIBERO Object    |          81% | [nvidia/gr00t17-lerobot-libero_object-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_object-640)   |
+| LIBERO Goal      |          97% | [nvidia/gr00t17-lerobot-libero_goal-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_goal-640)       |
+| LIBERO 10 (Long) |          84% | [nvidia/gr00t17-lerobot-libero_10-640](https://huggingface.co/nvidia/gr00t17-lerobot-libero_10-640)           |
+| **Average**      |   **88.25%** |                                                                                                               |
 
 ```bash
 export MODEL_ID=your_trained_model_on_huggingface

src/lerobot/annotations/steerable_pipeline/config.py

@@ -65,14 +65,6 @@ 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
 
@@ -168,11 +160,6 @@ 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,16 +413,7 @@ 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()
+    font = ImageFont.load_default()
     label = f"{timestamp:06.2f}s"
     left, top, right, bottom = draw.textbbox((0, 0), label, font=font)
     text_w, text_h = right - left, bottom - top

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

@@ -116,8 +116,6 @@ 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:
@@ -511,51 +509,6 @@ 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,23 +22,12 @@ 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.
-
-    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
+    """Read prompt template ``name.txt`` from the ``prompts/`` directory."""
     path = _DIR / f"{name}.txt"
     return path.read_text(encoding="utf-8")

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

@@ -1,35 +0,0 @@
-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,68 +1,112 @@
-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}"
+You are labeling a teleoperated robot demonstration.
 
-{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.
+The user originally asked: "{episode_task}"
 
-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.
+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.
 
-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.
+{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.
 
-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.
+Authoring rules — Hi Robot atom granularity, pi0.7-style short prompts:
 
-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.
+- 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".
 
 Output strictly valid JSON of shape:
 
   {{
     "subtasks": [
-      {{"text": "<short imperative action label>", "start": <float>, "end": <float>}},
+      {{"text": "<short imperative verb phrase>", "start": <float>, "end": <float>}},
       ...
     ]
   }}

src/lerobot/annotations/steerable_pipeline/vlm_client.py

@@ -285,8 +285,6 @@ 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}
@@ -298,13 +296,7 @@ def _make_openai_client(config: VlmConfig) -> VlmClient:
             chosen = clients[rr_counter["i"] % len(clients)]
             rr_counter["i"] += 1
         response = chosen.chat.completions.create(**kwargs)
-        # 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 ""
+        return response.choices[0].message.content 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/default.py

@@ -93,9 +93,6 @@ class EvalConfig:
     recording_repo_id: str | None = None
     # Whether the pushed recording repositories should be private.
     recording_private: bool = False
-    # Whether to save the policy's imagined/predicted video (world-model policies only) as mp4s.
-    # Requests intermediate predictions from the policy each step; policies that produce none are unaffected.
-    save_predicted_video: bool = False
 
     def __post_init__(self) -> None:
         if self.recording_repo_id is not None and not self.recording:

src/lerobot/datasets/compute_stats.py

@@ -519,6 +519,13 @@ def compute_episode_stats(
         if features[key]["dtype"] in {"string", "language"}:
             continue
 
+        # Features with zero-width shapes are skipped (no data to compute stats on)
+        if any(d == 0 for d in features[key].get("shape", ())):
+            logging.debug(
+                f"Skipping statistics computation for feature '{key}' with a zero-width shape {features[key]['shape']}."
+            )
+            continue
+
         if features[key]["dtype"] in ["image", "video"]:
             ep_ft_array = sample_images(data)
             axes_to_reduce = (0, 2, 3)

src/lerobot/datasets/feature_utils.py

@@ -67,9 +67,9 @@ def get_hf_features_from_features(features: dict) -> datasets.Features:
         elif ft["shape"] == (1,):
             hf_features[key] = datasets.Value(dtype=ft["dtype"])
         elif len(ft["shape"]) == 1:
-            hf_features[key] = datasets.Sequence(
-                length=ft["shape"][0], feature=datasets.Value(dtype=ft["dtype"])
-            )
+            # pyarrow rejects fixed-size lists of length 0, so use a variable length list instead
+            length = ft["shape"][0] if ft["shape"][0] > 0 else -1
+            hf_features[key] = datasets.Sequence(length=length, feature=datasets.Value(dtype=ft["dtype"]))
         elif len(ft["shape"]) == 2:
             hf_features[key] = datasets.Array2D(shape=ft["shape"], dtype=ft["dtype"])
         elif len(ft["shape"]) == 3:

src/lerobot/policies/evo1/internvl3_embedder.py

@@ -27,11 +27,9 @@ 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)
@@ -137,13 +135,9 @@ 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 is_flash_attn_2_available()) else "eager"
-        )
+        attn_implementation = "flash_attention_2" if (use_flash_attn and _flash_attn_available()) else "eager"
         if use_flash_attn and attn_implementation == "eager":
-            logger.warning(
-                "Flash Attention 2 is unavailable on this runtime. Falling back to eager attention."
-            )
+            logger.warning("flash_attn is not installed. Falling back to eager attention.")
 
         self.model = AutoModel.from_pretrained(
             model_name,
@@ -365,3 +359,11 @@ 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/factory.py

@@ -29,10 +29,8 @@ from lerobot.configs import FeatureType, PreTrainedConfig
 from lerobot.envs import EnvConfig, env_to_policy_features
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
-    NormalizerProcessorStep,
     PolicyProcessorPipeline,
     RelativeActionsProcessorStep,
-    UnnormalizerProcessorStep,
     batch_to_transition,
     policy_action_to_transition,
     transition_to_batch,
@@ -86,52 +84,6 @@ def _reconnect_relative_absolute_steps(
             step.relative_step = relative_step
 
 
-def _ensure_relative_actions(
-    preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline, policy_cfg
-) -> None:
-    """Enable (or inject) the relative/absolute action steps in a loaded pipeline.
-
-    When loading from a pretrained checkpoint, the saved processor is authoritative. If the base
-    predates the relative-action feature (e.g. FastWAM/LingBot bases) its pipeline has no
-    RelativeActionsProcessorStep, so lerobot-train's override cannot enable one — those override
-    keys are popped before `from_pretrained` (else it raises) and we reconstruct the steps here.
-    Bases that DO ship the (disabled) steps (e.g. pi0/pi05) are simply flipped on. No-op unless
-    ``policy_cfg.use_relative_actions`` is set, so non-relative runs are untouched.
-    """
-    if not getattr(policy_cfg, "use_relative_actions", False):
-        return
-
-    exclude_joints = list(getattr(policy_cfg, "relative_exclude_joints", []) or [])
-    action_names = getattr(policy_cfg, "action_feature_names", None)
-
-    pre_steps = list(preprocessor.steps)
-    relative_step = next((s for s in pre_steps if isinstance(s, RelativeActionsProcessorStep)), None)
-    if relative_step is None:
-        relative_step = RelativeActionsProcessorStep(
-            enabled=True, exclude_joints=exclude_joints, action_names=action_names
-        )
-        # Insert right before the normalizer (raw -> relative -> normalize); fall back to the front.
-        idx = next((i for i, s in enumerate(pre_steps) if isinstance(s, NormalizerProcessorStep)), 0)
-        pre_steps.insert(idx, relative_step)
-        preprocessor.steps = pre_steps
-    else:
-        relative_step.enabled = True
-        relative_step.exclude_joints = exclude_joints
-        relative_step.action_names = action_names
-
-    post_steps = list(postprocessor.steps)
-    absolute_step = next((s for s in post_steps if isinstance(s, AbsoluteActionsProcessorStep)), None)
-    if absolute_step is None:
-        absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
-        # Insert right after the unnormalizer (unnormalize -> absolute); fall back to the front.
-        idx = next((i for i, s in enumerate(post_steps) if isinstance(s, UnnormalizerProcessorStep)), -1)
-        post_steps.insert(idx + 1, absolute_step)
-        postprocessor.steps = post_steps
-    else:
-        absolute_step.enabled = True
-        absolute_step.relative_step = relative_step
-
-
 def get_policy_class(name: str) -> type[PreTrainedPolicy]:
     """
     Retrieves a policy class by its registered name.
@@ -368,20 +320,12 @@ def make_pre_post_processors(
                 ),
             )
 
-        # The relative/absolute override keys only match if the saved base already contains those
-        # steps (e.g. pi0/pi05). For bases that predate the feature (FastWAM/LingBot) they would
-        # raise "Override keys ... do not match any step". Pop them here and let
-        # _ensure_relative_actions() enable-or-inject the steps after loading (handles both cases).
-        pre_overrides = dict(kwargs.get("preprocessor_overrides") or {})
-        post_overrides = dict(kwargs.get("postprocessor_overrides") or {})
-        pre_overrides.pop("relative_actions_processor", None)
-        post_overrides.pop("absolute_actions_processor", None)
         preprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_model_name_or_path=pretrained_path,
             config_filename=kwargs.get(
                 "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
             ),
-            overrides=pre_overrides,
+            overrides=kwargs.get("preprocessor_overrides", {}),
             to_transition=batch_to_transition,
             to_output=transition_to_batch,
             revision=pretrained_revision,
@@ -391,12 +335,11 @@ def make_pre_post_processors(
             config_filename=kwargs.get(
                 "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"
             ),
-            overrides=post_overrides,
+            overrides=kwargs.get("postprocessor_overrides", {}),
             to_transition=policy_action_to_transition,
             to_output=transition_to_policy_action,
             revision=pretrained_revision,
         )
-        _ensure_relative_actions(preprocessor, postprocessor, policy_cfg)
         _reconnect_relative_absolute_steps(preprocessor, postprocessor)
         if isinstance(policy_cfg, Evo1Config):
             from .evo1.processor_evo1 import reconcile_evo1_processors

src/lerobot/policies/fastwam/configuration_fastwam.py

@@ -27,8 +27,6 @@ from lerobot.configs import (
 from lerobot.optim import AdamWConfig
 from lerobot.utils.constants import ACTION, OBS_STATE
 
-from ..rtc.configuration_rtc import RTCConfig
-
 WAN22_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B"
 WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 FASTWAM_BASE_MODEL_ID = "lerobot/fastwam_base"
@@ -190,25 +188,12 @@ class FastWAMConfig(PreTrainedConfig):
     action_video_freq_ratio: int = 4
     image_size: tuple[int, int] = (224, 448)
     context_len: int = 128
-
-    # Relative actions: converts absolute actions to relative (action -= state) during
-    # preprocessing, and reverses it at postprocessing. Requires `proprio_dim` (OBS_STATE).
-    use_relative_actions: bool = False
-    # Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
-    relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
-    # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
-    action_feature_names: list[str] | None = None
     model_id: str = WAN22_MODEL_ID
     tokenizer_model_id: str = WAN_T5_TOKENIZER_ID
     text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID
     base_model_id: str | None = FASTWAM_BASE_MODEL_ID
     tokenizer_max_len: int = 128
     load_text_encoder: bool = True
-    # Device for the frozen ~11GB UMT5-XXL text encoder. `None` keeps it on the main
-    # policy `device` (default). Set to e.g. "cpu" to keep it off the GPU and save VRAM;
-    # prompts are then encoded on that device and the resulting embeddings moved to the
-    # policy device. Trades GPU memory for slower (CPU) text encoding.
-    text_encoder_device: str | None = None
     mot_checkpoint_mixed_attn: bool = False
     torch_dtype: str = "bfloat16"
     prompt_template: str = (
@@ -216,11 +201,6 @@ class FastWAMConfig(PreTrainedConfig):
     )
     num_inference_steps: int = 10
     inference_seed: int | None = 42
-    # Real-Time Chunking (RTC): async chunk generation with prefix guidance so a new chunk
-    # inpaints smoothly onto the still-executing tail of the previous one. `None` disables it
-    # (default synchronous inference). Consumed by `RTCInferenceEngine`, which calls
-    # `predict_action_chunk(..., inference_delay=, prev_chunk_left_over=)`.
-    rtc_config: RTCConfig | None = None
     rand_device: str = "cpu"
     text_cfg_scale: float = 1.0
     negative_prompt: str = ""
@@ -299,10 +279,6 @@ class FastWAMConfig(PreTrainedConfig):
         finally:
             self.pretrained_path = pretrained_path
 
-    @property
-    def chunk_size(self) -> int:
-        return self.action_horizon
-
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay)
 

src/lerobot/policies/fastwam/modeling_fastwam.py

@@ -22,7 +22,6 @@ import torch
 from torch import Tensor
 
 from lerobot.policies.pretrained import PreTrainedPolicy
-from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.import_utils import require_package
 
@@ -86,29 +85,8 @@ class FastWAMPolicy(PreTrainedPolicy):
                 for layer in mot.layers:
                     if "video" in layer.blocks:
                         layer.blocks["video"].requires_grad_(False)
-        self.init_rtc_processor()
         self.reset()
 
-    def init_rtc_processor(self) -> None:
-        """Attach a Real-Time Chunking processor to the core model when configured.
-
-        Mirrors the PI0/PI05 pattern: the policy owns the `RTCProcessor` and hands it to
-        the core `FastWAM` model, which consults it inside `infer_action`'s denoising loop.
-        Must stay public and named exactly `init_rtc_processor`: the rollout loader
-        (`lerobot.rollout.context`) sets `policy.config.rtc_config = cfg.inference.rtc` and
-        then calls `policy.init_rtc_processor()` to (re)build the processor after load, so
-        `--inference.type=rtc` alone is enough to enable guidance — no separate policy-side
-        `rtc_config` needed. A private/renamed method would be silently skipped (guidance
-        off), degrading RTC to unguided async chunk-swapping.
-        """
-        self.rtc_processor = None
-        if self.config.rtc_config is not None:
-            self.rtc_processor = RTCProcessor(self.config.rtc_config)
-            self.model.rtc_processor = self.rtc_processor
-
-    def _rtc_enabled(self) -> bool:
-        return self.config.rtc_config is not None and self.config.rtc_config.enabled
-
     @classmethod
     def _load_as_safetensor(cls, model, model_file: str, map_location: str, strict: bool):
         """Shape-aware load that supports cross-embodiment fine-tuning.
@@ -172,24 +150,6 @@ class FastWAMPolicy(PreTrainedPolicy):
 
     def reset(self) -> None:
         self._action_queue: deque[Tensor] = deque([], maxlen=self.config.n_action_steps)
-        # Per-episode text-embedding cache (mirrors LingBot-VA's `_prompt_embeds`). The task
-        # is fixed for an episode, so the ~11GB UMT5 encoder runs once on the first chunk and
-        # the resulting context is reused for every subsequent chunk. Cleared here on reset so
-        # a new episode's (possibly different) task is re-encoded. Proprio is still appended
-        # fresh each chunk downstream, so only the text-only context is cached.
-        self._cached_prompt: Any = None
-        self._cached_context: Tensor | None = None
-        self._cached_context_mask: Tensor | None = None
-
-    def _encode_prompt_cached(self, prompt: Any) -> tuple[Tensor, Tensor]:
-        """Encode `prompt` to `(context, context_mask)`, reusing the cache when the prompt is
-        unchanged so UMT5 runs at most once per episode (per distinct task)."""
-        if self._cached_context is None or self._cached_prompt != prompt:
-            context, context_mask = self.model.encode_prompt(prompt)
-            self._cached_prompt = prompt
-            self._cached_context = context
-            self._cached_context_mask = context_mask
-        return self._cached_context, self._cached_context_mask
 
     def _batch_to_training_sample(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
         """Adapt a standard LeRobot batch to the FastWAM-native sample that
@@ -223,9 +183,7 @@ class FastWAMPolicy(PreTrainedPolicy):
                 sample["proprio"] = state.unsqueeze(1) if state.ndim == 2 else state
         return sample
 
-    def forward(
-        self, batch: dict[str, Tensor], reduction: str = "mean"
-    ) -> tuple[Tensor, dict[str, Any]]:
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
         """Compute FastWAM training loss for a LeRobot batch.
 
         Args:
@@ -233,42 +191,24 @@ class FastWAMPolicy(PreTrainedPolicy):
                 (`video`, `action`, `context`, `context_mask`) or LeRobot keys
                 that can be adapted (`observation.images.*`, `observation.state`,
                 `action`, `action_is_pad`).
-            reduction (str): "mean" returns the scalar loss (default, backward
-                compatible); "none" returns per-sample losses of shape (batch_size,)
-                for sample weighting (RA-BC).
 
         Returns:
-            tuple[Tensor, dict[str, Any]]: The loss to backprop (scalar for "mean",
-            per-sample (B,) for "none"), and a dict of logging metrics (e.g.
-            `loss_video`, `loss_action`) — the `(loss, output_dict)` contract the
-            LeRobot training loop expects.
+            tuple[Tensor, dict[str, Any]]: The scalar loss to backprop, and a dict of
+            logging metrics (e.g. `loss_video`, `loss_action`) — the `(loss, output_dict)`
+            contract the LeRobot training loop expects.
         """
 
         sample = self._batch_to_training_sample(batch)
-        loss, metrics = self.model.training_loss(sample, reduction=reduction)
+        loss, metrics = self.model.training_loss(sample)
         return loss, dict(metrics or {})
 
     @torch.no_grad()
-    def predict_action_chunk(
-        self,
-        batch: dict[str, Tensor],
-        inference_delay: int | None = None,
-        prev_chunk_left_over: Tensor | None = None,
-        execution_horizon: int | None = None,
-        **_: Any,
-    ) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **_: Any) -> Tensor:
         """Predict a chunk of actions from the current FastWAM observation.
 
         Args:
             batch (dict[str, Tensor]): Inference batch with `input_image` or
                 image observation keys, plus `context/context_mask` or `prompt`.
-            inference_delay (int | None): RTC — number of prefix steps assumed already
-                executed by the time this chunk lands (from measured inference latency).
-            prev_chunk_left_over (Tensor | None): RTC — the previous chunk's unexecuted
-                action tail `[T_prev, action_dim]` in model space; guides denoising so the
-                new chunk inpaints onto it. `None` (default) = plain synchronous inference.
-            execution_horizon (int | None): RTC — override for the prefix-weight horizon;
-                `None` falls back to `rtc_config.execution_horizon`.
 
         Returns:
             Tensor: Action chunk with shape `[B, action_horizon, action_dim]`.
@@ -276,20 +216,6 @@ class FastWAMPolicy(PreTrainedPolicy):
 
         self.eval()
         infer_kwargs = _batch_to_infer_kwargs(batch=batch, config=self.config)
-        # Encode the task once per episode and reuse it (LingBot-VA parity): swap the raw
-        # `prompt` for the cached `context`/`context_mask` so `infer_action` skips `encode_prompt`
-        # and the text encoder isn't re-run every chunk. Skipped when the caller supplies its own
-        # precomputed `context` (the two are mutually exclusive downstream).
-        if infer_kwargs.get("context") is None and infer_kwargs.get("prompt") is not None:
-            context, context_mask = self._encode_prompt_cached(infer_kwargs["prompt"])
-            infer_kwargs["prompt"] = None
-            infer_kwargs["context"] = context
-            infer_kwargs["context_mask"] = context_mask
-        # RTC guidance args flow straight to `infer_action`; they are inert unless an
-        # RTCProcessor is attached, enabled, and `prev_chunk_left_over` is provided.
-        infer_kwargs["inference_delay"] = inference_delay
-        infer_kwargs["prev_chunk_left_over"] = prev_chunk_left_over
-        infer_kwargs["execution_horizon"] = execution_horizon
         batch_size = _infer_kwargs_batch_size(infer_kwargs)
         if batch_size == 1:
             action = _action_from_model_output(self.model.infer_action(**infer_kwargs))
@@ -334,10 +260,9 @@ class FastWAMPolicy(PreTrainedPolicy):
             mixtures={"video": video_expert, "action": action_expert},
             mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn,
         )
-        text_encoder_device = config.text_encoder_device or device
         text_encoder = (
             load_pretrained_wan_text_encoder(
-                model_id=config.text_encoder_model_id, torch_dtype=dtype, device=text_encoder_device
+                model_id=config.text_encoder_model_id, torch_dtype=dtype, device=device
             )
             if config.load_text_encoder
             else None
@@ -348,7 +273,6 @@ class FastWAMPolicy(PreTrainedPolicy):
             mot=mot,
             vae=load_pretrained_wan_vae(torch_dtype=dtype, device=device),
             text_encoder=text_encoder,
-            text_encoder_device=config.text_encoder_device,
             tokenizer=build_wan_tokenizer(
                 model_id=config.tokenizer_model_id, tokenizer_max_len=config.tokenizer_max_len
             ),

src/lerobot/policies/fastwam/processor_fastwam.py

@@ -21,16 +21,13 @@ import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
-    AbsoluteActionsProcessorStep,
     ActionProcessorStep,
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
     NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
-    ProcessorStep,
     ProcessorStepRegistry,
-    RelativeActionsProcessorStep,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
     policy_action_to_transition,
@@ -108,20 +105,10 @@ 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.
 
-    # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
-    # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
-    # below so its cached raw state (set during preprocessing) flows to postprocessing.
-    relative_step = RelativeActionsProcessorStep(
-        enabled=config.use_relative_actions,
-        exclude_joints=getattr(config, "relative_exclude_joints", []),
-        action_names=getattr(config, "action_feature_names", None),
-    )
-
-    input_steps: list[ProcessorStep] = [
+    input_steps = [
         RenameObservationsProcessorStep(rename_map={}),
         AddBatchDimensionProcessorStep(),
         DeviceProcessorStep(device=config.device),
-        relative_step,
         NormalizerProcessorStep(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
@@ -129,13 +116,12 @@ def make_fastwam_pre_post_processors(
             device=config.device,
         ),
     ]
-    output_steps: list[ProcessorStep] = [
+    output_steps = [
         UnnormalizerProcessorStep(
             features=config.output_features,
             norm_map=config.normalization_mapping,
             stats=normalization_stats,
         ),
-        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
     ]
     if config.toggle_action_dimensions:
         output_steps.append(

src/lerobot/policies/fastwam/wan/modular.py

@@ -839,7 +839,6 @@ class FastWAM(torch.nn.Module):
         text_dim: int | None = None,
         proprio_dim: int | None = None,
         device: str = "cpu",
-        text_encoder_device: str | torch.device | None = None,
         torch_dtype: torch.dtype = torch.float32,
         video_train_shift: float = 5.0,
         video_infer_shift: float = 5.0,
@@ -908,27 +907,12 @@ class FastWAM(torch.nn.Module):
         self.train_scheduler = self.train_video_scheduler
         self.infer_scheduler = self.infer_video_scheduler
 
-        # Optional Real-Time Chunking processor (set by the policy wrapper). When present and
-        # enabled it guides the action denoising loop in `infer_action` so a freshly generated
-        # chunk inpaints onto the previous chunk's unexecuted tail. Plain attribute (not an
-        # nn.Module) — carries no parameters and stays out of state_dict / device moves.
-        self.rtc_processor = None
         self.device = torch.device(device)
-        # When pinned (e.g. "cpu"), the frozen text encoder stays on this device instead
-        # of following the model onto the GPU — `_apply` skips it and `encode_prompt` runs
-        # it here, moving embeddings back to `self.device`. `None` = follow `self.device`.
-        self._text_encoder_device = (
-            torch.device(text_encoder_device) if text_encoder_device is not None else None
-        )
         self.torch_dtype = torch_dtype
         self.loss_lambda_video = float(loss_lambda_video)
         self.loss_lambda_action = float(loss_lambda_action)
 
         self.to(self.device)
-        # `self.to` above (via `_apply`) skips a pinned text encoder; make sure it actually
-        # sits on the pinned device (it was loaded there, but this is a cheap safety net).
-        if self.text_encoder is not None and self._text_encoder_device is not None:
-            self.text_encoder._apply(lambda t: t.to(self._text_encoder_device))
 
     @classmethod
     def from_wan22_pretrained(
@@ -1019,8 +1003,7 @@ class FastWAM(torch.nn.Module):
         # while staying out of `state_dict()` / `parameters()`.
         super()._apply(fn, *args, **kwargs)
         self.vae._apply(fn)
-        # A pinned text encoder (e.g. on CPU) must NOT follow device moves — leave it put.
-        if self.text_encoder is not None and self._text_encoder_device is None:
+        if self.text_encoder is not None:
             self.text_encoder._apply(fn)
         return self
 
@@ -1041,12 +1024,9 @@ class FastWAM(torch.nn.Module):
                 "Prompt encoding requires loaded text encoder/tokenizer. "
                 "Set `load_text_encoder=true` or provide precomputed `context/context_mask`."
             )
-        # Run the encoder on its own device (may be pinned to CPU to save VRAM), then
-        # move the resulting embeddings/mask to the model device for the DiT.
-        te_device = self._text_encoder_device or self.device
         ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True)
-        ids = ids.to(te_device)
-        mask = mask.to(te_device, dtype=torch.bool)
+        ids = ids.to(self.device)
+        mask = mask.to(self.device, dtype=torch.bool)
         prompt_emb = self.text_encoder(ids, mask)
         seq_lens = mask.gt(0).sum(dim=1).long()
         for i, v in enumerate(seq_lens):
@@ -1054,7 +1034,7 @@ class FastWAM(torch.nn.Module):
         # Match FastWAM/Wan2.2 context semantics: padding embeddings are zeroed,
         # while cross-attention still sees a fixed-length context.
         mask = torch.ones_like(mask)
-        return prompt_emb.to(device=self.device), mask.to(device=self.device)
+        return prompt_emb.to(device=self.device), mask
 
     def _append_proprio_to_context(
         self,
@@ -1379,9 +1359,7 @@ class FastWAM(torch.nn.Module):
         pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre)
         return pred_video, pred_action
 
-    def _compute_training_video_loss(
-        self, inputs, pred_video, target_video, timestep_video, reduction: str = "mean"
-    ):
+    def _compute_training_video_loss(self, inputs, pred_video, target_video, timestep_video):
         include_initial_video_step = inputs["first_frame_latents"] is None
         if inputs["first_frame_latents"] is not None:
             pred_video = pred_video[:, :, 1:]
@@ -1396,13 +1374,9 @@ class FastWAM(torch.nn.Module):
             loss_video_per_sample.device,
             dtype=loss_video_per_sample.dtype,
         )
-        weighted = loss_video_per_sample * video_weight
-        # reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC).
-        return weighted if reduction == "none" else weighted.mean()
+        return (loss_video_per_sample * video_weight).mean()
 
-    def _compute_training_action_loss(
-        self, inputs, pred_action, target_action, timestep_action, reduction: str = "mean"
-    ):
+    def _compute_training_action_loss(self, inputs, pred_action, target_action, timestep_action):
         action_loss_token = functional.mse_loss(
             pred_action.float(), target_action.float(), reduction="none"
         ).mean(dim=2)
@@ -1419,11 +1393,9 @@ class FastWAM(torch.nn.Module):
             action_loss_per_sample.device,
             dtype=action_loss_per_sample.dtype,
         )
-        weighted = action_loss_per_sample * action_weight
-        # reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC).
-        return weighted if reduction == "none" else weighted.mean()
+        return (action_loss_per_sample * action_weight).mean()
 
-    def training_loss(self, sample, tiled: bool = False, reduction: str = "mean"):
+    def training_loss(self, sample, tiled: bool = False):
         inputs = self.build_inputs(sample, tiled=tiled)
         targets = self._sample_training_targets(inputs)
         pred_video, pred_action = self._run_training_mot(inputs=inputs, targets=targets)
@@ -1432,20 +1404,17 @@ class FastWAM(torch.nn.Module):
             pred_video=pred_video,
             target_video=targets["target_video"],
             timestep_video=targets["timestep_video"],
-            reduction=reduction,
         )
         loss_action = self._compute_training_action_loss(
             inputs=inputs,
             pred_action=pred_action,
             target_action=targets["target_action"],
             timestep_action=targets["timestep_action"],
-            reduction=reduction,
         )
-        # With reduction="none" both terms are (B,), so loss_total is the per-sample loss (B,).
         loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action
         loss_dict = {
-            "loss_video": self.loss_lambda_video * float(loss_video.detach().mean().item()),
-            "loss_action": self.loss_lambda_action * float(loss_action.detach().mean().item()),
+            "loss_video": self.loss_lambda_video * float(loss_video.detach().item()),
+            "loss_action": self.loss_lambda_action * float(loss_action.detach().item()),
         }
         return loss_total, loss_dict
 
@@ -1830,9 +1799,6 @@ class FastWAM(torch.nn.Module):
         seed: int | None = None,
         rand_device: str = "cpu",
         tiled: bool = False,
-        inference_delay: int | None = None,
-        prev_chunk_left_over: torch.Tensor | None = None,
-        execution_horizon: int | None = None,
     ) -> dict[str, Any]:
         self.eval()
         if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal":
@@ -1885,43 +1851,18 @@ class FastWAM(torch.nn.Module):
             dtype=latents_action.dtype,
             shift_override=sigma_shift,
         )
-        rtc_active = (
-            self.rtc_processor is not None
-            and getattr(self.rtc_processor.rtc_config, "enabled", False)
-            and prev_chunk_left_over is not None
-        )
-        num_train_timesteps = float(self.infer_action_scheduler.num_train_timesteps)
         for step_t_action, step_delta_action in zip(infer_timesteps_action, infer_deltas_action, strict=True):
             timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device)
 
-            def denoise(x_t, ts=timestep_action):
-                return self._predict_action_noise_with_cache(
-                    latents_action=x_t,
-                    timestep_action=ts,
-                    context=context,
-                    context_mask=context_mask,
-                    video_kv_cache=video_kv_cache,
-                    attention_mask=attention_mask,
-                    video_seq_len=video_seq_len,
-                )
-
-            if rtc_active:
-                # `time` is the flow-matching noise level sigma in [0, 1]: FastWAM's model
-                # predicts velocity v = noise - clean, so the clean-action estimate is
-                # x1 = x_t - sigma * v — exactly RTC's `x1_t = x_t - time * v_t`.
-                sigma = float(step_t_action.item()) / num_train_timesteps
-                pred_action = self.rtc_processor.denoise_step(
-                    x_t=latents_action,
-                    prev_chunk_left_over=prev_chunk_left_over.to(
-                        device=latents_action.device, dtype=latents_action.dtype
-                    ),
-                    inference_delay=inference_delay or 0,
-                    time=sigma,
-                    original_denoise_step_partial=denoise,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                pred_action = denoise(latents_action)
+            pred_action = self._predict_action_noise_with_cache(
+                latents_action=latents_action,
+                timestep_action=timestep_action,
+                context=context,
+                context_mask=context_mask,
+                video_kv_cache=video_kv_cache,
+                attention_mask=attention_mask,
+                video_seq_len=video_seq_len,
+            )
 
             latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action)
 

src/lerobot/policies/lingbot_va/configuration_lingbot_va.py

@@ -28,11 +28,7 @@ from dataclasses import dataclass, field
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import (
-    ConstantWithWarmupSchedulerConfig,
-    CosineAnnealingWithWarmupSchedulerConfig,
-    LRSchedulerConfig,
-)
+from lerobot.optim.schedulers import ConstantWithWarmupSchedulerConfig, LRSchedulerConfig
 from lerobot.utils.constants import ACTION
 
 
@@ -96,15 +92,6 @@ class LingBotVAConfig(PreTrainedConfig):
     # (un)normalization quantiles live in the checkpoint's ``policy_postprocessor.json``, not here.
     used_action_channel_ids: list[int] = field(default_factory=lambda: list(range(7)))
 
-    # Relative actions: converts absolute actions to relative (action -= state) during
-    # preprocessing, and reverses it at postprocessing. Requires the dataset to provide
-    # observation.state whose leading dims align 1:1 with the used action channels.
-    use_relative_actions: bool = False
-    # Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
-    relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
-    # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
-    action_feature_names: list[str] | None = None
-
     # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s.
     save_predicted_video: bool = False
 
@@ -125,11 +112,6 @@ class LingBotVAConfig(PreTrainedConfig):
     optimizer_weight_decay: float = 1e-4
     optimizer_grad_clip_norm: float = 1.0
     scheduler_warmup_steps: int = 1000
-    # Scheduler after warmup. "constant_with_warmup" (upstream default: warmup then flat peak LR)
-    # or "cosine_annealing_with_warmup" (warmup then cosine anneal peak->0 over the remaining steps).
-    # Cosine tightens the loss tail and often nudges final loss down; it does NOT reduce the
-    # flow-matching estimator's step-to-step noise (that's metric variance, LR-independent).
-    scheduler_type: str = "constant_with_warmup"
 
     def __post_init__(self):
         super().__post_init__()
@@ -168,10 +150,7 @@ class LingBotVAConfig(PreTrainedConfig):
         )
 
     def get_scheduler_preset(self) -> LRSchedulerConfig | None:
-        # Default (upstream): linear warmup then constant LR (warmup_constant_lambda).
-        # Optionally cosine-anneal peak->0 over the remaining steps via scheduler_type.
-        if self.scheduler_type == "cosine_annealing_with_warmup":
-            return CosineAnnealingWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps)
+        # Upstream uses a linear warmup followed by a constant LR (warmup_constant_lambda).
         return ConstantWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps)
 
     @property

src/lerobot/policies/lingbot_va/modeling_lingbot_va.py

@@ -38,7 +38,7 @@ import torch.nn.functional as F  # noqa: N812
 from einops import rearrange
 from torch import Tensor
 
-from lerobot.policies.pretrained import PreTrainedPolicy, unpack_action_output
+from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.utils.constants import ACTION
 from lerobot.utils.import_utils import require_package
 
@@ -99,6 +99,8 @@ class LingBotVAPolicy(PreTrainedPolicy):
         # from ``config.wan_pretrained_path`` the first time inference runs.
         self._frozen: dict = {}
 
+        self.last_predicted_frames: Tensor | None = None
+        self.last_predicted_latents: Tensor | None = None
         self.reset()
 
     # Frozen-module lazy loading (VAE + UMT5 + tokenizer)
@@ -168,6 +170,8 @@ class LingBotVAPolicy(PreTrainedPolicy):
         self._prompt: str | None = None
         self._prompt_embeds = None
         self._negative_prompt_embeds = None
+        self.last_predicted_frames = None
+        self.last_predicted_latents = None
         self._use_cfg = (cfg.guidance_scale > 1) or (cfg.action_guidance_scale > 1)
         # Two independent flow-matching schedulers (video latent + action streams).
         self._scheduler = FlowMatchScheduler(shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True)
@@ -253,12 +257,8 @@ class LingBotVAPolicy(PreTrainedPolicy):
             "grid_id": grid_id,
         }
 
-    def _flow_matching_loss(self, input_dict, pred, reduction: str = "mean"):
-        """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``).
-
-        ``reduction="mean"`` returns scalar (latent_loss, action_loss); ``"none"`` returns
-        per-sample vectors of shape ``(B,)`` each (averaged over latent frames) for RA-BC.
-        """
+    def _flow_matching_loss(self, input_dict, pred):
+        """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``)."""
         latent_pred, action_pred = pred
         ld, ad = input_dict["latent_dict"], input_dict["action_dict"]
         action_pred = rearrange(action_pred, "b (f n) c -> b c f n 1", f=ad["targets"].shape[-3])
@@ -278,8 +278,7 @@ class LingBotVAPolicy(PreTrainedPolicy):
         latent_loss = (
             (latent_loss * lw[:, None, :, None, None]).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
         )
-        # per (batch*frame) mean over spatial/channel -> (B*F,)
-        latent_loss = latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)
+        latent_loss = (latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)).mean()
 
         amask = ad["actions_mask"].float()
         action_loss = F.mse_loss(action_pred.float(), ad["targets"].float().detach(), reduction="none")
@@ -287,14 +286,10 @@ class LingBotVAPolicy(PreTrainedPolicy):
             (action_loss * aw[:, None, :, None, None] * amask).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
         )
         amask_f = amask.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
-        action_loss = action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)
+        action_loss = (action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)).mean()
+        return latent_loss, action_loss
 
-        if reduction == "none":
-            # (B*F,) -> (B, F) -> (B,): per-sample losses for RA-BC weighting.
-            return latent_loss.reshape(bn, fn).mean(dim=1), action_loss.reshape(bn, fn).mean(dim=1)
-        return latent_loss.mean(), action_loss.mean()
-
-    def training_loss_from_streams(self, latents, actions, actions_mask, text_emb, reduction: str = "mean"):
+    def training_loss_from_streams(self, latents, actions, actions_mask, text_emb):
         """Core dual-stream training loss given prepared latents / actions / text embeddings.
 
         ``latents``: ``[B, in_channels, F, h, w]`` (normalized video latents).
@@ -323,24 +318,20 @@ class LingBotVAPolicy(PreTrainedPolicy):
             "window_size": int(torch.randint(4, 65, (1,)).item()),
         }
         pred = self.transformer(input_dict, train_mode=True)
-        latent_loss, action_loss = self._flow_matching_loss(input_dict, pred, reduction)
-        # reduction="none": latent_loss/action_loss are (B,) -> loss is per-sample (B,).
+        latent_loss, action_loss = self._flow_matching_loss(input_dict, pred)
         loss = latent_loss + action_loss
-        return loss, {"latent_loss": latent_loss.mean().item(), "action_loss": action_loss.mean().item()}
+        return loss, {"latent_loss": latent_loss.item(), "action_loss": action_loss.item()}
 
-    def forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict | None]:
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
         """Training forward: dual-stream flow-matching loss.
 
         Builds the (video-latent, action, text) training streams from a LeRobot batch
         (VAE-encoding the camera frames and UMT5-encoding the task), then runs the flow-matching
         dual-stream loss. Requires the policy to be built with ``attn_mode='flex'``.
-
-        ``reduction="mean"`` returns the scalar loss (default); ``"none"`` returns per-sample
-        losses of shape ``(B,)`` for sample weighting (RA-BC).
         """
         self._ensure_frozen_modules()
         latents, actions, actions_mask, text_emb = self._build_training_streams(batch)
-        return self.training_loss_from_streams(latents, actions, actions_mask, text_emb, reduction=reduction)
+        return self.training_loss_from_streams(latents, actions, actions_mask, text_emb)
 
     @torch.no_grad()
     def _build_training_streams(self, batch):
@@ -409,31 +400,22 @@ class LingBotVAPolicy(PreTrainedPolicy):
         return torch.cat(per_cam, dim=-1).to(self.config.device)
 
     @torch.no_grad()
-    def select_action(
-        self, batch: dict[str, Tensor], return_intermediate_predictions: bool = False, **kwargs
-    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
         """Return one action, refilling the chunk (and feeding back observed keyframes) as needed.
 
         Mirrors the upstream LIBERO client loop (``evaluation/libero/client.py``): the first obs is
         the conditioning frame; every observation produced afterwards is buffered as a keyframe and,
         once the chunk's actions are exhausted, the buffered frames + executed actions are fed back
         into the KV cache before the next chunk is predicted.
-
-        When ``return_intermediate_predictions=True`` returns ``(action, predictions)``. Predictions
-        are produced only on the ticks that predict a fresh chunk (first tick and each chunk refill);
-        on the intermediate ticks that just pop a cached action, ``predictions`` is an empty dict.
         """
         self.eval()
         self._ensure_frozen_modules()
         self._maybe_init_prompt(batch)
 
-        predictions: dict[str, Tensor] = {}
         if not self._started:
             # First call: this observation conditions the first chunk (it is *not* a keyframe).
             self._started = True
-            actions, predictions = unpack_action_output(
-                self.predict_action_chunk(batch, return_intermediate_predictions=return_intermediate_predictions)
-            )  # [B, chunk_size, n_used]
+            actions = self.predict_action_chunk(batch)  # [B, chunk_size, n_used]
             self._action_queue.extend(actions.transpose(0, 1))  # [chunk_size, B, n_used]
             self._obs_buffer = []
             self._exec_step = 0
@@ -445,31 +427,17 @@ class LingBotVAPolicy(PreTrainedPolicy):
             if len(self._action_queue) == 0:
                 # All actions for the current chunk have been executed; feed the observed
                 # keyframes + executed actions back and predict the next chunk.
-                actions, predictions = unpack_action_output(
-                    self.predict_action_chunk(
-                        None, return_intermediate_predictions=return_intermediate_predictions
-                    )
-                )
+                actions = self.predict_action_chunk(None)
                 self._action_queue.extend(actions.transpose(0, 1))
                 self._exec_step = 0
 
         self._prev_j = self._exec_step % self.config.action_per_frame
         self._exec_step += 1
-        action = self._action_queue.popleft()
-        if return_intermediate_predictions:
-            return action, predictions
-        return action
+        return self._action_queue.popleft()
 
     @torch.no_grad()
-    def predict_action_chunk(
-        self, batch: dict[str, Tensor], return_intermediate_predictions: bool = False, **kwargs
-    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
-        """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized).
-
-        When ``return_intermediate_predictions=True`` returns ``(actions, predictions)`` where
-        ``predictions`` holds this chunk's VAE-decoded imagined video under ``"images.predicted"``
-        (``[T, H, W, 3]`` uint8 on CPU).
-        """
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+        """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized)."""
         self.eval()
         self._ensure_frozen_modules()
         self._maybe_init_prompt(batch)
@@ -491,6 +459,12 @@ class LingBotVAPolicy(PreTrainedPolicy):
         # actions: [B, action_dim, F, action_per_frame, 1] (model-normalized). Keep for KV feedback.
         self._executed_actions = actions
 
+        if self.config.save_predicted_video:
+            # Match upstream LingBot-VA visualization: collect chunk latents and decode the
+            # concatenated latent sequence once after the rollout finishes.
+            self.last_predicted_frames = None
+            self.last_predicted_latents = latents.detach().to("cpu")
+
         # On the first chunk, frame 0 is the conditioning frame (already "known"): the upstream
         # LIBERO client skips it (start_idx=1), so we drop the first frame's actions here.
         used = self.config.used_action_channel_ids
@@ -499,15 +473,7 @@ class LingBotVAPolicy(PreTrainedPolicy):
             a = a[:, :, 1:]  # drop frame 0 -> (F-1) frames of actions
         a = a.squeeze(-1).flatten(2)  # [B, n_used, n_steps]
         a = a.transpose(1, 2).contiguous()  # [B, n_steps, n_used]
-        a = a.to(torch.float32)
-
-        if return_intermediate_predictions:
-            # Decode this chunk's imagined video for visualization / eval. Per-chunk decode (the VAE
-            # has no streaming decoder) may differ slightly at chunk boundaries from a single decode
-            # over the whole concatenated latent sequence; acceptable for monitoring/inspection.
-            frames = self._decode_predicted_video(latents)  # [T, H, W, 3] uint8, CPU
-            return a, {"images.predicted": frames}
-        return a
+        return a.to(torch.float32)
 
     # Prompt / text encoding
     def _maybe_init_prompt(self, batch):
@@ -868,6 +834,11 @@ class LingBotVAPolicy(PreTrainedPolicy):
         return actions, latents
 
     # Predicted-video decoding (opt-in)
+    @torch.no_grad()
+    def decode_predicted_latents(self, latents) -> Tensor:
+        """Decode a concatenated predicted-latent sequence into ``[T, H, W, 3]`` uint8 frames."""
+        return self._decode_predicted_video(latents)
+
     @torch.no_grad()
     def _decode_predicted_video(self, latents) -> Tensor:
         """VAE-decode predicted latents into a uint8 frame stack ``[T, H, W, 3]`` on CPU."""

src/lerobot/policies/lingbot_va/processor_lingbot_va.py

@@ -25,14 +25,12 @@ import torch
 
 from lerobot.configs.types import FeatureType, NormalizationMode
 from lerobot.processor import (
-    AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
     NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
-    RelativeActionsProcessorStep,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
 )
@@ -54,19 +52,9 @@ def make_lingbot_va_pre_post_processors(
 ]:
     """Build the pre/post processor pipelines for LingBot-VA."""
 
-    # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
-    # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
-    # below so its cached raw state (set during preprocessing) flows to postprocessing.
-    relative_step = RelativeActionsProcessorStep(
-        enabled=config.use_relative_actions,
-        exclude_joints=getattr(config, "relative_exclude_joints", []),
-        action_names=getattr(config, "action_feature_names", None),
-    )
-
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),
         AddBatchDimensionProcessorStep(),
-        relative_step,
         NormalizerProcessorStep(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
@@ -75,16 +63,13 @@ def make_lingbot_va_pre_post_processors(
         DeviceProcessorStep(device=config.device),
     ]
 
-    # Unnormalize actions back to physical units. Config-driven norm_map (was hardcoded QUANTILES)
-    # so it stays symmetric with the preprocessor's NormalizerProcessorStep — required for
-    # use_relative_actions with ACTION=IDENTITY (and unchanged for QUANTILES runs).
+    # Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
     output_steps: list[ProcessorStep] = [
         UnnormalizerProcessorStep(
             features=config.output_features,
-            norm_map=config.normalization_mapping,
+            norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
             stats=dataset_stats,
         ),
-        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
         DeviceProcessorStep(device="cpu"),
     ]
 

src/lerobot/policies/pretrained.py

@@ -23,13 +23,12 @@ 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
-import packaging.version
-import safetensors
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
-import torch
 from torch import Tensor, nn
 
 from lerobot.__version__ import __version__
@@ -94,18 +93,6 @@ def _build_card_context(
 
 class ActionSelectKwargs(TypedDict, total=False):
     noise: Tensor | None
-    return_intermediate_predictions: bool
-
-
-def unpack_action_output(out: Tensor | tuple[Tensor, dict[str, Tensor]]) -> tuple[Tensor, dict[str, Tensor]]:
-    """Normalize a ``select_action`` / ``predict_action_chunk`` return to ``(action, predictions)``.
-
-    These methods return a bare action ``Tensor`` by default, or a ``(action, predictions)`` tuple when
-    called with ``return_intermediate_predictions=True``. A bare tensor becomes ``(tensor, {})``.
-    """
-    if isinstance(out, tuple):
-        return out[0], out[1]
-    return out, {}
 
 
 class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
@@ -234,14 +221,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        # safetensors' load_file 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.
-        if map_location == "cuda" and torch.cuda.is_available():
-            map_location = f"cuda:{torch.cuda.current_device()}"
-
         # Create base kwargs
         kwargs = {"strict": strict}
 
@@ -252,6 +231,16 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         # 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."
+            )
+            model.to(map_location)
         return model
 
     @abc.abstractmethod
@@ -284,34 +273,20 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def predict_action_chunk(
-        self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]
-    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
         """Returns the action chunk (for action chunking policies) for a given observation, potentially in batch mode.
 
         Child classes using action chunking should use this method within `select_action` to form the action chunk
         cached for selection.
-
-        By default returns just the action `Tensor`. If `return_intermediate_predictions=True`,
-        returns `(action, predictions)` where `predictions` is a (possibly empty) `dict[str, Tensor]`
-        of additional model predictions a policy may expose (e.g. world-model predicted frames).
-        Policies that produce nothing extra may ignore the kwarg.
         """
         raise NotImplementedError
 
     @abc.abstractmethod
-    def select_action(
-        self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]
-    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
+    def select_action(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
         """Return one action to run in the environment (potentially in batch mode).
 
         When the model uses a history of observations, or outputs a sequence of actions, this method deals
         with caching.
-
-        By default returns just the action `Tensor`. If `return_intermediate_predictions=True`,
-        returns `(action, predictions)` where `predictions` is a (possibly empty) `dict[str, Tensor]`
-        of additional model predictions a policy may expose (e.g. world-model predicted frames).
-        Policies that produce nothing extra may ignore the kwarg.
         """
         raise NotImplementedError
 

src/lerobot/policies/rtc/action_queue.py

@@ -236,21 +236,11 @@ class ActionQueue:
         if action_index_before_inference is not None:
             indexes_diff = max(0, self.last_index - action_index_before_inference)
             if indexes_diff != real_delay:
-                # The latency estimate (`real_delay`) and the number of actions the robot
-                # actually consumed during inference (`indexes_diff`) disagree. This happens
-                # when the queue starved (robot idle) or on the first chunk (nothing consumed
-                # yet). Discarding `real_delay` here would drop actions the arm never executed
-                # and splice the queue `real_delay` steps ahead of the physical pose — a hard
-                # jump/slam, worst on slow policies where `real_delay` is large. Never discard
-                # more than was actually consumed.
-                resolved = min(real_delay, indexes_diff)
                 logger.warning(
-                    "Indexes diff != real delay (indexes_diff=%d, real_delay=%d); "
-                    "clamping discard to %d to avoid a queue-splice jump.",
+                    "Indexes diff is not equal to real delay. indexes_diff=%d, real_delay=%d",
                     indexes_diff,
                     real_delay,
-                    resolved,
                 )
-                return resolved
+                return real_delay
 
         return effective_delay

src/lerobot/policies/vla_jepa/action_head.py

@@ -282,7 +282,6 @@ class VLAJEPAActionHead(nn.Module):
         actions: torch.Tensor,
         state: torch.Tensor | None = None,
         action_is_pad: torch.Tensor | None = None,
-        reduction: str = "mean",
     ) -> torch.Tensor:
         noise = torch.randn_like(actions)
         t = self.sample_time(actions.shape[0], actions.device, actions.dtype)
@@ -303,10 +302,6 @@ class VLAJEPAActionHead(nn.Module):
 
         loss = F.mse_loss(pred_actions, velocity, reduction="none")  # [B, T, action_dim]
         valid_mask = ~action_is_pad.unsqueeze(-1)  # [B, T, 1]
-        if reduction == "none":
-            # Per-sample loss (B,) for sample weighting (RA-BC): mask-average over T and action_dim.
-            per_sample_valid = valid_mask.sum(dim=(1, 2)) * loss.shape[-1]  # [B]
-            return (loss * valid_mask).sum(dim=(1, 2)) / per_sample_valid.clamp_min(1)
         num_valid = valid_mask.sum() * loss.shape[-1]
         return (loss * valid_mask).sum() / num_valid.clamp_min(1)
 

src/lerobot/policies/vla_jepa/configuration_vla_jepa.py

@@ -56,14 +56,6 @@ class VLAJEPAConfig(PreTrainedConfig):
     action_dim: int = 7
     state_dim: int = 8
 
-    # Relative actions: converts absolute actions to relative (action -= state) during
-    # preprocessing, and reverses it at postprocessing. Requires `state_dim` (OBS_STATE).
-    use_relative_actions: bool = False
-    # Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
-    relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
-    # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
-    action_feature_names: list[str] | None = None
-
     num_action_tokens_per_timestep: int = 8
     num_embodied_action_tokens_per_instruction: int = 32
     num_inference_timesteps: int = 4

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -194,12 +194,8 @@ class VLAJEPAModel(nn.Module):
             )
         return embodied_action_tokens, action_tokens
 
-    def _world_model_loss(self, videos: Tensor, action_tokens: Tensor, reduction: str = "mean") -> Tensor:
-        """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1].
-
-        `reduction="none"` returns a per-sample loss (B,) for sample weighting (RA-BC);
-        "mean" returns the scalar loss.
-        """
+    def _world_model_loss(self, videos: Tensor, action_tokens: Tensor) -> Tensor:
+        """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1]."""
         # Match the world model's expected view count: pad with the first view, or trim extras.
         num_views = self.config.jepa_tubelet_size
         if videos.shape[1] < num_views:
@@ -227,8 +223,7 @@ class VLAJEPAModel(nn.Module):
         # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression
         t_enc_total = self.config.num_video_frames // tubelet_size
         if t_enc_total < 2:
-            zero_shape = (video_embeddings.shape[0],) if reduction == "none" else ()
-            return torch.zeros(zero_shape, device=video_embeddings.device)
+            return torch.zeros((), device=video_embeddings.device)
 
         # Shift-by-one JEPA split: input_states = positions 0..T-2, gt_states = positions 1..T-1
         t_enc_ctx = t_enc_total - 1
@@ -244,10 +239,6 @@ class VLAJEPAModel(nn.Module):
         predicted_states = self.video_predictor(
             input_states.float(), action_tokens[:, :expected_actions].float()
         )
-        if reduction == "none":
-            # Per-sample loss (B,): mean over all non-batch dims (tokens, feature).
-            l = F.l1_loss(predicted_states, gt_states.float(), reduction="none")
-            return l.mean(dim=tuple(range(1, l.ndim)))
         return F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
 
     def _action_loss(
@@ -256,27 +247,17 @@ class VLAJEPAModel(nn.Module):
         actions: Tensor,
         state: Tensor | None,
         action_is_pad: Tensor | None,
-        reduction: str = "mean",
     ) -> Tensor:
-        """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`.
-
-        `reduction="none"` returns a per-sample loss (B,) — the `repeated_diffusion_steps`
-        independent noise draws are averaged back per original sample — for RA-BC weighting.
-        """
+        """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`."""
         device_type = next(self.parameters()).device.type
         with torch.autocast(device_type=device_type, dtype=torch.float32):
             r = self.config.repeated_diffusion_steps
             horizon = self.config.chunk_size
-            b = embodied_action_tokens.shape[0]
             actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)
             embodied = embodied_action_tokens.repeat(r, 1, 1)
             state_rep = state.to(embodied_action_tokens.dtype).repeat(r, 1, 1) if state is not None else None
             pad_rep = action_is_pad[:, -horizon:].repeat(r, 1) if action_is_pad is not None else None
-            loss = self.action_model(embodied, actions_target, state_rep, pad_rep, reduction=reduction)
-            if reduction == "none":
-                # `.repeat(r, 1, 1)` tiles as [rep0(b0..b_{B-1}), rep1(...), ...] → (r, B); mean over reps.
-                return loss.view(r, b).mean(dim=0)
-            return loss
+            return self.action_model(embodied, actions_target, state_rep, pad_rep)
 
     def forward(
         self,
@@ -286,29 +267,21 @@ class VLAJEPAModel(nn.Module):
         actions: Tensor | None = None,
         state: Tensor | None = None,
         action_is_pad: Tensor | None = None,
-        reduction: str = "mean",
     ) -> dict[str, Tensor]:
-        """Native forward: Qwen encode → optional world-model loss → optional action-head loss.
-
-        `reduction="none"` makes both loss terms per-sample (B,) for RA-BC weighting; "mean"
-        returns scalar losses.
-        """
+        """Native forward: Qwen encode → optional world-model loss → optional action-head loss."""
         embodied_action_tokens, action_tokens = self._encode_qwen(
             images, instructions, need_action_tokens=self.config.enable_world_model
         )
 
         if self.config.enable_world_model and videos is not None:
-            wm_loss = self._world_model_loss(videos, action_tokens, reduction=reduction)
+            wm_loss = self._world_model_loss(videos, action_tokens)
         else:
-            zero_shape = (embodied_action_tokens.shape[0],) if reduction == "none" else ()
-            wm_loss = torch.zeros(zero_shape, device=embodied_action_tokens.device)
+            wm_loss = torch.zeros((), device=embodied_action_tokens.device)
 
         if actions is None:
             return {"wm_loss": wm_loss}
 
-        action_loss = self._action_loss(
-            embodied_action_tokens, actions, state, action_is_pad, reduction=reduction
-        )
+        action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad)
         return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight}
 
     # ---- Native predict_action (follows original VLA_JEPA.predict_action) ----
@@ -394,19 +367,12 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         batch_size = batch[image_keys[0]].shape[0]
 
         # Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages.
-        # Resize to config.resize_images_to (as predict_action does) so training and inference feed
-        # Qwen the same resolution. Critical for memory: native camera frames (e.g. 720x1280) would
-        # otherwise blow up the Qwen3-VL vision-tower attention (patch count grows with resolution).
-        resize_hw = tuple(self.config.resize_images_to) if self.config.resize_images_to else None
         frames = []
         for key in image_keys:
             t = batch[key]
             if t.ndim == 5:  # [B, T, C, H, W] -> current observation (delta=0)
                 t = t[:, 0]
-            px = self.model.qwen.to_pixel_values(t)  # [B, C, H, W]
-            if resize_hw is not None and tuple(px.shape[-2:]) != resize_hw:
-                px = F.interpolate(px.float(), size=resize_hw, mode="area")
-            frames.append(px)
+            frames.append(self.model.qwen.to_pixel_values(t))
         images = [[frame[b] for frame in frames] for b in range(batch_size)]
 
         tasks = batch.get("task")
@@ -422,26 +388,7 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         # Videos [B, V, T, C, H, W] - only assembled during training when the world model consumes them.
         if self.model.config.enable_world_model and training:
             views = [batch[k].unsqueeze(1) if batch[k].ndim == 4 else batch[k] for k in image_keys]
-            # The world model consumes a SINGLE stacked [B, V, T, C, H, W] tensor, so all camera
-            # views must share a spatial size. Cameras can differ (e.g. base 480x640 vs wrist
-            # 720x1280), so resize each view to a common size before stacking — config.resize_images_to
-            # if set (same target predict_action uses), else the first view's size (a no-op when all
-            # views already match, preserving behavior for single-resolution datasets). The vjepa video
-            # processor does the final resize to the encoder resolution downstream.
-            cfg = self.model.config
-            target_hw = tuple(cfg.resize_images_to) if cfg.resize_images_to else tuple(views[0].shape[-2:])
-            resized = []
-            for v in views:
-                if tuple(v.shape[-2:]) != target_hw:
-                    b, t, c = v.shape[0], v.shape[1], v.shape[2]
-                    v = F.interpolate(
-                        v.reshape(b * t, c, v.shape[3], v.shape[4]).float(),
-                        size=target_hw,
-                        mode="bilinear",
-                        align_corners=False,
-                    ).reshape(b, t, c, target_hw[0], target_hw[1])
-                resized.append(v)
-            inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(resized, dim=1))
+            inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(views, dim=1))
 
         actions = batch.get(ACTION)
         if actions is not None:
@@ -459,17 +406,15 @@ class VLAJEPAPolicy(PreTrainedPolicy):
 
     # ---- LeRobot Policy Interface ----
 
-    def forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict]:
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """LeRobot train forward: convert → native forward → aggregate losses."""
-        native_output = self.model.forward(
-            **self._prepare_model_inputs(batch, training=True), reduction=reduction
-        )
+        native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True))
 
         ref = next(iter(native_output.values()))
-        zero = torch.zeros_like(ref)
+        zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
         total_loss = native_output.get("action_loss", zero) + native_output.get("wm_loss", zero)
-        logs = {k: v.detach().mean().item() for k, v in native_output.items()}
-        logs["loss"] = total_loss.detach().mean().item()
+        logs = {k: v.detach().item() for k, v in native_output.items()}
+        logs["loss"] = total_loss.detach().item()
         return total_loss, logs
 
     def get_optim_params(self) -> dict:

src/lerobot/policies/vla_jepa/processor_vla_jepa.py

@@ -20,7 +20,6 @@ import torch
 
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from lerobot.processor import (
-    AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
     EnvTransition,
@@ -29,7 +28,6 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
-    RelativeActionsProcessorStep,
     RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
@@ -114,21 +112,10 @@ def make_vla_jepa_pre_post_processors(
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
 ]:
     features = {**config.input_features, **config.output_features}
-
-    # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
-    # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
-    # below so its cached raw state (set during preprocessing) flows to postprocessing.
-    relative_step = RelativeActionsProcessorStep(
-        enabled=config.use_relative_actions,
-        exclude_joints=getattr(config, "relative_exclude_joints", []),
-        action_names=getattr(config, "action_feature_names", None),
-    )
-
     input_steps = [
         RenameObservationsProcessorStep(rename_map={}),
         AddBatchDimensionProcessorStep(),
         DeviceProcessorStep(device=config.device),
-        relative_step,
         NormalizerProcessorStep(
             features=features,
             norm_map=config.normalization_mapping,
@@ -149,11 +136,6 @@ def make_vla_jepa_pre_post_processors(
             stats=dataset_stats,
         )
     )
-    # Reverse the relative conversion on the unnormalized action, before gripper binarization.
-    # gripper is kept absolute by relative_exclude_joints, so the two steps touch disjoint dims.
-    output_steps.append(
-        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step)
-    )
     if config.binarize_gripper_action:
         output_steps.append(
             BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)

src/lerobot/processor/relative_action_processor.py

@@ -51,12 +51,6 @@ def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) ->
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
     if state.device != actions.device or state.dtype != actions.dtype:
         state = state.to(device=actions.device, dtype=actions.dtype)
-    # When the observation is temporally stacked (e.g. LingBot-VA loads several obs steps via
-    # observation_delta_indices, giving state shape (B, T_obs, state_dim)), the relative reference
-    # is the CURRENT frame (delta == 0, i.e. index 0). Collapse to it so the offset broadcasts over
-    # the action horizon. pi0/pi05 pass a 2D (B, state_dim) state and are unaffected.
-    if state.ndim == 3:
-        state = state[:, 0]
     state_offset = state[..., :dims] * mask_t
     if actions.ndim == 3:
         state_offset = state_offset.unsqueeze(-2)
@@ -79,10 +73,6 @@ def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) ->
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
     if state.device != actions.device or state.dtype != actions.dtype:
         state = state.to(device=actions.device, dtype=actions.dtype)
-    # Mirror to_relative_actions: collapse a temporally-stacked (B, T_obs, state_dim) state to the
-    # current frame (index 0) so the round-trip stays symmetric with the relative conversion.
-    if state.ndim == 3:
-        state = state[:, 0]
     state_offset = state[..., :dims] * mask_t
     if actions.ndim == 3:
         state_offset = state_offset.unsqueeze(-2)
@@ -136,7 +126,7 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
-        # Always cache state for the paired AbsoluteActionsProcessorStep.
+        # Always cache state for the paired AbsoluteActionsProcessorStep
         if state is not None:
             self._last_state = state
 
@@ -156,11 +146,6 @@ class RelativeActionsProcessorStep(ProcessorStep):
         """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
         return self._last_state
 
-    def set_cached_state(self, state: torch.Tensor | None) -> None:
-        """Override the cached anchor state, e.g. to re-pin a chunk's anchor after the
-        per-tick pipeline overwrote it (see ``SyncInferenceEngine``)."""
-        self._last_state = state
-
     def get_config(self) -> dict[str, Any]:
         return {
             "enabled": self.enabled,

src/lerobot/rewards/pretrained.py

@@ -21,6 +21,8 @@ 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
@@ -127,13 +129,29 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        missing_keys, unexpected_keys = load_model_as_safetensor(
-            model, model_file, strict=strict, device=map_location
-        )
+        # 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)
         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/rollout/configs.py

@@ -223,22 +223,9 @@ class RolloutConfig:
     fps: float = 30.0
     duration: float = 0.0  # 0 = infinite (24/7 mode)
     interpolation_multiplier: int = 1
-    # Safety net (opt-in): if any commanded joint's target differs from the robot's
-    # currently-measured position by more than this many units in a single control
-    # step, the action is treated as an unsafe jump — the robot is NOT commanded and
-    # the rollout stops. Units are the robot's raw `.pos` values (degrees for the
-    # SO-arms). Gripper joints are excluded (different unit/range). `None` disables it.
-    # Guards against chunk-splice / bad-chunk slams; a normal per-tick move at fps=30
-    # is only a few degrees, so a threshold like 20-30 catches slams without tripping
-    # on legitimate fast motion.
-    max_action_jump_deg: float | None = None
     device: str | None = None
     task: str = ""
     display_data: bool = False
-    # Also visualize model "extras" (e.g. a world model's imagined video) alongside observations.
-    # Off by default: requesting predictions forces per-chunk decoding on the control thread and only
-    # world-model policies produce anything. Implies display_data. Sync inference only.
-    display_extra_data: bool = False
     # Visualization backend used when display_data is True: "rerun" or "foxglove".
     display_mode: str = "rerun"
     # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket
@@ -268,26 +255,6 @@ class RolloutConfig:
 
     def __post_init__(self):
         """Validate config invariants and load the policy config from ``--policy.path``."""
-        # --- Visualization validation ---
-        # Extra-data visualization piggybacks on the display_data path (backend init + telemetry
-        # logging are both gated on display_data), so enabling it implies display_data.
-        if self.display_extra_data and not self.display_data:
-            logger.info("display_extra_data=True implies display_data=True; enabling display_data")
-            self.display_data = True
-        # Only the sync engine surfaces intermediate predictions (RTC runs the policy in a background
-        # thread); warn and let it be ignored rather than fail.
-        if self.display_extra_data and not isinstance(self.inference, SyncInferenceConfig):
-            logger.warning(
-                "display_extra_data is only supported with sync inference (--inference.type=sync); "
-                "it will be ignored for inference type '%s'",
-                self.inference.type,
-            )
-
-        if self.max_action_jump_deg is not None and self.max_action_jump_deg <= 0:
-            raise ValueError(
-                f"max_action_jump_deg must be positive when set, got {self.max_action_jump_deg}"
-            )
-
         # --- Strategy-specific validation ---
         if isinstance(self.strategy, DAggerStrategyConfig) and self.teleop is None:
             raise ValueError("DAgger strategy requires --teleop.type to be set")

src/lerobot/rollout/context.py

@@ -43,6 +43,7 @@ from lerobot.processor import (
     make_default_processors,
     rename_stats,
 )
+from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
 from lerobot.robots import make_robot_from_config
 from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
 from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
@@ -51,6 +52,7 @@ from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
 from .inference import (
     InferenceEngine,
     RTCInferenceConfig,
+    SyncInferenceConfig,
     create_inference_engine,
 )
 from .robot_wrapper import ThreadSafeRobot
@@ -397,6 +399,15 @@ def build_rollout_context(
         },
     )
 
+    if isinstance(cfg.inference, SyncInferenceConfig) and any(
+        isinstance(step, RelativeActionsProcessorStep) and step.enabled
+        for step in getattr(preprocessor, "steps", ())
+    ):
+        raise NotImplementedError(
+            "SyncInferenceEngine does not support policies with relative actions for now."
+            "Use --inference.type=rtc or remove relative action processor steps from the policy pipeline."
+        )
+
     # --- 7. Inference strategy (needs policy + pre/post + hardware) --
     logger.info(
         "Creating inference engine (type=%s)...",
@@ -418,7 +429,6 @@ def build_rollout_context(
         use_torch_compile=cfg.use_torch_compile,
         compile_warmup_inferences=cfg.compile_warmup_inferences,
         shutdown_event=shutdown_event,
-        visualize_predictions=cfg.display_extra_data,
     )
 
     # --- 8. Assemble ---------------------------------------------------

src/lerobot/rollout/inference/base.py

@@ -69,15 +69,6 @@ class InferenceEngine(abc.ABC):
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
         """Return the next action tensor, or ``None`` if unavailable."""
 
-    def get_intermediate_predictions(self) -> dict | None:
-        """Extra display-ready model outputs to visualize this tick, or ``None``.
-
-        Lets a backend surface a world model's intermediate predictions (e.g. imagined video
-        frames) into the rollout visualization path, keyed by ``"<datatype>.<name>"`` (mirroring
-        observation feature keys). Default: nothing extra.
-        """
-        return None
-
     def notify_observation(self, obs: dict) -> None:  # noqa: B027
         """Publish the latest processed observation.  Default: no-op."""
 

src/lerobot/rollout/inference/factory.py

@@ -95,7 +95,6 @@ def create_inference_engine(
     use_torch_compile: bool = False,
     compile_warmup_inferences: int = 2,
     shutdown_event: Event | None = None,
-    visualize_predictions: bool = False,
 ) -> InferenceEngine:
     """Instantiate the appropriate inference engine from a config object."""
     logger.info("Creating inference engine: %s", config.type)
@@ -109,7 +108,6 @@ def create_inference_engine(
             task=task,
             device=device,
             robot_type=robot_wrapper.robot_type,
-            visualize_predictions=visualize_predictions,
         )
     if isinstance(config, RTCInferenceConfig):
         return RTCInferenceEngine(
@@ -118,8 +116,7 @@ def create_inference_engine(
             postprocessor=postprocessor,
             robot_wrapper=robot_wrapper,
             rtc_config=config.rtc,
-            dataset_features=dataset_features,
-            ordered_action_keys=ordered_action_keys,
+            hw_features=hw_features,
             task=task,
             fps=fps,
             device=device,

src/lerobot/rollout/inference/rtc.py

@@ -34,13 +34,12 @@ import torch
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.rtc import ActionQueue, LatencyTracker, reanchor_relative_rtc_prefix
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
+from lerobot.policies.utils import prepare_observation_for_inference
 from lerobot.processor import (
     NormalizerProcessorStep,
     PolicyProcessorPipeline,
     RelativeActionsProcessorStep,
 )
-from lerobot.utils.constants import ACTION
 from lerobot.utils.feature_utils import build_dataset_frame
 
 from ..robot_wrapper import ThreadSafeRobot
@@ -64,28 +63,17 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0
 
 
 def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
-    """Pad or truncate RTC prefix actions to a fixed length for stable compiled inference.
-
-    Padding repeats the last real action ("hold") rather than filling with zeros. The RTC
-    guidance pulls the new chunk toward this prefix at the padded indices (they fall inside
-    the weighted region when the real leftover is shorter than ``target_steps``). A zero in
-    the model's normalized action space decodes to the dataset *mean* action — a nonzero
-    offset that yanks the spliced action toward a mean/neutral pose for one step, producing
-    an intermittent seam (e.g. 95 -> 103 -> 95). Holding the last real action keeps the
-    padded targets continuous with the prefix, so no fake target enters the guided region.
-    The fixed output length is preserved so ``torch.compile`` policies keep stable shapes.
-    """
+    """Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
     if prev_actions.ndim != 2:
         raise ValueError(f"Expected 2D [T, A] tensor, got shape={tuple(prev_actions.shape)}")
-    steps, _ = prev_actions.shape
+    steps, action_dim = prev_actions.shape
     if steps == target_steps:
         return prev_actions
     if steps > target_steps:
         return prev_actions[:target_steps]
-    if steps == 0:
-        raise ValueError("Cannot pad an empty prefix: no last action to hold.")
-    hold = prev_actions[-1:].expand(target_steps - steps, -1)
-    return torch.cat([prev_actions, hold], dim=0)
+    padded = torch.zeros((target_steps, action_dim), dtype=prev_actions.dtype, device=prev_actions.device)
+    padded[:steps] = prev_actions
+    return padded
 
 
 # ---------------------------------------------------------------------------
@@ -109,8 +97,7 @@ class RTCInferenceEngine(InferenceEngine):
         postprocessor: PolicyProcessorPipeline,
         robot_wrapper: ThreadSafeRobot,
         rtc_config: RTCConfig,
-        dataset_features: dict,
-        ordered_action_keys: list[str],
+        hw_features: dict,
         task: str,
         fps: float,
         device: str | None,
@@ -124,31 +111,7 @@ class RTCInferenceEngine(InferenceEngine):
         self._postprocessor = postprocessor
         self._robot = robot_wrapper
         self._rtc_config = rtc_config
-        # Build observations with the SAME feature spec sync uses (post
-        # `robot_observation_processor`), not the raw-hardware spec. `build_dataset_frame`
-        # orders `observation.state` by this spec's `names`; using the raw-hardware order
-        # here (as before) desynced the state vector from sync whenever the observation
-        # processor reorders/renames state keys, corrupting both normalization and the
-        # relative-action anchor. The `prefix="observation"` filter ignores the action
-        # entries in the combined dict.
-        self._obs_features = dataset_features
-        # The model emits actions in `dataset_features[ACTION]` order (the order it was
-        # trained on); the robot expects them in `ordered_action_keys` order. Sync remaps
-        # by NAME via `make_robot_action` + reindex (sync.py) before returning; RTC must do
-        # the SAME, otherwise the engine-agnostic strategy (`send_next_action`) maps the raw
-        # model-order tensor onto `ordered_action_keys` positionally and mis-assigns joints
-        # whenever the two orders differ — a per-joint permutation that drives the arm wrong.
-        self._ordered_action_keys = ordered_action_keys
-        state_ft = dataset_features.get("observation.state")
-        if state_ft is not None:
-            logger.info("RTC observation.state layout: %s", state_ft.get("names"))
-        action_ft = dataset_features.get(ACTION)
-        if action_ft is not None:
-            logger.info(
-                "RTC action layout: model/dataset=%s -> robot=%s",
-                action_ft.get("names"),
-                self._ordered_action_keys,
-            )
+        self._hw_features = hw_features
         self._task = task
         self._fps = fps
         self._device = device or "cpu"
@@ -267,19 +230,10 @@ class RTCInferenceEngine(InferenceEngine):
     # ------------------------------------------------------------------
 
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
-        """Pop the next action from the RTC queue (ignores ``obs_frame``).
-
-        The queued action is in the model's ``dataset_features[ACTION]`` order; remap it
-        by NAME into ``ordered_action_keys`` order before returning, so the engine-agnostic
-        strategy maps values onto the correct joints. Mirrors ``SyncInferenceEngine.get_action``.
-        """
+        """Pop the next action from the RTC queue (ignores ``obs_frame``)."""
         if self._action_queue is None:
             return None
-        action = self._action_queue.get()
-        if action is None:
-            return None
-        action_dict = make_robot_action(action, self._obs_features)
-        return torch.tensor([action_dict[k] for k in self._ordered_action_keys])
+        return self._action_queue.get()
 
     def notify_observation(self, obs: dict) -> None:
         """Publish the latest observation for the RTC thread to consume."""
@@ -298,8 +252,6 @@ class RTCInferenceEngine(InferenceEngine):
             policy_device = torch.device(self._device)
 
             warmup_required = max(1, self._compile_warmup_inferences) if self._use_torch_compile else 0
-            # exclude the first N inferences from the latency tracker to avoid cold-start spikes
-            latency_warmup_required = max(1, warmup_required)
             inference_count = 0
             consecutive_errors = 0
 
@@ -321,10 +273,10 @@ class RTCInferenceEngine(InferenceEngine):
                         idx_before = queue.get_action_index()
                         prev_actions = queue.get_left_over()
 
-                        latency = latency_tracker.p95()
+                        latency = latency_tracker.max()
                         delay = math.ceil(latency / time_per_chunk) if latency else 0
 
-                        obs_batch = build_dataset_frame(self._obs_features, obs, prefix="observation")
+                        obs_batch = build_dataset_frame(self._hw_features, obs, prefix="observation")
                         obs_batch = prepare_observation_for_inference(
                             obs_batch, policy_device, self._task, self._robot.robot_type
                         )
@@ -364,8 +316,7 @@ class RTCInferenceEngine(InferenceEngine):
                         inference_count += 1
                         consecutive_errors = 0
                         is_warmup = self._use_torch_compile and inference_count <= warmup_required
-                        # Ignore the first N inferences for latency tracking to avoid cold-start spikes
-                        if inference_count <= latency_warmup_required:
+                        if is_warmup:
                             latency_tracker.reset()
                         else:
                             latency_tracker.add(new_latency)

src/lerobot/rollout/inference/sync.py

@@ -22,23 +22,28 @@ from copy import copy
 
 import torch
 
-from lerobot.policies.pretrained import PreTrainedPolicy, unpack_action_output
+from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
-from lerobot.processor import PolicyProcessorPipeline, RelativeActionsProcessorStep
+from lerobot.processor import PolicyProcessorPipeline
 
 from .base import InferenceEngine
 
 logger = logging.getLogger(__name__)
 
 
-# Relative-action support: a predicted chunk of offsets is anchored to the robot
-# state at prediction time, but the sync engine reruns the pre/post pipeline every
-# tick, so ``RelativeActionsProcessorStep`` would re-anchor cached actions to the
-# current (moved) state and drift through the chunk. We pin the anchor per chunk:
-# a probe on the policy's public ``predict_action_chunk`` flags the ticks that
-# predict a fresh chunk; on the others the engine restores the anchor the relative
-# step overwrote. ``select_action`` stays on the hot path, so per-tick side effects
-# (e.g. LingBot-VA keyframe feedback) are preserved.
+# TODO(Steven): support relative-action policies.  The per-tick flow refreshes
+# ``RelativeActionsProcessorStep._last_state`` every call, so cached chunk
+# actions popped on later ticks get reanchored to the *current* robot state and
+# absolute targets drift through the chunk.  Relative-action policies are
+# rejected at context-build time today; RTC postprocesses the whole chunk and
+# is unaffected.
+#
+# Candidate fix: drive the policy via ``predict_action_chunk`` and serve a
+# local FIFO of postprocessed actions.  Eliminates drift by construction and
+# saves per-tick pre/post work, but bypasses ``select_action`` — needs
+# fallbacks for SAC (raises), ACT temporal ensembling (ensembler lives in
+# ``select_action``), and Diffusion-family (obs-history queues populated as a
+# side effect of ``select_action``).
 
 
 class SyncInferenceEngine(InferenceEngine):
@@ -59,7 +64,6 @@ class SyncInferenceEngine(InferenceEngine):
         task: str,
         device: str | None,
         robot_type: str,
-        visualize_predictions: bool = False,
     ) -> None:
         self._policy = policy
         self._preprocessor = preprocessor
@@ -69,45 +73,10 @@ class SyncInferenceEngine(InferenceEngine):
         self._task = task
         self._device = torch.device(device or "cpu")
         self._robot_type = robot_type
-
-        # Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk
-        # (see module comment), mirroring the RTC engine.
-        self._relative_step = next(
-            (
-                s
-                for s in getattr(preprocessor, "steps", ())
-                if isinstance(s, RelativeActionsProcessorStep) and s.enabled
-            ),
-            None,
-        )
-        # Set by the probe for the current tick / ever, respectively.
-        self._chunk_predicted = False
-        self._ever_predicted_chunk = False
-        self._original_predict_action_chunk = None  # set while the probe is installed
-        if self._relative_step is not None:
-            # ``action_names`` is optional on the step; fill it lazily from the
-            # policy/dataset so the relative<->absolute mask is built correctly. This is
-            # a deliberate engine->step side effect (the step is configured by its consumer).
-            if self._relative_step.action_names is None:
-                cfg_names = getattr(policy.config, "action_feature_names", None)
-                self._relative_step.action_names = list(cfg_names) if cfg_names else list(ordered_action_keys)
-            self._install_chunk_probe()
-            logger.info("Relative actions enabled: chunk anchor pinned per predicted chunk")
-
-        # Intermediate-prediction visualization (e.g. a world model's imagined video). When on,
-        # ``get_action`` requests predictions and keeps the current chunk's frame stacks; a playhead
-        # (``get_intermediate_predictions``) advances one step per tick, paced across the chunk's tick
-        # span so the imagined clip stays wall-clock aligned with execution.
-        self._visualize_predictions = visualize_predictions
-        self._pred_stacks: dict = {}  # key -> [T, H, W, 3] frame stack for the current chunk
-        self._pred_cursor = 0  # ticks elapsed since the current chunk's frames arrived
-        self._ticks_per_chunk = getattr(getattr(policy, "config", None), "chunk_size", None)
-
         logger.info(
-            "SyncInferenceEngine initialized (device=%s, action_keys=%d, visualize_predictions=%s)",
+            "SyncInferenceEngine initialized (device=%s, action_keys=%d)",
             self._device,
             len(ordered_action_keys),
-            self._visualize_predictions,
         )
 
     def start(self) -> None:
@@ -116,11 +85,6 @@ class SyncInferenceEngine(InferenceEngine):
 
     def stop(self) -> None:
         """No background resources to stop."""
-        # Undo the probe so the policy object isn't left permanently patched
-        # (it may outlive this engine or be reused by another).
-        if self._original_predict_action_chunk is not None:
-            self._policy.predict_action_chunk = self._original_predict_action_chunk
-            self._original_predict_action_chunk = None
         logger.info("SyncInferenceEngine stopped")
 
     def reset(self) -> None:
@@ -129,54 +93,6 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
-        # New episode: the next tick predicts a fresh chunk and re-anchors.
-        self._chunk_predicted = False
-        self._ever_predicted_chunk = False
-        self._pred_stacks = {}
-        self._pred_cursor = 0
-
-    def _install_chunk_probe(self) -> None:
-        """Wrap the policy's public ``predict_action_chunk`` so we learn which ticks
-        predict a fresh chunk (when the anchor must advance) without introspecting any
-        private action queue. Chunking policies call it from ``select_action``.
-
-        Wraps whatever callable is currently bound (e.g. an already-``torch.compile``d
-        one, since ``build_rollout_context`` compiles before building the engine); undone
-        in ``stop()``."""
-        self._original_predict_action_chunk = self._policy.predict_action_chunk
-        inner = self._original_predict_action_chunk
-
-        def probe(*args, **kwargs):
-            self._chunk_predicted = True
-            self._ever_predicted_chunk = True
-            return inner(*args, **kwargs)
-
-        self._policy.predict_action_chunk = probe
-
-    def get_intermediate_predictions(self) -> dict | None:
-        """Serve one imagined frame per key for this tick, advancing the playhead.
-
-        Maps the current chunk's ``T`` decoded frames onto its ``ticks_per_chunk`` control ticks so
-        the imagined video plays back in step with execution (falls back to one frame/tick, clamped,
-        when the chunk's tick span is unknown). Returns ``None`` until a chunk with frames arrives.
-        """
-        if not self._pred_stacks:
-            return None
-        tick = self._pred_cursor
-        span = self._ticks_per_chunk
-        out: dict = {}
-        for key, stack in self._pred_stacks.items():
-            n = len(stack)
-            if n == 0:
-                continue
-            idx = round(tick / (span - 1) * (n - 1)) if span and span > 1 else tick
-            idx = min(max(idx, 0), n - 1)
-            frame = stack[idx]
-            if hasattr(frame, "detach"):
-                frame = frame.detach().cpu().numpy()
-            out[key] = frame
-        self._pred_cursor += 1
-        return out or None
 
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
         """Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
@@ -191,34 +107,12 @@ class SyncInferenceEngine(InferenceEngine):
             if self._device.type == "cuda" and self._policy.config.use_amp
             else nullcontext()
         )
-        # Snapshot the chunk anchor before the preprocessor overwrites it with this
-        # tick's state; restore it below if this tick only served a cached action.
-        # ``clone`` so the snapshot survives even if the cached tensor is ever mutated
-        # in place (today it is only rebound, but the copy is cheap for a state vector).
-        anchor_before = None
-        if self._relative_step is not None:
-            cached = self._relative_step.get_cached_state()
-            anchor_before = cached.clone() if cached is not None else None
-        self._chunk_predicted = False
         with torch.inference_mode(), autocast_ctx:
             observation = prepare_observation_for_inference(
                 observation, self._device, self._task, self._robot_type
             )
             observation = self._preprocessor(observation)
-            if self._visualize_predictions:
-                action, predictions = unpack_action_output(
-                    self._policy.select_action(observation, return_intermediate_predictions=True)
-                )
-                if predictions:
-                    # A fresh chunk was predicted this tick — store its frame stacks and restart the playhead.
-                    self._pred_stacks = predictions
-                    self._pred_cursor = 0
-            else:
-                action = self._policy.select_action(observation)
-            # Hold the anchor only for a chunking policy serving a cached action this
-            # tick; policies that never chunk or that recomputed keep refreshing.
-            if self._relative_step is not None and self._ever_predicted_chunk and not self._chunk_predicted:
-                self._relative_step.set_cached_state(anchor_before)
+            action = self._policy.select_action(observation)
             action = self._postprocessor(action)
         action_tensor = action.squeeze(0).cpu()
 

src/lerobot/rollout/strategies/core.py

@@ -156,8 +156,8 @@ class RolloutStrategy(abc.ABC):
         except Exception as e:
             logger.warning("Could not return to initial position: %s", e)
 
+    @staticmethod
     def _log_telemetry(
-        self,
         obs_processed: dict | None,
         action_dict: dict | None,
         runtime_ctx: RuntimeContext,
@@ -166,16 +166,10 @@ class RolloutStrategy(abc.ABC):
         cfg = runtime_ctx.cfg
         if not cfg.display_data:
             return
-        # When extra-data visualization is on, pull any display-ready model predictions from the
-        # engine (e.g. a world model's imagined video) and log them on the dedicated prediction channel.
-        prediction = None
-        if cfg.display_extra_data and self._engine is not None:
-            prediction = self._engine.get_intermediate_predictions()
         log_visualization_data(
             cfg.display_mode,
             observation=obs_processed,
             action=action_dict,
-            prediction=prediction,
             compress_images=cfg.display_compressed_images,
         )
 
@@ -273,36 +267,6 @@ def estimate_max_episode_seconds(
 # ---------------------------------------------------------------------------
 
 
-def detect_unsafe_action_jumps(
-    action_dict: dict,
-    obs_raw: dict,
-    max_jump: float,
-    exclude_substrings: tuple[str, ...] = ("gripper",),
-) -> dict[str, float]:
-    """Return ``{joint: jump}`` for joints whose commanded target exceeds the current
-    measured position by more than ``max_jump`` in a single control step; empty if safe.
-
-    Only keys present in both ``action_dict`` and ``obs_raw`` and not matching an
-    excluded substring are checked (grippers use a different unit/range). A large
-    single-step jump is the signature of a chunk-splice slam: a well-behaved trajectory
-    advances a joint by only a few units per tick.
-    """
-    violations: dict[str, float] = {}
-    for key, target in action_dict.items():
-        if any(token in key for token in exclude_substrings):
-            continue
-        current = obs_raw.get(key)
-        if current is None:
-            continue
-        try:
-            jump = abs(float(target) - float(current))
-        except (TypeError, ValueError):
-            continue
-        if jump > max_jump:
-            violations[key] = jump
-    return violations
-
-
 def send_next_action(
     obs_processed: dict,
     obs_raw: dict,
@@ -336,25 +300,6 @@ def send_next_action(
     if len(interp) != len(ordered_keys):
         raise ValueError(f"Interpolated tensor length ({len(interp)}) != action keys ({len(ordered_keys)})")
     action_dict = {k: interp[i].item() for i, k in enumerate(ordered_keys)}
-
-    # Safety net: refuse to command a single-step joint jump larger than the configured
-    # limit and stop the rollout. Catches chunk-splice / bad-chunk slams before they reach
-    # the motors. Compares the commanded target against the freshly-measured pose in
-    # `obs_raw` (same `.pos` units); no-op when the limit is unset.
-    max_jump = ctx.runtime.cfg.max_action_jump_deg
-    if max_jump is not None:
-        unsafe = detect_unsafe_action_jumps(action_dict, obs_raw, max_jump)
-        if unsafe:
-            detail = ", ".join(f"{k}={v:.1f}" for k, v in sorted(unsafe.items(), key=lambda kv: -kv[1]))
-            logger.error(
-                "SAFETY STOP: commanded joint jump exceeds %.1f in one step (%s); "
-                "refusing to send action and signalling shutdown.",
-                max_jump,
-                detail,
-            )
-            ctx.runtime.shutdown_event.set()
-            return None
-
     processed = ctx.processors.robot_action_processor((action_dict, obs_raw))
     ctx.hardware.robot_wrapper.send_action(processed)
     return action_dict

src/lerobot/scripts/lerobot_annotate.py

@@ -127,9 +127,6 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
     """
     from huggingface_hub import HfApi  # noqa: PLC0415
 
-    from lerobot.datasets.io_utils import load_info  # noqa: PLC0415
-    from lerobot.datasets.utils import create_lerobot_dataset_card  # noqa: PLC0415
-
     repo_id = cfg.new_repo_id or cfg.repo_id
     commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
     api = HfApi()
@@ -146,17 +143,10 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         repo_id=repo_id,
         repo_type="dataset",
         commit_message=commit_message,
-        # 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"],
+        ignore_patterns=[".annotate_staging/**", "**/.DS_Store"],
     )
     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
@@ -165,9 +155,21 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
     # actually wrote (no accidental drift if the codebase floor moves).
     from lerobot.datasets.dataset_metadata import CODEBASE_VERSION  # noqa: PLC0415
 
-    version_tag = (
-        dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
-    )
+    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,
+            )
     revision = getattr(commit_info, "oid", None)
     tag_kwargs = {
         "repo_id": repo_id,

src/lerobot/scripts/lerobot_eval.py

@@ -83,7 +83,6 @@ from lerobot.envs import (
     preprocess_observation,
 )
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
-from lerobot.policies.pretrained import unpack_action_output
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.types import PolicyAction
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
@@ -170,7 +169,7 @@ def rollout(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
-    save_predicted_video: bool = False,
+    predicted_latents_callback: Callable[[PreTrainedPolicy], None] | None = None,
 ) -> dict:
     """Run a batched policy rollout once through a batch of environments.
 
@@ -200,10 +199,9 @@ def rollout(
             are returned optionally because they typically take more memory to cache. Defaults to False.
         render_callback: Optional rendering callback to be used after the environments are reset, and after
             every step.
-        save_predicted_video: When True, request intermediate predictions from the policy each step
-            (``select_action(..., return_intermediate_predictions=True)``) and collect any imagined
-            video frames a world-model policy returns. Collected per image key in
-            ``ret["predicted_frames"]`` as a list of ``[T, H, W, 3]`` uint8 chunk stacks.
+        predicted_latents_callback: Optional callback invoked after every ``select_action`` with the policy
+            itself. World-model policies (e.g. LingBot-VA) stash predicted video latents on
+            ``policy.last_predicted_latents``; this lets the caller concatenate chunks and decode once.
     Returns:
         The dictionary described above.
     """
@@ -247,9 +245,6 @@ def rollout(
     all_rewards = []
     all_successes = []
     all_dones = []
-    # Imagined-video frames returned by world-model policies, collected per image key. Each entry is
-    # a chunk stack [T, H, W, 3] uint8; concatenated on the time axis by the caller.
-    predicted_frames: dict[str, list] = {}
 
     step = 0
     # Keep track of which environments are done.
@@ -284,13 +279,9 @@ def rollout(
 
             observation = preprocessor(observation)
             with torch.inference_mode():
-                extra = {"return_intermediate_predictions": True} if save_predicted_video else {}
-                action, predictions = unpack_action_output(policy.select_action(observation, **extra))
-            # World-model policies return imagined frames only on chunk-boundary ticks; collect them.
-            for key, frames in predictions.items():
-                if hasattr(frames, "detach"):
-                    frames = frames.detach().to("cpu")
-                predicted_frames.setdefault(key, []).append(frames)
+                action = policy.select_action(observation)
+            if predicted_latents_callback is not None:
+                predicted_latents_callback(policy)
             action = postprocessor(action)
 
             action_transition = {ACTION: action}
@@ -403,9 +394,6 @@ def rollout(
             stacked_observations[key] = torch.stack([obs[key] for obs in all_observations], dim=1)
         ret[OBS_STR] = stacked_observations
 
-    if save_predicted_video:
-        ret["predicted_frames"] = predicted_frames
-
     if hasattr(policy, "use_original_modules"):
         policy.use_original_modules()
 
@@ -447,6 +435,11 @@ def eval_policy(
     if max_episodes_rendered > 0 and not videos_dir:
         raise ValueError("If max_episodes_rendered > 0, videos_dir must be provided.")
 
+    # World-model policies (e.g. LingBot-VA) opt into predicted-video saving via their config.
+    save_predicted_video = save_predicted_video or bool(
+        getattr(getattr(policy, "config", None), "save_predicted_video", False)
+    )
+
     if not isinstance(policy, PreTrainedPolicy):
         exc = ValueError(
             f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
@@ -496,6 +489,16 @@ def eval_policy(
         predicted_video_paths: list[str] = []
         n_predicted_rendered = 0
 
+    # Collect predicted-video latents across a rollout (world-model policies only). The latents are
+    # concatenated and decoded once after the rollout, matching upstream LingBot-VA's visualization path.
+    def collect_predicted_latents(policy: PreTrainedPolicy):
+        latents = getattr(policy, "last_predicted_latents", None)
+        if latents is not None:
+            pred_latents.append(
+                latents.detach().to("cpu") if hasattr(latents, "detach") else torch.as_tensor(latents).cpu()
+            )
+            policy.last_predicted_latents = None
+
     if return_episode_data:
         episode_data: dict | None = None
 
@@ -507,6 +510,9 @@ def eval_policy(
         if max_episodes_rendered > 0:
             ep_frames: list[np.ndarray] = []
 
+        if save_predicted_video:
+            pred_latents: list[torch.Tensor] = []
+
         if start_seed is None:
             seeds = None
         else:
@@ -527,7 +533,7 @@ def eval_policy(
             env_features=env_features,
             recording_repo_id=recording_repo_id,
             recording_private=recording_private,
-            save_predicted_video=save_predicted_video,
+            predicted_latents_callback=collect_predicted_latents if save_predicted_video else None,
         )
 
         # Figure out where in each rollout sequence the first done condition was encountered (results after
@@ -593,33 +599,33 @@ def eval_policy(
                 threads.append(thread)
                 n_episodes_rendered += 1
 
-        # Maybe save the policy's predicted (imagined) video for this batch's rollout. The policy
-        # returns display-ready frame stacks per image key; concatenate them on the time axis and
-        # write one mp4 per key (no decoding here — the policy already decoded).
-        pred_frames = rollout_data.get("predicted_frames", {}) if save_predicted_video else {}
-        if save_predicted_video and any(len(stacks) > 0 for stacks in pred_frames.values()):
+        # Maybe save the policy's predicted (imagined) video for this batch's rollout.
+        if save_predicted_video and len(pred_latents) > 0:
+            predicted_latent = torch.cat(pred_latents, dim=2)
+            decoder = getattr(policy, "decode_predicted_latents", None) or getattr(
+                policy, "_decode_predicted_video", None
+            )
+            if decoder is None:
+                raise AttributeError(
+                    "Policy config requested predicted-video saving, but the policy does not expose "
+                    "`decode_predicted_latents` or `_decode_predicted_video`."
+                )
+            predicted_video = decoder(predicted_latent)
+            if hasattr(predicted_video, "detach"):
+                predicted_video = predicted_video.detach().to("cpu").numpy()
             videos_dir.mkdir(parents=True, exist_ok=True)
-            multi_key = len(pred_frames) > 1
-            for key, stacks in pred_frames.items():
-                if len(stacks) == 0:
-                    continue
-                predicted_video = torch.cat(
-                    [s if hasattr(s, "dim") else torch.as_tensor(s) for s in stacks], dim=0
-                )
-                predicted_video = predicted_video.detach().to("cpu").numpy()  # [T, H, W, 3] uint8
-                suffix = f"_{key.replace('.', '_')}" if multi_key else ""
-                predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}{suffix}.mp4"
-                predicted_video_paths.append(str(predicted_video_path))
-                thread = threading.Thread(
-                    target=write_video,
-                    args=(
-                        str(predicted_video_path),
-                        predicted_video,
-                        env.unwrapped.metadata["render_fps"],
-                    ),
-                )
-                thread.start()
-                threads.append(thread)
+            predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}.mp4"
+            predicted_video_paths.append(str(predicted_video_path))
+            thread = threading.Thread(
+                target=write_video,
+                args=(
+                    str(predicted_video_path),
+                    predicted_video,
+                    env.unwrapped.metadata["render_fps"],
+                ),
+            )
+            thread.start()
+            threads.append(thread)
             n_predicted_rendered += 1
 
         progbar.set_postfix(
@@ -765,11 +771,6 @@ def eval_main(cfg: EvalPipelineConfig):
     recording_dir = Path(cfg.output_dir) / "recordings" if cfg.eval.recording else None
     max_episodes_rendered = 0 if cfg.eval.recording else 10
     videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos"
-    # Predicted-video saving needs a directory to write mp4s into; recording mode leaves videos_dir
-    # unset, so provide one explicitly.
-    save_predicted_video = cfg.eval.save_predicted_video
-    if save_predicted_video and videos_dir is None:
-        videos_dir = Path(cfg.output_dir) / "videos"
 
     with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
         info = eval_policy_all(
@@ -789,7 +790,6 @@ def eval_main(cfg: EvalPipelineConfig):
             env_features=cfg.env.features if cfg.eval.recording else None,
             recording_repo_id=cfg.eval.recording_repo_id,
             recording_private=cfg.eval.recording_private,
-            save_predicted_video=save_predicted_video,
         )
         print("Overall Aggregated Metrics:")
         print(info["overall"])
@@ -837,7 +837,6 @@ def eval_one(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
-    save_predicted_video: bool = False,
 ) -> TaskMetrics:
     """Evaluates one task_id of one suite using the provided vec env."""
 
@@ -859,7 +858,6 @@ def eval_one(
         env_features=env_features,
         recording_repo_id=recording_repo_id,
         recording_private=recording_private,
-        save_predicted_video=save_predicted_video,
     )
 
     per_episode = task_result["per_episode"]
@@ -891,7 +889,6 @@ def run_one(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
-    save_predicted_video: bool = False,
 ):
     """
     Run eval_one for a single (task_group, task_id, env).
@@ -926,7 +923,6 @@ def run_one(
         env_features=env_features,
         recording_repo_id=task_repo_id,
         recording_private=recording_private,
-        save_predicted_video=save_predicted_video,
     )
 
     if max_episodes_rendered > 0:
@@ -953,7 +949,6 @@ def eval_policy_all(
     return_episode_data: bool = False,
     start_seed: int | None = None,
     max_parallel_tasks: int = 1,
-    save_predicted_video: bool = False,
 ) -> dict:
     """
     Evaluate a nested `envs` dict: {task_group: {task_id: vec_env}}.
@@ -1013,7 +1008,6 @@ def eval_policy_all(
         env_features=env_features,
         recording_repo_id=recording_repo_id,
         recording_private=recording_private,
-        save_predicted_video=save_predicted_video,
     )
 
     if max_parallel_tasks <= 1:

src/lerobot/scripts/lerobot_train.py

@@ -20,7 +20,6 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
-import os
 import sys
 import time
 from contextlib import nullcontext
@@ -172,9 +171,6 @@ 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
 
 
@@ -465,14 +461,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # declares language columns; otherwise stay on PyTorch's default
     # collate so non-language training runs are unaffected.
     collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
-    # Multi-node fork-OOM mitigation: on EFA nodes (vm.overcommit_memory=0, no swap),
-    # forking dataloader workers from a multi-GB rank reserve-charges the rank's full virtual
-    # footprint, so 8 ranks x num_workers forking at once trips OSError(ENOMEM) despite free
-    # RAM. Honor LEROBOT_DATALOADER_MP_CONTEXT (e.g. "forkserver"/"spawn") to spawn workers
-    # from a clean context instead of fork(). Only meaningful when workers are used.
-    dataloader_mp_context = os.environ.get("LEROBOT_DATALOADER_MP_CONTEXT") or None
-    if cfg.num_workers == 0:
-        dataloader_mp_context = None
     dataloader = torch.utils.data.DataLoader(
         dataset,
         num_workers=cfg.num_workers,
@@ -484,7 +472,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         collate_fn=collate_fn,
         prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
         persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
-        multiprocessing_context=dataloader_mp_context,
     )
 
     # Build eval dataloader if a held-out split exists
@@ -512,7 +499,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             collate_fn=eval_collate_fn,
             prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
             persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
-            multiprocessing_context=dataloader_mp_context,
         )
 
     # Prepare everything with accelerator
@@ -586,7 +572,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        train_tracker, _ = update_policy(
+        train_tracker, output_dict = update_policy(
             train_tracker,
             policy,
             batch,
@@ -619,10 +605,9 @@ 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/constants.py

@@ -30,9 +30,6 @@ OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
 OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
 OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
 
-PREDICTION_STR = "prediction"
-PREDICTION_PREFIX = PREDICTION_STR + "."
-
 ACTION = "action"
 ACTION_PREFIX = ACTION + "."
 ACTION_TOKENS = ACTION + ".tokens"

src/lerobot/utils/foxglove_visualization.py

@@ -37,8 +37,6 @@ from .constants import (
     OBS_PREFIX,
     OBS_STATE,
     OBS_STR,
-    PREDICTION_PREFIX,
-    PREDICTION_STR,
     REWARD,
     SUCCESS,
     TRUNCATED,
@@ -285,11 +283,10 @@ def _log_foxglove_image(
 def log_foxglove_data(
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
-    prediction: dict | None = None,
     compress_images: bool = False,
 ) -> None:
     """
-    Logs observation, action and prediction data to a Foxglove WebSocket server for real-time visualization.
+    Logs observation and action data to a Foxglove WebSocket server for real-time visualization.
 
     Mirrors ``log_rerun_data`` but emits Foxglove messages over the server started by
     :func:`init_foxglove`. Data is mapped as follows:
@@ -305,8 +302,6 @@ def log_foxglove_data(
     Args:
         observation: An optional dictionary containing observation data to log.
         action: An optional dictionary containing action data to log.
-        prediction: An optional dictionary of display-ready model outputs (e.g. a world model's
-            imagined video), keyed "<datatype>.<name>", logged on ``/prediction/...`` topics.
         compress_images: Whether to JPEG-compress images before logging to save bandwidth in exchange
             for CPU and quality.
     """
@@ -339,30 +334,6 @@ def log_foxglove_data(
                     )
         _log_foxglove_scalars(_foxglove_topic(OBS_STATE), obs_scalars, log_time=now)
 
-    if prediction:
-        # Predicted outputs are keyed "<datatype>.<name>" (e.g. "images.predicted"); route images to
-        # /prediction/images/<name> and any scalars to an aggregate /prediction/state topic.
-        pred_scalars: dict[str, float] = {}
-        for k, v in prediction.items():
-            if v is None:
-                continue
-            key = k[len(PREDICTION_PREFIX) :] if str(k).startswith(PREDICTION_PREFIX) else str(k)
-            if _is_scalar(v):
-                pred_scalars[key] = float(v)
-            elif isinstance(v, np.ndarray):
-                if v.ndim == 1:
-                    pred_scalars.update(_labeled_scalars(key, v))
-                else:
-                    name = key[len("images.") :] if key.startswith("images.") else key
-                    _log_foxglove_image(
-                        f"/{PREDICTION_STR}/images/{_foxglove_safe_name(name)}",
-                        name,
-                        v,
-                        compress_images=compress_images,
-                        log_time=now,
-                    )
-        _log_foxglove_scalars(f"/{PREDICTION_STR}/state", pred_scalars, log_time=now)
-
     if action:
         action_scalars: dict[str, float] = {}
         for k, v in action.items():

src/lerobot/utils/logging_utils.py

@@ -104,7 +104,6 @@ class MetricsTracker:
         "episodes",
         "epochs",
         "accelerator",
-        "_caller_metrics",
     ]
 
     def __init__(
@@ -130,9 +129,6 @@ 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__:
@@ -160,21 +156,6 @@ 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"``

src/lerobot/utils/rerun_visualization.py

@@ -27,7 +27,7 @@ import numpy as np
 from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit
 from lerobot.types import RobotAction, RobotObservation
 
-from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, PREDICTION_PREFIX
+from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
 from .import_utils import require_package
 
 
@@ -37,43 +37,6 @@ def _is_scalar(x):
     )
 
 
-def _log_scalar_or_image_mapping(rr, data, prefix, scalar_paths, image_paths, compress_images):
-    """Log a mapping of scalars/images (observation- or prediction-style) under ``prefix``.
-
-    Scalars and 1D arrays go to ``scalar_paths`` (time-series); 2D/3D arrays are treated as images
-    (CHW->HWC as needed, depth for single-channel) and go to ``image_paths`` (spatial views).
-    """
-    for k, v in data.items():
-        if v is None:
-            continue
-        key = str(k) if str(k).startswith(prefix) else f"{prefix}{k}"
-
-        if _is_scalar(v):
-            rr.log(key, rr.Scalars(float(v)))
-            scalar_paths.add(key)
-        elif isinstance(v, np.ndarray):
-            arr = v
-            # Convert CHW -> HWC when needed
-            if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
-                arr = np.transpose(arr, (1, 2, 0))
-            if arr.ndim == 1:
-                rr.log(key, rr.Scalars(arr.astype(float)))
-                scalar_paths.add(key)
-            else:
-                if arr.shape[-1] == 1:
-                    # At record time, the depth unit is inferred from the frame type.
-                    depth_unit = infer_depth_unit(arr.dtype)
-                    img_entity = rr.DepthImage(
-                        arr,
-                        meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0,
-                        colormap=rr.components.Colormap.Viridis,
-                    )
-                else:
-                    img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
-                rr.log(key, entity=img_entity, static=True)
-                image_paths.add(key)
-
-
 def init_rerun(
     session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
 ) -> None:
@@ -110,16 +73,10 @@ def shutdown_rerun() -> None:
     rr.rerun_shutdown()
 
 
-def _build_blueprint(
-    observation_paths: set[str],
-    action_paths: set[str],
-    image_paths: set[str],
-    prediction_paths: set[str],
-):
-    """Build a Rerun blueprint laying out camera/predicted images and scalar series in separate views.
+def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]):
+    """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views.
 
-    Images (observation and prediction) each get a spatial view; observation, action, and prediction
-    scalars each get their own time-series view. All arranged in a grid.
+    Camera images, observation and action scalars are arranged in a grid.
     """
 
     # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
@@ -131,29 +88,22 @@ def _build_blueprint(
         views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths)))
     if action_paths:
         views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths)))
-    if prediction_paths:
-        views.append(rrb.TimeSeriesView(name="prediction", contents=sorted(prediction_paths)))
 
     return rrb.Blueprint(rrb.Grid(*views))
 
 
-def _ensure_blueprint(
-    observation_paths: set[str],
-    action_paths: set[str],
-    image_paths: set[str],
-    prediction_paths: set[str],
-) -> None:
-    """Build and send the blueprint once, from the first observation/action/prediction data."""
+def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None:
+    """Build and send the blueprint once, from the first observation and action data."""
     if getattr(log_rerun_data, "blueprint", None) is not None:
         return
 
-    if not (observation_paths or action_paths or image_paths or prediction_paths):
+    if not (observation_paths or action_paths or image_paths):
         return
 
     # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
     import rerun as rr
 
-    blueprint = _build_blueprint(observation_paths, action_paths, image_paths, prediction_paths)
+    blueprint = _build_blueprint(observation_paths, action_paths, image_paths)
     log_rerun_data.blueprint = blueprint
     rr.send_blueprint(blueprint)
 
@@ -161,11 +111,10 @@ def _ensure_blueprint(
 def log_rerun_data(
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
-    prediction: dict | None = None,
     compress_images: bool = False,
 ) -> None:
     """
-    Logs observation, action and prediction data to Rerun for real-time visualization.
+    Logs observation and action data to Rerun for real-time visualization.
 
     This function iterates through the provided observation and action dictionaries and sends their contents
     to the Rerun viewer. It handles different data types appropriately:
@@ -184,8 +133,6 @@ def log_rerun_data(
     Args:
         observation: An optional dictionary containing observation data to log.
         action: An optional dictionary containing action data to log.
-        prediction: An optional dictionary of display-ready model outputs (e.g. a world model's
-            imagined video), keyed "<datatype>.<name>", logged on a dedicated "prediction." channel.
         compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality.
     """
 
@@ -195,19 +142,37 @@ def log_rerun_data(
     observation_paths: set[str] = set()
     action_paths: set[str] = set()
     image_paths: set[str] = set()
-    prediction_paths: set[str] = set()
 
     if observation:
-        _log_scalar_or_image_mapping(
-            rr, observation, OBS_PREFIX, observation_paths, image_paths, compress_images
-        )
+        for k, v in observation.items():
+            if v is None:
+                continue
+            key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}"
 
-    if prediction:
-        # Predicted images share the spatial-view set (their "prediction." names keep them distinct);
-        # predicted scalars get their own time-series view.
-        _log_scalar_or_image_mapping(
-            rr, prediction, PREDICTION_PREFIX, prediction_paths, image_paths, compress_images
-        )
+            if _is_scalar(v):
+                rr.log(key, rr.Scalars(float(v)))
+                observation_paths.add(key)
+            elif isinstance(v, np.ndarray):
+                arr = v
+                # Convert CHW -> HWC when needed
+                if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
+                    arr = np.transpose(arr, (1, 2, 0))
+                if arr.ndim == 1:
+                    rr.log(key, rr.Scalars(arr.astype(float)))
+                    observation_paths.add(key)
+                else:
+                    if arr.shape[-1] == 1:
+                        # At record time, the depth unit is inferred from the frame type.
+                        depth_unit = infer_depth_unit(arr.dtype)
+                        img_entity = rr.DepthImage(
+                            arr,
+                            meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0,
+                            colormap=rr.components.Colormap.Viridis,
+                        )
+                    else:
+                        img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
+                    rr.log(key, entity=img_entity, static=True)
+                    image_paths.add(key)
 
     if action:
         for k, v in action.items():
@@ -223,4 +188,4 @@ def log_rerun_data(
                 rr.log(key, rr.Scalars(v.reshape(-1).astype(float)))
                 action_paths.add(key)
 
-    _ensure_blueprint(observation_paths, action_paths, image_paths, prediction_paths)
+    _ensure_blueprint(observation_paths, action_paths, image_paths)

src/lerobot/utils/visualization_utils.py

@@ -56,19 +56,14 @@ def log_visualization_data(
     display_mode: str,
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
-    prediction: dict | None = None,
     compress_images: bool = False,
 ) -> None:
-    """Logs observation/action/prediction data to the backend selected by ``display_mode``."""
+    """Logs observation/action data to the backend selected by ``display_mode``."""
 
     if display_mode == "rerun":
-        log_rerun_data(
-            observation=observation, action=action, prediction=prediction, compress_images=compress_images
-        )
+        log_rerun_data(observation=observation, action=action, compress_images=compress_images)
     elif display_mode == "foxglove":
-        log_foxglove_data(
-            observation=observation, action=action, prediction=prediction, compress_images=compress_images
-        )
+        log_foxglove_data(observation=observation, action=action, compress_images=compress_images)
     else:
         raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.")
 

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(
         {
-            "COMPLETED manipulation events": {
+            "atomic subtasks": {
                 "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(
         {
-            "COMPLETED manipulation events": {
+            "atomic subtasks": {
                 "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 "COMPLETED manipulation events" in _prompt_text(m)]
+    subtask_calls = [m for m in captured if "atomic subtasks" 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/datasets/test_compute_stats.py

@@ -13,6 +13,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import logging
 from unittest.mock import patch
 
 import numpy as np
@@ -687,6 +688,28 @@ def test_compute_episode_stats_string_features_skipped():
     assert "q01" in stats["action"]
 
 
+def test_compute_episode_stats_zero_width_features_skipped(caplog):
+    """Test that features with a zero-width dim (e.g. shape=(0,)) are skipped with a debug log."""
+    episode_data = {
+        "empty": np.zeros((100, 0), dtype=np.float32),  # Zero-width feature
+        "action": np.random.normal(0, 1, (100, 5)),
+    }
+    features = {
+        "empty": {"dtype": "float32", "shape": (0,)},
+        "action": {"dtype": "float32", "shape": (5,)},
+    }
+
+    with caplog.at_level(logging.DEBUG):
+        stats = compute_episode_stats(episode_data, features)
+
+    # Zero-width features should be skipped with a debug log, others computed as usual
+    assert "empty" not in stats
+    assert "empty" in caplog.text
+    assert "action" in stats
+    assert "q01" in stats["action"]
+    assert stats["action"]["mean"].shape == (5,)
+
+
 def test_aggregate_feature_stats_with_quantiles():
     """Test aggregating feature stats that include quantiles."""
     stats_ft_list = [

tests/datasets/test_datasets.py

@@ -1804,3 +1804,11 @@ def test_episode_filter_unknown_key_raises(tmp_path, lerobot_dataset_factory):
             root=dataset.root,
             episode_filter=lambda ep: ep["not_a_real_field"] > 0,
         )
+
+
+def test_get_hf_features_zero_width_feature_does_not_raise_on_from_dict():
+    import datasets
+
+    features = {"empty": {"dtype": "float32", "shape": (0,), "names": ["empty"]}}
+    hf_features = get_hf_features_from_features(features)
+    datasets.Dataset.from_dict({"empty": [[], []]}, features=hf_features)

tests/policies/test_relative_actions.py

@@ -346,10 +346,3 @@ def test_state_not_modified_by_relative_processor(dataset, action_dim):
 
     result_state = result[TransitionKey.OBSERVATION][OBS_STATE]
     torch.testing.assert_close(result_state, original_state)
-
-
-def test_cached_anchor_not_in_config():
-    """The cached anchor is ephemeral runtime state and must not leak into the config."""
-    step = RelativeActionsProcessorStep(enabled=True)
-    step.set_cached_state(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
-    assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}

tests/scripts/test_lerobot_annotate.py

@@ -30,9 +30,7 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
 
     root = tmp_path / "dataset"
     (root / "meta").mkdir(parents=True)
-    (root / "meta" / "info.json").write_text(
-        json.dumps({"codebase_version": "v3.0", "fps": 30, "features": {}})
-    )
+    (root / "meta" / "info.json").write_text(json.dumps({"codebase_version": "v3.0"}))
 
     calls = {}
 
@@ -57,11 +55,6 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
 
     monkeypatch.setattr("huggingface_hub.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",
         new_repo_id="annotated/dataset",
@@ -78,13 +71,6 @@ 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/test_rollout.py

@@ -348,200 +348,3 @@ def test_rollout_context_fields():
 
     field_names = {f.name for f in dataclasses.fields(RolloutContext)}
     assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
-
-
-# ---------------------------------------------------------------------------
-# Sync engine: relative-action anchoring (drift-free chunk execution)
-# ---------------------------------------------------------------------------
-
-_REL_ACTION_NAMES = ["j0.pos", "j1.pos", "j2.pos", "gripper.pos"]
-_REL_ACTION_DIM = len(_REL_ACTION_NAMES)
-
-
-def _relative_pre_post(exclude_joints=None):
-    """Pre/post processors wrapping the real relative (caches anchor) and absolute
-    (relative + cached state) steps, mirroring what the sync engine feeds them."""
-    from lerobot.processor import (
-        AbsoluteActionsProcessorStep,
-        RelativeActionsProcessorStep,
-        TransitionKey,
-        create_transition,
-    )
-    from lerobot.utils.constants import OBS_STATE
-
-    relative_step = RelativeActionsProcessorStep(
-        enabled=True, exclude_joints=exclude_joints or [], action_names=list(_REL_ACTION_NAMES)
-    )
-    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
-
-    class _Pre:
-        steps = [relative_step]
-
-        def __call__(self, observation):
-            # Run the relative step so it caches the anchor, then pass the batch through.
-            transition = create_transition(observation={OBS_STATE: observation[OBS_STATE]})
-            relative_step(transition)
-            return observation
-
-        def reset(self):
-            pass
-
-    class _Post:
-        def __call__(self, action):
-            transition = create_transition(action=action)
-            return absolute_step(transition)[TransitionKey.ACTION]
-
-        def reset(self):
-            pass
-
-    return _Pre(), _Post(), relative_step
-
-
-def _fake_relative_policy(chunk_rel, n_action_steps, chunking=True):
-    """Fake relative-action policy for the sync engine.
-
-    ``chunking=True`` buffers a chunk and serves it one action per tick, calling the
-    public ``predict_action_chunk`` only on refill (pi0/fastwam/lingbot). ``False``
-    returns an action directly and never calls it. The engine's anchor probe keys off
-    that public call, so the fake routes through it rather than any private queue.
-    """
-    from collections import deque
-
-    policy = MagicMock()
-    policy.config.use_amp = False
-    policy.config.action_feature_names = list(_REL_ACTION_NAMES)
-    state = {"predict_calls": 0}
-    queue = deque(maxlen=n_action_steps)
-
-    def predict_action_chunk(_batch=None, **_kwargs):
-        state["predict_calls"] += 1
-        return chunk_rel.unsqueeze(0)  # [B=1, n, dim]
-
-    def select_action(_observation):
-        if not chunking:
-            return chunk_rel[0].unsqueeze(0)
-        if len(queue) == 0:
-            actions = policy.predict_action_chunk(_observation)
-            queue.extend(actions.transpose(0, 1))  # [n, 1, dim]
-        return queue.popleft()
-
-    policy.predict_action_chunk.side_effect = predict_action_chunk
-    policy.select_action.side_effect = select_action
-    policy.reset.side_effect = queue.clear
-    policy._predict_state = state
-    return policy
-
-
-def _build_sync_engine(policy, pre, post):
-    from lerobot.rollout import SyncInferenceEngine
-
-    return SyncInferenceEngine(
-        policy=policy,
-        preprocessor=pre,
-        postprocessor=post,
-        dataset_features={"action": {"names": list(_REL_ACTION_NAMES)}},
-        ordered_action_keys=list(_REL_ACTION_NAMES),
-        task="test",
-        device="cpu",
-        robot_type="mock",
-    )
-
-
-def _obs_frame(state_values):
-    import numpy as np
-
-    return {"observation.state": np.asarray(state_values, dtype=np.float32)}
-
-
-def test_sync_relative_holds_anchor_across_chunk():
-    """Every action popped within a chunk must anchor to the tick-0 state (no drift)."""
-    n = 4
-    # A distinct relative offset per chunk step so a wrong anchor would be visible.
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
-    pre, post, relative_step = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
-    engine = _build_sync_engine(policy, pre, post)
-
-    assert engine._relative_step is relative_step  # introspection wired the step
-
-    s0 = [1.0, 2.0, 3.0, 4.0]
-    outputs = []
-    for tick in range(n):
-        # Feed a *different* state each tick; a drifting anchor would use it.
-        state = [v + tick for v in s0]
-        outputs.append(engine.get_action(_obs_frame(state)))
-
-    # Exactly one chunk was predicted across the n ticks.
-    assert policy._predict_state["predict_calls"] == 1
-    for tick in range(n):
-        expected = torch.tensor(s0) + chunk_rel[tick]
-        torch.testing.assert_close(outputs[tick], expected)
-
-    # Next tick empties the queue -> fresh chunk -> anchor advances to the new state.
-    s_next = [10.0, 20.0, 30.0, 40.0]
-    out = engine.get_action(_obs_frame(s_next))
-    assert policy._predict_state["predict_calls"] == 2
-    torch.testing.assert_close(out, torch.tensor(s_next) + chunk_rel[0])
-    # The anchor now reflects the fresh-chunk state, not the held one.
-    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next]))
-
-
-def test_sync_relative_reset_reanchors_new_episode():
-    """After ``reset()`` the first tick of the next episode anchors to the new state."""
-    n = 3
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
-    pre, post, relative_step = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
-    engine = _build_sync_engine(policy, pre, post)
-
-    # Episode 1: one tick anchors to s0 and leaves cached actions in the queue.
-    engine.get_action(_obs_frame([1.0, 1.0, 1.0, 1.0]))
-    assert policy._predict_state["predict_calls"] == 1
-
-    engine.reset()  # clears the queue and the per-episode chunk flags
-
-    # Episode 2: a fresh state must produce a fresh chunk anchored to that state,
-    # not carry over the previous episode's anchor.
-    s_new = [7.0, 8.0, 9.0, 10.0]
-    out = engine.get_action(_obs_frame(s_new))
-    assert policy._predict_state["predict_calls"] == 2
-    torch.testing.assert_close(out, torch.tensor(s_new) + chunk_rel[0])
-    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_new]))
-
-
-def test_sync_relative_non_chunking_policy_refreshes_every_tick():
-    """A policy that never calls ``predict_action_chunk`` must not freeze the anchor."""
-    n = 3
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.5) for _ in range(n)])
-    pre, post, _ = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n, chunking=False)
-    engine = _build_sync_engine(policy, pre, post)
-
-    s0 = [1.0, 1.0, 1.0, 1.0]
-    for tick in range(3):
-        state = [v + tick for v in s0]
-        out = engine.get_action(_obs_frame(state))
-        # Anchor must track the current state every tick (no chunk => no hold).
-        torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
-    assert policy._predict_state["predict_calls"] == 0
-
-
-def test_sync_engine_no_relative_step_is_none():
-    """Without an enabled relative step, the engine takes the plain select_action path."""
-    policy = MagicMock()
-    policy.config.use_amp = False
-    engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock())
-    assert engine._relative_step is None
-
-
-def test_sync_relative_stop_restores_policy_method():
-    """``stop()`` un-patches the probe so the policy object isn't permanently modified."""
-    n = 3
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
-    pre, post, _ = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
-    original = policy.predict_action_chunk
-    engine = _build_sync_engine(policy, pre, post)
-    assert policy.predict_action_chunk is not original  # probe installed
-    engine.stop()
-    assert policy.predict_action_chunk is original  # restored

tests/utils/test_logging_utils.py

@@ -233,37 +233,3 @@ 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)