chore(config): add libs note

Per @jashshah999, mention that spawn workers re-import modules and so
add some startup time vs fork. Also trim the failure-mode dump from
the inline comment — the linked issue covers the symptoms in detail.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

README.md

@@ -83,7 +83,7 @@ episode_index=0
 print(f"{dataset[episode_index]['action'].shape=}\n")
 ```
 
-Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3)
+Learn more about it in the [LeRobotDataset Documentation](https://huggingface.co/docs/lerobot/lerobot-dataset-v3).
 
 ## SoTA Models
 
@@ -109,7 +109,7 @@ lerobot-train \
 | **World Models**           | [VLA-JEPA](./docs/source/vla_jepa.mdx), [LingBot-VA](./docs/source/lingbot_va.mdx), [FastWAM](./docs/source/fastwam.mdx)                                                                                                                                                                                                                                                                   |
 | **Reward Models**          | [SARM](./docs/source/sarm.mdx), [TOPReward](./docs/source/topreward.mdx), [Robometer](./docs/source/robometer.mdx)                                                                                                                                                                                                                                                                         |
 
-Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub
+Similarly to the hardware, you can easily implement your own policy & leverage LeRobot's data collection, training, and visualization tools, and share your model to the HF Hub.
 
 For detailed policy setup guides, see the [Policy Documentation](https://huggingface.co/docs/lerobot/bring_your_own_policies). For GPU/RAM requirements and expected training time per policy, see the [Compute Hardware Guide](https://huggingface.co/docs/lerobot/hardware_guide).
 
@@ -126,7 +126,7 @@ lerobot-eval \
   --eval.n_episodes=10
 ```
 
-Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub)
+Learn how to implement your own simulation environment or benchmark and distribute it from the HF Hub by following the [EnvHub Documentation](https://huggingface.co/docs/lerobot/envhub).
 
 ## Resources
 

docs/source/annotation_pipeline.mdx

@@ -89,8 +89,8 @@ 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)
-for the production settings (single camera, timestamped contact sheets,
+[Running on Hugging Face Jobs](#running-on-hugging-face-jobs) for the
+production settings (single camera, timestamped contact sheets,
 auto-windowed subtask generation).
 
 ### Tools
@@ -110,28 +110,67 @@ not-yet-implemented.
 
 ## Running on Hugging Face Jobs
 
-Annotation runs on [Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs).
-The repo ships a launcher script you copy and tweak for your dataset:
+Annotating a real dataset needs a GPU big enough to serve the VLM, so
+`lerobot-annotate` can dispatch itself to
+[Hugging Face Jobs](https://huggingface.co/docs/hub/en/jobs) — same as
+`lerobot-train`. Add `--job.target=<flavor>` to the exact command you'd
+run locally and it runs on that hardware instead:
 
 ```bash
-HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
+hf auth login   # once
+
+uv run lerobot-annotate \
+    --repo_id=user/my_dataset \
+    --new_repo_id=user/my_dataset_annotated \
+    --push_to_hub=true \
+    --vlm.model_id=Qwen/Qwen3.6-27B \
+    --vlm.num_gpus=1 \
+    --vlm.serve_command="vllm serve Qwen/Qwen3.6-27B --tensor-parallel-size 1 \
+        --max-model-len 32768 --gpu-memory-utilization 0.8 \
+        --uvicorn-log-level warning --port {port}" \
+    --vlm.serve_ready_timeout_s=1800 \
+    --vlm.chat_template_kwargs='{"enable_thinking": false}' \
+    --job.target=h200
 ```
 
-[`run_hf_job.py`](https://github.com/huggingface/lerobot/blob/main/examples/annotations/run_hf_job.py)
-starts a single-GPU `h200` job (bump it to `h200x4` for big datasets)
-that:
+That submits a single-GPU `h200` job that:
 
-1. installs `lerobot` (from `main`) plus the annotation extras,
-2. boots one vLLM server per GPU (using the `vllm/vllm-openai` image) and
-   drives it over the OpenAI-compatible API,
-3. runs the `plan` / `interjections` / `vqa` modules across the dataset
-   with `lerobot-annotate`,
+1. starts from the `vllm/vllm-openai` image and installs `lerobot` on top,
+2. boots one vLLM server per GPU and drives it over the OpenAI-compatible API,
+3. runs the `plan` / `interjections` / `vqa` modules across the dataset,
 4. with `--push_to_hub=true`, uploads the result to `--new_repo_id` (or
    back to `--repo_id` in place if you leave that unset).
 
-To use a different dataset, model, or hub repo, edit the `CMD` block in
-the script. Every flag there maps directly to a `lerobot-annotate` flag
-(run `lerobot-annotate --help` for the full list).
+The command streams the job's logs; `Ctrl-C` detaches without cancelling
+it. List the available flavors and their pricing with `hf jobs hardware`.
+
+<Tip warning={true}>
+
+Qwen3.6 ships with thinking enabled, which eats the token budget the
+annotator needs for its JSON answer — `--vlm.chat_template_kwargs='{"enable_thinking": false}'`
+turns it off. Without `--push_to_hub=true` the annotated dataset is
+discarded when the pod exits.
+
+</Tip>
+
+### Job options
+
+| Flag                | Default                   | What it does                                                                    |
+| ------------------- | ------------------------- | ------------------------------------------------------------------------------- |
+| `--job.target`      | `local`                   | HF Jobs flavor to run on (e.g. `h200`, `h200x4`). Omitted/`local` runs here.    |
+| `--job.image`       | `vllm/vllm-openai:latest` | Runtime image for the pod.                                                      |
+| `--job.timeout`     | `2h`                      | Wall-clock cap. Raise it for large datasets.                                    |
+| `--job.detach`      | `false`                   | Submit and exit instead of streaming logs.                                      |
+| `--job.lerobot_ref` | `main`                    | Git ref of lerobot installed on the pod — point it at a branch to test changes. |
+| `--job.tags`        | `[]`                      | Extra tags on the job and on any dataset it pushes (`lerobot` is always added). |
+
+For a bigger dataset, scale to `h200x4` and raise
+`--vlm.parallel_servers` / `--vlm.num_gpus` to match, and give the job
+more headroom with e.g. `--job.timeout=8h`.
+
+Remote runs need `--repo_id` (the pod pulls the dataset from the Hub;
+`--root` names a directory only your machine has). A dataset that exists
+only in your local cache is pushed to a **private** repo first.
 
 ## Key options
 

docs/source/omx.mdx

@@ -1,3 +1,11 @@
+# OMX
+
+<img
+  src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/lerobot/omx_mainimage.png"
+  alt="OMX"
+  width=600
+/>
+
 ## Order and Assemble the parts
 
 First, assemble the OMX hardware following the official assembly guide.

examples/annotations/run_hf_job.py

@@ -1,80 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-"""Launch ``lerobot-annotate`` on a Hugging Face job (vllm + Qwen3.6-27B VLM).
-
-Spawns one single-GPU ``h200`` job that:
-
-  1. installs ``lerobot`` from ``main`` plus the annotation extras,
-  2. boots one vllm server with Qwen3.6-27B (dense VLM),
-  3. runs the plan / interjections / vqa modules across the dataset
-     in free-form mode (each episode generates its own subtasks +
-     memory),
-  4. uploads the annotated dataset to ``--new_repo_id`` (when set)
-     or back to ``--repo_id``.
-
-Usage:
-
-    HF_TOKEN=hf_... uv run python examples/annotations/run_hf_job.py
-
-Adjust ``CMD`` (dataset, model, hub repo) and ``flavor`` below for your
-run. For larger datasets, scale to ``h200x4`` and raise
-``--vlm.parallel_servers`` / ``--vlm.num_gpus`` to match.
-"""
-
-import os
-
-from huggingface_hub import get_token, run_job
-
-token = os.environ.get("HF_TOKEN") or get_token()
-if not token:
-    raise RuntimeError("No HF token. Run `huggingface-cli login` or `export HF_TOKEN=hf_...`")
-
-CMD = (
-    "apt-get update -qq && apt-get install -y -qq git ffmpeg && "
-    "pip install --no-deps "
-    "'lerobot @ git+https://github.com/huggingface/lerobot.git@main' && "
-    # Pins mirror pyproject.toml — unpinned installs pull av 18 / datasets 5 /
-    # draccus 0.11, which break lerobot at import time.
-    "pip install --upgrade-strategy only-if-needed "
-    "'datasets>=4.7.0,<5.0.0' 'pyarrow>=21.0.0,<30.0.0' 'av>=15.0.0,<16.0.0' 'draccus==0.10.0' "
-    "'pandas>=2.0.0,<3.0.0' jsonlines gymnasium torchcodec mergedeep pyyaml-include toml typing-inspect "
-    "openai && "
-    "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
-    "export VLLM_VIDEO_BACKEND=pyav && "
-    "lerobot-annotate "
-    "--repo_id=pepijn223/robocasa_pretrain_human300_v4 "
-    "--new_repo_id=pepijn223/robocasa_pretrain_human300_v4_annotated "
-    "--push_to_hub=true "
-    "--vlm.backend=openai "
-    "--vlm.model_id=Qwen/Qwen3.6-27B "
-    "--vlm.num_gpus=1 "
-    '--vlm.serve_command="vllm serve Qwen/Qwen3.6-27B '
-    "--tensor-parallel-size 1 --max-model-len 32768 "
-    '--gpu-memory-utilization 0.8 --uvicorn-log-level warning --port {port}" '
-    "--vlm.serve_ready_timeout_s=1800 "
-    # Qwen3.6 ships with thinking on; annotation wants plain JSON answers.
-    "--vlm.chat_template_kwargs='{\"enable_thinking\": false}'"
-)
-
-job = run_job(
-    image="vllm/vllm-openai:latest",
-    command=["bash", "-c", CMD],
-    flavor="h200",
-    secrets={"HF_TOKEN": token},
-    timeout="2h",
-)
-print(f"Job URL: {job.url}")
-print(f"Job ID:  {job.id}")

src/lerobot/annotations/steerable_pipeline/config.py

@@ -20,6 +20,29 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any
 
+from lerobot.configs.default import JobConfig
+
+# The annotation pipeline boots its own vLLM server, so the pod starts from the
+# official vLLM runtime rather than the prebuilt `lerobot-gpu` training image;
+# `lerobot` is pip-installed on top (see `lerobot.jobs.annotate`).
+DEFAULT_ANNOTATE_JOB_IMAGE = "vllm/vllm-openai:latest"
+
+
+@dataclass
+class AnnotationJobConfig(JobConfig):
+    """`JobConfig` with the annotation runtime's defaults.
+
+    Adds `lerobot_ref` because the vLLM image ships no lerobot: the pod installs
+    it from git, and the ref decides which code actually annotates. Point it at a
+    branch/tag/SHA to try unmerged changes remotely.
+    """
+
+    image: str = DEFAULT_ANNOTATE_JOB_IMAGE
+    # Annotation is a bounded pass over a dataset; a tighter cap than training's
+    # "2d" keeps a wedged vLLM server from burning a day of GPU time.
+    timeout: str | None = "2h"
+    lerobot_ref: str = "main"
+
 
 @dataclass
 class PlanConfig:
@@ -207,6 +230,11 @@ class AnnotationPipelineConfig:
     vlm: VlmConfig = field(default_factory=VlmConfig)
     executor: ExecutorConfig = field(default_factory=ExecutorConfig)
 
+    # Where the annotation runs: omitted / "local" annotates on this machine, any
+    # other value is an HF Jobs flavor (e.g. "h200") and submits the run there.
+    # List flavors + pricing with `hf jobs hardware`.
+    job: AnnotationJobConfig = field(default_factory=AnnotationJobConfig)
+
     skip_validation: bool = False
     only_episodes: tuple[int, ...] | None = None
 

src/lerobot/annotations/steerable_pipeline/executor.py

@@ -30,8 +30,8 @@ Phase 3 is why the ``plan`` module must be re-entered after the
 timestamps.
 
 Distributed execution is provided by Hugging Face Jobs (see
-``examples/annotations/run_hf_job.py``); the runner inside the job
-invokes ``lerobot-annotate`` which uses this in-process executor.
+``lerobot.jobs.annotate``, reached via ``--job.target=<flavor>``); the pod
+inside the job invokes ``lerobot-annotate`` which uses this in-process executor.
 Episode-level concurrency is controlled by
 ``ExecutorConfig.episode_parallelism``.
 """

src/lerobot/annotations/steerable_pipeline/vlm_client.py

@@ -194,12 +194,13 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
     """Build the shared VLM client.
 
     Only the ``openai`` backend is supported for now. The shipped workflow
-    is Hugging Face Jobs (``examples/annotations/run_hf_job.py``): it boots
-    a vLLM server inside the ``vllm/vllm-openai`` image and the pipeline
-    talks to it over the OpenAI-compatible API (``--vlm.backend=openai``,
-    optionally auto-spawning the server via ``auto_serve`` /
-    ``serve_command``). The former in-process ``vllm`` / ``transformers``
-    backends were removed to keep the support surface to the HF Jobs path.
+    is Hugging Face Jobs (``lerobot-annotate --job.target=<flavor>``): it
+    boots a vLLM server inside the ``vllm/vllm-openai`` image and the
+    pipeline talks to it over the OpenAI-compatible API
+    (``--vlm.backend=openai``, optionally auto-spawning the server via
+    ``auto_serve`` / ``serve_command``). The former in-process ``vllm`` /
+    ``transformers`` backends were removed to keep the support surface to
+    the HF Jobs path.
 
     For ``stub``, construct :class:`StubVlmClient` directly with a responder
     callable; it is rejected here to make accidental misuse obvious.
@@ -213,8 +214,8 @@ def make_vlm_client(config: VlmConfig) -> VlmClient:
     if config.backend in {"vllm", "transformers"}:
         raise ValueError(
             f"backend={config.backend!r} (in-process local model) is not supported for now — "
-            "only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline via "
-            "examples/annotations/run_hf_job.py, which serves the model with vLLM in the "
+            "only backend='openai' (the Hugging Face Jobs flow) is. Run the pipeline with "
+            "`lerobot-annotate --job.target=<flavor>`, which serves the model with vLLM in the "
             "vllm/vllm-openai image and talks to it over the OpenAI-compatible API."
         )
     raise ValueError(f"Unknown VLM backend: {config.backend!r}")

src/lerobot/configs/train.py

@@ -14,6 +14,7 @@
 import builtins
 import datetime as dt
 import json
+import multiprocessing
 import os
 import tempfile
 from dataclasses import dataclass, field
@@ -101,6 +102,12 @@ class TrainPipelineConfig(HubMixin):
     batch_size: int = 8
     prefetch_factor: int = 4
     persistent_workers: bool = True
+    # DataLoader worker start method. "spawn" is safer than "fork" with
+    # non-fork-safe libs (PyAV / torchcodec / ffmpeg), but adds some
+    # worker-startup time per run since workers re-import modules instead
+    # of inheriting parent state. Override with `--dataloader_multiprocessing_context=fork`
+    # when appropriate, or set it to `null` to use Python's platform default.
+    dataloader_multiprocessing_context: str | None = "spawn"
     steps: int = 100_000
     # Run policy in the simulation environment every N steps to measure reward/success (0 = disabled).
     env_eval_freq: int = 20_000
@@ -212,6 +219,17 @@ class TrainPipelineConfig(HubMixin):
             self.reward_model.pretrained_path = str(policy_dir)
 
     def validate(self) -> None:
+        available_contexts = multiprocessing.get_all_start_methods()
+        if (
+            self.dataloader_multiprocessing_context is not None
+            and self.dataloader_multiprocessing_context not in available_contexts
+        ):
+            raise ValueError(
+                "`dataloader_multiprocessing_context` must be None or one of "
+                f"{available_contexts} on this platform, got "
+                f"{self.dataloader_multiprocessing_context!r}."
+            )
+
         self._resolve_pretrained_from_cli()
 
         if self.policy is None and self.reward_model is None:

src/lerobot/datasets/lerobot_dataset.py

@@ -478,18 +478,19 @@ class LeRobotDataset(torch.utils.data.Dataset):
         """Return the number of frames in the selected episodes."""
         return self.num_frames
 
-    def __getitem__(self, idx) -> dict:
-        """Return a single frame by index, with all transforms applied.
+    def __getitem__(self, idx: int | slice) -> dict | list[dict]:
+        """Return one frame or a slice of frames, with all transforms applied.
 
         Loads the frame from the underlying HF dataset, expands delta-timestamp
         windows, decodes video frames, and applies image transforms. Delegates
-        the core logic to :meth:`DatasetReader.get_item`.
+        the core logic to :class:`DatasetReader`.
 
         Args:
-            idx: Index into the (possibly episode-filtered) dataset.
+            idx: Integer index or slice into the possibly episode-filtered dataset.
 
         Returns:
-            Dict mapping feature names to their tensor values for this frame.
+            A frame dictionary for an integer index, or a list of frame
+            dictionaries for a slice.
 
         Raises:
             RuntimeError: If the dataset is currently being recorded and
@@ -499,6 +500,9 @@ class LeRobotDataset(torch.utils.data.Dataset):
             raise RuntimeError(
                 "Cannot read from a dataset that is being recorded. Call finalize() first, then access items."
             )
+        if isinstance(idx, slice):
+            return [self[item_idx] for item_idx in range(*idx.indices(len(self)))]
+
         reader = self._ensure_reader()
         if reader.hf_dataset is None:
             # One-shot load after finalize()

src/lerobot/envs/configs.py

@@ -322,7 +322,7 @@ class HILSerlRobotEnvConfig(EnvConfig):
 class LiberoEnv(EnvConfig):
     task: str = "libero_10"  # can also choose libero_spatial, libero_object, etc.
     task_ids: list[int] | None = None
-    fps: int = 30
+    fps: int = 20  # Must match robosuite's default control_freq (20 Hz)
     episode_length: int | None = None
     obs_type: str = "pixels_agent_pos"
     render_mode: str = "rgb_array"
@@ -354,6 +354,9 @@ class LiberoEnv(EnvConfig):
     control_mode: str = "relative"  # or "absolute"
 
     def __post_init__(self):
+        if self.fps <= 0:
+            raise ValueError(f"fps must be positive, got {self.fps}")
+
         if self.obs_type == "pixels":
             self.features[LIBERO_KEY_PIXELS_AGENTVIEW] = PolicyFeature(
                 type=FeatureType.VISUAL, shape=(self.observation_height, self.observation_width, 3)
@@ -412,6 +415,7 @@ class LiberoEnv(EnvConfig):
             "render_mode": self.render_mode,
             "observation_height": self.observation_height,
             "observation_width": self.observation_width,
+            "control_freq": self.fps,
         }
         if self.task_ids is not None:
             kwargs["task_ids"] = self.task_ids

src/lerobot/envs/libero.py

@@ -125,10 +125,13 @@ class LiberoEnv(gym.Env):
         n_envs: int = 1,
         camera_name_mapping: dict[str, str] | None = None,
         num_steps_wait: int = 10,
+        control_freq: int = 20,
         control_mode: str = "relative",
         is_libero_plus: bool = False,
     ):
         super().__init__()
+        if control_freq <= 0:
+            raise ValueError(f"control_freq must be positive, got {control_freq}")
         self.task_id = task_id
         self.is_libero_plus = is_libero_plus
         self.obs_type = obs_type
@@ -154,6 +157,7 @@ class LiberoEnv(gym.Env):
             }
         self.camera_name_mapping = camera_name_mapping
         self.num_steps_wait = num_steps_wait
+        self.control_freq = control_freq
         self.episode_index = episode_index
         self.episode_length = episode_length
         # Load once and keep
@@ -260,6 +264,7 @@ class LiberoEnv(gym.Env):
             bddl_file_name=self._task_bddl_file,
             camera_heights=self.observation_height,
             camera_widths=self.observation_width,
+            control_freq=self.control_freq,
         )
         env.reset()
         self._env = env

src/lerobot/jobs/__init__.py

@@ -18,6 +18,7 @@ from lerobot.utils.import_utils import require_package
 # guard the optional dependency here so importing this package fails loudly if it's missing.
 require_package("datasets", extra="dataset")
 
+from .annotate import submit_annotate_to_hf
 from .hf import submit_to_hf
 
-__all__ = ["submit_to_hf"]
+__all__ = ["submit_annotate_to_hf", "submit_to_hf"]

src/lerobot/jobs/annotate.py

@@ -0,0 +1,176 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Run ``lerobot-annotate`` on HF Jobs (HuggingFace GPUs).
+
+Same shape as the training submitter in ``hf.py``, with one difference: the
+annotation pipeline serves its own VLM, so the pod starts from the official
+``vllm/vllm-openai`` image (which has no lerobot) instead of the prebuilt
+``lerobot-gpu`` image, and installs lerobot on top before running.
+
+Because there is no config repo to stage, the pod replays the user's own CLI
+flags — everything except the client-only ``--job.*`` and the host-local
+``--root``, which is replaced by ``--repo_id`` so the pod pulls the dataset
+from the Hub.
+"""
+
+from __future__ import annotations
+
+import shlex
+import sys
+from dataclasses import is_dataclass
+from typing import TYPE_CHECKING
+
+from huggingface_hub import HfApi, get_token, run_job
+
+from .dataset import ensure_dataset_available
+
+# Package-internal reuse of the training submitter's job plumbing: following a
+# submitted job and forwarding argv are identical for annotation runs.
+from .hf import _pod_forwarded_args, follow_job, resolve_job_tags
+
+if TYPE_CHECKING:
+    from lerobot.annotations.steerable_pipeline.config import AnnotationPipelineConfig
+
+LEROBOT_GIT_URL = "https://github.com/huggingface/lerobot.git"
+
+# Mirrors the pins in pyproject.toml. The vLLM image resolves dependencies on its
+# own otherwise, and pulls av 18 / datasets 5 / draccus 0.11 — each of which breaks
+# lerobot at import time. `--upgrade-strategy only-if-needed` keeps vLLM's own
+# (torch, transformers, ...) pins intact.
+_RUNTIME_REQUIREMENTS = (
+    "'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"
+)
+
+# Flags the submitter resolves itself instead of forwarding verbatim: `--root`
+# names a directory only this machine has, `--repo_id` is re-emitted from the
+# config, and the config-file args name local files (rejected up front by
+# `submit_annotate_to_hf`). `--job.*` is dropped separately, by prefix; bare
+# `--job` is not, hence its entry here — it is the one arg that could smuggle a
+# remote `target` onto the pod and have the job recursively submit itself.
+_SUBMITTER_OWNED_ARGS = ("--root", "--repo_id", "--config_path", "--job")
+
+
+def _local_config_file_args(cfg: AnnotationPipelineConfig) -> list[str]:
+    """The CLI args that name a config file on the client's disk.
+
+    draccus exposes ``--config_path`` for the whole config plus a ``--<field>``
+    for every nested dataclass (``--vlm``, ``--plan``, ``--job``, ...). The pod has
+    none of those files, so a remote run has to reject them rather than silently
+    drop the settings they carry.
+    """
+    return ["--config_path", *(f"--{name}" for name in vars(cfg) if is_dataclass(getattr(cfg, name)))]
+
+
+def build_pod_setup(lerobot_ref: str) -> str:
+    """Shell prelude that turns the vLLM image into a ``lerobot-annotate`` runtime."""
+    spec = f"lerobot @ git+{LEROBOT_GIT_URL}@{lerobot_ref}"
+    return (
+        # git to install from the repo, ffmpeg to decode the dataset's videos.
+        "apt-get update -qq && apt-get install -y -qq git ffmpeg && "
+        f"pip install --no-deps {shlex.quote(spec)} && "
+        f"pip install --upgrade-strategy only-if-needed {_RUNTIME_REQUIREMENTS} && "
+        # vLLM's cudagraph memory estimate over-reserves and starves the KV cache;
+        # PyAV is the video backend the server can decode our frames with.
+        "export VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0 && "
+        "export VLLM_VIDEO_BACKEND=pyav"
+    )
+
+
+def build_pod_command(repo_id: str, lerobot_ref: str, argv: list[str]) -> list[str]:
+    """Build the ``bash -c`` command the pod runs: setup prelude, then annotation.
+
+    ``argv`` is the user's CLI (``sys.argv[1:]``) minus the flags in
+    ``_SUBMITTER_OWNED_ARGS``; ``--repo_id`` is re-added from the config so the pod
+    always annotates the dataset we just made sure is reachable on the Hub.
+    ``--job.target=local`` stops the pod from re-dispatching to itself.
+    """
+    forwarded = _pod_forwarded_args(argv, drop_names=_SUBMITTER_OWNED_ARGS, drop_prefixes=("--job.",))
+    annotate = shlex.join(["lerobot-annotate", f"--repo_id={repo_id}", *forwarded, "--job.target=local"])
+    return ["bash", "-c", f"{build_pod_setup(lerobot_ref)} && {annotate}"]
+
+
+def submit_annotate_to_hf(cfg: AnnotationPipelineConfig) -> None:
+    """Submit an annotation run to HF Jobs infrastructure.
+
+    Resolves credentials, makes sure the source dataset is reachable from the pod,
+    submits the job, then tails its logs until the job reaches a terminal stage —
+    or returns immediately with ``--job.detach``. Ctrl-C detaches without
+    cancelling the remote job.
+    """
+    token = get_token()
+    if not token:
+        raise RuntimeError("Not logged in to Hugging Face. Run `hf auth login` first.")
+
+    if cfg.repo_id is None:
+        raise ValueError(
+            "Remote annotation requires --repo_id: the pod downloads the dataset from the Hub, "
+            "and --root only names a directory on this machine."
+        )
+
+    argv = sys.argv[1:]
+    passed = {tok.split("=", 1)[0] for tok in argv}
+    used_config_files = sorted(passed.intersection(_local_config_file_args(cfg)))
+    if used_config_files:
+        raise ValueError(
+            f"{', '.join(used_config_files)} cannot be used with a remote --job.target: the pod "
+            "cannot read config files from this machine. Pass the settings as CLI flags instead."
+        )
+
+    if not cfg.push_to_hub:
+        # The pod's filesystem is discarded when the job ends, so without a push the
+        # run produces nothing. Warn rather than fail: a smoke test over
+        # --only_episodes that only inspects the logs is a legitimate use.
+        print(
+            "WARNING: --push_to_hub is off. The annotated dataset lives only on the pod and is "
+            "discarded when the job ends. Pass --push_to_hub=true to keep the result."
+        )
+
+    api = HfApi(token=token)
+    tags = resolve_job_tags(cfg.job.tags)
+    ensure_dataset_available(cfg.repo_id, api=api, tags=tags)
+
+    command = build_pod_command(cfg.repo_id, cfg.job.lerobot_ref, argv)
+
+    print(f"Submitting job to HF Jobs (flavor={cfg.job.target}, image={cfg.job.image}) ...")
+    job_info = run_job(
+        image=cfg.job.image,
+        command=command,
+        flavor=cfg.job.target,
+        secrets={"HF_TOKEN": token},
+        timeout=cfg.job.timeout,
+        # HF Jobs labels are key/value; expose each tag as a queryable label.
+        labels=dict.fromkeys(tags, "true"),
+    )
+    job_id = job_info.id
+    job_url = getattr(job_info, "url", None)
+    print(f"Job submitted: {job_id}")
+    if job_url:
+        print(f"  Job page:     {job_url}")
+    target_repo_id = cfg.new_repo_id or cfg.repo_id
+    if cfg.push_to_hub:
+        print(f"  Dataset repo: https://huggingface.co/datasets/{target_repo_id}")
+    print(f"  Monitor:      hf jobs logs {job_id}")
+    print(f"  Cancel:       hf jobs cancel {job_id}")
+
+    # No success marker: `lerobot-annotate` keeps working after the upload log line
+    # (dataset card, version tag), so completion has to be stage-based.
+    if not follow_job(job_id, detach=cfg.job.detach):
+        return
+
+    if cfg.push_to_hub:
+        print(f"\nAnnotation complete — dataset pushed to https://huggingface.co/datasets/{target_repo_id}")
+    else:
+        print("\nAnnotation complete. Note: --push_to_hub was off, so the result stayed on the pod.")

src/lerobot/jobs/hf.py

@@ -223,6 +223,74 @@ def _poll_until_done(
     return None
 
 
+def follow_job(job_id: str, *, detach: bool = False, success_marker: str | None = None) -> bool:
+    """Watch a submitted job to the end, streaming its logs to stdout.
+
+    Returns True when the job finished successfully and False when we stopped watching
+    without a verdict — `detach`, or the user pressing Ctrl-C, which detaches rather than
+    cancelling the remote job. Raises RuntimeError when the job reaches a terminal stage
+    other than COMPLETED.
+
+    `success_marker` finishes as soon as that string appears in the logs instead of waiting
+    out the platform's post-run finalization (~30s). Callers that have a log line meaning
+    "the artifact is on the Hub" should pass it; without one, completion is stage-based.
+    """
+    if detach:
+        return False
+
+    done = threading.Event()
+    detached = threading.Event()
+    marker_seen = threading.Event()
+    stage_holder: dict[str, str | None] = {}
+
+    def _poll() -> None:
+        stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
+
+    poll_thread = threading.Thread(target=_poll, daemon=True)
+    poll_thread.start()
+    log_thread = threading.Thread(
+        target=_tail_logs, args=(job_id, done, success_marker, marker_seen), daemon=True
+    )
+    log_thread.start()
+
+    def _detach(sig, frame):
+        detached.set()
+        done.set()
+        print("\nDetached. Job is still running.")
+        print(f"  Monitor: hf jobs logs {job_id}")
+        print(f"  Cancel:  hf jobs cancel {job_id}")
+
+    # signal.signal only works on the main thread; when called from a worker thread
+    # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
+    # handler rather than crashing with ValueError.
+    install_sigint = threading.current_thread() is threading.main_thread()
+    original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
+    if install_sigint:
+        signal.signal(signal.SIGINT, _detach)
+    try:
+        # Timeout-based join so SIGINT is delivered to the main thread promptly.
+        while poll_thread.is_alive():
+            poll_thread.join(timeout=0.5)
+        log_thread.join(timeout=5)
+    finally:
+        if install_sigint:
+            signal.signal(signal.SIGINT, original_sigint)
+
+    if detached.is_set():
+        return False
+    if marker_seen.is_set():
+        return True
+
+    stage = stage_holder.get("stage")
+    if stage != "COMPLETED":
+        message = stage_holder.get("message")
+        detail = f" ({message})" if message else ""
+        raise RuntimeError(
+            f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
+        )
+    return True
+
+
 def _pod_forwarded_args(
     argv: list[str], drop_names: tuple[str, ...] = (), drop_prefixes: tuple[str, ...] = ()
 ) -> list[str]:
@@ -362,64 +430,11 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None:
     print(f"  Monitor:    hf jobs logs {job_id}")
     print(f"  Cancel:     hf jobs cancel {job_id}")
 
-    if cfg.job.detach:
-        return
-
-    done = threading.Event()
-    detached = threading.Event()
-    pushed_ok = threading.Event()
-    stage_holder: dict[str, str | None] = {}
-
-    def _poll() -> None:
-        stage_holder["stage"] = _poll_until_done(job_id, done, status_holder=stage_holder)
-
-    poll_thread = threading.Thread(target=_poll, daemon=True)
-    poll_thread.start()
     # Finish as soon as the model is pushed, rather than waiting out the platform's
     # post-run finalization before the job stage flips to COMPLETED. This matches the
     # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay
     # in sync. If it ever stops matching we just fall back to stage-based completion
     # (~30s slower), so the contract is an optimization, not a correctness requirement.
     success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
-    log_thread = threading.Thread(
-        target=_tail_logs, args=(job_id, done, success_marker, pushed_ok), daemon=True
-    )
-    log_thread.start()
-
-    def _detach(sig, frame):
-        detached.set()
-        done.set()
-        print("\nDetached. Job is still running.")
-        print(f"  Monitor: hf jobs logs {job_id}")
-        print(f"  Cancel:  hf jobs cancel {job_id}")
-
-    # signal.signal only works on the main thread; when called from a worker thread
-    # (e.g. an orchestration framework) skip the Ctrl-C-detaches-instead-of-cancels
-    # handler rather than crashing with ValueError.
-    install_sigint = threading.current_thread() is threading.main_thread()
-    original_sigint = signal.getsignal(signal.SIGINT) if install_sigint else None
-    if install_sigint:
-        signal.signal(signal.SIGINT, _detach)
-    try:
-        # Timeout-based join so SIGINT is delivered to the main thread promptly.
-        while poll_thread.is_alive():
-            poll_thread.join(timeout=0.5)
-        log_thread.join(timeout=5)
-    finally:
-        if install_sigint:
-            signal.signal(signal.SIGINT, original_sigint)
-
-    if detached.is_set():
-        return
-
-    if pushed_ok.is_set():
+    if follow_job(job_id, detach=cfg.job.detach, success_marker=success_marker):
         print(f"\nTraining complete — model pushed to https://huggingface.co/{repo_id}")
-        return
-
-    stage = stage_holder.get("stage")
-    if stage != "COMPLETED":
-        message = stage_holder.get("message")
-        detail = f" ({message})" if message else ""
-        raise RuntimeError(
-            f"Job {job_id} ended with stage={stage}{detail}. Check logs: hf jobs logs {job_id}"
-        )

src/lerobot/motors/damiao/damiao.py

@@ -20,7 +20,6 @@ import logging
 import time
 from contextlib import contextmanager
 from copy import deepcopy
-from functools import cached_property
 from typing import TYPE_CHECKING, Any, TypedDict
 
 from lerobot.utils.decorators import check_if_already_connected, check_if_not_connected
@@ -854,7 +853,7 @@ class DamiaoMotorsBus(MotorsBusBase):
         else:
             raise ValueError(f"Motor {motor_obj} doesn't have a valid recv_id (None).")
 
-    @cached_property
+    @property
     def is_calibrated(self) -> bool:
         """Check if motors are calibrated."""
         return bool(self.calibration)

src/lerobot/motors/motors_bus.py

@@ -23,6 +23,7 @@ from __future__ import annotations
 
 import abc
 import logging
+import time
 from collections.abc import Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
@@ -818,13 +819,13 @@ class SerialMotorsBus(MotorsBusBase):
         """
         motor_names = self._get_motors_list(motors)
 
-        start_positions = self.sync_read("Present_Position", motor_names, normalize=False)
+        start_positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
         mins = start_positions.copy()
         maxes = start_positions.copy()
 
         user_pressed_enter = False
         while not user_pressed_enter:
-            positions = self.sync_read("Present_Position", motor_names, normalize=False)
+            positions = self.sync_read("Present_Position", motor_names, normalize=False, num_retry=5)
             mins = {motor: min(positions[motor], min_) for motor, min_ in mins.items()}
             maxes = {motor: max(positions[motor], max_) for motor, max_ in maxes.items()}
 
@@ -837,9 +838,12 @@ class SerialMotorsBus(MotorsBusBase):
             if enter_pressed():
                 user_pressed_enter = True
 
-            if display_values and not user_pressed_enter:
-                # Move cursor up to overwrite the previous output
-                move_cursor_up(len(motor_names) + 3)
+            if not user_pressed_enter:
+                if display_values:
+                    # Move cursor up to overwrite the previous output
+                    move_cursor_up(len(motor_names) + 3)
+                # Throttle reads even when the live table is disabled.
+                time.sleep(0.02)
 
         same_min_max = [motor for motor in motor_names if mins[motor] == maxes[motor]]
         if same_min_max:

src/lerobot/policies/diffusion/configuration_diffusion.py

@@ -79,6 +79,8 @@ class DiffusionConfig(PreTrainedConfig):
         use_film_scale_modulation: FiLM (https://huggingface.co/papers/1709.07871) is used for the Unet conditioning.
             Bias modulation is used be default, while this parameter indicates whether to also use scale
             modulation.
+        gradient_checkpointing: Whether to checkpoint the Unet residual blocks during training. This reduces
+            activation memory at the cost of recomputing those blocks during the backward pass.
         noise_scheduler_type: Name of the noise scheduler to use. Supported options: ["DDPM", "DDIM"].
         num_train_timesteps: Number of diffusion steps for the forward diffusion schedule.
         beta_schedule: Name of the diffusion beta schedule as per DDPMScheduler from Hugging Face diffusers.
@@ -132,6 +134,7 @@ class DiffusionConfig(PreTrainedConfig):
     n_groups: int = 8
     diffusion_step_embed_dim: int = 128
     use_film_scale_modulation: bool = True
+    gradient_checkpointing: bool = False
     # Noise scheduler.
     noise_scheduler_type: str = "DDPM"
     num_train_timesteps: int = 100

src/lerobot/policies/diffusion/modeling_diffusion.py

@@ -31,6 +31,7 @@ import torch
 import torch.nn.functional as F  # noqa: N812
 import torchvision
 from torch import Tensor, nn
+from torch.utils.checkpoint import checkpoint
 
 from lerobot.utils.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
 from lerobot.utils.import_utils import _diffusers_available, require_package
@@ -727,22 +728,35 @@ class DiffusionConditionalUnet1d(nn.Module):
         else:
             global_feature = timesteps_embed
 
+        use_gc = self.config.gradient_checkpointing and self.training
+
         # Run encoder, keeping track of skip features to pass to the decoder.
         encoder_skip_features: list[Tensor] = []
         for resnet, resnet2, downsample in self.down_modules:
-            x = resnet(x, global_feature)
-            x = resnet2(x, global_feature)
+            if use_gc:
+                x = checkpoint(resnet, x, global_feature, use_reentrant=False)
+                x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
+            else:
+                x = resnet(x, global_feature)
+                x = resnet2(x, global_feature)
             encoder_skip_features.append(x)
             x = downsample(x)
 
         for mid_module in self.mid_modules:
-            x = mid_module(x, global_feature)
+            if use_gc:
+                x = checkpoint(mid_module, x, global_feature, use_reentrant=False)
+            else:
+                x = mid_module(x, global_feature)
 
         # Run decoder, using the skip features from the encoder.
         for resnet, resnet2, upsample in self.up_modules:
             x = torch.cat((x, encoder_skip_features.pop()), dim=1)
-            x = resnet(x, global_feature)
-            x = resnet2(x, global_feature)
+            if use_gc:
+                x = checkpoint(resnet, x, global_feature, use_reentrant=False)
+                x = checkpoint(resnet2, x, global_feature, use_reentrant=False)
+            else:
+                x = resnet(x, global_feature)
+                x = resnet2(x, global_feature)
             x = upsample(x)
 
         x = self.final_conv(x)

src/lerobot/policies/pi05/modeling_pi05.py

@@ -524,8 +524,6 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
     def embed_suffix(self, noisy_actions, timestep):
         """Embed noisy_actions, timestep to prepare for Expert Gemma processing."""
-        embs = []
-        pad_masks = []
         att_masks = []
 
         # Embed timestep using sine-cosine positional encoding
@@ -551,23 +549,17 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
             return F.silu(x)
 
         time_emb = self._apply_checkpoint(time_mlp_func, time_emb)
-        action_time_emb = action_emb
         adarms_cond = time_emb
 
-        embs.append(action_time_emb)
-        bsize, action_time_dim = action_time_emb.shape[:2]
-        action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
-        pad_masks.append(action_time_mask)
+        bsize, action_time_dim = action_emb.shape[:2]
+        pad_masks = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=timestep.device)
 
         # Set attention masks so that image, language and state inputs do not attend to action tokens
         att_masks += [1] + ([0] * (self.config.chunk_size - 1))
-
-        embs = torch.cat(embs, dim=1)
-        pad_masks = torch.cat(pad_masks, dim=1)
-        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)
+        att_masks = torch.tensor(att_masks, dtype=action_emb.dtype, device=action_emb.device)
         att_masks = att_masks[None, :].expand(bsize, len(att_masks))
 
-        return embs, pad_masks, att_masks, adarms_cond
+        return action_emb, pad_masks, att_masks, adarms_cond
 
     def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
         """Do a full training forward pass and compute the loss."""

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -61,9 +61,15 @@ import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
 from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
-from lerobot.utils.device_utils import get_safe_dtype
 from lerobot.utils.import_utils import require_package
 
+from ..common.flow_matching import euler_integrate, sample_noise, sample_time_beta
+from ..common.vla_utils import (
+    create_sinusoidal_pos_embedding,
+    make_att_2d_masks,
+    pad_vector,
+    resize_with_pad,
+)
 from ..pretrained import PreTrainedPolicy
 from ..rtc.modeling_rtc import RTCProcessor
 from ..utils import (
@@ -79,96 +85,6 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
-def create_sinusoidal_pos_embedding(
-    time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
-) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
-    if dimension % 2 != 0:
-        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
-
-    if time.ndim != 1:
-        raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.")
-
-    dtype = get_safe_dtype(torch.float64, device.type)
-    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
-    period = min_period * (max_period / min_period) ** fraction
-
-    # Compute the outer product
-    scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
-    return pos_emb
-
-
-def make_att_2d_masks(pad_masks, att_masks):
-    """Copied from big_vision.
-
-    Tokens can attend to valid inputs tokens which have a cumulative mask_ar
-    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to
-    setup several types of attention, for example:
-
-      [[1 1 1 1 1 1]]: pure causal attention.
-
-      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between
-          themselves and the last 3 tokens have a causal attention. The first
-          entry could also be a 1 without changing behaviour.
-
-      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a
-          block can attend all previous blocks and all tokens on the same block.
-
-    Args:
-      input_mask: bool[B, N] true if its part of the input, false if padding.
-      mask_ar: int32[B, N] mask that's 1 where previous tokens cannot depend on
-        it and 0 where it shares the same attention mask as the previous token.
-    """
-    if att_masks.ndim != 2:
-        raise ValueError(att_masks.ndim)
-    if pad_masks.ndim != 2:
-        raise ValueError(pad_masks.ndim)
-
-    cumsum = torch.cumsum(att_masks, dim=1)
-    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
-    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
-    att_2d_masks = att_2d_masks & pad_2d_masks
-    return att_2d_masks
-
-
-def resize_with_pad(img, width, height, pad_value=-1):
-    # assume no-op when width height fits already
-    if img.ndim != 4:
-        raise ValueError(f"(b,c,h,w) expected, but {img.shape}")
-
-    cur_height, cur_width = img.shape[2:]
-
-    ratio = max(cur_width / width, cur_height / height)
-    resized_height = int(cur_height / ratio)
-    resized_width = int(cur_width / ratio)
-    resized_img = F.interpolate(
-        img, size=(resized_height, resized_width), mode="bilinear", align_corners=False
-    )
-
-    pad_height = max(0, int(height - resized_height))
-    pad_width = max(0, int(width - resized_width))
-
-    # pad on left and top of image
-    padded_img = F.pad(resized_img, (pad_width, 0, pad_height, 0), value=pad_value)
-    return padded_img
-
-
-def pad_vector(vector, new_dim):
-    """Can be (batch_size x sequence_length x features_dimension)
-    or (batch_size x features_dimension)
-    """
-    if vector.shape[-1] == new_dim:
-        return vector
-    shape = list(vector.shape)
-    current_dim = shape[-1]
-    shape[-1] = new_dim
-    new_vector = torch.zeros(*shape, dtype=vector.dtype, device=vector.device)
-    new_vector[..., :current_dim] = vector
-    return new_vector
-
-
 def normalize(x, min_val, max_val):
     return (x - min_val) / (max_val - min_val)
 
@@ -429,7 +345,13 @@ class SmolVLAPolicy(PreTrainedPolicy):
         for key in present_img_keys:
             img = batch[key][:, -1, :, :, :] if batch[key].ndim == 5 else batch[key]
             if self.config.resize_imgs_with_padding is not None:
-                img = resize_with_pad(img, *self.config.resize_imgs_with_padding, pad_value=0)
+                # SmolVLA stores the target as (width, height); the shared helper expects (height, width).
+                img = resize_with_pad(
+                    img,
+                    self.config.resize_imgs_with_padding[1],
+                    self.config.resize_imgs_with_padding[0],
+                    pad_value=0,
+                )
 
             # Normalize from range [0,1] to [-1,1] as expacted by siglip
             img = img * 2.0 - 1.0
@@ -619,20 +541,10 @@ class VLAFlowMatching(nn.Module):
             params.requires_grad = self.config.train_state_proj
 
     def sample_noise(self, shape, device):
-        noise = torch.normal(
-            mean=0.0,
-            std=1.0,
-            size=shape,
-            dtype=torch.float32,
-            device=device,
-        )
-        return noise
+        return sample_noise(shape, device)
 
     def sample_time(self, bsize, device):
-        beta_dist = torch.distributions.Beta(concentration1=1.5, concentration0=1.0)
-        time_beta = beta_dist.sample((bsize,)).to(device=device, dtype=torch.float32)
-        time = time_beta * 0.999 + 0.001
-        return time
+        return sample_time_beta(bsize, device, alpha=1.5, beta=1.0, scale=0.999, offset=0.001)
 
     def embed_prefix(
         self, images, img_masks, lang_tokens, lang_masks, state: torch.Tensor = None
@@ -800,7 +712,6 @@ class VLAFlowMatching(nn.Module):
             past_key_values=None,
             inputs_embeds=[prefix_embs, suffix_embs],
             use_cache=False,
-            fill_kv_cache=False,
         )
         suffix_out = suffix_out[:, -self.config.chunk_size :]
         # Original openpi code, upcast attention output
@@ -839,46 +750,24 @@ class VLAFlowMatching(nn.Module):
             past_key_values=None,
             inputs_embeds=[prefix_embs, None],
             use_cache=self.config.use_cache,
-            fill_kv_cache=True,
         )
         num_steps = self.config.num_steps
-        dt = -1.0 / num_steps
 
-        x_t = noise
-        for step in range(num_steps):
-            time = 1.0 + step * dt
-            time_tensor = torch.tensor(time, dtype=torch.float32, device=device).expand(bsize)
-
-            def denoise_step_partial_call(input_x_t, current_timestep=time_tensor):
-                return self.denoise_step(
-                    x_t=input_x_t,
-                    prefix_pad_masks=prefix_pad_masks,
-                    past_key_values=past_key_values,
-                    timestep=current_timestep,
-                )
-
-            if self._rtc_enabled():
-                inference_delay = kwargs.get("inference_delay")
-                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
-                execution_horizon = kwargs.get("execution_horizon")
-
-                v_t = self.rtc_processor.denoise_step(
-                    x_t=x_t,
-                    prev_chunk_left_over=prev_chunk_left_over,
-                    inference_delay=inference_delay,
-                    time=time,
-                    original_denoise_step_partial=denoise_step_partial_call,
-                    execution_horizon=execution_horizon,
-                )
-            else:
-                v_t = denoise_step_partial_call(x_t)
-
-            x_t = x_t + dt * v_t
-
-            if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled():
-                self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t)
-
-        return x_t
+        return euler_integrate(
+            lambda input_x_t, current_timestep: self.denoise_step(
+                x_t=input_x_t,
+                prefix_pad_masks=prefix_pad_masks,
+                past_key_values=past_key_values,
+                timestep=current_timestep,
+            ),
+            noise,
+            num_steps,
+            rtc_processor=self.rtc_processor,
+            rtc_enabled=self._rtc_enabled(),
+            inference_delay=kwargs.get("inference_delay"),
+            prev_chunk_left_over=kwargs.get("prev_chunk_left_over"),
+            execution_horizon=kwargs.get("execution_horizon"),
+        )
 
     def denoise_step(
         self,
@@ -907,8 +796,10 @@ class VLAFlowMatching(nn.Module):
             past_key_values=past_key_values,
             inputs_embeds=[None, suffix_embs],
             use_cache=self.config.use_cache,
-            fill_kv_cache=False,
         )
+        if past_key_values is not None:
+            # Self-attention layers append suffix K/V in place; restore the prefix for the next step.
+            past_key_values.crop(prefix_len)
         suffix_out = outputs_embeds[1]
         suffix_out = suffix_out[:, -self.config.chunk_size :]
         suffix_out = suffix_out.to(dtype=torch.float32)

src/lerobot/policies/smolvla/smolvlm_with_expert.py

@@ -26,6 +26,7 @@ if TYPE_CHECKING or _transformers_available:
         AutoModel,
         AutoModelForImageTextToText,
         AutoProcessor,
+        DynamicCache,
         SmolVLMForConditionalGeneration,
     )
 else:
@@ -33,6 +34,7 @@ else:
     AutoModel = None
     AutoModelForImageTextToText = None
     AutoProcessor = None
+    DynamicCache = None
     SmolVLMForConditionalGeneration = None
 
 
@@ -216,9 +218,8 @@ class SmolVLMWithExpertModel(nn.Module):
         batch_size,
         head_dim,
         use_cache: bool = True,
-        fill_kv_cache: bool = True,
-        past_key_values=None,
-    ) -> list[torch.Tensor]:
+        past_key_values: "DynamicCache | None" = None,
+    ) -> "tuple[list[torch.Tensor], DynamicCache | None]":
         query_states = []
         key_states = []
         value_states = []
@@ -259,22 +260,16 @@ class SmolVLMWithExpertModel(nn.Module):
         query_states = apply_rope(query_states, position_ids_)
         key_states = apply_rope(key_states, position_ids_)
 
-        if use_cache and past_key_values is None:
-            past_key_values = {}
-
         if use_cache:
-            if fill_kv_cache:
-                past_key_values[layer_idx] = {
-                    "key_states": key_states,
-                    "value_states": value_states,
-                }
-            else:
-                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
-                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
-                # the max len, then we (for instance) double the cache size. This implementation already exists
-                # in `transformers`. (molbap)
-                key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
-                value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1)
+            # `DynamicCache` stores tensors as [batch, heads, seq, head_dim]; this module works with
+            # [batch, seq, heads, head_dim]. During prefix prefill this stores the (post-RoPE) K/V and
+            # returns them unchanged; during denoising it appends the suffix K/V and returns
+            # [prefix; suffix], exactly like the previous hand-rolled dict cache.
+            key_states, value_states = past_key_values.update(
+                key_states.transpose(1, 2), value_states.transpose(1, 2), layer_idx
+            )
+            key_states = key_states.transpose(1, 2)
+            value_states = value_states.transpose(1, 2)
 
         attention_interface = self.get_attention_interface()
 
@@ -293,13 +288,12 @@ class SmolVLMWithExpertModel(nn.Module):
         batch_size,
         head_dim,
         use_cache: bool = True,
-        fill_kv_cache: bool = True,
-        past_key_values=None,
-    ) -> list[torch.Tensor]:
+        past_key_values: "DynamicCache | None" = None,
+    ) -> "tuple[list[torch.Tensor], DynamicCache | None]":
         attention_interface = self.get_attention_interface()
 
         att_outputs = []
-        assert len(inputs_embeds) == 2 or (use_cache and past_key_values is not None and not fill_kv_cache), (
+        assert len(inputs_embeds) == 2 or (use_cache and past_key_values is not None), (
             f"Both len(inputs_embeds) == {len(inputs_embeds)} and past_key_values is {past_key_values}"
         )
 
@@ -332,22 +326,13 @@ class SmolVLMWithExpertModel(nn.Module):
         else:
             expert_position_id = position_ids
 
-        if use_cache and past_key_values is None:
-            past_key_values = {}
-
-        if use_cache:
-            if fill_kv_cache:
-                past_key_values[layer_idx] = {
-                    "key_states": key_states,
-                    "value_states": value_states,
-                }
-            else:
-                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
-                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
-                # the max len, then we (for instance) double the cache size. This implementation already exists
-                # in `transformers`. (molbap)
-                key_states = past_key_values[layer_idx]["key_states"]
-                value_states = past_key_values[layer_idx]["value_states"]
+        if use_cache and past_key_values is not None:
+            # Cross-attention layers never fill the cache themselves: during the prefix prefill every
+            # layer goes through `forward_attn_layer`, which stores the (post-RoPE) VLM K/V for this
+            # layer index. Here we only read them back (no concatenation: the expert cross-attends to
+            # the fixed prefix). `DynamicCache` stores [batch, heads, seq, head_dim]; transpose back.
+            key_states = past_key_values.layers[layer_idx].keys.transpose(1, 2)
+            value_states = past_key_values.layers[layer_idx].values.transpose(1, 2)
 
         # Expert
         expert_layer = model_layers[1][layer_idx]
@@ -360,14 +345,15 @@ class SmolVLMWithExpertModel(nn.Module):
             expert_hidden_states = expert_hidden_states.to(dtype=expert_layer.self_attn.q_proj.weight.dtype)
             expert_query_state = expert_layer.self_attn.q_proj(expert_hidden_states).view(expert_hidden_shape)
 
-            _key_states = key_states.to(dtype=expert_layer.self_attn.k_proj.weight.dtype).view(
+            # reshape (not view): K/V read back from the cache are transposed, hence non-contiguous
+            _key_states = key_states.to(dtype=expert_layer.self_attn.k_proj.weight.dtype).reshape(
                 *key_states.shape[:2], -1
             )
             expert_key_states = expert_layer.self_attn.k_proj(_key_states).view(
                 *_key_states.shape[:-1], -1, expert_layer.self_attn.head_dim
             )  # k_proj should have same dim as kv
 
-            _value_states = value_states.to(dtype=expert_layer.self_attn.v_proj.weight.dtype).view(
+            _value_states = value_states.to(dtype=expert_layer.self_attn.v_proj.weight.dtype).reshape(
                 *value_states.shape[:2], -1
             )
             expert_value_states = expert_layer.self_attn.v_proj(_value_states).view(
@@ -416,10 +402,9 @@ class SmolVLMWithExpertModel(nn.Module):
         self,
         attention_mask: torch.Tensor | None = None,
         position_ids: torch.LongTensor | None = None,
-        past_key_values: list[torch.FloatTensor] | None = None,
+        past_key_values: "DynamicCache | None" = None,
         inputs_embeds: list[torch.FloatTensor] = None,
         use_cache: bool | None = None,
-        fill_kv_cache: bool | None = None,
     ):
         models = [self.get_vlm_model().text_model, self.lm_expert]
         model_layers = self.get_model_layers(models)
@@ -431,6 +416,13 @@ class SmolVLMWithExpertModel(nn.Module):
                 continue
             batch_size = hidden_states.shape[0]
 
+        # Prefix prefill: no cache was passed, so create one and fill it (every layer runs
+        # self-attention over the prefix). When a filled cache is passed (denoising), layers
+        # read from it instead.
+        fill_kv_cache = use_cache and past_key_values is None
+        if fill_kv_cache:
+            past_key_values = DynamicCache()
+
         # RMSNorm
         num_layers = self.num_vlm_layers
         head_dim = self.vlm.config.text_config.head_dim
@@ -449,7 +441,6 @@ class SmolVLMWithExpertModel(nn.Module):
                     batch_size,
                     head_dim,
                     use_cache=use_cache,
-                    fill_kv_cache=fill_kv_cache,
                     past_key_values=past_key_values,
                 )
             else:
@@ -462,7 +453,6 @@ class SmolVLMWithExpertModel(nn.Module):
                     batch_size,
                     head_dim,
                     use_cache=use_cache,
-                    fill_kv_cache=fill_kv_cache,
                     past_key_values=past_key_values,
                 )
             outputs_embeds = []

src/lerobot/policies/vla_jepa/action_head.py

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

src/lerobot/policies/vla_jepa/configuration_vla_jepa.py

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

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

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

src/lerobot/policies/vla_jepa/processor_vla_jepa.py

@@ -17,19 +17,14 @@ from __future__ import annotations
 from typing import Any
 
 import torch
-import torch.nn.functional as F  # noqa: N812
 
-from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from lerobot.processor import (
-    AbsoluteActionsProcessorStep,
     EnvTransition,
-    ObservationProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
-    RelativeActionsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
     make_default_policy_processor_steps,
@@ -37,69 +32,6 @@ from lerobot.processor import (
 )
 
 
-@ProcessorStepRegistry.register(name="vla_jepa_image_prep")
-class ImagePrepProcessorStep(ObservationProcessorStep):
-    """Prepares image observations for the VLA-JEPA model: float cast, 1->3 channel expand, resize.
-
-    This makes explicit (in the serialized pipeline) the image prep the model used to do
-    internally. The model keeps the same operations as idempotent guards, so:
-      - checkpoints saved WITHOUT this step (older uploads) are unaffected — the model still
-        does the prep;
-      - checkpoints saved WITH this step get it done here, and the model-side guards no-op.
-
-    Mirrors `Qwen3VLInterface.to_pixel_values` + the `F.interpolate(mode="area")` resize in
-    `VLAJEPAPolicy._prepare_model_inputs`/`predict_action`. Deliberately does NOT clamp (the
-    model path doesn't), so values stay bit-identical. Handles [C,H,W], [B,C,H,W]/[T,C,H,W]
-    and [B,T,C,H,W] image tensors.
-    """
-
-    def __init__(self, resize_to: tuple[int, int] | None = None, expand_channels: bool = True):
-        self.resize_to = tuple(resize_to) if resize_to is not None else None
-        self.expand_channels = expand_channels
-
-    def observation(self, observation: dict) -> dict:
-        new_observation = dict(observation)
-        for key in observation:
-            if "image" not in key:
-                continue
-            image = observation[key].float()
-            if self.expand_channels and image.shape[-3] == 1:
-                repeats = [1] * image.ndim
-                repeats[-3] = 3
-                image = image.repeat(*repeats)
-            if self.resize_to is not None and tuple(image.shape[-2:]) != self.resize_to:
-                device = image.device
-                # NOTE: no "area" kernel on mps; resize on cpu then move back.
-                if device.type == "mps":
-                    image = image.cpu()
-                lead = image.shape[:-3]
-                c, h, w = image.shape[-3:]
-                flat = image.reshape(-1, c, h, w)
-                flat = F.interpolate(flat, size=self.resize_to, mode="area")
-                image = flat.reshape(*lead, c, *self.resize_to).to(device)
-            new_observation[key] = image
-        return new_observation
-
-    def get_config(self) -> dict[str, Any]:
-        return {
-            "resize_to": list(self.resize_to) if self.resize_to is not None else None,
-            "expand_channels": self.expand_channels,
-        }
-
-    def transform_features(self, features):
-        for key in features[PipelineFeatureType.OBSERVATION]:
-            if "image" not in key:
-                continue
-            feat = features[PipelineFeatureType.OBSERVATION][key]
-            # Match `to_pixel_values`: only a single channel is expanded to 3.
-            nb_channel = 3 if (self.expand_channels and feat.shape[0] == 1) else feat.shape[0]
-            spatial = self.resize_to if self.resize_to is not None else tuple(feat.shape[1:])
-            features[PipelineFeatureType.OBSERVATION][key] = PolicyFeature(
-                type=feat.type, shape=(nb_channel, *spatial)
-            )
-        return features
-
-
 @ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
 class ClipActionsProcessorStep(ProcessorStep):
     """Clips action tensor to [-1, 1] before unnormalization."""
@@ -177,24 +109,10 @@ def make_vla_jepa_pre_post_processors(
 ]:
     features = {**config.input_features, **config.output_features}
     steps = make_default_policy_processor_steps(config, dataset_stats)
-
-    # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
-    # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
-    # below so its cached raw state (set during preprocessing) flows to postprocessing.
-    relative_step = RelativeActionsProcessorStep(
-        enabled=config.use_relative_actions,
-        exclude_joints=getattr(config, "relative_exclude_joints", []),
-        action_names=getattr(config, "action_feature_names", None),
-    )
-
     input_steps = [
         steps.rename_observations,
         steps.add_batch_dim,
         steps.to_device,
-        ImagePrepProcessorStep(
-            resize_to=tuple(config.resize_images_to) if config.resize_images_to else None,
-        ),
-        relative_step,
         steps.normalize,
     ]
     output_steps: list[ProcessorStep] = []
@@ -213,11 +131,6 @@ def make_vla_jepa_pre_post_processors(
             stats=dataset_stats,
         )
     )
-    # Reverse the relative conversion on the unnormalized action, before gripper binarization.
-    # gripper is kept absolute by relative_exclude_joints, so the two steps touch disjoint dims.
-    output_steps.append(
-        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step)
-    )
     if config.binarize_gripper_action:
         output_steps.append(
             BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)

src/lerobot/robots/openarm_follower/openarm_follower.py

@@ -150,9 +150,6 @@ class OpenArmFollower(Robot):
 
         self.configure()
 
-        if self.is_calibrated:
-            self.bus.set_zero_position()
-
         self.bus.enable_torque()
 
         logger.info(f"{self} connected.")

src/lerobot/scripts/lerobot_annotate.py

@@ -24,7 +24,14 @@ Example:
       --root=/path/to/dataset \\
       --vlm.model_id=Qwen/Qwen2.5-VL-7B-Instruct
 
-For distributed runs, see ``examples/annotations/run_hf_job.py``.
+Pass ``--job.target=<flavor>`` to run the same command on a Hugging Face
+Jobs GPU instead of this machine (see ``lerobot.jobs.annotate``):
+
+  uv run lerobot-annotate \\
+      --repo_id=user/dataset \\
+      --new_repo_id=user/dataset_annotated \\
+      --push_to_hub=true \\
+      --job.target=h200
 """
 
 import logging
@@ -69,6 +76,14 @@ def _resolve_root(cfg: AnnotationPipelineConfig) -> Path:
 def annotate(cfg: AnnotationPipelineConfig) -> None:
     """Run the steerable annotation pipeline against a dataset."""
     logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+    if cfg.job.is_remote:
+        # Imported lazily: the submitter pulls in LeRobotDataset (the `dataset`
+        # extra), which a local annotation run over --root doesn't need.
+        from lerobot.jobs.annotate import submit_annotate_to_hf
+
+        return submit_annotate_to_hf(cfg)
+
     root = _resolve_root(cfg)
     logger.info("annotate: root=%s", root)
 

src/lerobot/scripts/lerobot_setup_motors.py

@@ -51,19 +51,7 @@ from lerobot.teleoperators import (  # noqa: F401
     rebot_102_leader,
     so_leader,
 )
-
-COMPATIBLE_DEVICES = [
-    "koch_follower",
-    "koch_leader",
-    "omx_follower",
-    "omx_leader",
-    "openarm_mini",
-    "so100_follower",
-    "so100_leader",
-    "so101_follower",
-    "so101_leader",
-    "lekiwi",
-]
+from lerobot.utils.import_utils import register_third_party_plugins
 
 
 @dataclass
@@ -80,18 +68,19 @@ class SetupConfig:
 
 @draccus.wrap()
 def setup_motors(cfg: SetupConfig):
-    if cfg.device.type not in COMPATIBLE_DEVICES:
-        raise NotImplementedError
-
     if isinstance(cfg.device, RobotConfig):
         device = make_robot_from_config(cfg.device)
     else:
         device = make_teleoperator_from_config(cfg.device)
 
-    device.setup_motors()
+    setup = getattr(device, "setup_motors", None)
+    if not callable(setup):
+        raise NotImplementedError(f"Device type '{cfg.device.type}' does not support motor setup.")
+    setup()
 
 
 def main():
+    register_third_party_plugins()
     setup_motors()
 
 

src/lerobot/scripts/lerobot_train.py

@@ -71,6 +71,16 @@ from lerobot.utils.utils import (
 from .lerobot_eval import eval_policy_all
 
 
+def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]:
+    """Return worker-only DataLoader options, disabling them for single-process loading."""
+    workers_enabled = cfg.num_workers > 0
+    return {
+        "prefetch_factor": cfg.prefetch_factor if workers_enabled else None,
+        "persistent_workers": cfg.persistent_workers and workers_enabled,
+        "multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None,
+    }
+
+
 def update_policy(
     train_metrics: MetricsTracker,
     policy: PreTrainedPolicy,
@@ -473,8 +483,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         pin_memory=device.type == "cuda",
         drop_last=False,
         collate_fn=collate_fn,
-        prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
-        persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
+        **_dataloader_worker_kwargs(cfg),
     )
 
     # Build eval dataloader if a held-out split exists
@@ -500,8 +509,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             pin_memory=device.type == "cuda",
             drop_last=False,
             collate_fn=eval_collate_fn,
-            prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
-            persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
+            **_dataloader_worker_kwargs(cfg),
         )
 
     # Prepare everything with accelerator

src/lerobot/teleoperators/gamepad/configuration_gamepad.py

@@ -23,3 +23,5 @@ from ..config import TeleoperatorConfig
 @dataclass
 class GamepadTeleopConfig(TeleoperatorConfig):
     use_gripper: bool = True
+    # Use hidapi instead of pygame for controllers that pygame cannot detect reliably.
+    hidapi_fallback: bool = False

src/lerobot/teleoperators/gamepad/teleop_gamepad.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import sys
 from enum import IntEnum
 from typing import Any
@@ -27,6 +28,8 @@ from ..teleoperator import Teleoperator
 from ..utils import TeleopEvents
 from .configuration_gamepad import GamepadTeleopConfig
 
+logger = logging.getLogger(__name__)
+
 
 class GripperAction(IntEnum):
     CLOSE = 0
@@ -56,6 +59,13 @@ class GamepadTeleop(Teleoperator):
 
         self.gamepad = None
 
+        self.hidapi_fallback = config.hidapi_fallback
+        if sys.platform == "darwin" and not self.hidapi_fallback:
+            logger.warning(
+                "On macOS, pygame may not reliably detect input from some controllers. "
+                "If you experience issues, set `hidapi_fallback=true`."
+            )
+
     @property
     def action_features(self) -> dict:
         if self.config.use_gripper:
@@ -76,9 +86,7 @@ class GamepadTeleop(Teleoperator):
         return {}
 
     def connect(self) -> None:
-        # use HidApi for macos
-        if sys.platform == "darwin":
-            # NOTE: On macOS, pygame doesn’t reliably detect input from some controllers so we fall back to hidapi
+        if self.hidapi_fallback:
             from .gamepad_utils import GamepadControllerHID as Gamepad
         else:
             from .gamepad_utils import GamepadController as Gamepad

src/lerobot/utils/device_utils.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import logging
-from contextlib import nullcontext
 
 import torch
 
@@ -122,25 +121,3 @@ def is_amp_available(device: str):
         return False
     else:
         raise ValueError(f"Unknown device '{device}.")
-
-
-def get_autocast_context(device_type: str, dtype: torch.dtype = torch.bfloat16):
-    """Return a device-safe autocast context manager.
-
-    Hardcoding `torch.autocast(dtype=torch.bfloat16)` breaks on backends without AMP
-    (MPS) and silently misbehaves on pre-Ampere CUDA GPUs that lack bf16 support. This
-    picks a safe context per device:
-      - no AMP support (e.g. mps): `nullcontext()` (run in the tensors' native dtype)
-      - CUDA requesting bf16 on compute capability < 8.0 (pre-Ampere): fall back to fp16
-      - otherwise: `torch.autocast(device_type, dtype)`
-    """
-    if not is_amp_available(device_type):
-        return nullcontext()
-    if (
-        device_type == "cuda"
-        and dtype == torch.bfloat16
-        and torch.cuda.is_available()
-        and torch.cuda.get_device_capability()[0] < 8
-    ):
-        dtype = torch.float16
-    return torch.autocast(device_type=device_type, dtype=dtype)

tests/datasets/test_datasets.py

@@ -114,6 +114,20 @@ def test_dataset_initialization(tmp_path, lerobot_dataset_factory):
     assert dataset.num_frames == len(dataset)
 
 
+def test_dataset_slice(tmp_path, lerobot_dataset_factory):
+    dataset = lerobot_dataset_factory(
+        root=tmp_path / "test", total_episodes=3, total_frames=30, use_videos=False
+    )
+
+    assert len(dataset[:5]) == 5
+    assert len(dataset[::2]) == (len(dataset) + 1) // 2
+    assert [item["index"].item() for item in dataset[4::-1]] == [4, 3, 2, 1, 0]
+    assert [item["index"].item() for item in dataset[-3:]] == list(range(len(dataset) - 3, len(dataset)))
+    assert dataset[len(dataset) :] == []
+    assert isinstance(dataset[0], dict)
+    assert dataset[:1][0].keys() == dataset[0].keys()
+
+
 # TODO(rcadene, aliberts): do not run LeRobotDataset.create, instead refactor LeRobotDatasetMetadata.create
 # and test the small resulting function that validates the features
 def test_dataset_feature_with_forward_slash_raises_error():
@@ -1741,6 +1755,38 @@ def test_delta_timestamps_query_returns_correct_values(tmp_path, empty_lerobot_d
     assert is_pad == [True, False], f"Expected [True, False], got {is_pad}"
 
 
+def test_dataset_slice_with_delta_timestamps(tmp_path, empty_lerobot_dataset_factory):
+    features = {
+        "observation.state": {"dtype": "float32", "shape": (1,), "names": ["x"]},
+    }
+    dataset = empty_lerobot_dataset_factory(
+        root=tmp_path / "test_slice_delta", features=features, use_videos=False, fps=10
+    )
+
+    for frame_idx in range(5):
+        dataset.add_frame(
+            {
+                "observation.state": torch.tensor([frame_idx], dtype=torch.float32),
+                "task": "task_0",
+            }
+        )
+    dataset.save_episode()
+    dataset.finalize()
+
+    sliced_dataset = LeRobotDataset(
+        dataset.repo_id,
+        root=dataset.root,
+        delta_timestamps={"observation.state": [-0.1, 0.0]},
+        tolerance_s=0.04,
+    )
+
+    items = sliced_dataset[:2]
+
+    assert items[0]["observation.state"].tolist() == [0.0, 0.0]
+    assert items[0]["observation.state_is_pad"].tolist() == [True, False]
+    assert items[1]["observation.state"].tolist() == [0.0, 1.0]
+
+
 def test_episode_filter_filters_dataset(tmp_path, lerobot_dataset_factory):
     """episode_filter on LeRobotDataset narrows the loaded dataset to matching episodes."""
     dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=8, total_frames=200)

tests/envs/test_dispatch.py

@@ -35,6 +35,17 @@ def test_unknown_type():
         make_env_config("nonexistent")
 
 
+def test_libero_fps_controls_simulator_frequency():
+    cfg = LiberoEnv(fps=17)
+
+    assert cfg.gym_kwargs["control_freq"] == 17
+
+
+def test_libero_rejects_nonpositive_fps():
+    with pytest.raises(ValueError, match="fps must be positive"):
+        LiberoEnv(fps=0)
+
+
 def test_identity_processors():
     """Base class get_env_processors() returns identity pipelines."""
     cfg = make_env_config("aloha")

tests/jobs/test_annotate.py

@@ -0,0 +1,245 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import shlex
+import sys
+from unittest.mock import MagicMock
+
+import draccus
+import pytest
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.annotations.steerable_pipeline.config import (
+    DEFAULT_ANNOTATE_JOB_IMAGE,
+    AnnotationJobConfig,
+    AnnotationPipelineConfig,
+)
+from lerobot.jobs.annotate import build_pod_command, build_pod_setup, submit_annotate_to_hf
+
+
+def _parse(*args):
+    return draccus.parse(AnnotationPipelineConfig, args=list(args))
+
+
+def _set_argv(monkeypatch, *args):
+    monkeypatch.setattr(sys, "argv", ["lerobot-annotate", *args])
+
+
+# --- config ----------------------------------------------------------------
+
+
+def test_annotation_job_defaults_are_local_with_vllm_image():
+    cfg = AnnotationJobConfig()
+    assert cfg.target is None
+    assert cfg.is_remote is False
+    assert cfg.image == DEFAULT_ANNOTATE_JOB_IMAGE
+    assert cfg.timeout == "2h"
+    assert cfg.lerobot_ref == "main"
+
+
+def test_annotation_config_parses_job_target():
+    cfg = _parse("--repo_id", "u/d", "--job.target", "h200")
+    assert cfg.job.target == "h200"
+    assert cfg.job.is_remote is True
+
+
+def test_annotation_config_defaults_to_local():
+    assert _parse("--repo_id", "u/d").job.is_remote is False
+
+
+# --- pod command -----------------------------------------------------------
+
+
+def test_pod_setup_installs_requested_ref():
+    setup = build_pod_setup("my-branch")
+    assert "git+https://github.com/huggingface/lerobot.git@my-branch" in setup
+    # The vLLM image has neither ffmpeg (video decode) nor lerobot's pinned deps.
+    assert "ffmpeg" in setup
+    assert "'draccus==0.10.0'" in setup
+
+
+def _annotate_argv(command):
+    """Extract the `lerobot-annotate ...` argv from a `bash -c` pod command."""
+    assert command[:2] == ["bash", "-c"]
+    _setup, _, annotate = command[2].rpartition(" && ")
+    return shlex.split(annotate)
+
+
+def test_pod_command_forwards_user_flags_and_pins_local_target():
+    command = build_pod_command(
+        "u/d",
+        "main",
+        ["--repo_id=u/d", "--new_repo_id=u/d_annotated", "--push_to_hub=true", "--job.target=h200"],
+    )
+    argv = _annotate_argv(command)
+    assert argv[0] == "lerobot-annotate"
+    # --job.* is client-side orchestration; the pod must not re-dispatch itself.
+    assert not any(a.startswith("--job.") for a in argv[1:-1])
+    assert argv[-1] == "--job.target=local"
+    assert "--new_repo_id=u/d_annotated" in argv
+    assert "--push_to_hub=true" in argv
+
+
+def test_pod_command_replaces_host_local_root_with_repo_id():
+    """--root points at a directory only the client has; the pod resolves by repo_id."""
+    command = build_pod_command("u/d", "main", ["--root", "/home/me/datasets/d", "--seed=7"])
+    argv = _annotate_argv(command)
+    assert "--root" not in argv
+    assert "/home/me/datasets/d" not in argv
+    assert argv.count("--repo_id=u/d") == 1
+    assert "--seed=7" in argv
+
+
+def test_pod_command_does_not_duplicate_repo_id():
+    command = build_pod_command("u/d", "main", ["--repo_id", "u/d"])
+    assert _annotate_argv(command).count("--repo_id=u/d") == 1
+
+
+def test_pod_command_quotes_flags_containing_spaces_and_json():
+    """serve_command and chat_template_kwargs must survive the trip through `bash -c`."""
+    serve = "--vlm.serve_command=vllm serve Qwen/Qwen3.6-27B --max-model-len 32768 --port {port}"
+    kwargs = '--vlm.chat_template_kwargs={"enable_thinking": false}'
+    command = build_pod_command("u/d", "main", [serve, kwargs])
+    argv = _annotate_argv(command)
+    assert serve in argv
+    assert kwargs in argv
+
+
+# --- submission ------------------------------------------------------------
+
+
+def test_submit_requires_login(monkeypatch):
+    monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: None)
+    with pytest.raises(RuntimeError, match="hf auth login"):
+        submit_annotate_to_hf(_parse("--repo_id", "u/d", "--job.target", "h200"))
+
+
+def test_submit_requires_repo_id(monkeypatch):
+    """A remote run over --root alone can't work: the pod can't see the client's disk."""
+    monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
+    cfg = _parse("--root", "/tmp/d", "--job.target", "h200")
+    with pytest.raises(ValueError, match="--repo_id"):
+        submit_annotate_to_hf(cfg)
+
+
+@pytest.mark.parametrize("arg", ["--config_path=annotate.yaml", "--vlm=vlm.yaml", "--job=job.yaml"])
+def test_submit_rejects_local_config_files(monkeypatch, arg):
+    """draccus takes a config file for the whole config and for each nested one; the
+    pod can read none of them, so a remote run must refuse rather than drop them."""
+    monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
+    _set_argv(monkeypatch, arg, "--job.target=h200")
+    cfg = _parse("--repo_id", "u/d", "--job.target", "h200")
+    with pytest.raises(ValueError, match="cannot read config files"):
+        submit_annotate_to_hf(cfg)
+
+
+def test_pod_command_drops_bare_job_config_file_arg():
+    """`--job` isn't caught by the `--job.` prefix, and could carry a remote target
+    that would make the pod submit a job of its own — recursively."""
+    argv = _annotate_argv(build_pod_command("u/d", "main", ["--job", "job.yaml", "--seed=7"]))
+    assert "--job" not in argv
+    assert "job.yaml" not in argv
+    assert argv[-1] == "--job.target=local"
+
+
+def test_submit_dispatches_job(monkeypatch):
+    monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
+    monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
+    monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
+
+    run_job_calls = []
+
+    def fake_run_job(**kwargs):
+        run_job_calls.append(kwargs)
+        return MagicMock(id="job-123")
+
+    monkeypatch.setattr("lerobot.jobs.annotate.run_job", fake_run_job)
+    _set_argv(monkeypatch, "--repo_id=u/d", "--push_to_hub=true", "--job.target=h200", "--job.detach=true")
+
+    cfg = _parse("--repo_id", "u/d", "--push_to_hub", "true", "--job.target", "h200", "--job.detach", "true")
+    submit_annotate_to_hf(cfg)
+
+    assert len(run_job_calls) == 1
+    call = run_job_calls[0]
+    assert call["flavor"] == "h200"
+    assert call["image"] == DEFAULT_ANNOTATE_JOB_IMAGE
+    assert call["timeout"] == "2h"
+    # The Hub token is forwarded so the pod can pull a private dataset and push the result.
+    assert call["secrets"]["HF_TOKEN"] == "tok"
+    assert call["labels"].get("lerobot") == "true"
+    argv = _annotate_argv(call["command"])
+    assert argv[0] == "lerobot-annotate"
+    assert "--push_to_hub=true" in argv
+
+
+@pytest.mark.timeout(15)
+def test_submit_follows_job_to_completion(monkeypatch, capsys):
+    """Non-detach path must stream logs and RETURN (not hang) once the job is terminal.
+
+    Exercises the `follow_job` helper shared with the training submitter from the
+    annotation side, which is why the job-state patches target `lerobot.jobs.hf`.
+    Asserting on the completion message and not merely on "didn't hang" is what makes
+    this fail if `follow_job` ever reports detached-without-a-verdict instead.
+    """
+    monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
+    monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
+    monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
+    monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1", url="http://x"))
+    monkeypatch.setattr(
+        "lerobot.jobs.hf.inspect_job",
+        lambda job_id: MagicMock(status=MagicMock(stage=MagicMock(value="COMPLETED"), message=None)),
+    )
+    monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
+    _set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200")
+
+    submit_annotate_to_hf(_parse("--repo_id", "u/d", "--push_to_hub", "true", "--job.target", "h200"))
+    assert "Annotation complete" in capsys.readouterr().out
+
+
+@pytest.mark.timeout(15)
+def test_submit_raises_when_job_fails(monkeypatch):
+    """A job that ends in a non-COMPLETED stage must surface as an error, not a silent return."""
+    monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
+    monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
+    monkeypatch.setattr("lerobot.jobs.annotate.ensure_dataset_available", lambda *a, **kw: None)
+    monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1", url=None))
+    monkeypatch.setattr(
+        "lerobot.jobs.hf.inspect_job",
+        lambda job_id: MagicMock(status=MagicMock(stage=MagicMock(value="ERROR"), message="Job timeout")),
+    )
+    monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", lambda job_id, follow=True: iter(()))
+    _set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200")
+
+    with pytest.raises(RuntimeError, match="stage=ERROR .Job timeout."):
+        submit_annotate_to_hf(_parse("--repo_id", "u/d", "--job.target", "h200"))
+
+
+def test_submit_ensures_dataset_is_on_the_hub(monkeypatch):
+    """A local-only dataset is pushed (privately) before the job can reach it by repo_id."""
+    monkeypatch.setattr("lerobot.jobs.annotate.get_token", lambda: "tok")
+    monkeypatch.setattr("lerobot.jobs.annotate.HfApi", lambda token=None: MagicMock())
+    monkeypatch.setattr("lerobot.jobs.annotate.run_job", lambda **kw: MagicMock(id="job-1"))
+
+    seen = []
+    monkeypatch.setattr(
+        "lerobot.jobs.annotate.ensure_dataset_available",
+        lambda repo_id, *, api, tags=None: seen.append((repo_id, tags)),
+    )
+    _set_argv(monkeypatch, "--repo_id=u/d", "--job.target=h200", "--job.detach=true")
+
+    submit_annotate_to_hf(
+        _parse("--repo_id", "u/d", "--job.target", "h200", "--job.detach", "true", "--job.tags", '["lelab"]')
+    )
+    assert seen == [("u/d", ["lerobot", "lelab"])]

tests/jobs/test_hf.py

@@ -29,12 +29,26 @@ from lerobot.jobs.hf import (
     _poll_until_done,
     build_remote_config_file,
     build_repo_id,
+    follow_job,
     resolve_job_tags,
     resolve_wandb_api_key,
     submit_to_hf,
 )
 
 
+def test_follow_job_detach_returns_without_watching(monkeypatch):
+    """`detach` must short-circuit before any polling or log streaming starts."""
+
+    def _boom(*a, **kw):
+        raise AssertionError("detach must not touch the job")
+
+    monkeypatch.setattr("lerobot.jobs.hf.inspect_job", _boom)
+    monkeypatch.setattr("lerobot.jobs.hf.fetch_job_logs", _boom)
+    # False = "stopped watching without a verdict", so callers stay quiet rather than
+    # claiming success for a job that is still running.
+    assert follow_job("job-1", detach=True) is False
+
+
 def test_resolve_job_tags_always_includes_lerobot_and_dedups():
     assert resolve_job_tags(None) == ["lerobot"]
     assert resolve_job_tags([]) == ["lerobot"]

tests/motors/test_dynamixel.py

@@ -405,12 +405,18 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors):
     read_pos_stub = mock_motors.build_sequential_sync_read_stub(
         *X_SERIES_CONTROL_TABLE["Present_Position"], positions
     )
-    with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]):
-        bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
-        bus.connect(handshake=False)
+    bus = DynamixelMotorsBus(port=mock_motors.port, motors=dummy_motors)
+    bus.connect(handshake=False)
 
+    with (
+        patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]),
+        patch("lerobot.motors.motors_bus.time.sleep") as mock_sleep,
+        patch.object(bus, "sync_read", wraps=bus.sync_read) as mock_sync_read,
+    ):
         mins, maxes = bus.record_ranges_of_motion(display_values=False)
 
     assert mock_motors.stubs[read_pos_stub].calls == 3
+    assert all(call.kwargs["num_retry"] == 5 for call in mock_sync_read.call_args_list)
+    mock_sleep.assert_called_once_with(0.02)
     assert mins == expected_mins
     assert maxes == expected_maxes

tests/motors/test_feetech.py

@@ -509,12 +509,18 @@ def test_record_ranges_of_motion(mock_motors, dummy_motors):
     stub = mock_motors.build_sequential_sync_read_stub(
         *STS_SMS_SERIES_CONTROL_TABLE["Present_Position"], positions
     )
-    with patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]):
-        bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
-        bus.connect(handshake=False)
+    bus = FeetechMotorsBus(port=mock_motors.port, motors=dummy_motors)
+    bus.connect(handshake=False)
 
+    with (
+        patch("lerobot.motors.motors_bus.enter_pressed", side_effect=[False, True]),
+        patch("lerobot.motors.motors_bus.time.sleep") as mock_sleep,
+        patch.object(bus, "sync_read", wraps=bus.sync_read) as mock_sync_read,
+    ):
         mins, maxes = bus.record_ranges_of_motion(display_values=False)
 
     assert mock_motors.stubs[stub].calls == 3
+    assert all(call.kwargs["num_retry"] == 5 for call in mock_sync_read.call_args_list)
+    mock_sleep.assert_called_once_with(0.02)
     assert mins == expected_mins
     assert maxes == expected_maxes

tests/policies/vla_jepa/test_processor_vla_jepa.py

@@ -1,182 +0,0 @@
-#!/usr/bin/env python
-
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-"""Tests for the VLA-JEPA image-prep processor step and its back-compat contract with the model.
-
-The step moves image resize + 1->3 channel-expand out of the model into the (serialized)
-preprocessor. The model keeps the same ops as idempotent guards, so:
-  - old checkpoints (JSON without the step) are unaffected — the model still does the prep;
-  - new checkpoints (JSON with the step) get it done in the step, and the model guards no-op.
-These tests pin the step's numerics (bit-identical to the model's F.interpolate(area)) and the
-equivalence of the two paths on the Qwen image path.
-"""
-
-from __future__ import annotations
-
-from copy import deepcopy
-
-import pytest
-import torch
-import torch.nn.functional as F  # noqa: N812
-
-pytest.importorskip("transformers")
-pytest.importorskip("diffusers")
-
-from conftest import (  # noqa: E402
-    BATCH_SIZE,
-    IMAGE_SIZE,
-    make_config,
-    make_inference_batch,
-    make_train_batch,
-)
-
-from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature  # noqa: E402
-from lerobot.policies.vla_jepa.modeling_vla_jepa import VLAJEPAPolicy  # noqa: E402
-from lerobot.policies.vla_jepa.processor_vla_jepa import (  # noqa: E402
-    ImagePrepProcessorStep,
-    make_vla_jepa_pre_post_processors,
-)
-from lerobot.processor import ProcessorStepRegistry  # noqa: E402
-from lerobot.utils.constants import OBS_IMAGES, OBS_STATE  # noqa: E402
-
-RESIZE = (IMAGE_SIZE // 2, IMAGE_SIZE // 2)  # (4, 4)
-IMG_KEY = f"{OBS_IMAGES}.laptop"
-
-
-# ---------------------------------------------------------------------------
-# Step numerics / shape handling
-# ---------------------------------------------------------------------------
-
-
-@pytest.mark.parametrize(
-    "shape",
-    [
-        (3, IMAGE_SIZE, IMAGE_SIZE),  # [C, H, W] (raw single-sample inference)
-        (BATCH_SIZE, 3, IMAGE_SIZE, IMAGE_SIZE),  # [B, C, H, W]
-        (BATCH_SIZE, 2, 3, IMAGE_SIZE, IMAGE_SIZE),  # [B, T, C, H, W] (video stack)
-    ],
-)
-def test_image_prep_resize_shapes_and_area_numerics(shape: tuple[int, ...]) -> None:
-    step = ImagePrepProcessorStep(resize_to=RESIZE)
-    x = torch.rand(*shape)
-    out = step.observation({IMG_KEY: x})[IMG_KEY]
-
-    assert out.shape[:-2] == x.shape[:-2]  # leading + channel dims unchanged
-    assert tuple(out.shape[-2:]) == RESIZE
-    assert out.dtype == torch.float32
-
-    # bit-identical to the model-side F.interpolate(mode="area"), no clamp
-    ref = F.interpolate(x.float().reshape(-1, *x.shape[-3:]), size=RESIZE, mode="area").reshape(
-        *x.shape[:-2], *RESIZE
-    )
-    assert torch.equal(out, ref)
-
-
-def test_image_prep_channel_expand() -> None:
-    step = ImagePrepProcessorStep(resize_to=None, expand_channels=True)
-    x = torch.rand(BATCH_SIZE, 1, IMAGE_SIZE, IMAGE_SIZE)
-    out = step.observation({IMG_KEY: x})[IMG_KEY]
-    assert out.shape[1] == 3
-    # all three channels are copies of the single input channel
-    assert torch.equal(out[:, 0], x[:, 0]) and torch.equal(out[:, 1], x[:, 0])
-
-
-def test_image_prep_resize_skip_when_already_target_size() -> None:
-    step = ImagePrepProcessorStep(resize_to=RESIZE)
-    x = torch.rand(BATCH_SIZE, 3, *RESIZE)
-    out = step.observation({IMG_KEY: x})[IMG_KEY]
-    # size already matches -> only the float cast happens, values preserved exactly
-    assert torch.equal(out, x)
-
-
-def test_image_prep_leaves_non_image_keys_untouched() -> None:
-    step = ImagePrepProcessorStep(resize_to=RESIZE)
-    state = torch.randn(BATCH_SIZE, 4)
-    out = step.observation({IMG_KEY: torch.rand(BATCH_SIZE, 3, IMAGE_SIZE, IMAGE_SIZE), OBS_STATE: state})
-    assert torch.equal(out[OBS_STATE], state)
-
-
-def test_image_prep_config_roundtrip_via_registry() -> None:
-    step = ImagePrepProcessorStep(resize_to=RESIZE, expand_channels=True)
-    cfg = step.get_config()
-    assert cfg == {"resize_to": [RESIZE[0], RESIZE[1]], "expand_channels": True}
-    rebuilt = ProcessorStepRegistry.get("vla_jepa_image_prep")(**cfg)
-    assert rebuilt.resize_to == RESIZE
-    assert rebuilt.expand_channels is True
-
-
-def test_image_prep_transform_features() -> None:
-    step = ImagePrepProcessorStep(resize_to=RESIZE, expand_channels=True)
-    features = {
-        PipelineFeatureType.OBSERVATION: {
-            IMG_KEY: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_SIZE, IMAGE_SIZE)),
-            "observation.images.depth": PolicyFeature(
-                type=FeatureType.VISUAL, shape=(1, IMAGE_SIZE, IMAGE_SIZE)
-            ),
-            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(4,)),
-        }
-    }
-    out = step.transform_features(features)[PipelineFeatureType.OBSERVATION]
-    assert out[IMG_KEY].shape == (3, *RESIZE)  # already 3-channel, only resized
-    assert out["observation.images.depth"].shape == (3, *RESIZE)  # 1->3 expanded
-    assert out[OBS_STATE].shape == (4,)  # non-image untouched
-
-
-# ---------------------------------------------------------------------------
-# Pipeline wiring + back-compat with the model
-# ---------------------------------------------------------------------------
-
-
-def test_image_prep_step_wired_into_preprocessor() -> None:
-    cfg = make_config()
-    cfg.resize_images_to = RESIZE
-    preprocessor, _ = make_vla_jepa_pre_post_processors(cfg, dataset_stats=None)
-    prep_steps = [s for s in preprocessor.steps if isinstance(s, ImagePrepProcessorStep)]
-    assert len(prep_steps) == 1
-    assert prep_steps[0].resize_to == RESIZE
-
-
-@torch.no_grad()
-@pytest.mark.parametrize("batch_fn", [make_inference_batch, make_train_batch])
-def test_image_prep_matches_model_qwen_path(patch_vla_jepa_external_models: None, batch_fn) -> None:
-    """The Qwen image path is identical whether the step resized (new ckpt) or the model does (old ckpt).
-
-    Both use F.interpolate(mode="area"), so pre-resizing in the step then letting the model's
-    size guard no-op yields byte-identical Qwen inputs to the pure model path. This is the
-    contract that keeps already-uploaded checkpoints correct.
-    """
-    cfg = make_config()
-    cfg.resize_images_to = RESIZE
-    policy = VLAJEPAPolicy(cfg)
-    policy.eval()
-    training = batch_fn is make_train_batch
-
-    batch = batch_fn()
-
-    # Path A (old checkpoint, no processor step): the model resizes internally.
-    imgs_a = policy._prepare_model_inputs(deepcopy(batch), training=training)["images"]
-
-    # Path B (new checkpoint): the step resizes first; the model's guard becomes a no-op.
-    step = ImagePrepProcessorStep(resize_to=RESIZE)
-    resized = step.observation({IMG_KEY: batch[IMG_KEY]})
-    batch_b = deepcopy(batch)
-    batch_b[IMG_KEY] = resized[IMG_KEY]
-    imgs_b = policy._prepare_model_inputs(batch_b, training=training)["images"]
-
-    assert len(imgs_a) == len(imgs_b) == BATCH_SIZE
-    for views_a, views_b in zip(imgs_a, imgs_b, strict=True):
-        for a, b in zip(views_a, views_b, strict=True):
-            assert torch.equal(a, b)

tests/scripts/test_setup_motors.py

@@ -0,0 +1,49 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+import lerobot.scripts.lerobot_setup_motors as motors_module
+
+
+def test_main_registers_plugins_before_parsing(monkeypatch):
+    calls = []
+    monkeypatch.setattr(motors_module, "register_third_party_plugins", lambda: calls.append("register"))
+    monkeypatch.setattr(motors_module, "setup_motors", lambda: calls.append("setup"))
+
+    motors_module.main()
+
+    assert calls == ["register", "setup"]
+
+
+def test_setup_motors_accepts_third_party_device(monkeypatch):
+    device = MagicMock()
+    monkeypatch.setattr(motors_module, "make_teleoperator_from_config", lambda _: device)
+    cfg = SimpleNamespace(device=SimpleNamespace(type="third_party"))
+
+    motors_module.setup_motors.__wrapped__(cfg)
+
+    device.setup_motors.assert_called_once_with()
+
+
+def test_setup_motors_reports_unsupported_device(monkeypatch):
+    device = object()
+    monkeypatch.setattr(motors_module, "make_teleoperator_from_config", lambda _: device)
+    cfg = SimpleNamespace(device=SimpleNamespace(type="third_party"))
+
+    with pytest.raises(NotImplementedError, match="third_party"):
+        motors_module.setup_motors.__wrapped__(cfg)