fix typo

Co-authored-by: Martino Russi <77496684+nepyope@users.noreply.github.com>

docs/source/annotation_pipeline.mdx

@@ -81,6 +81,12 @@ merged. Both prompts also carry a causal **event-boundary** definition (a
 new event starts when an object becomes held / is released / reaches a new
 location / a lid changes state / contents move) to sharpen where cuts land.
 
+Optionally, a third **seeded-relabel** pass (`--plan.subtask_seeded_relabel`)
+revisits each span with its previous/current/next segment contact sheets and
+minimally corrects the label, using the first label as a prior — it keeps the
+boundaries fixed and only sharpens wording, at the cost of one extra call per
+subtask.
+
 The resulting spans are then stitched into a gap-free, full-episode
 cover, so **every frame has exactly one active subtask**. See
 [`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
@@ -157,30 +163,33 @@ 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.                                                               |
+| 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. |
 
 ### 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.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.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`).    |
 
 ### Interjections + VQA
 

examples/annotations/run_hf_job.py

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

src/lerobot/annotations/steerable_pipeline/config.py

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

src/lerobot/annotations/steerable_pipeline/frames.py

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

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

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

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

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

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

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

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

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

src/lerobot/annotations/steerable_pipeline/vlm_client.py

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

src/lerobot/datasets/compute_stats.py

@@ -18,9 +18,8 @@ from __future__ import annotations
 import logging
 
 import numpy as np
-import torch
 
-from lerobot.processor import RelativeActionsProcessorStep, to_relative_actions
+from lerobot.processor import RelativeActionsProcessorStep
 from lerobot.utils.constants import ACTION, OBS_STATE
 
 from .io_utils import load_image_as_numpy
@@ -661,8 +660,6 @@ def _compute_relative_chunk_batch(
     all_states: np.ndarray,
     chunk_size: int,
     relative_mask: np.ndarray,
-    pose_representation: str = "componentwise",
-    se3_pose_groups: list[list[int]] | None = None,
 ) -> np.ndarray:
     """Vectorised relative-action computation for a batch of start indices.
 
@@ -674,18 +671,6 @@ def _compute_relative_chunk_batch(
     frame_idx = start_indices[:, None] + offsets[None, :]
     chunks = all_actions[frame_idx].copy()
     states = all_states[start_indices]
-    if pose_representation == "se3":
-        return (
-            to_relative_actions(
-                torch.from_numpy(chunks),
-                torch.from_numpy(states),
-                relative_mask.astype(bool).tolist(),
-                pose_representation=pose_representation,
-                se3_pose_groups=se3_pose_groups,
-            )
-            .numpy()
-            .reshape(-1, all_actions.shape[1])
-        )
     mask_dim = len(relative_mask)
     chunks[:, :, :mask_dim] -= states[:, None, :mask_dim] * relative_mask[None, None, :]
     return chunks.reshape(-1, all_actions.shape[1])
@@ -697,9 +682,6 @@ def compute_relative_action_stats(
     chunk_size: int,
     exclude_joints: list[str] | None = None,
     num_workers: int = 0,
-    state_from_action: bool = False,
-    pose_representation: str = "componentwise",
-    se3_pose_groups: list[list[int]] | None = None,
 ) -> dict[str, np.ndarray]:
     """Compute normalization statistics for relative actions over the full dataset.
 
@@ -718,9 +700,6 @@ def compute_relative_action_stats(
         num_workers: Number of parallel threads for computation. Values ≤1
             mean single-threaded. Numpy releases the GIL so threads give
             real parallelism here.
-        state_from_action: Use the current absolute action as state. This is
-            intended for state-less pose datasets where each action row is the
-            synchronized measured robot pose.
 
     Returns:
         Statistics dict with keys "mean", "std", "min", "max", "q01", …, "q99".
@@ -743,7 +722,7 @@ def compute_relative_action_stats(
 
     logging.info("Loading action/state data for relative action stats...")
     all_actions = np.array(hf_dataset[ACTION], dtype=np.float32)
-    all_states = all_actions if state_from_action else np.array(hf_dataset[OBS_STATE], dtype=np.float32)
+    all_states = np.array(hf_dataset[OBS_STATE], dtype=np.float32)
     episode_indices = np.array(hf_dataset["episode_index"])
 
     valid_starts = _get_valid_chunk_starts(episode_indices, chunk_size)
@@ -775,8 +754,6 @@ def compute_relative_action_stats(
                     all_states,
                     chunk_size,
                     relative_mask,
-                    pose_representation,
-                    se3_pose_groups,
                 )
                 for batch in batches
             ]
@@ -785,15 +762,7 @@ def compute_relative_action_stats(
     else:
         for batch in batches:
             running_stats.update(
-                _compute_relative_chunk_batch(
-                    batch,
-                    all_actions,
-                    all_states,
-                    chunk_size,
-                    relative_mask,
-                    pose_representation,
-                    se3_pose_groups,
-                )
+                _compute_relative_chunk_batch(batch, all_actions, all_states, chunk_size, relative_mask)
             )
 
     stats = running_stats.get_statistics()
@@ -808,58 +777,3 @@ def compute_relative_action_stats(
     )
 
     return stats
-
-
-def compute_state_history_stats(
-    hf_dataset,
-    features: dict,
-    history_steps: int,
-    exclude_joints: list[str] | None = None,
-    relative: bool = False,
-    pose_representation: str = "componentwise",
-    se3_pose_groups: list[list[int]] | None = None,
-) -> dict[str, np.ndarray]:
-    """Compute stats for flattened state history synthesized from absolute actions.
-
-    History is left-padded with the first action of each episode, matching dataset
-    boundary padding. When ``relative`` is enabled, every history pose is expressed
-    relative to its newest pose while excluded dimensions remain absolute.
-    """
-    if history_steps < 1:
-        raise ValueError("history_steps must be at least 1")
-    if exclude_joints is None:
-        exclude_joints = []
-
-    actions = np.asarray(hf_dataset[ACTION], dtype=np.float32)
-    episode_indices = np.asarray(hf_dataset["episode_index"])
-    sample_indices = np.arange(len(actions))
-    episode_starts = np.maximum.accumulate(
-        np.where(
-            np.concatenate(([True], episode_indices[1:] != episode_indices[:-1])),
-            sample_indices,
-            0,
-        )
-    )
-    offsets = np.arange(-(history_steps - 1), 1)
-    history_indices = np.maximum(sample_indices[:, None] + offsets[None, :], episode_starts[:, None])
-    history = actions[history_indices].copy()
-
-    if relative:
-        state_dim = actions.shape[-1]
-        names = features.get(ACTION, {}).get("names")
-        mask_step = RelativeActionsProcessorStep(
-            enabled=True,
-            exclude_joints=exclude_joints,
-            action_names=names,
-        )
-        mask = mask_step._build_mask(state_dim)
-        history = to_relative_actions(
-            torch.from_numpy(history),
-            torch.from_numpy(history[:, -1].copy()),
-            mask,
-            pose_representation=pose_representation,
-            se3_pose_groups=se3_pose_groups,
-        ).numpy()
-
-    flattened = history.reshape(len(history), -1)
-    return get_feature_stats(flattened, axis=0, keepdims=False)

src/lerobot/datasets/dataset_tools.py

@@ -54,7 +54,6 @@ from .compute_stats import (
     aggregate_stats,
     compute_episode_stats,
     compute_relative_action_stats,
-    compute_state_history_stats,
 )
 from .dataset_metadata import LeRobotDatasetMetadata
 from .image_writer import write_image
@@ -1567,12 +1566,6 @@ def recompute_stats(
     relative_exclude_joints: list[str] | None = None,
     chunk_size: int = 50,
     num_workers: int = 0,
-    state_from_action: bool = False,
-    state_history_steps: int = 1,
-    relative_state_history: bool = False,
-    relative_state_exclude_joints: list[str] | None = None,
-    relative_pose_representation: str = "componentwise",
-    relative_se3_pose_groups: list[list[int]] | None = None,
 ) -> LeRobotDataset:
     """Recompute stats.json from scratch by iterating all episodes.
 
@@ -1590,15 +1583,6 @@ def recompute_stats(
             ``policy.chunk_size``. Only used when ``relative_action=True``.
         num_workers: Number of parallel threads for relative action stats computation.
             Values ≤1 mean single-threaded. Only used when ``relative_action=True``.
-        state_from_action: Use absolute action rows as synthetic state while
-            computing relative-action stats, and write their absolute statistics
-            under ``observation.state``.
-        state_history_steps: Number of consecutive synthesized state samples.
-        relative_state_history: Express state history relative to its newest pose.
-        relative_state_exclude_joints: State dimensions to retain as absolute.
-        relative_pose_representation: ``componentwise`` for legacy subtraction or
-            ``se3`` for ``inv(T_current) @ T_target`` pose composition.
-        relative_se3_pose_groups: Six-index xyz+rotation-vector pose groups.
 
     Returns:
         The same dataset with updated stats.
@@ -1622,21 +1606,7 @@ def recompute_stats(
     # (matching what the model sees during training) and skip action in the
     # per-episode pass below.
     relative_action_stats = None
-    synthetic_state_stats = None
-    if state_from_action:
-        if ACTION not in features:
-            raise ValueError("state_from_action requires an action feature")
-        synthetic_state_stats = compute_state_history_stats(
-            dataset.hf_dataset,
-            features,
-            history_steps=state_history_steps,
-            exclude_joints=relative_state_exclude_joints,
-            relative=relative_state_history,
-            pose_representation=relative_pose_representation,
-            se3_pose_groups=relative_se3_pose_groups,
-        )
-
-    if relative_action and ACTION in features and (OBS_STATE in features or state_from_action):
+    if relative_action and ACTION in features and OBS_STATE in features:
         if relative_exclude_joints is None:
             relative_exclude_joints = ["gripper"]
         relative_action_stats = compute_relative_action_stats(
@@ -1645,9 +1615,6 @@ def recompute_stats(
             chunk_size=chunk_size,
             exclude_joints=relative_exclude_joints,
             num_workers=num_workers,
-            state_from_action=state_from_action,
-            pose_representation=relative_pose_representation,
-            se3_pose_groups=relative_se3_pose_groups,
         )
         features_to_compute.pop(ACTION, None)
 
@@ -1687,8 +1654,6 @@ def recompute_stats(
 
     if relative_action_stats is not None:
         new_stats[ACTION] = relative_action_stats
-    if synthetic_state_stats is not None:
-        new_stats[OBS_STATE] = synthetic_state_stats
 
     # Merge: keep existing stats for features we didn't recompute
     if dataset.meta.stats:

src/lerobot/policies/evo1/internvl3_embedder.py

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

src/lerobot/policies/pi05/configuration_pi05.py

@@ -55,18 +55,6 @@ class PI05Config(PreTrainedConfig):
     relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
     # Populated at runtime from dataset metadata by make_policy.
     action_feature_names: list[str] | None = None
-    # ``se3`` uses inv(T_current) @ T_target for each xyz+rotation-vector pose group.
-    # ``componentwise`` preserves the legacy action - state behavior.
-    relative_pose_representation: str = "componentwise"
-    relative_se3_pose_groups: list[list[int]] = field(default_factory=lambda: [list(range(6))])
-
-    # Build proprioception from absolute action samples when the dataset has no
-    # observation.state. With history_steps=2, training samples request t-1 as
-    # well as the normal t..t+chunk_size-1 action targets.
-    state_from_action: bool = False
-    proprioception_history_steps: int = 1
-    use_relative_state_history: bool = False
-    relative_state_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
 
     # Real-Time Chunking (RTC) configuration
     rtc_config: RTCConfig | None = None
@@ -133,20 +121,6 @@ class PI05Config(PreTrainedConfig):
         if self.dtype not in ["bfloat16", "float32"]:
             raise ValueError(f"Invalid dtype: {self.dtype}")
 
-        if self.proprioception_history_steps < 1:
-            raise ValueError("proprioception_history_steps must be at least 1")
-
-        if self.relative_pose_representation not in {"componentwise", "se3"}:
-            raise ValueError(
-                "relative_pose_representation must be either 'componentwise' or 'se3', got "
-                f"{self.relative_pose_representation!r}"
-            )
-        for group in self.relative_se3_pose_groups:
-            if len(group) != 6 or len(set(group)) != 6 or any(index < 0 for index in group):
-                raise ValueError(f"Invalid six-index SE(3) pose group: {group}")
-        if self.relative_pose_representation == "se3" and not self.relative_se3_pose_groups:
-            raise ValueError("relative_pose_representation='se3' requires relative_se3_pose_groups")
-
     def validate_features(self) -> None:
         """Validate and set up input/output features."""
         for i in range(self.empty_cameras):
@@ -157,6 +131,13 @@ class PI05Config(PreTrainedConfig):
             )
             self.input_features[key] = empty_camera
 
+        if OBS_STATE not in self.input_features:
+            state_feature = PolicyFeature(
+                type=FeatureType.STATE,
+                shape=(self.max_state_dim,),  # Padded to max_state_dim
+            )
+            self.input_features[OBS_STATE] = state_feature
+
         if ACTION not in self.output_features:
             action_feature = PolicyFeature(
                 type=FeatureType.ACTION,
@@ -164,25 +145,6 @@ class PI05Config(PreTrainedConfig):
             )
             self.output_features[ACTION] = action_feature
 
-        if OBS_STATE not in self.input_features:
-            state_shape = (self.max_state_dim,)
-            if self.state_from_action and ACTION in self.output_features:
-                state_shape = self.output_features[ACTION].shape
-            state_feature = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=state_shape,
-            )
-            self.input_features[OBS_STATE] = state_feature
-
-        state_dim = self.input_features[OBS_STATE].shape[-1]
-        history_state_dim = state_dim * self.proprioception_history_steps
-        if history_state_dim > self.max_state_dim:
-            raise ValueError(
-                "Flattened proprioception history exceeds max_state_dim: "
-                f"{state_dim} * {self.proprioception_history_steps} = {history_state_dim} > "
-                f"{self.max_state_dim}"
-            )
-
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(
             lr=self.optimizer_lr,
@@ -206,8 +168,7 @@ class PI05Config(PreTrainedConfig):
 
     @property
     def action_delta_indices(self) -> list:
-        history_prefix = self.proprioception_history_steps - 1 if self.state_from_action else 0
-        return list(range(-history_prefix, self.chunk_size))
+        return list(range(self.chunk_size))
 
     @property
     def reward_delta_indices(self) -> None:

src/lerobot/policies/pi05/processor_pi05.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 
 from copy import deepcopy
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Any
 
 import numpy as np
@@ -36,7 +36,6 @@ from lerobot.processor import (
     TokenizerProcessorStep,
     UnnormalizerProcessorStep,
     policy_action_to_transition,
-    to_relative_actions,
     transition_to_policy_action,
 )
 from lerobot.types import EnvTransition, TransitionKey
@@ -49,150 +48,6 @@ from lerobot.utils.constants import (
 from .configuration_pi05 import PI05Config
 
 
-@ProcessorStepRegistry.register(name="pi05_state_from_action_processor_step")
-@dataclass
-class Pi05StateFromActionProcessorStep(ProcessorStep):
-    """Synthesize proprioception from absolute actions in state-less datasets.
-
-    The dataset loader supplies ``history_steps - 1`` actions before the normal
-    target chunk. Those leading samples and action(t) become state history; only
-    the leading samples are then removed from the action targets.
-    """
-
-    enabled: bool = False
-    history_steps: int = 1
-    _inference_history: torch.Tensor | None = field(default=None, init=False, repr=False)
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        if not self.enabled:
-            return transition
-
-        observation = transition.get(TransitionKey.OBSERVATION, {})
-        observed_state = observation.get(OBS_STATE)
-        if observed_state is not None:
-            # At inference the robot normally provides only the current state and
-            # there is no action target. Build a rolling history in the processor.
-            if transition.get(TransitionKey.ACTION) is None and observed_state.ndim == 2:
-                if self._inference_history is None:
-                    self._inference_history = observed_state.unsqueeze(1).repeat(1, self.history_steps, 1)
-                else:
-                    self._inference_history = torch.cat(
-                        [self._inference_history[:, 1:], observed_state.unsqueeze(1)], dim=1
-                    )
-                new_transition = transition.copy()
-                new_observation = dict(observation)
-                new_observation[OBS_STATE] = self._inference_history.clone()
-                new_transition[TransitionKey.OBSERVATION] = new_observation
-                return new_transition
-            return transition
-
-        action = transition.get(TransitionKey.ACTION)
-        if action is None:
-            raise ValueError("Cannot synthesize PI0.5 state without action")
-        if action.ndim != 3:
-            raise ValueError(f"Expected batched action chunks with shape (B, T, D), got {action.shape}")
-        if action.shape[1] < self.history_steps:
-            raise ValueError(
-                f"Action chunk has {action.shape[1]} steps, fewer than history_steps={self.history_steps}"
-            )
-
-        new_transition = transition.copy()
-        new_observation = dict(observation)
-        state = action[:, : self.history_steps].clone()
-        if self.history_steps == 1:
-            state = state[:, 0]
-        new_observation[OBS_STATE] = state
-        new_transition[TransitionKey.OBSERVATION] = new_observation
-        new_transition[TransitionKey.ACTION] = action[:, self.history_steps - 1 :]
-        return new_transition
-
-    def get_config(self) -> dict[str, Any]:
-        return {"enabled": self.enabled, "history_steps": self.history_steps}
-
-    def reset(self) -> None:
-        self._inference_history = None
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
-
-
-@ProcessorStepRegistry.register(name="pi05_flatten_state_history_processor_step")
-@dataclass
-class Pi05FlattenStateHistoryProcessorStep(ProcessorStep):
-    """Optionally relativize raw state history, then flatten it for PI0.5."""
-
-    history_steps: int = 1
-    max_state_dim: int = 32
-    relative: bool = False
-    exclude_joints: list[str] = field(default_factory=list)
-    state_names: list[str] | None = None
-    pose_representation: str = "componentwise"
-    se3_pose_groups: list[list[int]] = field(default_factory=list)
-
-    def __call__(self, transition: EnvTransition) -> EnvTransition:
-        observation = transition.get(TransitionKey.OBSERVATION, {})
-        state = observation.get(OBS_STATE)
-        if state is None:
-            raise ValueError("State is required for PI05")
-        if self.history_steps == 1 and state.ndim == 2:
-            state = state.unsqueeze(1)
-        if state.ndim != 3 or state.shape[1] != self.history_steps:
-            raise ValueError(
-                f"Expected state history with shape (B, {self.history_steps}, D), got {state.shape}"
-            )
-
-        flattened_dim = state.shape[1] * state.shape[2]
-        if flattened_dim > self.max_state_dim:
-            raise ValueError(
-                f"Flattened state history has {flattened_dim} dimensions, above max_state_dim={self.max_state_dim}"
-            )
-
-        processed_state = state.clone()
-        if self.relative:
-            mask_step = RelativeActionsProcessorStep(
-                enabled=True,
-                exclude_joints=self.exclude_joints,
-                action_names=self.state_names,
-            )
-            processed_state = to_relative_actions(
-                state,
-                state[:, -1],
-                mask_step._build_mask(state.shape[-1]),
-                pose_representation=self.pose_representation,
-                se3_pose_groups=self.se3_pose_groups,
-            )
-
-        new_transition = transition.copy()
-        new_observation = dict(observation)
-        new_observation[OBS_STATE] = processed_state.flatten(start_dim=1)
-        new_transition[TransitionKey.OBSERVATION] = new_observation
-        return new_transition
-
-    def get_config(self) -> dict[str, Any]:
-        return {
-            "history_steps": self.history_steps,
-            "max_state_dim": self.max_state_dim,
-            "relative": self.relative,
-            "exclude_joints": self.exclude_joints,
-            "state_names": self.state_names,
-            "pose_representation": self.pose_representation,
-            "se3_pose_groups": self.se3_pose_groups,
-        }
-
-    def transform_features(
-        self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
-    ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        transformed = deepcopy(features)
-        for feature_group in transformed.values():
-            state_feature = feature_group.get(OBS_STATE)
-            if state_feature is not None and self.history_steps > 1:
-                state_dim = state_feature.shape[-1] * self.history_steps
-                feature_group[OBS_STATE] = PolicyFeature(type=state_feature.type, shape=(state_dim,))
-        return transformed
-
-
 @ProcessorStepRegistry.register(name="pi05_prepare_state_tokenizer_processor_step")
 @dataclass
 class Pi05PrepareStateTokenizerProcessorStep(ProcessorStep):
@@ -278,28 +133,13 @@ def make_pi05_pre_post_processors(
         enabled=config.use_relative_actions,
         exclude_joints=getattr(config, "relative_exclude_joints", []),
         action_names=getattr(config, "action_feature_names", None),
-        pose_representation=config.relative_pose_representation,
-        se3_pose_groups=config.relative_se3_pose_groups,
     )
 
     # OpenPI order: raw → relative → normalize → model → unnormalize → absolute
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
         AddBatchDimensionProcessorStep(),
-        Pi05StateFromActionProcessorStep(
-            enabled=config.state_from_action,
-            history_steps=config.proprioception_history_steps,
-        ),
         relative_step,
-        Pi05FlattenStateHistoryProcessorStep(
-            history_steps=config.proprioception_history_steps,
-            max_state_dim=config.max_state_dim,
-            relative=config.use_relative_state_history,
-            exclude_joints=config.relative_state_exclude_joints,
-            state_names=config.action_feature_names,
-            pose_representation=config.relative_pose_representation,
-            se3_pose_groups=config.relative_se3_pose_groups,
-        ),
         # NOTE: NormalizerProcessorStep MUST come before Pi05PrepareStateTokenizerProcessorStep
         # because the tokenizer step expects normalized state in [-1, 1] range for discretization
         NormalizerProcessorStep(

src/lerobot/policies/pretrained.py

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

src/lerobot/policies/xvla/modeling_xvla.py

@@ -154,7 +154,7 @@ class XVLAModel(nn.Module):
         # Freeze or unfreeze policy transformer
         if not self.config.train_policy_transformer:
             for name, param in self.transformer.named_parameters():
-                if "soft_prompts" not in name:
+                if "soft_prompt" not in name:
                     param.requires_grad = False
 
         # Freeze or unfreeze soft prompts

src/lerobot/processor/__init__.py

@@ -90,9 +90,7 @@ from .relative_action_processor import (
     AbsoluteActionsProcessorStep,
     RelativeActionsProcessorStep,
     to_absolute_actions,
-    to_absolute_se3_pose,
     to_relative_actions,
-    to_relative_se3_pose,
 )
 from .rename_processor import RenameObservationsProcessorStep, rename_stats
 from .tokenizer_processor import ActionTokenizerProcessorStep, TokenizerProcessorStep
@@ -137,10 +135,6 @@ __all__ = [
     "make_default_robot_observation_processor",
     "AbsoluteActionsProcessorStep",
     "RelativeActionsProcessorStep",
-    "to_absolute_actions",
-    "to_absolute_se3_pose",
-    "to_relative_actions",
-    "to_relative_se3_pose",
     "MapDeltaActionToRobotActionStep",
     "MapTensorToDeltaActionDictStep",
     "NewLineTaskProcessorStep",
@@ -174,6 +168,8 @@ __all__ = [
     "transition_to_batch",
     "TransitionKey",
     "TruncatedProcessorStep",
+    "to_absolute_actions",
+    "to_relative_actions",
     "UnnormalizerProcessorStep",
     "VanillaObservationProcessorStep",
 ]

src/lerobot/processor/relative_action_processor.py

@@ -34,206 +34,57 @@ __all__ = [
     "AbsoluteActionsProcessorStep",
     "to_relative_actions",
     "to_absolute_actions",
-    "to_relative_se3_pose",
-    "to_absolute_se3_pose",
 ]
 
 
-def _rotvec_to_quaternion(rotvec: Tensor) -> Tensor:
-    angle = torch.linalg.vector_norm(rotvec, dim=-1, keepdim=True)
-    angle_sq = angle.square()
-    small_scale = 0.5 - angle_sq / 48.0 + angle_sq.square() / 3840.0
-    scale = torch.where(angle > 1e-6, torch.sin(angle / 2.0) / angle.clamp_min(1e-12), small_scale)
-    return torch.cat((torch.cos(angle / 2.0), rotvec * scale), dim=-1)
-
-
-def _quaternion_to_rotvec(quaternion: Tensor) -> Tensor:
-    quaternion = quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
-    quaternion = quaternion * torch.where(quaternion[..., :1] < 0, -1.0, 1.0)
-    vector = quaternion[..., 1:]
-    sin_half_angle = torch.linalg.vector_norm(vector, dim=-1, keepdim=True)
-    angle = 2.0 * torch.atan2(sin_half_angle, quaternion[..., :1].clamp_min(0.0))
-    small_scale = 2.0 + sin_half_angle.square() / 3.0
-    scale = torch.where(
-        sin_half_angle > 1e-6,
-        angle / sin_half_angle.clamp_min(1e-12),
-        small_scale,
-    )
-    return vector * scale
-
-
-def _quaternion_multiply(left: Tensor, right: Tensor) -> Tensor:
-    left_w, left_xyz = left[..., :1], left[..., 1:]
-    right_w, right_xyz = right[..., :1], right[..., 1:]
-    return torch.cat(
-        (
-            left_w * right_w - (left_xyz * right_xyz).sum(dim=-1, keepdim=True),
-            left_w * right_xyz + right_w * left_xyz + torch.linalg.cross(left_xyz, right_xyz, dim=-1),
-        ),
-        dim=-1,
-    )
-
-
-def _quaternion_conjugate(quaternion: Tensor) -> Tensor:
-    return torch.cat((quaternion[..., :1], -quaternion[..., 1:]), dim=-1)
-
-
-def _quaternion_rotate(quaternion: Tensor, vector: Tensor) -> Tensor:
-    quaternion_xyz = quaternion[..., 1:]
-    uv = torch.linalg.cross(quaternion_xyz, vector, dim=-1)
-    uuv = torch.linalg.cross(quaternion_xyz, uv, dim=-1)
-    return vector + 2.0 * (quaternion[..., :1] * uv + uuv)
-
-
-def to_relative_se3_pose(target_pose: Tensor, reference_pose: Tensor) -> Tensor:
-    """Encode a pose as ``inv(T_reference) @ T_target``.
-
-    Poses use ``[x, y, z, rx, ry, rz]`` with an axis-angle rotation vector.
-    The relative translation is therefore expressed in the reference EE frame.
-    """
-    if target_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
-        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
-    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
-    target_quaternion = _rotvec_to_quaternion(target_pose[..., 3:])
-    inverse_reference_quaternion = _quaternion_conjugate(reference_quaternion)
-    relative_translation = _quaternion_rotate(
-        inverse_reference_quaternion, target_pose[..., :3] - reference_pose[..., :3]
-    )
-    relative_quaternion = _quaternion_multiply(inverse_reference_quaternion, target_quaternion)
-    return torch.cat((relative_translation, _quaternion_to_rotvec(relative_quaternion)), dim=-1)
-
-
-def to_absolute_se3_pose(relative_pose: Tensor, reference_pose: Tensor) -> Tensor:
-    """Decode a pose with ``T_target = T_reference @ T_relative``."""
-    if relative_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
-        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
-    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
-    relative_quaternion = _rotvec_to_quaternion(relative_pose[..., 3:])
-    target_translation = reference_pose[..., :3] + _quaternion_rotate(
-        reference_quaternion, relative_pose[..., :3]
-    )
-    target_quaternion = _quaternion_multiply(reference_quaternion, relative_quaternion)
-    return torch.cat((target_translation, _quaternion_to_rotvec(target_quaternion)), dim=-1)
-
-
-def _broadcast_reference(actions: Tensor, state: Tensor) -> Tensor:
-    if state.device != actions.device or state.dtype != actions.dtype:
-        state = state.to(device=actions.device, dtype=actions.dtype)
-    if actions.ndim == state.ndim + 1:
-        state = state.unsqueeze(-2)
-    return state
-
-
-def _validate_se3_pose_groups(
-    pose_representation: str,
-    se3_pose_groups: Sequence[Sequence[int]] | None,
-    mask: Sequence[bool],
-    action_dim: int,
-) -> list[list[int]]:
-    if pose_representation not in {"componentwise", "se3"}:
-        raise ValueError(
-            f"Unsupported pose_representation={pose_representation!r}; expected 'componentwise' or 'se3'"
-        )
-    if pose_representation == "componentwise":
-        return []
-    if not se3_pose_groups:
-        raise ValueError("pose_representation='se3' requires at least one six-index se3_pose_group")
-
-    normalized_groups: list[list[int]] = []
-    used_indices: set[int] = set()
-    for raw_group in se3_pose_groups:
-        group = [int(index) for index in raw_group]
-        if len(group) != 6:
-            raise ValueError(f"Each SE(3) pose group must contain six indices, got {group}")
-        if len(set(group)) != 6 or any(index < 0 or index >= action_dim for index in group):
-            raise ValueError(f"Invalid SE(3) pose group for action_dim={action_dim}: {group}")
-        if any(index >= len(mask) for index in group):
-            raise ValueError(f"SE(3) pose group lies outside the relative mask: {group}")
-        if used_indices.intersection(group):
-            raise ValueError(f"SE(3) pose groups must not overlap: {group}")
-        group_mask = [bool(mask[index]) for index in group]
-        if any(group_mask) and not all(group_mask):
-            raise ValueError(f"An SE(3) pose group must be wholly relative or wholly absolute: {group}")
-        used_indices.update(group)
-        if all(group_mask):
-            normalized_groups.append(group)
-    return normalized_groups
-
-
-def to_relative_actions(
-    actions: Tensor,
-    state: Tensor,
-    mask: Sequence[bool],
-    *,
-    pose_representation: str = "componentwise",
-    se3_pose_groups: Sequence[Sequence[int]] | None = None,
-) -> Tensor:
-    """Convert absolute actions to a configured relative representation.
-
-    Component-wise mode computes ``action - state``. SE(3) mode computes
-    ``inv(T_state) @ T_action`` for each configured pose group.
+def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
+    """Convert absolute actions to relative: relative = action - state (for masked dims).
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
-    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, actions.shape[-1])
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    state = _broadcast_reference(actions, state)
-    component_mask = mask_t.clone()
-    for group in groups:
-        component_mask[group] = 0
-    state_offset = state[..., :dims] * component_mask
+    if state.device != actions.device or state.dtype != actions.dtype:
+        state = state.to(device=actions.device, dtype=actions.dtype)
+    state_offset = state[..., :dims] * mask_t
+    if actions.ndim == 3:
+        state_offset = state_offset.unsqueeze(-2)
     actions = actions.clone()
     actions[..., :dims] -= state_offset
-    for group in groups:
-        actions[..., group] = to_relative_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
-def to_absolute_actions(
-    actions: Tensor,
-    state: Tensor,
-    mask: Sequence[bool],
-    *,
-    pose_representation: str = "componentwise",
-    se3_pose_groups: Sequence[Sequence[int]] | None = None,
-) -> Tensor:
-    """Convert relative actions back to absolute actions.
-
-    Component-wise mode computes ``relative + state``. SE(3) mode computes
-    ``T_state @ T_relative`` for each configured pose group.
+def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
+    """Convert relative actions back to absolute: absolute = relative + state (for masked dims).
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
-    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, actions.shape[-1])
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    state = _broadcast_reference(actions, state)
-    component_mask = mask_t.clone()
-    for group in groups:
-        component_mask[group] = 0
-    state_offset = state[..., :dims] * component_mask
+    if state.device != actions.device or state.dtype != actions.dtype:
+        state = state.to(device=actions.device, dtype=actions.dtype)
+    state_offset = state[..., :dims] * mask_t
+    if actions.ndim == 3:
+        state_offset = state_offset.unsqueeze(-2)
     actions = actions.clone()
     actions[..., :dims] += state_offset
-    for group in groups:
-        actions[..., group] = to_absolute_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
 @ProcessorStepRegistry.register("relative_actions_processor")
 @dataclass
 class RelativeActionsProcessorStep(ProcessorStep):
-    """Converts absolute actions to the configured relative representation.
+    """Converts absolute actions to relative actions (action -= state) for masked dimensions.
 
     Mirrors OpenPI's DeltaActions transform. Applied during preprocessing so the model
     trains on relative offsets instead of absolute positions.
@@ -250,8 +101,6 @@ class RelativeActionsProcessorStep(ProcessorStep):
     enabled: bool = False
     exclude_joints: list[str] = field(default_factory=list)
     action_names: list[str] | None = None
-    pose_representation: str = "componentwise"
-    se3_pose_groups: list[list[int]] = field(default_factory=list)
     _last_state: torch.Tensor | None = field(default=None, init=False, repr=False)
 
     def _build_mask(self, action_dim: int) -> list[bool]:
@@ -277,30 +126,20 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
-        # State history has shape (B, H, D). Relative actions are referenced to
-        # the newest proprioceptive state, not the whole history tensor.
-        reference_state = state[:, -1] if state is not None and state.ndim == 3 else state
-
         # Always cache state for the paired AbsoluteActionsProcessorStep
-        if reference_state is not None:
-            self._last_state = reference_state
+        if state is not None:
+            self._last_state = state
 
         if not self.enabled:
             return transition
 
         new_transition = transition.copy()
         action = new_transition.get(TransitionKey.ACTION)
-        if action is None or reference_state is None:
+        if action is None or state is None:
             return new_transition
 
         mask = self._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_relative_actions(
-            action,
-            reference_state,
-            mask,
-            pose_representation=self.pose_representation,
-            se3_pose_groups=self.se3_pose_groups,
-        )
+        new_transition[TransitionKey.ACTION] = to_relative_actions(action, state, mask)
         return new_transition
 
     def get_cached_state(self) -> torch.Tensor | None:
@@ -312,8 +151,6 @@ class RelativeActionsProcessorStep(ProcessorStep):
             "enabled": self.enabled,
             "exclude_joints": self.exclude_joints,
             "action_names": self.action_names,
-            "pose_representation": self.pose_representation,
-            "se3_pose_groups": self.se3_pose_groups,
         }
 
     def transform_features(
@@ -362,13 +199,7 @@ class AbsoluteActionsProcessorStep(ProcessorStep):
             return new_transition
 
         mask = self.relative_step._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_absolute_actions(
-            action,
-            cached_state,
-            mask,
-            pose_representation=self.relative_step.pose_representation,
-            se3_pose_groups=self.relative_step.se3_pose_groups,
-        )
+        new_transition[TransitionKey.ACTION] = to_absolute_actions(action, cached_state, mask)
         return new_transition
 
     def get_config(self) -> dict[str, Any]:

src/lerobot/rewards/pretrained.py

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

src/lerobot/scripts/lerobot_annotate.py

@@ -28,7 +28,12 @@ For distributed runs, see ``examples/annotations/run_hf_job.py``.
 """
 
 import logging
+from contextlib import suppress
 from pathlib import Path
+from typing import TYPE_CHECKING
+
+from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub.errors import RevisionNotFoundError
 
 from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
 from lerobot.annotations.steerable_pipeline.executor import Executor
@@ -42,6 +47,12 @@ from lerobot.annotations.steerable_pipeline.validator import StagingValidator
 from lerobot.annotations.steerable_pipeline.vlm_client import make_vlm_client
 from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
 from lerobot.configs import parser
+from lerobot.utils.import_utils import _datasets_available, require_package
+
+if TYPE_CHECKING or _datasets_available:
+    from lerobot.datasets.dataset_metadata import CODEBASE_VERSION
+    from lerobot.datasets.io_utils import load_info
+    from lerobot.datasets.utils import create_lerobot_dataset_card
 
 logger = logging.getLogger(__name__)
 
@@ -50,8 +61,6 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
     if cfg.root is not None:
         return Path(cfg.root)
     if cfg.repo_id is not None:
-        from huggingface_hub import snapshot_download
-
         return Path(snapshot_download(repo_id=cfg.repo_id, repo_type="dataset"))
     raise ValueError("Either --root or --repo_id must be provided.")
 
@@ -125,7 +134,7 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
 
     Pushes to ``cfg.new_repo_id`` when set, otherwise back to ``cfg.repo_id``.
     """
-    from huggingface_hub import HfApi  # noqa: PLC0415
+    require_package("datasets", "dataset")
 
     repo_id = cfg.new_repo_id or cfg.repo_id
     commit_message = cfg.push_commit_message or "Add steerable annotations (lerobot-annotate)"
@@ -143,33 +152,26 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         repo_id=repo_id,
         repo_type="dataset",
         commit_message=commit_message,
-        ignore_patterns=[".annotate_staging/**", "**/.DS_Store"],
+        # README.md is excluded because when pushing to ``new_repo_id`` the
+        # source card's links (e.g. the visualize badge) would keep pointing
+        # at the source dataset; a fresh card is generated below instead.
+        ignore_patterns=[".annotate_staging/**", "**/.DS_Store", "README.md"],
     )
     print(f"[lerobot-annotate] uploaded to https://huggingface.co/datasets/{repo_id}", flush=True)
 
+    dataset_info = load_info(root)
+    card = create_lerobot_dataset_card(dataset_info=dataset_info, license="apache-2.0", repo_id=repo_id)
+    card.push_to_hub(repo_id=repo_id, repo_type="dataset")
+
     # Tag the upload with the codebase version. ``LeRobotDatasetMetadata``
     # resolves the dataset revision via ``get_safe_version`` which scans
     # for tags like ``v3.0``; without a tag it raises
     # ``RevisionNotFoundError``. Read the version straight from the
     # dataset's own ``meta/info.json`` so we tag whatever the writer
     # actually wrote (no accidental drift if the codebase floor moves).
-    from lerobot.datasets.dataset_metadata import CODEBASE_VERSION  # noqa: PLC0415
-
-    info_path = root / "meta" / "info.json"
-    version_tag = CODEBASE_VERSION
-    if info_path.exists():
-        try:
-            from lerobot.utils.io_utils import load_json  # noqa: PLC0415
-
-            info = load_json(info_path)
-            ds_version = info.get("codebase_version")
-            if isinstance(ds_version, str) and ds_version.startswith("v"):
-                version_tag = ds_version
-        except Exception as exc:  # noqa: BLE001
-            print(
-                f"[lerobot-annotate] could not read codebase_version from info.json ({exc}); falling back to {version_tag}",
-                flush=True,
-            )
+    version_tag = (
+        dataset_info.codebase_version if dataset_info.codebase_version.startswith("v") else CODEBASE_VERSION
+    )
     revision = getattr(commit_info, "oid", None)
     tag_kwargs = {
         "repo_id": repo_id,
@@ -180,10 +182,6 @@ def _push_to_hub(root: Path, cfg: AnnotationPipelineConfig) -> None:
         tag_kwargs["revision"] = revision
 
     try:
-        from contextlib import suppress  # noqa: PLC0415
-
-        from huggingface_hub.errors import RevisionNotFoundError  # noqa: PLC0415
-
         with suppress(RevisionNotFoundError):
             api.delete_tag(repo_id, tag=version_tag, repo_type="dataset")
         api.create_tag(**tag_kwargs)

src/lerobot/scripts/lerobot_edit_dataset.py

@@ -325,8 +325,6 @@ class RecomputeStatsConfig(OperationConfig):
     relative_exclude_joints: list[str] | None = None
     chunk_size: int = 50
     num_workers: int = 0
-    relative_pose_representation: str = "componentwise"
-    relative_se3_pose_groups: list[list[int]] | None = None
     overwrite: bool = False
 
 
@@ -700,8 +698,6 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
         relative_exclude_joints=cfg.operation.relative_exclude_joints,
         chunk_size=cfg.operation.chunk_size,
         num_workers=cfg.operation.num_workers,
-        relative_pose_representation=cfg.operation.relative_pose_representation,
-        relative_se3_pose_groups=cfg.operation.relative_se3_pose_groups,
     )
 
     logging.info(f"Stats written to {dataset.root}")

src/lerobot/scripts/lerobot_train.py

@@ -171,6 +171,9 @@ def update_policy(
     train_metrics.update_s = time.perf_counter() - start_time
     if torch.cuda.is_available():
         train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
+    # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
+    if output_dict:
+        train_metrics.update_metrics(output_dict)
     return train_metrics, output_dict
 
 
@@ -572,7 +575,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        train_tracker, output_dict = update_policy(
+        train_tracker, _ = update_policy(
             train_tracker,
             policy,
             batch,
@@ -605,9 +608,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     train_tracker.samples_per_s = effective_batch_size / step_time
                 logging.info(train_tracker)
                 if wandb_logger:
+                    # Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
+                    # tracker by update_policy, so to_dict() already carries their windowed,
+                    # rank-reduced averages — no per-step output_dict passthrough needed.
                     wandb_log_dict = train_tracker.to_dict()
-                    if output_dict:
-                        wandb_log_dict.update(output_dict)
                     # Log sample weighting statistics if enabled
                     if sample_weighter is not None:
                         weighter_stats = sample_weighter.get_stats()

src/lerobot/utils/device_utils.py

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

src/lerobot/utils/logging_utils.py

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

tests/annotations/test_modules.py

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

tests/processor/test_pi05_state_from_action.py

@@ -1,227 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from math import pi
-
-import numpy as np
-import pytest
-import torch
-
-pytest.importorskip("transformers")
-
-from lerobot.datasets.compute_stats import (  # noqa: E402
-    compute_relative_action_stats,
-    compute_state_history_stats,
-)
-from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
-from lerobot.policies.pi05.processor_pi05 import (  # noqa: E402
-    Pi05FlattenStateHistoryProcessorStep,
-    Pi05StateFromActionProcessorStep,
-)
-from lerobot.processor.relative_action_processor import (  # noqa: E402
-    AbsoluteActionsProcessorStep,
-    RelativeActionsProcessorStep,
-)
-from lerobot.types import TransitionKey  # noqa: E402
-from lerobot.utils.constants import OBS_STATE  # noqa: E402
-
-
-def _transition(action: torch.Tensor | None, state: torch.Tensor | None = None) -> dict:
-    observation = {} if state is None else {OBS_STATE: state}
-    return {
-        TransitionKey.OBSERVATION: observation,
-        TransitionKey.ACTION: action,
-        TransitionKey.REWARD: None,
-        TransitionKey.DONE: None,
-        TransitionKey.TRUNCATED: None,
-        TransitionKey.COMPLEMENTARY_DATA: {},
-    }
-
-
-def test_pi05_config_requests_action_history_prefix():
-    config = PI05Config(
-        device="cpu",
-        chunk_size=4,
-        n_action_steps=4,
-        state_from_action=True,
-        proprioception_history_steps=2,
-    )
-
-    assert config.action_delta_indices == [-1, 0, 1, 2, 3]
-
-
-def test_state_from_action_extracts_history_and_preserves_target_horizon():
-    action = torch.arange(2 * 5 * 3, dtype=torch.float32).reshape(2, 5, 3)
-    step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
-
-    result = step(_transition(action))
-
-    torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], action[:, :2])
-    torch.testing.assert_close(result[TransitionKey.ACTION], action[:, 1:])
-
-
-def test_relative_actions_use_newest_state_in_history_and_roundtrip():
-    state_history = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
-    absolute = torch.tensor([[[3.0, 30.0], [4.0, 40.0]]])
-    relative_step = RelativeActionsProcessorStep(enabled=True)
-    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
-
-    relative = relative_step(_transition(absolute, state_history))
-    expected = torch.tensor([[[1.0, 10.0], [2.0, 20.0]]])
-    torch.testing.assert_close(relative[TransitionKey.ACTION], expected)
-
-    recovered = absolute_step(_transition(relative[TransitionKey.ACTION]))
-    torch.testing.assert_close(recovered[TransitionKey.ACTION], absolute)
-
-
-def test_flatten_state_history_preserves_chronological_order():
-    state_history = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
-    step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=4)
-
-    result = step(_transition(torch.zeros(1, 2, 2), state_history))
-
-    torch.testing.assert_close(
-        result[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[1.0, 2.0, 3.0, 4.0]])
-    )
-
-
-def test_state_history_can_be_relative_with_absolute_gripper():
-    state_history = torch.tensor([[[1.0, 10.0, 0.2], [3.0, 20.0, 0.4]]])
-    step = Pi05FlattenStateHistoryProcessorStep(
-        history_steps=2,
-        max_state_dim=6,
-        relative=True,
-        exclude_joints=["gripper"],
-        state_names=["x", "y", "gripper_width"],
-    )
-
-    result = step(_transition(torch.zeros(1, 2, 3), state_history))
-
-    torch.testing.assert_close(
-        result[TransitionKey.OBSERVATION][OBS_STATE],
-        torch.tensor([[-2.0, -10.0, 0.2, 0.0, 0.0, 0.4]]),
-    )
-
-
-def test_state_history_can_use_se3_composition_with_absolute_gripper():
-    state_history = torch.tensor(
-        [[[0.0, 1.0, 0.0, 0.0, 0.0, pi / 2, 0.2], [0.0, 0.0, 0.0, 0.0, 0.0, pi / 2, 0.4]]]
-    )
-    step = Pi05FlattenStateHistoryProcessorStep(
-        history_steps=2,
-        max_state_dim=14,
-        relative=True,
-        exclude_joints=["gripper"],
-        state_names=["x", "y", "z", "rx", "ry", "rz", "gripper_width"],
-        pose_representation="se3",
-        se3_pose_groups=[list(range(6))],
-    )
-
-    result = step(_transition(torch.zeros(1, 2, 7), state_history))
-
-    expected = torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.2, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.4]])
-    torch.testing.assert_close(result[TransitionKey.OBSERVATION][OBS_STATE], expected, atol=1e-6, rtol=1e-6)
-
-
-def test_inference_state_history_is_rolled_and_reset():
-    step = Pi05StateFromActionProcessorStep(enabled=True, history_steps=2)
-
-    first = step(_transition(None, torch.tensor([[1.0, 2.0]])))
-    second = step(_transition(None, torch.tensor([[3.0, 4.0]])))
-    torch.testing.assert_close(
-        first[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [1.0, 2.0]]])
-    )
-    torch.testing.assert_close(
-        second[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[1.0, 2.0], [3.0, 4.0]]])
-    )
-
-    step.reset()
-    reset = step(_transition(None, torch.tensor([[5.0, 6.0]])))
-    torch.testing.assert_close(
-        reset[TransitionKey.OBSERVATION][OBS_STATE], torch.tensor([[[5.0, 6.0], [5.0, 6.0]]])
-    )
-
-
-def test_flatten_state_history_checks_max_state_dim():
-    step = Pi05FlattenStateHistoryProcessorStep(history_steps=2, max_state_dim=3)
-
-    with pytest.raises(ValueError, match="above max_state_dim"):
-        step(_transition(torch.zeros(1, 2, 2), torch.zeros(1, 2, 2)))
-
-
-def test_relative_stats_can_use_absolute_action_as_state():
-    actions = np.asarray([[0.0, 0.0], [1.0, 2.0], [2.0, 4.0], [3.0, 6.0]], dtype=np.float32)
-    dataset = {"action": actions, "episode_index": np.zeros(4, dtype=np.int64)}
-    features = {"action": {"shape": [2], "names": ["x", "y"]}}
-
-    stats = compute_relative_action_stats(
-        dataset,
-        features,
-        chunk_size=2,
-        state_from_action=True,
-    )
-
-    np.testing.assert_allclose(stats["mean"], [0.5, 1.0])
-
-
-def test_relative_state_history_stats_match_processor_representation():
-    actions = np.asarray(
-        [[0.0, 0.1], [1.0, 0.2], [3.0, 0.3]],
-        dtype=np.float32,
-    )
-    dataset = {"action": actions, "episode_index": np.zeros(3, dtype=np.int64)}
-    features = {"action": {"shape": [2], "names": ["x", "gripper_width"]}}
-
-    stats = compute_state_history_stats(
-        dataset,
-        features,
-        history_steps=2,
-        exclude_joints=["gripper"],
-        relative=True,
-    )
-
-    expected = np.asarray([[0.0, 0.1, 0.0, 0.1], [-1.0, 0.1, 0.0, 0.2], [-2.0, 0.2, 0.0, 0.3]])
-    np.testing.assert_allclose(stats["mean"], expected.mean(axis=0))
-
-
-def test_se3_relative_action_stats_use_reference_frame():
-    actions = np.asarray(
-        [
-            [0.0, 0.0, 0.0, 0.0, 0.0, pi / 2, 0.2],
-            [0.0, 1.0, 0.0, 0.0, 0.0, pi / 2, 0.3],
-        ],
-        dtype=np.float32,
-    )
-    dataset = {"action": actions, "episode_index": np.zeros(2, dtype=np.int64)}
-    features = {
-        "action": {
-            "shape": [7],
-            "names": ["x", "y", "z", "rx", "ry", "rz", "gripper_width"],
-        }
-    }
-
-    stats = compute_relative_action_stats(
-        dataset,
-        features,
-        chunk_size=2,
-        exclude_joints=["gripper"],
-        state_from_action=True,
-        pose_representation="se3",
-        se3_pose_groups=[list(range(6))],
-    )
-
-    np.testing.assert_allclose(stats["mean"][:3], [0.5, 0.0, 0.0], atol=1e-6)
-    np.testing.assert_allclose(stats["mean"][6], 0.25, atol=1e-6)

tests/processor/test_se3_relative_actions.py

@@ -1,70 +0,0 @@
-import math
-
-import pytest
-import torch
-
-from lerobot.processor.relative_action_processor import (
-    to_absolute_actions,
-    to_absolute_se3_pose,
-    to_relative_actions,
-    to_relative_se3_pose,
-)
-
-POSE_GROUP = [list(range(6))]
-
-
-def test_se3_translation_is_expressed_in_reference_frame():
-    reference = torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, math.pi / 2]])
-    target = torch.tensor([[1.0, 3.0, 3.0, 0.0, 0.0, math.pi / 2]])
-
-    relative = to_relative_se3_pose(target, reference)
-
-    torch.testing.assert_close(relative, torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]]), atol=1e-6, rtol=1e-6)
-
-
-def test_se3_pose_roundtrip_for_batched_chunks():
-    torch.manual_seed(0)
-    reference = torch.randn(4, 6)
-    reference[:, 3:] *= 0.8
-    target = torch.randn(4, 11, 6)
-    target[..., 3:] *= 0.8
-
-    relative = to_relative_se3_pose(target, reference.unsqueeze(1))
-    recovered = to_absolute_se3_pose(relative, reference.unsqueeze(1))
-
-    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
-
-
-def test_mixed_se3_pose_and_absolute_gripper_roundtrip():
-    reference = torch.tensor([[0.2, -0.1, 0.4, 0.1, 0.2, -0.3, 0.06]])
-    target = torch.tensor([[[0.3, 0.2, 0.5, -0.2, 0.1, 0.4, 0.03], [0.1, -0.3, 0.2, 0.5, -0.1, 0.2, 0.05]]])
-    mask = [True, True, True, True, True, True, False]
-
-    relative = to_relative_actions(
-        target,
-        reference,
-        mask,
-        pose_representation="se3",
-        se3_pose_groups=POSE_GROUP,
-    )
-    recovered = to_absolute_actions(
-        relative,
-        reference,
-        mask,
-        pose_representation="se3",
-        se3_pose_groups=POSE_GROUP,
-    )
-
-    torch.testing.assert_close(relative[..., 6], target[..., 6])
-    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
-
-
-def test_se3_pose_group_cannot_be_partially_relative():
-    with pytest.raises(ValueError, match="wholly relative or wholly absolute"):
-        to_relative_actions(
-            torch.zeros(1, 7),
-            torch.zeros(1, 7),
-            [True, True, True, False, False, False, False],
-            pose_representation="se3",
-            se3_pose_groups=POSE_GROUP,
-        )

tests/scripts/test_lerobot_annotate.py

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

tests/utils/test_logging_utils.py

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