precommit

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/inference.mdx

@@ -39,11 +39,11 @@ lerobot-rollout \
     --duration=60
 ```
 
-| Flag             | Description                                            |
-| ---------------- | ------------------------------------------------------ |
-| `--duration`     | Run time in seconds (0 = infinite)                     |
-| `--task`         | Task description passed to the policy                  |
-| `--display_data` | Stream observations/actions to Rerun for visualization |
+| Flag             | Description                                                                                                                                 |
+| ---------------- | ------------------------------------------------------------------------------------------------------------------------------------------- |
+| `--duration`     | Run time in seconds (0 = infinite)                                                                                                          |
+| `--task`         | Task description passed to the policy                                                                                                       |
+| `--display_data` | Stream observations/actions to the visualization backend (`--display_mode`; add `--display_extra_data` for a policy's imagined predictions) |
 
 ### Sentry (`--strategy.type=sentry`)
 
@@ -241,22 +241,25 @@ See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.
 
 ## Common Flags
 
-| Flag                              | Description                                                       | Default |
-| --------------------------------- | ----------------------------------------------------------------- | ------- |
-| `--policy.path`                   | **Required.** HF Hub model ID or local checkpoint path            | --      |
-| `--robot.type`                    | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`) | --      |
-| `--robot.port`                    | Serial port for the robot                                         | --      |
-| `--robot.cameras`                 | Camera configuration (JSON dict)                                  | --      |
-| `--fps`                           | Control loop frequency                                            | 30      |
-| `--duration`                      | Run time in seconds (0 = infinite)                                | 0       |
-| `--device`                        | Torch device (`cpu`, `cuda`, `mps`)                               | auto    |
-| `--task`                          | Task description (used when no dataset is provided)               | --      |
-| `--display_data`                  | Stream telemetry to Rerun visualization                           | false   |
-| `--display_ip` / `--display_port` | Remote Rerun server address                                       | --      |
-| `--interpolation_multiplier`      | Action interpolation factor                                       | 1       |
-| `--use_torch_compile`             | Enable `torch.compile` for inference                              | false   |
-| `--resume`                        | Resume a previous recording session                               | false   |
-| `--play_sounds`                   | Vocal synthesis for events                                        | true    |
+| Flag                              | Description                                                                                                                                                 | Default |
+| --------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
+| `--policy.path`                   | **Required.** HF Hub model ID or local checkpoint path                                                                                                      | --      |
+| `--robot.type`                    | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`)                                                                                           | --      |
+| `--robot.port`                    | Serial port for the robot                                                                                                                                   | --      |
+| `--robot.cameras`                 | Camera configuration (JSON dict)                                                                                                                            | --      |
+| `--fps`                           | Control loop frequency                                                                                                                                      | 30      |
+| `--duration`                      | Run time in seconds (0 = infinite)                                                                                                                          | 0       |
+| `--device`                        | Torch device (`cpu`, `cuda`, `mps`)                                                                                                                         | auto    |
+| `--task`                          | Task description (used when no dataset is provided)                                                                                                         | --      |
+| `--display_data`                  | Stream telemetry to the visualization backend                                                                                                               | false   |
+| `--display_mode`                  | Visualization backend: `rerun` or `foxglove`                                                                                                                | rerun   |
+| `--display_extra_data`            | Also stream a policy's intermediate predictions (e.g. a world model's imagined video) on a dedicated channel; implies `--display_data`, sync inference only | false   |
+| `--display_compressed_images`     | JPEG-compress images before streaming (less bandwidth, more CPU)                                                                                            | false   |
+| `--display_ip` / `--display_port` | Remote Rerun server address (or bind interface/port for foxglove)                                                                                           | --      |
+| `--interpolation_multiplier`      | Action interpolation factor (upsamples the control rate)                                                                                                    | 1       |
+| `--use_torch_compile`             | Enable `torch.compile` for inference                                                                                                                        | false   |
+| `--resume`                        | Resume a previous recording session                                                                                                                         | false   |
+| `--play_sounds`                   | Vocal synthesis for events                                                                                                                                  | true    |
 
 ---
 

docs/source/lingbot_va.mdx

@@ -111,26 +111,17 @@ Requirements:
 - The block-causal masks use PyTorch **flex-attention**, so build the policy with
   `--policy.attn_mode=flex` for training (the default `torch` SDPA is inference-only).
 - The full 5B DiT does not fit a single 24–32 GB GPU under AdamW; fine-tune with **LoRA**
-  (`--peft.method_type=LORA`) and/or optimizer offload. `get_optim_params` returns only the
-  trainable (e.g. adapter) parameters; the VAE + UMT5 text encoder stay frozen. Install the
-  `lerobot[peft]` extra to enable PEFT support (see the [PEFT training guide](./peft_training)).
+  (`--policy.use_peft=true`) and/or optimizer offload. `get_optim_params` returns only the
+  trainable (e.g. adapter) parameters; the VAE + UMT5 text encoder stay frozen.
 
 ```bash
 lerobot-train \
   --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \
-  --peft.method_type=LORA --peft.r=32 --peft.lora_alpha=32 \
-  --peft.target_modules='transformer\.blocks\.\d+\.attn[12]\.(to_q|to_v)' \
+  --policy.use_peft=true \
   --dataset.repo_id=<your LeRobot-format dataset> \
   --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va
 ```
 
-Unlike SmolVLA / π₀, LingBot-VA does not ship built-in default LoRA targets, so you must pass
-`--peft.target_modules` explicitly. Only `self.transformer` (the dual-stream Wan transformer) is
-trainable; the example above adapts the query/value projections of both its self-attention
-(`attn1`) and cross-attention (`attn2`) blocks — the standard LoRA target set. Broaden it (e.g.
-add `to_k`/`to_out`, or the `ffn` layers) if you need a higher-capacity adapter. Passing
-`--peft.method_type` implies PEFT, so `--policy.use_peft=true` is not required.
-
 The dataset must provide camera clips (a temporal window per camera, VAE-encoded to
 `frame_chunk_size` latent frames) and `frame_chunk_size * action_per_frame` action steps per item.
 

docs/source/peft_training.mdx

@@ -62,10 +62,3 @@ to the `--peft.full_training_modules` parameter:
 
 The learning rate and the scheduled target learning rate can usually be scaled by a factor of 10 compared to the
 learning rate used for full fine-tuning (e.g., 1e-4 normal, so 1e-3 using LoRA).
-
-## Other policies
-
-The same `--peft.*` flags work for any pre-trained policy. Some policies (SmolVLA, π₀, π₀.₅) ship
-built-in default `target_modules`, so `--peft.method_type=LORA` is enough. Others do not, and will
-ask you to pass `--peft.target_modules` explicitly — for example LingBot-VA, whose recommended
-targets are documented in its [dedicated guide](./lingbot_va#training--fine-tuning).

examples/annotations/run_hf_job.py

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

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,6 +93,9 @@ 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/configs/train.py

@@ -221,11 +221,6 @@ class TrainPipelineConfig(HubMixin):
             )
 
         active_cfg = self.trainable_config
-
-        # Keep the policy-level `use_peft` flag in sync with the presence of a `--peft.*` config.
-        if self.peft is not None and self.policy is not None:
-            self.policy.use_peft = True
-
         if self.rename_map and active_cfg.pretrained_path is None:
             raise ValueError(
                 "`rename_map` requires a pretrained policy checkpoint. "

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

@@ -513,37 +513,6 @@ def make_pre_post_processors(
     return processors
 
 
-def _has_peft_adapter_config(
-    pretrained_path: str,
-    revision: str | None = None,
-) -> bool:
-    """Return whether ``pretrained_path`` points to an existing PEFT adapter.
-
-    A PEFT adapter checkpoint always ships an ``adapter_config.json``.
-    A plain base-model checkpoint does not. This distinction lets us tell apart
-    two very different ``use_peft=True`` scenarios that both set ``pretrained_path``:
-
-    * loading/resuming a previously trained adapter (config lives at ``pretrained_path``)
-    * starting a *fresh* PEFT fine-tune on top of a base model
-
-    Works for both local directories and Hub repo ids.
-    """
-    import os
-
-    adapter_config_name = "adapter_config.json"
-
-    if os.path.isdir(pretrained_path):
-        return os.path.isfile(os.path.join(pretrained_path, adapter_config_name))
-
-    from huggingface_hub import file_exists
-    from huggingface_hub.errors import HfHubHTTPError
-
-    try:
-        return file_exists(pretrained_path, adapter_config_name, revision=revision)
-    except (HfHubHTTPError, OSError):
-        return False
-
-
 def make_policy(
     cfg: PreTrainedConfig,
     ds_meta: LeRobotDatasetMetadata | None = None,
@@ -638,24 +607,13 @@ def make_policy(
             "the PEFT config parameters to be set. For training with PEFT, see `lerobot_train.py` on how to do that."
         )
 
-    # When `use_peft=True` and a checkpoint is given, the checkpoint can be one of two things:
-    # 1. A base model checkpoint (e.g., a pretrained policy) on which we want to start a fresh PEFT fine-tune.
-    # 2. A PEFT adapter checkpoint (e.g., a previously trained PEFT adapter)
-    # We distinguish between these two cases
-    load_existing_adapter = (
-        cfg.pretrained_path
-        and cfg.use_peft
-        and _has_peft_adapter_config(str(cfg.pretrained_path), cfg.pretrained_revision)
-    )
-
-    if cfg.pretrained_path and not load_existing_adapter:
+    if cfg.pretrained_path and not cfg.use_peft:
         # Load a pretrained policy and override the config if needed (for example, if there are inference-time
-        # hyperparameters that we want to vary). This also covers starting a fresh PEFT fine-tune on top of a
-        # base model: the base weights are loaded here and `wrap_with_peft` builds the adapter afterwards.
+        # hyperparameters that we want to vary).
         kwargs["pretrained_name_or_path"] = cfg.pretrained_path
         kwargs["revision"] = cfg.pretrained_revision
         policy = policy_cls.from_pretrained(**kwargs)
-    elif load_existing_adapter:
+    elif cfg.pretrained_path and cfg.use_peft:
         # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
         # of the adapter and the adapter's config contains the path to the base policy. So we need the
         # adapter config first, then load the correct policy and then apply PEFT.

src/lerobot/policies/lingbot_va/configuration_lingbot_va.py

@@ -92,9 +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)))
 
-    # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s.
-    save_predicted_video: bool = False
-
     # Normalization: IDENTITY here; images are scaled + VAE-encoded and actions are
     # quantile-(un)normalized inside the policy / dedicated processor steps.
     normalization_mapping: dict[str, NormalizationMode] = field(

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
+from lerobot.policies.pretrained import PreTrainedPolicy, unpack_action_output
 from lerobot.utils.constants import ACTION
 from lerobot.utils.import_utils import require_package
 
@@ -99,8 +99,6 @@ 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)
@@ -170,8 +168,6 @@ 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)
@@ -400,22 +396,33 @@ 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], **kwargs) -> Tensor:
+    def select_action(
+        self, batch: dict[str, Tensor], return_intermediate_predictions: bool = False, **kwargs
+    ) -> Tensor | tuple[Tensor, dict[str, 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 = self.predict_action_chunk(batch)  # [B, chunk_size, n_used]
+            actions, predictions = unpack_action_output(
+                self.predict_action_chunk(
+                    batch, return_intermediate_predictions=return_intermediate_predictions
+                )
+            )  # [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
@@ -427,17 +434,31 @@ 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 = self.predict_action_chunk(None)
+                actions, predictions = unpack_action_output(
+                    self.predict_action_chunk(
+                        None, return_intermediate_predictions=return_intermediate_predictions
+                    )
+                )
                 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
-        return self._action_queue.popleft()
+        action = self._action_queue.popleft()
+        if return_intermediate_predictions:
+            return action, predictions
+        return action
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
-        """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized)."""
+    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).
+        """
         self.eval()
         self._ensure_frozen_modules()
         self._maybe_init_prompt(batch)
@@ -459,12 +480,6 @@ 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
@@ -473,7 +488,15 @@ 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]
-        return a.to(torch.float32)
+        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
 
     # Prompt / text encoding
     def _maybe_init_prompt(self, batch):
@@ -834,11 +857,6 @@ 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/pretrained.py

@@ -23,6 +23,8 @@ 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
@@ -32,7 +34,6 @@ from torch import Tensor, nn
 from lerobot.__version__ import __version__
 from lerobot.configs import PreTrainedConfig
 from lerobot.configs.train import TrainPipelineConfig
-from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
 
 from .utils import log_model_loading_keys
@@ -92,6 +93,18 @@ 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):
@@ -220,10 +233,26 @@ class PreTrainedPolicy(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=resolve_safetensors_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)
         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
@@ -256,20 +285,34 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
+    def predict_action_chunk(
+        self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]
+    ) -> Tensor | tuple[Tensor, dict[str, 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:
+    def select_action(
+        self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]
+    ) -> Tensor | tuple[Tensor, dict[str, 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/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
@@ -28,7 +30,6 @@ from safetensors.torch import load_model as load_model_as_safetensor, save_model
 from torch import Tensor, nn
 
 from lerobot.configs.rewards import RewardModelConfig
-from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
 
 if TYPE_CHECKING:
@@ -128,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=resolve_safetensors_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

@@ -226,6 +226,10 @@ class RolloutConfig:
     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
@@ -255,6 +259,21 @@ 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,
+            )
+
         # --- 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

@@ -429,6 +429,7 @@ 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,6 +69,15 @@ 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,6 +95,7 @@ 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)
@@ -108,6 +109,7 @@ def create_inference_engine(
             task=task,
             device=device,
             robot_type=robot_wrapper.robot_type,
+            visualize_predictions=visualize_predictions,
         )
     if isinstance(config, RTCInferenceConfig):
         return RTCInferenceEngine(

src/lerobot/rollout/inference/prediction_buffer.py

@@ -0,0 +1,79 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Paced playback of a policy's temporal predictions for visualization.
+
+Intermediate predictions (e.g. a world model's imagined clip) arrive as a ``{key: sequence}`` stack
+with time on axis 0, one fresh stack per predicted chunk. Advancement is kept separate from reading
+so the loop's two clocks don't get conflated: :meth:`advance` steps the playhead at policy cadence
+(one new prediction per chunk), while :meth:`current` is a side-effect-free read for the display,
+which may run faster (e.g. action-interpolated). Between policy ticks the viewer just holds the last
+frame. Modality-agnostic: a step may be an image, vector, or scalar — only axis 0 is assumed to be
+time.
+"""
+
+from __future__ import annotations
+
+
+class PacedPredictionBuffer:
+    """Maps a temporal prediction stack onto a chunk's policy ticks for paced visualization.
+
+    Drive from policy inference (:meth:`load` a fresh chunk, :meth:`advance` on its later ticks);
+    read from the display loop (:meth:`current`).
+    """
+
+    def __init__(self, ticks_per_chunk: int | None = None) -> None:
+        # Policy ticks per predicted chunk; the T steps of each stack are spread across this span.
+        # ``None`` falls back to one step per policy tick (clamped to the stack length).
+        self._ticks_per_chunk = ticks_per_chunk
+        self._stacks: dict = {}  # key -> sequence with time on axis 0, for the current chunk
+        self._cursor = 0  # policy ticks elapsed since the current chunk's stacks were loaded
+
+    def reset(self) -> None:
+        """Drop the current chunk's stacks and rewind the playhead."""
+        self._stacks = {}
+        self._cursor = 0
+
+    def load(self, stacks: dict) -> None:
+        """Store a freshly predicted chunk's stacks and rewind the playhead to its first step."""
+        self._stacks = stacks
+        self._cursor = 0
+
+    def advance(self) -> None:
+        """Move the playhead forward one policy tick (no-op until a chunk is loaded)."""
+        if self._stacks:
+            self._cursor += 1
+
+    def current(self) -> dict | None:
+        """Current playhead step per key (``None`` until a chunk is loaded); no side effects.
+
+        Maps each stack's ``T`` steps onto the ``ticks_per_chunk`` span (one step/tick, clamped, when
+        the span is unknown).
+        """
+        if not self._stacks:
+            return None
+        tick = self._cursor
+        span = self._ticks_per_chunk
+        out: dict = {}
+        for key, stack in self._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)
+            step = stack[idx]
+            if hasattr(step, "detach"):  # torch tensor -> numpy for the visualization backend
+                step = step.detach().cpu().numpy()
+            out[key] = step
+        return out or None

src/lerobot/rollout/inference/sync.py

@@ -22,11 +22,12 @@ from copy import copy
 
 import torch
 
-from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.pretrained import PreTrainedPolicy, unpack_action_output
 from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
 from lerobot.processor import PolicyProcessorPipeline
 
 from .base import InferenceEngine
+from .prediction_buffer import PacedPredictionBuffer
 
 logger = logging.getLogger(__name__)
 
@@ -64,6 +65,7 @@ class SyncInferenceEngine(InferenceEngine):
         task: str,
         device: str | None,
         robot_type: str,
+        visualize_predictions: bool = False,
     ) -> None:
         self._policy = policy
         self._preprocessor = preprocessor
@@ -73,10 +75,18 @@ class SyncInferenceEngine(InferenceEngine):
         self._task = task
         self._device = torch.device(device or "cpu")
         self._robot_type = robot_type
+
+        # Intermediate-prediction visualization (e.g. a world model's imagined future). ``get_action``
+        # (once per policy tick) loads/advances the buffer; ``get_intermediate_predictions`` reads it
+        # for display. Pacing on policy ticks keeps playback aligned with execution regardless of
+        # action-interpolation upsampling.
+        self._visualize_predictions = visualize_predictions
+        self._predictions = PacedPredictionBuffer(getattr(policy.config, "chunk_size", None))
         logger.info(
-            "SyncInferenceEngine initialized (device=%s, action_keys=%d)",
+            "SyncInferenceEngine initialized (device=%s, action_keys=%d, visualize_predictions=%s)",
             self._device,
             len(ordered_action_keys),
+            self._visualize_predictions,
         )
 
     def start(self) -> None:
@@ -93,6 +103,11 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
+        self._predictions.reset()
+
+    def get_intermediate_predictions(self) -> dict | None:
+        """Read the current chunk's paced prediction frame for display (``None`` if none)."""
+        return self._predictions.current()
 
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
         """Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
@@ -112,7 +127,18 @@ class SyncInferenceEngine(InferenceEngine):
                 observation, self._device, self._task, self._robot_type
             )
             observation = self._preprocessor(observation)
-            action = self._policy.select_action(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 — load its stacks and rewind the playhead.
+                    self._predictions.load(predictions)
+                else:
+                    # Same chunk, later policy tick — step the playhead through the imagined frames.
+                    self._predictions.advance()
+            else:
+                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,10 +166,16 @@ 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,
         )
 

src/lerobot/scripts/lerobot_annotate.py

@@ -28,12 +28,7 @@ For distributed runs, see ``examples/annotations/run_hf_job.py``.
 """
 
 import logging
-from contextlib import suppress
 from pathlib import Path
-from typing import TYPE_CHECKING
-
-from huggingface_hub import HfApi, snapshot_download
-from huggingface_hub.errors import RevisionNotFoundError
 
 from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
 from lerobot.annotations.steerable_pipeline.executor import Executor
@@ -47,12 +42,6 @@ from lerobot.annotations.steerable_pipeline.validator import StagingValidator
 from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
 from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
 from lerobot.configs import parser
-from lerobot.utils.import_utils import _datasets_available, require_package
-
-if TYPE_CHECKING or _datasets_available:
-    from lerobot.datasets.dataset_metadata import CODEBASE_VERSION
-    from lerobot.datasets.io_utils import load_info
-    from lerobot.datasets.utils import create_lerobot_dataset_card
 
 logger = logging.getLogger(__name__)
 
@@ -61,6 +50,8 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
     if cfg.root is not None:
         return Path(cfg.root)
     if cfg.repo_id is not None:
+        from huggingface_hub import snapshot_download
+
         return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset"))
     raise ValueError("Either --root or --repo_id must be provided.")
 
@@ -134,7 +125,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
 
     Pushes to ``cfg.new_repo_id`` when set, otherwise back to ``cfg.repo_id``.
     """
-    require_package("datasets", "dataset")
+    from huggingface_hub import HfApi  # noqa: PLC0415
 
     repo_id = cfg.new_repo_id or cfg.repo_id
     commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
@@ -152,26 +143,33 @@ 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
     # ``RevisionNotFoundError``. Read the version straight from the
     # dataset's own ``meta/info.json`` so we tag whatever the writer
     # actually wrote (no accidental drift if the codebase floor moves).
-    version_tag = (
-        dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
-    )
+    from lerobot.datasets.dataset_metadata import CODEBASE_VERSION  # noqa: PLC0415
+
+    info_path = root / "meta" / "info.json"
+    version_tag = CODEBASE_VERSION
+    if info_path.exists():
+        try:
+            from lerobot.utils.io_utils import load_json  # noqa: PLC0415
+
+            info = load_json(info_path)
+            ds_version = info.get("codebase_version")
+            if isinstance(ds_version, str) and ds_version.startswith("v"):
+                version_tag = ds_version
+        except Exception as exc:  # noqa: BLE001
+            print(
+                f"[lerobot-annotate] could not read codebase_version from info.json ({exc}); falling back to {version_tag}",
+                flush=True,
+            )
     revision = getattr(commit_info, "oid", None)
     tag_kwargs = {
         "repo_id": repo_id,
@@ -182,6 +180,10 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         tag_kwargs["revision"] = revision
 
     try:
+        from contextlib import suppress  # noqa: PLC0415
+
+        from huggingface_hub.errors import RevisionNotFoundError  # noqa: PLC0415
+
         with suppress(RevisionNotFoundError):
             api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
         api.create_tag(**tag_kwargs)

src/lerobot/scripts/lerobot_eval.py

@@ -83,6 +83,7 @@ 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
@@ -169,7 +170,7 @@ def rollout(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
-    predicted_latents_callback: Callable[[PreTrainedPolicy], None] | None = None,
+    save_predicted_video: bool = False,
 ) -> dict:
     """Run a batched policy rollout once through a batch of environments.
 
@@ -199,9 +200,10 @@ 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.
-        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.
+        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.
     Returns:
         The dictionary described above.
     """
@@ -245,6 +247,9 @@ 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.
@@ -279,9 +284,13 @@ def rollout(
 
             observation = preprocessor(observation)
             with torch.inference_mode():
-                action = policy.select_action(observation)
-            if predicted_latents_callback is not None:
-                predicted_latents_callback(policy)
+                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 = postprocessor(action)
 
             action_transition = {ACTION: action}
@@ -394,6 +403,9 @@ 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()
 
@@ -435,11 +447,6 @@ 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."
@@ -489,16 +496,6 @@ 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
 
@@ -510,9 +507,6 @@ 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:
@@ -533,7 +527,7 @@ def eval_policy(
             env_features=env_features,
             recording_repo_id=recording_repo_id,
             recording_private=recording_private,
-            predicted_latents_callback=collect_predicted_latents if save_predicted_video else None,
+            save_predicted_video=save_predicted_video,
         )
 
         # Figure out where in each rollout sequence the first done condition was encountered (results after
@@ -599,33 +593,33 @@ def eval_policy(
                 threads.append(thread)
                 n_episodes_rendered += 1
 
-        # 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()
+        # 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()):
             videos_dir.mkdir(parents=True, exist_ok=True)
-            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)
+            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)
             n_predicted_rendered += 1
 
         progbar.set_postfix(
@@ -771,6 +765,11 @@ 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(
@@ -790,6 +789,7 @@ 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,6 +837,7 @@ 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."""
 
@@ -858,6 +859,7 @@ 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"]
@@ -889,6 +891,7 @@ 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).
@@ -923,6 +926,7 @@ 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:
@@ -949,6 +953,7 @@ 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}}.
@@ -1008,6 +1013,7 @@ 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

@@ -171,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
 
 
@@ -575,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,
@@ -608,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,6 +30,9 @@ 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/device_utils.py

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

src/lerobot/utils/foxglove_visualization.py

@@ -37,6 +37,8 @@ from .constants import (
     OBS_PREFIX,
     OBS_STATE,
     OBS_STR,
+    PREDICTION_PREFIX,
+    PREDICTION_STR,
     REWARD,
     SUCCESS,
     TRUNCATED,
@@ -283,10 +285,11 @@ 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 and action data to a Foxglove WebSocket server for real-time visualization.
+    Logs observation, action and prediction 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:
@@ -302,6 +305,8 @@ 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.
     """
@@ -334,6 +339,30 @@ 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, OBS_STR
+from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, PREDICTION_PREFIX
 from .import_utils import require_package
 
 
@@ -37,6 +37,43 @@ 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:
@@ -73,10 +110,16 @@ def shutdown_rerun() -> None:
     rr.rerun_shutdown()
 
 
-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.
+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.
 
-    Camera images, observation and action scalars are arranged in a grid.
+    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.
     """
 
     # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
@@ -88,22 +131,29 @@ def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_
         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]) -> None:
-    """Build and send the blueprint once, from the first observation and action data."""
+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."""
     if getattr(log_rerun_data, "blueprint", None) is not None:
         return
 
-    if not (observation_paths or action_paths or image_paths):
+    if not (observation_paths or action_paths or image_paths or prediction_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)
+    blueprint = _build_blueprint(observation_paths, action_paths, image_paths, prediction_paths)
     log_rerun_data.blueprint = blueprint
     rr.send_blueprint(blueprint)
 
@@ -111,10 +161,11 @@ def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image
 def log_rerun_data(
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
+    prediction: dict | None = None,
     compress_images: bool = False,
 ) -> None:
     """
-    Logs observation and action data to Rerun for real-time visualization.
+    Logs observation, action and prediction 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:
@@ -133,6 +184,8 @@ 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.
     """
 
@@ -142,37 +195,19 @@ 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:
-        for k, v in observation.items():
-            if v is None:
-                continue
-            key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}"
+        _log_scalar_or_image_mapping(
+            rr, observation, OBS_PREFIX, observation_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 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 action:
         for k, v in action.items():
@@ -188,4 +223,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)
+    _ensure_blueprint(observation_paths, action_paths, image_paths, prediction_paths)

src/lerobot/utils/visualization_utils.py

@@ -56,14 +56,19 @@ 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 data to the backend selected by ``display_mode``."""
+    """Logs observation/action/prediction data to the backend selected by ``display_mode``."""
 
     if display_mode == "rerun":
-        log_rerun_data(observation=observation, action=action, compress_images=compress_images)
+        log_rerun_data(
+            observation=observation, action=action, prediction=prediction, compress_images=compress_images
+        )
     elif display_mode == "foxglove":
-        log_foxglove_data(observation=observation, action=action, compress_images=compress_images)
+        log_foxglove_data(
+            observation=observation, action=action, prediction=prediction, 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/policies/test_peft_base_model_loading.py

@@ -1,142 +0,0 @@
-# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Regression tests for PEFT loading when the checkpoint is a base model (see issue #3975).
-
-Starting a fresh LoRA/PEFT fine-tune points ``--policy.path`` at a *base model* (no
-``adapter_config.json``) while also setting ``use_peft=True``. This must NOT be mistaken for
-loading an existing PEFT adapter. These tests lock in the base-model vs. adapter distinction
-made by ``lerobot.policies.factory._has_peft_adapter_config`` and the branch it drives in
-``make_policy``. They are pure/fast (no network, no ``peft``), so they run in CI.
-"""
-
-import json
-from unittest.mock import MagicMock, patch
-
-import torch
-from huggingface_hub.errors import HfHubHTTPError
-
-from lerobot.policies.factory import _has_peft_adapter_config
-
-
-def test_local_base_model_dir_has_no_adapter_config(tmp_path):
-    # A base-model checkpoint directory (only model weights, no adapter config).
-    (tmp_path / "model.safetensors").write_bytes(b"")
-    (tmp_path / "config.json").write_text("{}")
-    assert _has_peft_adapter_config(str(tmp_path)) is False
-
-
-def test_local_adapter_dir_has_adapter_config(tmp_path):
-    (tmp_path / "adapter_config.json").write_text(json.dumps({"peft_type": "LORA"}))
-    (tmp_path / "adapter_model.safetensors").write_bytes(b"")
-    assert _has_peft_adapter_config(str(tmp_path)) is True
-
-
-def test_hub_base_model_repo_has_no_adapter_config():
-    with patch("huggingface_hub.file_exists", return_value=False) as mock_exists:
-        assert _has_peft_adapter_config("lerobot/lingbot_va_base") is False
-    mock_exists.assert_called_once()
-    assert mock_exists.call_args.args[1] == "adapter_config.json"
-
-
-def test_hub_adapter_repo_has_adapter_config():
-    with patch("huggingface_hub.file_exists", return_value=True):
-        assert _has_peft_adapter_config("some/adapter-repo", revision="main") is True
-
-
-def test_hub_lookup_error_falls_back_to_base_model():
-    # Offline / private / transient Hub errors must not crash; treat as "not an adapter".
-    with patch(
-        "huggingface_hub.file_exists",
-        side_effect=HfHubHTTPError("boom", response=MagicMock()),
-    ):
-        assert _has_peft_adapter_config("some/private-repo") is False
-    with patch("huggingface_hub.file_exists", side_effect=OSError("offline")):
-        assert _has_peft_adapter_config("some/repo") is False
-
-
-def _make_dummy_policy_cfg(pretrained_path, use_peft):
-    cfg = MagicMock()
-    cfg.type = "act"
-    cfg.device = "cpu"
-    cfg.pretrained_path = pretrained_path
-    cfg.pretrained_revision = None
-    cfg.use_peft = use_peft
-    cfg.input_features = {}
-    cfg.output_features = {}
-    return cfg
-
-
-@patch("lerobot.policies.factory.validate_visual_features_consistency")
-@patch("lerobot.policies.factory.env_to_policy_features", return_value={})
-@patch("lerobot.policies.factory.get_policy_class")
-def test_make_policy_base_model_with_use_peft_loads_base_not_adapter(
-    mock_get_cls, _mock_features, _mock_validate
-):
-    """`use_peft=True` on a base model must load the base weights, not a PEFT adapter.
-
-    Before the #3975 fix this went down the ``PeftConfig.from_pretrained`` path and failed
-    looking for a non-existent ``adapter_config.json``.
-    """
-    from lerobot.policies import factory
-
-    policy_cls = MagicMock()
-    loaded_policy = torch.nn.Linear(1, 1)  # a real nn.Module so make_policy's assert passes
-    policy_cls.from_pretrained.return_value = loaded_policy
-    mock_get_cls.return_value = policy_cls
-
-    cfg = _make_dummy_policy_cfg(pretrained_path="lerobot/lingbot_va_base", use_peft=True)
-    env_cfg = MagicMock()
-
-    with patch.object(factory, "_has_peft_adapter_config", return_value=False) as mock_has_adapter:
-        policy = factory.make_policy(cfg=cfg, env_cfg=env_cfg)
-
-    mock_has_adapter.assert_called_once()
-    # Base model is loaded via the normal pretrained path...
-    policy_cls.from_pretrained.assert_called_once()
-    assert policy_cls.from_pretrained.call_args.kwargs["pretrained_name_or_path"] == (
-        "lerobot/lingbot_va_base"
-    )
-    # ...and PEFT adapter loading is NOT attempted (would need peft + adapter_config.json).
-    assert policy is loaded_policy
-
-
-@patch("lerobot.policies.factory.validate_visual_features_consistency")
-@patch("lerobot.policies.factory.env_to_policy_features", return_value={})
-@patch("lerobot.policies.factory.get_policy_class")
-def test_make_policy_existing_adapter_uses_peft_loading(mock_get_cls, _mock_features, _mock_validate):
-    """A real adapter checkpoint (has ``adapter_config.json``) must go through PEFT loading."""
-    from lerobot.policies import factory
-
-    policy_cls = MagicMock()
-    mock_get_cls.return_value = policy_cls
-
-    cfg = _make_dummy_policy_cfg(pretrained_path="some/adapter-repo", use_peft=True)
-    env_cfg = MagicMock()
-
-    policy_cls.from_pretrained.return_value = torch.nn.Linear(1, 1)
-
-    fake_peft = MagicMock()
-    fake_peft_config = MagicMock()
-    fake_peft_config.base_model_name_or_path = "lerobot/lingbot_va_base"
-    fake_peft.PeftConfig.from_pretrained.return_value = fake_peft_config
-    fake_peft.PeftModel.from_pretrained.return_value = torch.nn.Linear(1, 1)
-
-    with (
-        patch.object(factory, "_has_peft_adapter_config", return_value=True),
-        patch.dict("sys.modules", {"peft": fake_peft}),
-    ):
-        factory.make_policy(cfg=cfg, env_cfg=env_cfg)
-
-    fake_peft.PeftConfig.from_pretrained.assert_called_once_with("some/adapter-repo")
-    fake_peft.PeftModel.from_pretrained.assert_called_once()

tests/scripts/test_lerobot_annotate.py

@@ -18,8 +18,6 @@ import json
 from types import SimpleNamespace
 
 import pytest
-import requests
-from huggingface_hub.errors import RevisionNotFoundError
 
 # ``lerobot.scripts.lerobot_annotate`` (and the ``_push_to_hub`` path it
 # exercises) imports ``lerobot.datasets``, which only ships under the
@@ -28,13 +26,11 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
 
 
 def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
-    from lerobot.scripts import lerobot_annotate
+    from lerobot.scripts.lerobot_annotate import _push_to_hub
 
     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 = {}
 
@@ -47,6 +43,9 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
             return SimpleNamespace(oid="abc123")
 
         def delete_tag(self, repo_id, **kwargs):
+            import requests
+            from huggingface_hub.errors import RevisionNotFoundError
+
             calls["delete_tag"] = {"repo_id": repo_id, **kwargs}
             # Simulate the common case: no stale tag to delete.
             raise RevisionNotFoundError("no such tag", response=requests.Response())
@@ -54,12 +53,7 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
         def create_tag(self, **kwargs):
             calls["create_tag"] = kwargs
 
-    monkeypatch.setattr(lerobot_annotate, "HfApi", FakeHfApi)
-
-    def fake_card_push(self, **kwargs):
-        calls["card_push"] = {"content": str(self), **kwargs}
-
-    monkeypatch.setattr("huggingface_hub.DatasetCard.push_to_hub", fake_card_push)
+    monkeypatch.setattr("huggingface_hub.HfApi", FakeHfApi)
 
     cfg = SimpleNamespace(
         repo_id="source/dataset",
@@ -68,7 +62,7 @@ def test_push_to_hub_tags_uploaded_dataset_revision(tmp_path, monkeypatch):
         push_commit_message=None,
     )
 
-    lerobot_annotate._push_to_hub(root, cfg)
+    _push_to_hub(root, cfg)
 
     assert calls["create_repo"] == {
         "repo_id": "annotated/dataset",
@@ -77,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

@@ -19,6 +19,7 @@ from __future__ import annotations
 import dataclasses
 from unittest.mock import MagicMock
 
+import numpy as np
 import pytest
 import torch
 
@@ -348,3 +349,62 @@ def test_rollout_context_fields():
 
     field_names = {f.name for f in dataclasses.fields(RolloutContext)}
     assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
+
+
+# ---------------------------------------------------------------------------
+# Paced prediction buffer
+# ---------------------------------------------------------------------------
+
+
+def _walk_policy_ticks(buffer, key, n_ticks):
+    """Read the buffer over ``n_ticks`` policy ticks (load counts as the first)."""
+    seen = [buffer.current()[key]]
+    for _ in range(n_ticks - 1):
+        buffer.advance()
+        seen.append(buffer.current()[key])
+    return seen
+
+
+def test_prediction_buffer_paces_over_chunk():
+    from lerobot.rollout.inference.prediction_buffer import PacedPredictionBuffer
+
+    buffer = PacedPredictionBuffer(ticks_per_chunk=4)
+    assert buffer.current() is None  # nothing until a chunk is loaded
+
+    # 2 steps spread over a 4-policy-tick chunk: first half -> step 0, second half -> step 1.
+    buffer.load({"x": ["a", "b"]})
+    assert _walk_policy_ticks(buffer, "x", 4) == ["a", "a", "b", "b"]
+
+    buffer.reset()
+    assert buffer.current() is None
+
+    # Unknown span -> one step per policy tick, clamped at the last step.
+    unpaced = PacedPredictionBuffer(ticks_per_chunk=None)
+    unpaced.load({"x": ["a", "b"]})
+    assert _walk_policy_ticks(unpaced, "x", 3) == ["a", "b", "b"]
+
+
+def test_prediction_buffer_read_is_independent_of_display_rate():
+    from lerobot.rollout.inference.prediction_buffer import PacedPredictionBuffer
+
+    # Pacing fix: the playhead advances per policy tick; the display may read it many times per tick
+    # (e.g. interpolation multiplier=3), each read returning the same frame with no side effects — so
+    # playback never races to the end and freezes.
+    buffer = PacedPredictionBuffer(ticks_per_chunk=2)
+    buffer.load({"x": ["a", "b"]})
+    assert [buffer.current()["x"] for _ in range(3)] == ["a", "a", "a"]  # policy tick 0
+    buffer.advance()
+    assert [buffer.current()["x"] for _ in range(3)] == ["b", "b", "b"]  # policy tick 1
+
+
+def test_prediction_buffer_is_modality_agnostic():
+    from lerobot.rollout.inference.prediction_buffer import PacedPredictionBuffer
+
+    # Non-image temporal predictions pass through; torch tensors are detached to numpy on the way out.
+    buffer = PacedPredictionBuffer(ticks_per_chunk=2)
+    states = torch.arange(6).reshape(2, 3)  # [T=2, D=3]
+    buffer.load({"observation.state": states})
+    out = buffer.current()["observation.state"]
+    assert isinstance(out, np.ndarray) and np.array_equal(out, states[0].numpy())
+    buffer.advance()
+    assert np.array_equal(buffer.current()["observation.state"], states[1].numpy())

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)