docs(rtc): use 15-step training delay

* feat(pi05): add continuous 6D SE3 actions

* fix(processors): align rotation 6D with UMI rows

docs/source/annotation_pipeline.mdx

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

docs/source/bring_your_own_policies.mdx

@@ -150,14 +150,14 @@ class MyPolicy(PreTrainedPolicy):
 
 The methods called by the train/eval loops:
 
-| Method                                                            | Used by           | What it does                                                                                                                                                                                                                                                                                 |
-| ----------------------------------------------------------------- | ----------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
-| `reset() -> None`                                                 | `lerobot-eval`    | Clear per-episode state at the start of each episode.                                                                                                                                                                                                                                        |
-| `select_action(batch, **kwargs) -> Tensor`                        | `lerobot-eval`    | Return the next action `(B, action_dim)`. Called every step.                                                                                                                                                                                                                                 |
-| `predict_action_chunk(batch, **kwargs) -> Tensor`                 | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk.                                                                                                                                       |
-| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train`   | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting.                                                                                                                                                                                         |
-| `get_optim_params() -> dict`                                      | the optimizer     | Return `self.parameters()` for simple policies; return a named parameter dict for multi-optimizer policies (see `get_optim_params` in [`modeling_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/modeling_act.py) for a per-group learning-rate example). |
-| `update() -> None` _(optional)_                                   | `lerobot-train`   | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this).                                                                                                                                                                                   |
+| Method                                                            | Used by           | What it does                                                                                                                                                                                                                                         |
+| ----------------------------------------------------------------- | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `reset() -> None`                                                 | `lerobot-eval`    | Clear per-episode state at the start of each episode.                                                                                                                                                                                                |
+| `select_action(batch, **kwargs) -> Tensor`                        | `lerobot-eval`    | Return the next action `(B, action_dim)`. Called every step.                                                                                                                                                                                         |
+| `predict_action_chunk(batch, **kwargs) -> Tensor`                 | the policy itself | Return an action chunk `(B, chunk_size, action_dim)`. Currently abstract on the base class — raise `NotImplementedError` if your policy doesn't chunk.                                                                                               |
+| `forward(batch, reduction="mean") -> tuple[Tensor, dict \| None]` | `lerobot-train`   | Return `(loss, output_dict)`. Accept `reduction="none"` if you want to support per-sample weighting.                                                                                                                                                 |
+| `get_optim_params() -> dict`                                      | the optimizer     | Return `self.parameters()` for simple policies; return a named parameter dict for [multi-optimizer policies](https://github.com/huggingface/lerobot/blob/ecd38c50d7d15b4184cf42649ff1185ee2e11eeb/src/lerobot/policies/sac/modeling_sac.py#L61-L73). |
+| `update() -> None` _(optional)_                                   | `lerobot-train`   | Called after each optimizer step _if defined_. Use for EMA, target nets, replay buffers (TDMPC uses this).                                                                                                                                           |
 
 Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.<motor>`), `OBS_IMAGES` (`observation.images.<camera>`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes.
 
@@ -295,10 +295,12 @@ The file names are load-bearing: the factory does lazy imports by name, and the
 
 ### Wiring
 
-Two places need to know about your policy. All by name.
+Four places need to know about your policy. All by name.
 
-1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
-2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
+1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast).
+2. **`factory.py:get_policy_class`** — add a branch returning `MyPolicy` from a lazy import.
+3. **`factory.py:make_policy_config`** and **`factory.py:make_pre_post_processors`** — same idea, two more branches.
+4. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page.
 
 Mirror an existing policy that's structurally similar to yours; the diff is small.
 
@@ -330,10 +332,6 @@ This way:
 
 Add a matching extra to [`pyproject.toml`](https://github.com/huggingface/lerobot/blob/main/pyproject.toml) `[project.optional-dependencies]` and include it in the `all` extra so `pip install 'lerobot[all]'` keeps installing everything.
 
-### Avoid copying a modeling file — subclass it
-
-If your policy needs to modify a backbone that already exists in `transformers` (custom conditioning, extra inputs, a swapped sub-module), **do not vendor a copy of its `modeling_*.py`**. Instead, subclass the smallest upstream unit and override only what changes. [`pi_gemma.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/pi_gemma.py) is the canonical reference: it injects AdaRMS conditioning into PaliGemma/Gemma in ~370 lines by subclassing `GemmaModel`/`PaliGemmaModel` and overriding the decoder-layer forward, instead of forking the ~2,000-line modeling file. Model surgery on a _loaded_ native model is also fine (layer truncation, tokenizer expansion, hidden-state capture — see `evo1/internvl3_embedder.py`, `eo1/modeling_eo1.py`, `groot/groot_n1_7.py` for working examples). Reviewers will ask for this pattern when a PR arrives with a copied modeling file; the only accepted exception is a model that does not exist in `transformers` at all.
-
 ### Benchmarks and a published checkpoint
 
 A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
@@ -369,7 +367,7 @@ If your policy is real-robot-only and no sim benchmark applies, swap the sim eva
 The general expectations are in [`CONTRIBUTING.md`](https://github.com/huggingface/lerobot/blob/main/CONTRIBUTING.md) and the [PR template](https://github.com/huggingface/lerobot/blob/main/.github/PULL_REQUEST_TEMPLATE.md). On top of those, reviewers will look for:
 
 - [ ] `MyPolicy` and `MyPolicyConfig` cover the surface above; `__init_subclass__` accepts the class.
-- [ ] `policies/__init__.py` re-exports the config (this registers the policy; the factory resolves modeling/processor by naming convention).
+- [ ] `factory.py` and `policies/__init__.py` are wired (lazy imports for modeling).
 - [ ] `make_my_policy_pre_post_processors` follows the naming convention.
 - [ ] Optional deps live behind a `[project.optional-dependencies]` extra and the `TYPE_CHECKING + require_package` guard.
 - [ ] `tests/policies/` updated; backward-compat artifact committed & policy-specific tests.

docs/source/inference.mdx

@@ -228,12 +228,13 @@ lerobot-rollout \
     --device=cuda
 ```
 
-| Flag                                        | Description                                                    |
-| ------------------------------------------- | -------------------------------------------------------------- |
-| `--inference.rtc.execution_horizon`         | Steps to blend with previous chunk (default: varies by policy) |
-| `--inference.rtc.max_guidance_weight`       | Consistency enforcement strength (default: varies by policy)   |
-| `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS`               |
-| `--inference.queue_threshold`               | Max queue size before backpressure (default: 30)               |
+| Flag                                        | Description                                                                     |
+| ------------------------------------------- | ------------------------------------------------------------------------------- |
+| `--inference.rtc.execution_horizon`         | Steps to blend with previous chunk (default: varies by policy)                  |
+| `--inference.rtc.mode`                      | `guided` (default) or trained-prefix `trained` for compatible Pi0.5 checkpoints |
+| `--inference.rtc.max_guidance_weight`       | Consistency enforcement strength (default: varies by policy)                    |
+| `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS`                                |
+| `--inference.queue_threshold`               | Backpressure threshold; trained RTC requires at least its maximum delay         |
 
 See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.
 

docs/source/rtc.mdx

@@ -1,6 +1,6 @@
 # Real-Time Chunking (RTC)
 
-Real-Time Chunking (RTC) is an inference-time method that allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency.
+Real-Time Chunking (RTC) allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency. LeRobot provides the original inference-time guided mode and, for compatible Pi0.5 checkpoints, training-time action conditioning with cheap hard-prefix inference.
 
 These policies generate chunks of future actions (e.g., 50 steps at a time) instead of single actions.
 Because the models are large, producing each chunk takes longer than the time it takes the robot to execute it.
@@ -92,6 +92,15 @@ for step in range(num_steps):
 
 `RTCConfig` has the following parameters to tune:
 
+**`mode`** selects the action-prefix conditioning method:
+
+- `guided` (default) applies the original Jacobian guidance during denoising and works with ordinary flow-matching checkpoints.
+- `trained` hard-inpaints the previous chunk's prefix with per-action flow timesteps. It currently requires a Pi0.5 checkpoint trained with `policy.rtc_training_max_delay > 0` and avoids the guidance backward pass.
+
+For trained mode, both `execution_horizon` and the rollout backend's
+`inference.queue_threshold` must be at least the checkpoint's
+`rtc_training_max_delay`; rollout validates this before connecting the robot.
+
 **`execution_horizon`**: How many timesteps from the previous chunk to maintain consistency with. Higher values mean smoother transitions but potentially less reactivity.
 
 Typical values: 8-12 steps
@@ -111,6 +120,27 @@ RTCConfig(execution_horizon=10)
 
 **`inference_delay`**: How many timesteps of inference latency your system has. This is passed to `predict_action_chunk()` rather than the config, since it may vary at runtime.
 
+## Training Pi0.5 for Trained RTC
+
+Set the maximum prefix delay when fine-tuning Pi0.5:
+
+```bash
+lerobot-train \
+    --policy.type=pi05 \
+    --policy.pretrained_path=lerobot/pi05_base \
+    --policy.rtc_training_max_delay=15 \
+    --dataset.repo_id=${HF_USERNAME}/dataset_repo_id \
+    --output_dir=outputs/pi05_training_rtc
+```
+
+Each example samples a clean prefix from zero through the configured maximum;
+the flow loss is computed only on the remaining postfix. Choose the maximum as
+approximately `ceil(p95 end-to-end inference latency * control frequency)` and
+keep it smaller than `policy.chunk_size`. Setting the value to zero preserves
+ordinary Pi0.5 training and existing checkpoints remain compatible with guided
+RTC. At a 50 Hz control rate, 15 steps cover up to 300 ms of end-to-end
+inference latency.
+
 ## Testing RTC Offline
 
 Before running on a real robot, test RTC with dataset samples to visualize how it works:
@@ -124,6 +154,10 @@ python examples/rtc/eval_dataset.py \
     --device=cuda
 ```
 
+Add `--rtc.mode=trained` when evaluating a compatible training-time RTC Pi0.5
+checkpoint. Unsupported policies reject trained mode instead of falling back to
+guided RTC.
+
 The script generates a visualization of the denoising process, comparing standard generation (left) with RTC (right). In the RTC plots, you can see how the first few steps (blue/purple lines) are guided to match the red ground truth trajectory (previous chunk's tail), ensuring a smooth transition between chunks.
 
 <p align="center">
@@ -141,6 +175,7 @@ lerobot-rollout \
     --strategy.type=base \
     --policy.path=${HF_USERNAME}/policy_repo_id \
     --inference.type=rtc \
+    --inference.rtc.mode=guided \
     --inference.rtc.execution_horizon=10 \
     --inference.rtc.max_guidance_weight=10.0 \
     --robot.type=so100_follower \
@@ -151,6 +186,25 @@ lerobot-rollout \
     --device=cuda
 ```
 
+For a training-time RTC Pi0.5 checkpoint, change the mode to `trained`. The
+checkpoint records its maximum supported delay, and rollout validates measured
+latency against it:
+
+```bash
+lerobot-rollout \
+    --strategy.type=base \
+    --policy.path=${HF_USERNAME}/pi05_training_rtc \
+    --inference.type=rtc \
+    --inference.rtc.mode=trained \
+    --inference.rtc.execution_horizon=15 \
+    --inference.queue_threshold=15 \
+    --robot.type=so100_follower \
+    --robot.port=/dev/tty.usbmodem58FA0834591 \
+    --task="Move green small object into the purple platform" \
+    --duration=120 \
+    --device=cuda
+```
+
 ## How It Differs from the Async Inference in LeRobot
 
 Both RTC and [async inference](./async) improve real-time robot control, but they solve different problems.
@@ -189,3 +243,5 @@ See `examples/rtc/eval_dataset.py` for a complete example of offline RTC visuali
 - [Smooth-As-Butter Robot Policies](https://alexander-soare.github.io/robotics/2025/08/05/smooth-as-butter-robot-policies.html) - Excellent technical explanation with real robot results
 - [Physical Intelligence - Real-Time Chunking](https://www.physicalintelligence.company/research/real_time_chunking) - Original paper and research
 - [Kinetix RTC Implementation](https://github.com/Physical-Intelligence/real-time-chunking-kinetix) - Reference implementation from Physical Intelligence
+- [Training-Time Action Conditioning](https://arxiv.org/abs/2512.05964) - Efficient RTC with clean-prefix conditioning during training
+- [RLDX-1](https://github.com/RLWRLD/RLDX-1) - PyTorch reference used for the training-time RTC integration

examples/annotations/run_hf_job.py

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

examples/rtc/eval_dataset.py

@@ -306,6 +306,7 @@ class RTCEvaluator:
         # Configure RTC
         rtc_config = RTCConfig(
             enabled=rtc_enabled,
+            mode=self.cfg.rtc.mode,
             execution_horizon=self.cfg.rtc.execution_horizon,
             max_guidance_weight=self.cfg.rtc.max_guidance_weight,
             prefix_attention_schedule=self.cfg.rtc.prefix_attention_schedule,

src/lerobot/annotations/steerable_pipeline/config.py

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

src/lerobot/annotations/steerable_pipeline/frames.py

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

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

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

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

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

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

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

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

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

src/lerobot/annotations/steerable_pipeline/vlm_client.py

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

src/lerobot/configs/policies.py

@@ -205,30 +205,24 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC):  # type: igno
                     f"{CONFIG_NAME} not found on the HuggingFace Hub in {model_id}"
                 ) from e
 
+        # HACK: Parse the original config to get the config subclass, so that we can
+        # apply cli overrides.
+        # This is very ugly, ideally we'd like to be able to do that natively with draccus
+        # something like --policy.path (in addition to --policy.type)
+        with draccus.config_type("json"):
+            orig_config = draccus.parse(cls, config_file, args=[])
+
         if config_file is None:
             raise FileNotFoundError(f"{CONFIG_NAME} not found in {model_id}")
 
         with open(config_file) as f:
             config = json.load(f)
 
-        # Resolve the concrete config subclass from the serialized "type" tag, then parse
-        # the config (with CLI overrides) directly for that class. The "type" key is
-        # stripped because draccus only consumes it when parsing the registry base class.
-        policy_type = config.pop("type", None)
-        if policy_type is None:
-            raise ValueError(f"Missing 'type' field in {CONFIG_NAME} of {model_id}")
-        try:
-            config_cls = cls.get_choice_class(policy_type)
-        except Exception as e:
-            raise ValueError(
-                f"Policy type '{policy_type}' (from {CONFIG_NAME} of {model_id}) is not registered. "
-                f"Available policy types: {cls.get_known_choices()}"
-            ) from e
-
+        config.pop("type")
         with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f:
             json.dump(config, f)
             config_file = f.name
 
         cli_overrides = policy_kwargs.pop("cli_overrides", [])
         with draccus.config_type("json"):
-            return draccus.parse(config_cls, config_file, args=cli_overrides)
+            return draccus.parse(orig_config.__class__, config_file, args=cli_overrides)

src/lerobot/datasets/compute_stats.py

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

src/lerobot/datasets/dataset_tools.py

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

src/lerobot/policies/__init__.py

@@ -32,7 +32,6 @@ from .pretrained import PreTrainedPolicy as PreTrainedPolicy
 from .smolvla.configuration_smolvla import SmolVLAConfig as SmolVLAConfig
 from .tdmpc.configuration_tdmpc import TDMPCConfig as TDMPCConfig
 from .utils import make_robot_action, prepare_observation_for_inference
-from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig as VLAJEPAConfig
 from .vqbet.configuration_vqbet import VQBeTConfig as VQBeTConfig
 from .wall_x.configuration_wall_x import WallXConfig as WallXConfig
 from .xvla.configuration_xvla import XVLAConfig as XVLAConfig
@@ -58,7 +57,6 @@ __all__ = [
     "PI05Config",
     "SmolVLAConfig",
     "TDMPCConfig",
-    "VLAJEPAConfig",
     "VQBeTConfig",
     "WallXConfig",
     "XVLAConfig",

src/lerobot/policies/act/processor_act.py

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

src/lerobot/policies/diffusion/processor_diffusion.py

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

src/lerobot/policies/eo1/processor_eo1.py

@@ -23,16 +23,24 @@ import torch
 
 from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
+    AddBatchDimensionProcessorStep,
     ComplementaryDataProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
-    make_default_policy_processor_steps,
-    make_policy_processor_pipelines,
+    RenameObservationsProcessorStep,
+    UnnormalizerProcessorStep,
 )
+from lerobot.processor.converters import policy_action_to_transition, transition_to_policy_action
 from lerobot.types import TransitionKey
-from lerobot.utils.constants import OBS_STATE
+from lerobot.utils.constants import (
+    OBS_STATE,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
 from lerobot.utils.import_utils import _transformers_available, require_package
 
 from .configuration_eo1 import EO1Config
@@ -234,12 +242,14 @@ def make_eo1_pre_post_processors(
 ]:
     """Build pre/post processor pipelines for EO1."""
 
-    steps = make_default_policy_processor_steps(config, dataset_stats)
-
     input_steps: list[ProcessorStep] = [
-        steps.rename_observations,
-        steps.add_batch_dim,
-        steps.normalize,
+        RenameObservationsProcessorStep(rename_map={}),
+        AddBatchDimensionProcessorStep(),
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
         EO1ConversationTemplateStep(input_features=config.input_features, chunk_size=config.chunk_size),
         EO1QwenProcessorStep(
             processor_name=config.vlm_base,
@@ -247,12 +257,27 @@ def make_eo1_pre_post_processors(
             image_max_pixels=config.image_max_pixels,
             use_fast_processor=config.use_fast_processor,
         ),
-        steps.to_device,
+        DeviceProcessorStep(device=config.device),
     ]
 
     output_steps: list[ProcessorStep] = [
-        steps.unnormalize,
-        steps.to_cpu,
+        UnnormalizerProcessorStep(
+            features=config.output_features,
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
+        DeviceProcessorStep(device="cpu"),
     ]
 
-    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/evo1/internvl3_embedder.py

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

src/lerobot/policies/factory.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import importlib
-import inspect
 import logging
 from typing import TYPE_CHECKING, Any, TypedDict, Unpack
 
@@ -45,10 +44,26 @@ from lerobot.utils.constants import (
 )
 from lerobot.utils.feature_utils import dataset_to_policy_features
 
+from .act.configuration_act import ACTConfig
+from .diffusion.configuration_diffusion import DiffusionConfig
+from .eo1.configuration_eo1 import EO1Config
 from .evo1.configuration_evo1 import Evo1Config
+from .fastwam.configuration_fastwam import FastWAMConfig
+from .gaussian_actor.configuration_gaussian_actor import GaussianActorConfig
 from .groot.configuration_groot import GrootConfig
+from .lingbot_va.configuration_lingbot_va import LingBotVAConfig
+from .molmoact2.configuration_molmoact2 import MolmoAct2Config
+from .multi_task_dit.configuration_multi_task_dit import MultiTaskDiTConfig
+from .pi0.configuration_pi0 import PI0Config
+from .pi05.configuration_pi05 import PI05Config
 from .pretrained import PreTrainedPolicy
+from .smolvla.configuration_smolvla import SmolVLAConfig
+from .tdmpc.configuration_tdmpc import TDMPCConfig
 from .utils import validate_visual_features_consistency
+from .vla_jepa.configuration_vla_jepa import VLAJEPAConfig
+from .vqbet.configuration_vqbet import VQBeTConfig
+from .wall_x.configuration_wall_x import WallXConfig
+from .xvla.configuration_xvla import XVLAConfig
 
 
 def _reconnect_relative_absolute_steps(
@@ -73,23 +88,100 @@ def get_policy_class(name: str) -> type[PreTrainedPolicy]:
     """
     Retrieves a policy class by its registered name.
 
-    Resolution is convention-based: the draccus-registered config class of ``name`` is
-    looked up, its ``configuration_*`` module path is rewritten to ``modeling_*``, and
-    the ``<X>Policy`` class is imported from there. The modeling module is only imported
-    at call time, keeping heavy optional dependencies lazy. This works for both built-in
-    policies and third-party lerobot plugins (anything registered via
-    ``@PreTrainedConfig.register_subclass``).
+    This function uses dynamic imports to avoid loading all policy classes into memory
+    at once, improving startup time and reducing dependencies.
 
     Args:
-        name: The registered name of the policy (e.g. "act", "diffusion", "pi0").
+        name: The name of the policy. Supported names are "tdmpc", "diffusion", "act",
+            "multi_task_dit", "vqbet", "pi0", "pi05", "gaussian_actor", "smolvla", "wall_x",
+            "molmoact2", "eo1", "evo1".
     Returns:
         The policy class corresponding to the given name.
 
     Raises:
-        ValueError: If the policy name is not registered.
-        ImportError: If the policy's optional dependencies are not installed.
+        NotImplementedError: If the policy name is not recognized.
     """
-    return _get_policy_cls_from_policy_name(name=name)
+    if name == "tdmpc":
+        from .tdmpc.modeling_tdmpc import TDMPCPolicy
+
+        return TDMPCPolicy
+    elif name == "diffusion":
+        from .diffusion.modeling_diffusion import DiffusionPolicy
+
+        return DiffusionPolicy
+    elif name == "act":
+        from .act.modeling_act import ACTPolicy
+
+        return ACTPolicy
+    elif name == "multi_task_dit":
+        from .multi_task_dit.modeling_multi_task_dit import MultiTaskDiTPolicy
+
+        return MultiTaskDiTPolicy
+    elif name == "vqbet":
+        from .vqbet.modeling_vqbet import VQBeTPolicy
+
+        return VQBeTPolicy
+    elif name == "pi0":
+        from .pi0.modeling_pi0 import PI0Policy
+
+        return PI0Policy
+    elif name == "pi0_fast":
+        from .pi0_fast.modeling_pi0_fast import PI0FastPolicy
+
+        return PI0FastPolicy
+    elif name == "pi05":
+        from .pi05.modeling_pi05 import PI05Policy
+
+        return PI05Policy
+    elif name == "gaussian_actor":
+        from .gaussian_actor.modeling_gaussian_actor import GaussianActorPolicy
+
+        return GaussianActorPolicy
+    elif name == "smolvla":
+        from .smolvla.modeling_smolvla import SmolVLAPolicy
+
+        return SmolVLAPolicy
+    elif name == "groot":
+        from .groot.modeling_groot import GrootPolicy
+
+        return GrootPolicy
+    elif name == "xvla":
+        from .xvla.modeling_xvla import XVLAPolicy
+
+        return XVLAPolicy
+    elif name == "wall_x":
+        from .wall_x.modeling_wall_x import WallXPolicy
+
+        return WallXPolicy
+    elif name == "eo1":
+        from .eo1.modeling_eo1 import EO1Policy
+
+        return EO1Policy
+    elif name == "molmoact2":
+        from .molmoact2.modeling_molmoact2 import MolmoAct2Policy
+
+        return MolmoAct2Policy
+    elif name == "vla_jepa":
+        from .vla_jepa.modeling_vla_jepa import VLAJEPAPolicy
+
+        return VLAJEPAPolicy
+    elif name == "lingbot_va":
+        from .lingbot_va.modeling_lingbot_va import LingBotVAPolicy
+
+        return LingBotVAPolicy
+    elif name == "fastwam":
+        from .fastwam.modeling_fastwam import FastWAMPolicy
+
+        return FastWAMPolicy
+    elif name == "evo1":
+        from .evo1.modeling_evo1 import Evo1Policy
+
+        return Evo1Policy
+    else:
+        try:
+            return _get_policy_cls_from_policy_name(name=name)
+        except Exception as e:
+            raise ValueError(f"Policy type '{name}' is not available.") from e
 
 
 def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
@@ -100,8 +192,9 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
     mapping a string identifier to the corresponding config class.
 
     Args:
-        policy_type: The registered type of the policy (any name registered via
-                     ``@PreTrainedConfig.register_subclass``, e.g. "act", "diffusion", "pi0").
+        policy_type: The type of the policy. Supported types include "tdmpc",
+                     "multi_task_dit", "diffusion", "act", "vqbet", "pi0", "pi05", "gaussian_actor",
+                     "smolvla", "wall_x", "molmoact2", "eo1", "evo1".
         **kwargs: Keyword arguments to be passed to the configuration class constructor.
 
     Returns:
@@ -110,11 +203,48 @@ def make_policy_config(policy_type: str, **kwargs) -> PreTrainedConfig:
     Raises:
         ValueError: If the `policy_type` is not recognized.
     """
-    try:
-        config_cls = PreTrainedConfig.get_choice_class(policy_type)
-    except Exception as e:
-        raise ValueError(f"Policy type '{policy_type}' is not available.") from e
-    return config_cls(**kwargs)
+    if policy_type == "tdmpc":
+        return TDMPCConfig(**kwargs)
+    elif policy_type == "diffusion":
+        return DiffusionConfig(**kwargs)
+    elif policy_type == "act":
+        return ACTConfig(**kwargs)
+    elif policy_type == "multi_task_dit":
+        return MultiTaskDiTConfig(**kwargs)
+    elif policy_type == "vqbet":
+        return VQBeTConfig(**kwargs)
+    elif policy_type == "pi0":
+        return PI0Config(**kwargs)
+    elif policy_type == "pi05":
+        return PI05Config(**kwargs)
+    elif policy_type == "gaussian_actor":
+        return GaussianActorConfig(**kwargs)
+    elif policy_type == "smolvla":
+        return SmolVLAConfig(**kwargs)
+    elif policy_type == "groot":
+        return GrootConfig(**kwargs)
+    elif policy_type == "xvla":
+        return XVLAConfig(**kwargs)
+    elif policy_type == "wall_x":
+        return WallXConfig(**kwargs)
+    elif policy_type == "eo1":
+        return EO1Config(**kwargs)
+    elif policy_type == "molmoact2":
+        return MolmoAct2Config(**kwargs)
+    elif policy_type == "vla_jepa":
+        return VLAJEPAConfig(**kwargs)
+    elif policy_type == "lingbot_va":
+        return LingBotVAConfig(**kwargs)
+    elif policy_type == "fastwam":
+        return FastWAMConfig(**kwargs)
+    elif policy_type == "evo1":
+        return Evo1Config(**kwargs)
+    else:
+        try:
+            config_cls = PreTrainedConfig.get_choice_class(policy_type)
+            return config_cls(**kwargs)
+        except Exception as e:
+            raise ValueError(f"Policy type '{policy_type}' is not available.") from e
 
 
 class ProcessorConfigKwargs(TypedDict, total=False):
@@ -168,7 +298,8 @@ def make_pre_post_processors(
         A tuple containing the input (pre-processor) and output (post-processor) pipelines.
 
     Raises:
-        ValueError: If no processor factory exists for the given policy configuration type.
+        NotImplementedError: If a processor factory is not implemented for the given
+            policy configuration type.
     """
     if pretrained_path:
         if isinstance(policy_cfg, GrootConfig):
@@ -220,13 +351,166 @@ def make_pre_post_processors(
             )
         return preprocessor, postprocessor
 
-    # Create new processors from the policy config, resolving the per-policy factory
-    # function by naming convention (lazy import keeps optional dependencies optional).
-    return _make_processors_from_policy_config(
-        config=policy_cfg,
-        dataset_stats=kwargs.get("dataset_stats"),
-        dataset_meta=kwargs.get("dataset_meta"),
-    )
+    # Create a new processor based on policy type
+    if isinstance(policy_cfg, TDMPCConfig):
+        from .tdmpc.processor_tdmpc import make_tdmpc_pre_post_processors
+
+        processors = make_tdmpc_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, DiffusionConfig):
+        from .diffusion.processor_diffusion import make_diffusion_pre_post_processors
+
+        processors = make_diffusion_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, ACTConfig):
+        from .act.processor_act import make_act_pre_post_processors
+
+        processors = make_act_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, MultiTaskDiTConfig):
+        from .multi_task_dit.processor_multi_task_dit import (
+            make_multi_task_dit_pre_post_processors,
+        )
+
+        processors = make_multi_task_dit_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, VQBeTConfig):
+        from .vqbet.processor_vqbet import make_vqbet_pre_post_processors
+
+        processors = make_vqbet_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, PI0Config):
+        from .pi0.processor_pi0 import make_pi0_pre_post_processors
+
+        processors = make_pi0_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, PI05Config):
+        from .pi05.processor_pi05 import make_pi05_pre_post_processors
+
+        processors = make_pi05_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, GaussianActorConfig):
+        from .gaussian_actor.processor_gaussian_actor import make_gaussian_actor_pre_post_processors
+
+        processors = make_gaussian_actor_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, SmolVLAConfig):
+        from .smolvla.processor_smolvla import make_smolvla_pre_post_processors
+
+        processors = make_smolvla_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, GrootConfig):
+        from .groot.processor_groot import make_groot_pre_post_processors
+
+        processors = make_groot_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+            dataset_meta=kwargs.get("dataset_meta"),
+        )
+
+    elif isinstance(policy_cfg, XVLAConfig):
+        from .xvla.processor_xvla import (
+            make_xvla_pre_post_processors,
+        )
+
+        processors = make_xvla_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, WallXConfig):
+        from .wall_x.processor_wall_x import make_wall_x_pre_post_processors
+
+        processors = make_wall_x_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, EO1Config):
+        from .eo1.processor_eo1 import make_eo1_pre_post_processors
+
+        processors = make_eo1_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+    elif isinstance(policy_cfg, Evo1Config):
+        from .evo1.processor_evo1 import make_evo1_pre_post_processors
+
+        processors = make_evo1_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, MolmoAct2Config):
+        from .molmoact2.processor_molmoact2 import make_molmoact2_pre_post_processors
+
+        processors = make_molmoact2_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+            dataset_meta=kwargs.get("dataset_meta"),
+        )
+
+    elif isinstance(policy_cfg, VLAJEPAConfig):
+        from .vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
+
+        processors = make_vla_jepa_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, LingBotVAConfig):
+        from .lingbot_va.processor_lingbot_va import make_lingbot_va_pre_post_processors
+
+        processors = make_lingbot_va_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    elif isinstance(policy_cfg, FastWAMConfig):
+        from .fastwam.processor_fastwam import make_fastwam_pre_post_processors
+
+        processors = make_fastwam_pre_post_processors(
+            config=policy_cfg,
+            dataset_stats=kwargs.get("dataset_stats"),
+        )
+
+    else:
+        try:
+            processors = _make_processors_from_policy_config(
+                config=policy_cfg,
+                dataset_stats=kwargs.get("dataset_stats"),
+            )
+        except Exception as e:
+            raise ValueError(f"Processor for policy type '{policy_cfg.type}' is not implemented.") from e
+
+    return processors
 
 
 def make_policy(
@@ -370,12 +654,10 @@ def make_policy(
     return policy
 
 
-def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
+def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedConfig]:
     """Get policy class from its registered name using dynamic imports.
 
-    Works for built-in policies and 3rd party lerobot plugins alike: the config class
-    registered under ``name`` is resolved via the draccus ChoiceRegistry, and the policy
-    class is imported from the sibling ``modeling_*`` module by naming convention.
+    This is used as a helper function to import policies from 3rd party lerobot plugins.
 
     Args:
         name: The name of the policy.
@@ -401,39 +683,22 @@ def _get_policy_cls_from_policy_name(name: str) -> type[PreTrainedPolicy]:
         "configuration_", "modeling_"
     )  # e.g., configuration_diffusion -> modeling_diffusion
 
-    try:
-        module = importlib.import_module(module_path)
-    except ModuleNotFoundError as e:
-        if e.name == module_path:
-            # The modeling_* module itself does not exist for this policy type. A missing
-            # optional dependency inside an existing module propagates unchanged instead,
-            # so its actionable install hint stays visible.
-            raise ValueError(f"Policy class for '{name}' is not implemented.") from e
-        raise
-    policy_cls = getattr(module, cls_name, None)
-    if policy_cls is None:
-        raise ValueError(
-            f"Policy class '{cls_name}' not found in '{module_path}'. "
-            f"Policies must expose '<Name>Policy' in the sibling 'modeling_*' module by naming convention."
-        )
+    module = importlib.import_module(module_path)
+    policy_cls = getattr(module, cls_name)
     return policy_cls
 
 
 def _make_processors_from_policy_config(
     config: PreTrainedConfig,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    dataset_meta: Any | None = None,
 ) -> tuple[Any, Any]:
     """Create pre- and post-processors from a policy configuration using dynamic imports.
 
-    Resolves ``make_{type}_pre_post_processors`` from the policy's ``processor_*`` module
-    by naming convention. Works for built-in policies and 3rd party lerobot plugins.
+    This is used as a helper function to import processor factories from 3rd party lerobot plugins.
 
     Args:
         config: The policy configuration object.
         dataset_stats: Dataset statistics for normalization.
-        dataset_meta: Dataset metadata, forwarded only to factories that declare a
-            ``dataset_meta`` parameter (e.g. groot, molmoact2).
     Returns:
         A tuple containing the input (pre-processor) and output (post-processor) pipelines.
     """
@@ -446,19 +711,6 @@ def _make_processors_from_policy_config(
     logging.debug(
         f"Instantiating pre/post processors using function '{function_name}' from module '{module_path}'"
     )
-    try:
-        module = importlib.import_module(module_path)
-    except ModuleNotFoundError as e:
-        if e.name == module_path:
-            # The processor_* module itself does not exist for this policy type. A missing
-            # optional dependency inside an existing module propagates unchanged instead,
-            # so its actionable install hint stays visible.
-            raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.") from e
-        raise
-    function = getattr(module, function_name, None)
-    if function is None:
-        raise ValueError(f"Processor for policy type '{policy_type}' is not implemented.")
-    call_kwargs: dict[str, Any] = {"dataset_stats": dataset_stats}
-    if "dataset_meta" in inspect.signature(function).parameters:
-        call_kwargs["dataset_meta"] = dataset_meta
-    return function(config, **call_kwargs)
+    module = importlib.import_module(module_path)
+    function = getattr(module, function_name)
+    return function(config, dataset_stats=dataset_stats)

src/lerobot/policies/fastwam/processor_fastwam.py

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

src/lerobot/policies/gaussian_actor/processor_gaussian_actor.py

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

src/lerobot/policies/lingbot_va/processor_lingbot_va.py

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

src/lerobot/policies/multi_task_dit/processor_multi_task_dit.py

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

src/lerobot/policies/pi0/processor_pi0.py

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

src/lerobot/policies/pi05/configuration_pi05.py

@@ -55,9 +55,25 @@ class PI05Config(PreTrainedConfig):
     relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
     # Populated at runtime from dataset metadata by make_policy.
     action_feature_names: list[str] | None = None
+    # ``se3`` uses inv(T_current) @ T_target for each xyz+rotation-vector pose group.
+    # ``se3_6d`` uses the same composition and expands each relative rotation
+    # vector to the continuous first-two-row 6-D rotation representation.
+    # ``componentwise`` preserves the legacy action - state behavior.
+    relative_pose_representation: str = "componentwise"
+    relative_se3_pose_groups: list[list[int]] = field(default_factory=lambda: [list(range(6))])
+
+    # Build proprioception from absolute action samples when the dataset has no
+    # observation.state. With history_steps=2, training samples request t-1 as
+    # well as the normal t..t+chunk_size-1 action targets.
+    state_from_action: bool = False
+    proprioception_history_steps: int = 1
+    use_relative_state_history: bool = False
+    relative_state_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
 
     # Real-Time Chunking (RTC) configuration
     rtc_config: RTCConfig | None = None
+    # Maximum clean action-prefix length sampled during training. Zero disables trained RTC.
+    rtc_training_max_delay: int = 0
 
     image_resolution: tuple[int, int] = (
         DEFAULT_IMAGE_SIZE,
@@ -111,6 +127,11 @@ class PI05Config(PreTrainedConfig):
             raise ValueError(
                 f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
             )
+        if not 0 <= self.rtc_training_max_delay < self.chunk_size:
+            raise ValueError(
+                "rtc_training_max_delay must satisfy "
+                f"0 <= delay < chunk_size ({self.chunk_size}), got {self.rtc_training_max_delay}"
+            )
 
         if self.paligemma_variant not in ["gemma_300m", "gemma_2b"]:
             raise ValueError(f"Invalid paligemma_variant: {self.paligemma_variant}")
@@ -121,6 +142,25 @@ class PI05Config(PreTrainedConfig):
         if self.dtype not in ["bfloat16", "float32"]:
             raise ValueError(f"Invalid dtype: {self.dtype}")
 
+        if self.proprioception_history_steps < 1:
+            raise ValueError("proprioception_history_steps must be at least 1")
+
+        if self.relative_pose_representation not in {"componentwise", "se3", "se3_6d"}:
+            raise ValueError(
+                "relative_pose_representation must be 'componentwise', 'se3', or 'se3_6d', got "
+                f"{self.relative_pose_representation!r}"
+            )
+        for group in self.relative_se3_pose_groups:
+            if len(group) != 6 or len(set(group)) != 6 or any(index < 0 for index in group):
+                raise ValueError(f"Invalid six-index SE(3) pose group: {group}")
+            if self.relative_pose_representation == "se3_6d" and group != list(range(group[0], group[0] + 6)):
+                raise ValueError("se3_6d pose groups must contain six contiguous ascending indices")
+        if self.relative_pose_representation in {"se3", "se3_6d"} and not self.relative_se3_pose_groups:
+            raise ValueError(
+                f"relative_pose_representation={self.relative_pose_representation!r} "
+                "requires relative_se3_pose_groups"
+            )
+
     def validate_features(self) -> None:
         """Validate and set up input/output features."""
         for i in range(self.empty_cameras):
@@ -131,19 +171,54 @@ class PI05Config(PreTrainedConfig):
             )
             self.input_features[key] = empty_camera
 
-        if OBS_STATE not in self.input_features:
-            state_feature = PolicyFeature(
-                type=FeatureType.STATE,
-                shape=(self.max_state_dim,),  # Padded to max_state_dim
-            )
-            self.input_features[OBS_STATE] = state_feature
-
         if ACTION not in self.output_features:
             action_feature = PolicyFeature(
                 type=FeatureType.ACTION,
                 shape=(self.max_action_dim,),  # Padded to max_action_dim
             )
             self.output_features[ACTION] = action_feature
+        elif self.relative_pose_representation == "se3_6d":
+            action_feature = self.output_features[ACTION]
+            source_dim = (
+                len(self.action_feature_names)
+                if self.action_feature_names is not None
+                else action_feature.shape[-1]
+            )
+            model_dim = source_dim + 3 * len(self.relative_se3_pose_groups)
+            if action_feature.shape[-1] == source_dim:
+                self.output_features[ACTION] = PolicyFeature(
+                    type=action_feature.type,
+                    shape=(model_dim,),
+                )
+            elif action_feature.shape[-1] != model_dim:
+                raise ValueError(
+                    "se3_6d action feature has incompatible width: "
+                    f"source={source_dim}, expected model width={model_dim}, "
+                    f"got={action_feature.shape[-1]}"
+                )
+            if model_dim > self.max_action_dim:
+                raise ValueError(
+                    f"se3_6d action width {model_dim} exceeds max_action_dim={self.max_action_dim}"
+                )
+
+        if OBS_STATE not in self.input_features:
+            state_shape = (self.max_state_dim,)
+            if self.state_from_action and ACTION in self.output_features:
+                state_shape = self.output_features[ACTION].shape
+            state_feature = PolicyFeature(
+                type=FeatureType.STATE,
+                shape=state_shape,
+            )
+            self.input_features[OBS_STATE] = state_feature
+
+        state_dim = self.input_features[OBS_STATE].shape[-1]
+        history_state_dim = state_dim * self.proprioception_history_steps
+        if history_state_dim > self.max_state_dim:
+            raise ValueError(
+                "Flattened proprioception history exceeds max_state_dim: "
+                f"{state_dim} * {self.proprioception_history_steps} = {history_state_dim} > "
+                f"{self.max_state_dim}"
+            )
 
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(
@@ -168,7 +243,8 @@ class PI05Config(PreTrainedConfig):
 
     @property
     def action_delta_indices(self) -> list:
-        return list(range(self.chunk_size))
+        history_prefix = self.proprioception_history_steps - 1 if self.state_from_action else 0
+        return list(range(-history_prefix, self.chunk_size))
 
     @property
     def reward_delta_indices(self) -> None:

src/lerobot/policies/pi05/modeling_pi05.py

@@ -66,6 +66,107 @@ class ActionSelectKwargs(TypedDict, total=False):
     execution_horizon: int | None
 
 
+def _prepare_trained_rtc_prefix(
+    x_t: Tensor,
+    prev_chunk_left_over: Tensor | None,
+    inference_delay: int,
+    training_max_delay: int,
+) -> tuple[Tensor | None, Tensor | None]:
+    """Pad and validate a hard prefix for training-time RTC inference."""
+    if prev_chunk_left_over is None or inference_delay <= 0:
+        return None, None
+    if training_max_delay <= 0:
+        raise ValueError(
+            "RTC mode='trained' requires a checkpoint trained with policy.rtc_training_max_delay > 0."
+        )
+    if inference_delay > training_max_delay:
+        raise ValueError(
+            f"Measured RTC inference delay ({inference_delay}) exceeds the checkpoint's "
+            f"rtc_training_max_delay ({training_max_delay})."
+        )
+    if inference_delay >= x_t.shape[1]:
+        raise ValueError(
+            f"RTC inference delay ({inference_delay}) must be smaller than chunk_size ({x_t.shape[1]})."
+        )
+
+    previous = prev_chunk_left_over.to(device=x_t.device, dtype=x_t.dtype)
+    if not torch.isfinite(previous).all():
+        raise ValueError("RTC prefix contains NaN or Inf values.")
+    if previous.ndim == 2:
+        previous = previous.unsqueeze(0)
+    if previous.ndim != 3:
+        raise ValueError(f"Expected RTC prefix shape (B, T, A), got {tuple(previous.shape)}")
+    if previous.shape[0] == 1 and x_t.shape[0] > 1:
+        previous = previous.expand(x_t.shape[0], -1, -1)
+    if previous.shape[0] != x_t.shape[0]:
+        raise ValueError(
+            f"RTC prefix batch size ({previous.shape[0]}) does not match policy batch ({x_t.shape[0]})."
+        )
+    if previous.shape[1] < inference_delay:
+        raise ValueError(f"RTC prefix has {previous.shape[1]} steps, but inference_delay={inference_delay}.")
+    if previous.shape[2] > x_t.shape[2]:
+        raise ValueError(
+            f"RTC prefix action dimension ({previous.shape[2]}) exceeds model dimension ({x_t.shape[2]})."
+        )
+
+    padded_prefix = torch.zeros_like(x_t)
+    padded_prefix[:, :inference_delay, : previous.shape[2]] = previous[:, :inference_delay]
+    prefix_mask = torch.arange(x_t.shape[1], device=x_t.device) < inference_delay
+    prefix_mask = prefix_mask[None, :, None].expand(x_t.shape[0], -1, x_t.shape[2])
+    return padded_prefix, prefix_mask
+
+
+def _sample_training_rtc_prefix_mask(
+    batch_size: int,
+    action_horizon: int,
+    max_delay: int,
+    device: torch.device,
+) -> Tensor | None:
+    """Sample a clean action-prefix length independently for each training example."""
+    if max_delay <= 0:
+        return None
+    delays = torch.randint(0, max_delay + 1, (batch_size,), device=device)
+    positions = torch.arange(action_horizon, device=device)
+    return positions.unsqueeze(0) < delays.unsqueeze(1)
+
+
+def _build_flow_matching_inputs(
+    actions: Tensor,
+    noise: Tensor,
+    time: Tensor,
+    prefix_mask: Tensor | None,
+) -> tuple[Tensor, Tensor]:
+    """Keep the sampled RTC prefix clean while noising the remaining action chunk."""
+    if prefix_mask is None:
+        model_time = time
+        expanded_time = time[:, None, None]
+    else:
+        model_time = time[:, None].expand_as(prefix_mask)
+        model_time = torch.where(prefix_mask, torch.zeros_like(model_time), model_time)
+        expanded_time = model_time.unsqueeze(-1)
+    x_t = expanded_time * noise + (1 - expanded_time) * actions
+    return x_t, model_time
+
+
+def _reduce_training_rtc_loss(
+    losses: Tensor,
+    prefix_mask: Tensor | None,
+    reduction: str,
+) -> Tensor:
+    """Average flow loss over predicted postfix actions, excluding the clean RTC prefix."""
+    if reduction not in {"mean", "none"}:
+        raise ValueError(f"Unsupported loss reduction: {reduction!r}")
+    if prefix_mask is None:
+        return losses.mean() if reduction == "mean" else losses.mean(dim=(1, 2))
+
+    postfix_mask = (~prefix_mask).unsqueeze(-1).expand_as(losses)
+    if reduction == "none":
+        numerator = (losses * postfix_mask).sum(dim=(1, 2))
+        denominator = postfix_mask.sum(dim=(1, 2))
+        return numerator / denominator.clamp(min=1)
+    return (losses * postfix_mask).sum() / postfix_mask.sum().clamp(min=1)
+
+
 def get_safe_dtype(target_dtype, device_type):
     """Get a safe dtype for the given device type."""
     if device_type == "mps" and target_dtype == torch.float64:
@@ -82,21 +183,20 @@ def get_safe_dtype(target_dtype, device_type):
 def create_sinusoidal_pos_embedding(  # see openpi `create_sinusoidal_pos_embedding` (exact copy)
     time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu"
 ) -> Tensor:
-    """Computes sine-cosine positional embedding vectors for scalar positions."""
+    """Compute sine-cosine embeddings for scalar or per-action 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, )`.")
+    if time.ndim not in (1, 2):
+        raise ValueError("The time tensor must have shape (batch_size,) or (batch_size, action_horizon).")
 
     dtype = get_safe_dtype(torch.float64, device.type)
     fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
     period = min_period * (max_period / min_period) ** fraction
 
-    # Compute the outer product
     scaling_factor = 1.0 / period * 2 * math.pi
-    sin_input = scaling_factor[None, :] * time[:, None]
-    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+    sin_input = time[..., None] * scaling_factor
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=-1)
 
 
 def sample_beta(alpha, beta, bsize, device):  # see openpi `sample_beta` (exact copy)
@@ -739,14 +839,23 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         return embs, pad_masks, att_masks, adarms_cond
 
-    def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
+    def forward(
+        self,
+        images,
+        img_masks,
+        tokens,
+        masks,
+        actions,
+        noise,
+        time,
+        prefix_mask: Tensor | None = None,
+    ) -> Tensor:
         """Do a full training forward pass and compute the loss."""
-        time_expanded = time[:, None, None]
-        x_t = time_expanded * noise + (1 - time_expanded) * actions
+        x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
         u_t = noise - actions
 
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, tokens, masks)
-        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, time)
+        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, model_time)
 
         if (
             self.paligemma_with_expert.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype
@@ -833,11 +942,35 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
         dt = -1.0 / num_steps
 
         x_t = noise
+        rtc_mode = "guided"
+        trained_prefix = trained_prefix_mask = None
+        if self._rtc_enabled():
+            rtc_mode = self.rtc_processor.rtc_config.mode
+            if rtc_mode == "trained":
+                training_max_delay = int(getattr(self.config, "rtc_training_max_delay", 0))
+                if training_max_delay <= 0:
+                    raise ValueError(
+                        "RTC mode='trained' requires a checkpoint trained with "
+                        "policy.rtc_training_max_delay > 0."
+                    )
+                trained_prefix, trained_prefix_mask = _prepare_trained_rtc_prefix(
+                    x_t,
+                    kwargs.get("prev_chunk_left_over"),
+                    int(kwargs.get("inference_delay") or 0),
+                    training_max_delay,
+                )
+
         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):
+            denoise_timestep = time_tensor
+            if trained_prefix is not None:
+                x_t = torch.where(trained_prefix_mask, trained_prefix, x_t)
+                denoise_timestep = time_tensor[:, None].expand(bsize, x_t.shape[1]).clone()
+                denoise_timestep[trained_prefix_mask[..., 0]] = 0.0
+
+            def denoise_step_partial_call(input_x_t, current_timestep=denoise_timestep):
                 return self.denoise_step(
                     prefix_pad_masks=prefix_pad_masks,
                     past_key_values=past_key_values,
@@ -845,7 +978,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
                     timestep=current_timestep,
                 )
 
-            if self._rtc_enabled():
+            if self._rtc_enabled() and rtc_mode == "guided":
                 inference_delay = kwargs.get("inference_delay")
                 prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
                 execution_horizon = kwargs.get("execution_horizon")
@@ -862,6 +995,8 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
                 v_t = denoise_step_partial_call(x_t)
 
             x_t = x_t + dt * v_t
+            if trained_prefix is not None:
+                x_t = torch.where(trained_prefix_mask, trained_prefix, x_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)
@@ -1137,7 +1272,10 @@ class PI05Policy(PreTrainedPolicy):
         # Create processor if config provided
         # If RTC is not enabled - we can still track the denoising data
         if self.config.rtc_config is not None:
-            self.rtc_processor = RTCProcessor(self.config.rtc_config)
+            self.rtc_processor = RTCProcessor(
+                self.config.rtc_config,
+                trained_mode_supported=int(getattr(self.config, "rtc_training_max_delay", 0)) > 0,
+            )
 
             model_value = getattr(self, "model", None)
             if model_value is not None:
@@ -1269,28 +1407,35 @@ class PI05Policy(PreTrainedPolicy):
 
         noise = self.model.sample_noise(actions.shape, actions.device)
         time = self.model.sample_time(actions.shape[0], actions.device)
+        prefix_mask = _sample_training_rtc_prefix_mask(
+            actions.shape[0],
+            actions.shape[1],
+            self.config.rtc_training_max_delay,
+            actions.device,
+        )
 
         # Compute loss (no separate state needed for PI05)
-        losses = self.model.forward(images, img_masks, tokens, masks, actions, noise, time)
+        losses = self.model.forward(images, img_masks, tokens, masks, actions, noise, time, prefix_mask)
 
         # Truncate losses to actual action dimensions
         original_action_dim = self.config.output_features[ACTION].shape[0]
         losses = losses[:, :, :original_action_dim]
 
-        loss_dict = {
-            "loss_per_dim": losses.mean(dim=[0, 1]).detach().cpu().numpy().tolist(),
-        }
+        if prefix_mask is None:
+            loss_per_dim = losses.mean(dim=(0, 1))
+        else:
+            postfix_mask = (~prefix_mask).unsqueeze(-1).expand_as(losses)
+            loss_per_dim = (losses * postfix_mask).sum(dim=(0, 1)) / postfix_mask.sum(dim=(0, 1)).clamp(min=1)
+        loss_dict = {"loss_per_dim": loss_per_dim.detach().cpu().numpy().tolist()}
 
         if reduction == "none":
-            # Return per-sample losses (B,) by averaging over time and action dims
-            per_sample_loss = losses.mean(dim=(1, 2))
+            per_sample_loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="none")
             loss_dict["loss"] = per_sample_loss.mean().item()
             return per_sample_loss, loss_dict
-        else:
-            # Default: return scalar mean loss
-            loss = losses.mean()
-            loss_dict["loss"] = loss.item()
-            return loss, loss_dict
+
+        loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean")
+        loss_dict["loss"] = loss.item()
+        return loss, loss_dict
 
     def _get_default_peft_targets(self) -> dict[str, any]:
         """Return default PEFT target modules for PI0.5 fine-tuning."""

src/lerobot/policies/pi05/processor_pi05.py

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

src/lerobot/policies/pi0_fast/processor_pi0_fast.py

@@ -25,17 +25,26 @@ from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
     ActionTokenizerProcessorStep,
+    AddBatchDimensionProcessorStep,
+    DeviceProcessorStep,
+    NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
     RelativeActionsProcessorStep,
+    RenameObservationsProcessorStep,
     TokenizerProcessorStep,
-    make_default_policy_processor_steps,
-    make_policy_processor_pipelines,
+    UnnormalizerProcessorStep,
+    policy_action_to_transition,
+    transition_to_policy_action,
 )
 from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import OBS_STATE
+from lerobot.utils.constants import (
+    OBS_STATE,
+    POLICY_POSTPROCESSOR_DEFAULT_NAME,
+    POLICY_PREPROCESSOR_DEFAULT_NAME,
+)
 
 from .configuration_pi0_fast import PI0FastConfig
 
@@ -126,8 +135,6 @@ def make_pi0_fast_pre_post_processors(
         action_names=getattr(config, "action_feature_names", None),
     )
 
-    steps = make_default_policy_processor_steps(config, dataset_stats)
-
     # Pi0Fast order: relative → normalize → tokenize → model → unnormalize → absolute
     # This matches pi0/pi0.5: RelativeActionsProcessorStep runs first on raw absolute actions,
     # caching the raw state. NormalizerProcessorStep then normalizes the raw relative actions,
@@ -137,10 +144,14 @@ def make_pi0_fast_pre_post_processors(
     # before Pi0FastPrepareStateAndLanguageTokenizerProcessorStep, so the state tokenizer
     # continues to receive normalized state in [-1, 1] as expected.
     input_steps: list[ProcessorStep] = [
-        steps.rename_observations,  # To mimic the same processor as pretrained one
-        steps.add_batch_dim,
+        RenameObservationsProcessorStep(rename_map={}),  # To mimic the same processor as pretrained one
+        AddBatchDimensionProcessorStep(),
         relative_step,
-        steps.normalize,
+        NormalizerProcessorStep(
+            features={**config.input_features, **config.output_features},
+            norm_map=config.normalization_mapping,
+            stats=dataset_stats,
+        ),
         Pi0FastPrepareStateAndLanguageTokenizerProcessorStep(max_state_dim=config.max_state_dim),
         TokenizerProcessorStep(
             tokenizer_name=config.text_tokenizer_name,
@@ -154,13 +165,26 @@ def make_pi0_fast_pre_post_processors(
             fast_skip_tokens=config.fast_skip_tokens,
             paligemma_tokenizer_name=config.text_tokenizer_name,
         ),
-        steps.to_device,
+        DeviceProcessorStep(device=config.device),
     ]
 
     output_steps: list[ProcessorStep] = [
-        steps.unnormalize,
+        UnnormalizerProcessorStep(
+            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
+        ),
         AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
-        steps.to_cpu,
+        DeviceProcessorStep(device="cpu"),
     ]
 
-    return make_policy_processor_pipelines(input_steps=input_steps, output_steps=output_steps)
+    return (
+        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
+            steps=input_steps,
+            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
+        ),
+        PolicyProcessorPipeline[PolicyAction, PolicyAction](
+            steps=output_steps,
+            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
+            to_transition=policy_action_to_transition,
+            to_output=transition_to_policy_action,
+        ),
+    )

src/lerobot/policies/pi_gemma.py

@@ -121,7 +121,10 @@ class PiGemmaRMSNorm(nn.Module):
         if cond.shape[-1] != self.cond_dim:
             raise ValueError(f"Expected cond dim {self.cond_dim}, got {cond.shape[-1]}")
         modulation = self.dense(cond)
-        if len(x.shape) == 3:
+        # Per-sample conditioning (B, D) is broadcast across the sequence.
+        # Training-time RTC supplies per-action conditioning (B, T, D), which
+        # is already aligned with x and must keep its token dimension intact.
+        if len(x.shape) == 3 and modulation.dim() == 2:
             modulation = modulation.unsqueeze(1)
         scale, shift, gate = modulation.chunk(3, dim=-1)
         normed = normed * (1 + scale.float()) + shift.float()

src/lerobot/policies/pretrained.py

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

src/lerobot/policies/rtc/configuration_rtc.py

@@ -37,6 +37,10 @@ class RTCConfig:
     # Infrastructure
     enabled: bool = True
 
+    # ``guided`` is the original inference-time Jacobian guidance. ``trained``
+    # hard-inpaints a prefix and requires a compatible training-time RTC checkpoint.
+    mode: str = "guided"
+
     # Core RTC settings
     # Todo change to exp
     prefix_attention_schedule: RTCAttentionSchedule = RTCAttentionSchedule.LINEAR
@@ -49,6 +53,8 @@ class RTCConfig:
 
     def __post_init__(self):
         """Validate RTC configuration parameters."""
+        if self.mode not in {"guided", "trained"}:
+            raise ValueError(f"mode must be 'guided' or 'trained', got {self.mode!r}")
         if self.max_guidance_weight <= 0:
             raise ValueError(f"max_guidance_weight must be positive, got {self.max_guidance_weight}")
         if self.debug_maxlen <= 0:

src/lerobot/policies/rtc/modeling_rtc.py

@@ -42,7 +42,12 @@ class RTCProcessor:
     prefix attention, and adaptive chunk processing.
     """
 
-    def __init__(self, rtc_config: RTCConfig):
+    def __init__(self, rtc_config: RTCConfig, *, trained_mode_supported: bool = False):
+        if rtc_config.enabled and rtc_config.mode == "trained" and not trained_mode_supported:
+            raise ValueError(
+                "RTC mode='trained' requires a PI05-compatible checkpoint trained with "
+                "rtc_training_max_delay > 0."
+            )
         self.rtc_config = rtc_config
 
         self.tracker = None

src/lerobot/policies/rtc/relative.py

@@ -49,7 +49,13 @@ def reanchor_relative_rtc_prefix(
 
     action_cpu = prev_actions_absolute.detach().cpu()
     mask = relative_step._build_mask(action_cpu.shape[-1])
-    relative_actions = to_relative_actions(action_cpu, state, mask)
+    relative_actions = to_relative_actions(
+        action_cpu,
+        state,
+        mask,
+        pose_representation=relative_step.pose_representation,
+        se3_pose_groups=relative_step.se3_pose_groups,
+    )
 
     transition = create_transition(action=relative_actions)
     if normalizer_step is not None:

src/lerobot/policies/smolvla/processor_smolvla.py

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

src/lerobot/policies/tdmpc/processor_tdmpc.py

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

src/lerobot/policies/vla_jepa/processor_vla_jepa.py

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

src/lerobot/policies/vqbet/processor_vqbet.py

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

src/lerobot/policies/wall_x/processor_wall_x.py

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

src/lerobot/policies/xvla/processor_xvla.py

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

src/lerobot/processor/__init__.py

@@ -42,14 +42,10 @@ from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorTo
 from .device_processor import DeviceProcessorStep
 from .env_processor import IsaaclabArenaProcessorStep, LiberoProcessorStep
 from .factory import (
-    DefaultPolicyProcessorSteps,
-    make_default_policy_processor_steps,
-    make_default_pre_post_processors,
     make_default_processors,
     make_default_robot_action_processor,
     make_default_robot_observation_processor,
     make_default_teleop_action_processor,
-    make_policy_processor_pipelines,
 )
 from .gym_action_processor import (
     Numpy2TorchActionProcessorStep,
@@ -93,8 +89,15 @@ from .policy_robot_bridge import (
 from .relative_action_processor import (
     AbsoluteActionsProcessorStep,
     RelativeActionsProcessorStep,
+    relative_action_output_dim,
+    rotation_6d_to_rotvec,
+    rotvec_to_rotation_6d,
     to_absolute_actions,
+    to_absolute_se3_pose,
+    to_absolute_se3_pose_6d,
     to_relative_actions,
+    to_relative_se3_pose,
+    to_relative_se3_pose_6d,
 )
 from .rename_processor import RenameObservationsProcessorStep, rename_stats
 from .tokenizer_processor import ActionTokenizerProcessorStep, TokenizerProcessorStep
@@ -133,16 +136,21 @@ __all__ = [
     "ImageCropResizeProcessorStep",
     "InfoProcessorStep",
     "InterventionActionProcessorStep",
-    "DefaultPolicyProcessorSteps",
-    "make_default_policy_processor_steps",
-    "make_default_pre_post_processors",
     "make_default_processors",
     "make_default_teleop_action_processor",
     "make_default_robot_action_processor",
     "make_default_robot_observation_processor",
-    "make_policy_processor_pipelines",
     "AbsoluteActionsProcessorStep",
     "RelativeActionsProcessorStep",
+    "relative_action_output_dim",
+    "rotation_6d_to_rotvec",
+    "rotvec_to_rotation_6d",
+    "to_absolute_actions",
+    "to_absolute_se3_pose",
+    "to_absolute_se3_pose_6d",
+    "to_relative_actions",
+    "to_relative_se3_pose",
+    "to_relative_se3_pose_6d",
     "MapDeltaActionToRobotActionStep",
     "MapTensorToDeltaActionDictStep",
     "NewLineTaskProcessorStep",
@@ -176,8 +184,6 @@ __all__ = [
     "transition_to_batch",
     "TransitionKey",
     "TruncatedProcessorStep",
-    "to_absolute_actions",
-    "to_relative_actions",
     "UnnormalizerProcessorStep",
     "VanillaObservationProcessorStep",
 ]

src/lerobot/processor/factory.py

@@ -14,33 +14,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from dataclasses import dataclass
-from typing import Any
+from lerobot.types import RobotAction, RobotObservation
 
-import torch
-
-from lerobot.configs.policies import PreTrainedConfig
-from lerobot.types import PolicyAction, RobotAction, RobotObservation
-from lerobot.utils.constants import POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME
-
-from .batch_processor import AddBatchDimensionProcessorStep
 from .converters import (
     observation_to_transition,
-    policy_action_to_transition,
     robot_action_observation_to_transition,
     transition_to_observation,
-    transition_to_policy_action,
     transition_to_robot_action,
 )
-from .device_processor import DeviceProcessorStep
-from .normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
-from .pipeline import (
-    IdentityProcessorStep,
-    PolicyProcessorPipeline,
-    ProcessorStep,
-    RobotProcessorPipeline,
-)
-from .rename_processor import RenameObservationsProcessorStep
+from .pipeline import IdentityProcessorStep, RobotProcessorPipeline
 
 
 def make_default_teleop_action_processor() -> RobotProcessorPipeline[
@@ -79,97 +61,3 @@ def make_default_processors():
     robot_action_processor = make_default_robot_action_processor()
     robot_observation_processor = make_default_robot_observation_processor()
     return (teleop_action_processor, robot_action_processor, robot_observation_processor)
-
-
-@dataclass
-class DefaultPolicyProcessorSteps:
-    """The canonical processor steps shared by most policies' pre/post pipelines.
-
-    Policies compose these in their own order (step ORDER is a Hub-serialized contract
-    and intentionally stays explicit per policy) and interleave their custom steps.
-    """
-
-    rename_observations: RenameObservationsProcessorStep
-    add_batch_dim: AddBatchDimensionProcessorStep
-    to_device: DeviceProcessorStep
-    normalize: NormalizerProcessorStep
-    unnormalize: UnnormalizerProcessorStep
-    to_cpu: DeviceProcessorStep
-
-
-def make_default_policy_processor_steps(
-    config: PreTrainedConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    *,
-    normalizer_device: torch.device | str | None = None,
-) -> DefaultPolicyProcessorSteps:
-    """Construct the canonical policy processor steps from a policy config.
-
-    Args:
-        config: A `PreTrainedConfig` providing `device`, `input_features`,
-            `output_features` and `normalization_mapping`.
-        dataset_stats: Dataset statistics used for (un)normalization.
-        normalizer_device: Device passed to `NormalizerProcessorStep` (some policies pin
-            their normalization stats to the policy device; most leave it unset).
-    """
-    return DefaultPolicyProcessorSteps(
-        rename_observations=RenameObservationsProcessorStep(rename_map={}),
-        add_batch_dim=AddBatchDimensionProcessorStep(),
-        to_device=DeviceProcessorStep(device=config.device),
-        normalize=NormalizerProcessorStep(
-            features={**config.input_features, **config.output_features},
-            norm_map=config.normalization_mapping,
-            stats=dataset_stats,
-            device=normalizer_device,
-        ),
-        unnormalize=UnnormalizerProcessorStep(
-            features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
-        ),
-        to_cpu=DeviceProcessorStep(device="cpu"),
-    )
-
-
-def make_policy_processor_pipelines(
-    input_steps: list[ProcessorStep],
-    output_steps: list[ProcessorStep],
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """Wrap pre/post step lists into the canonical policy pipeline pair.
-
-    Uses the standard pipeline names (which determine the serialized JSON filenames on
-    the Hub) and the standard policy-action converters on the postprocessor.
-    """
-    return (
-        PolicyProcessorPipeline[dict[str, Any], dict[str, Any]](
-            steps=input_steps,
-            name=POLICY_PREPROCESSOR_DEFAULT_NAME,
-        ),
-        PolicyProcessorPipeline[PolicyAction, PolicyAction](
-            steps=output_steps,
-            name=POLICY_POSTPROCESSOR_DEFAULT_NAME,
-            to_transition=policy_action_to_transition,
-            to_output=transition_to_policy_action,
-        ),
-    )
-
-
-def make_default_pre_post_processors(
-    config: PreTrainedConfig,
-    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
-    *,
-    normalizer_device: torch.device | str | None = None,
-) -> tuple[
-    PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
-    PolicyProcessorPipeline[PolicyAction, PolicyAction],
-]:
-    """The pure-scaffold policy pipeline pair: Rename -> Batch -> Device -> Normalize,
-    and Unnormalize -> Device(cpu). Policies with custom steps or a different step order
-    compose `make_default_policy_processor_steps` themselves instead.
-    """
-    s = make_default_policy_processor_steps(config, dataset_stats, normalizer_device=normalizer_device)
-    return make_policy_processor_pipelines(
-        input_steps=[s.rename_observations, s.add_batch_dim, s.to_device, s.normalize],
-        output_steps=[s.unnormalize, s.to_cpu],
-    )

src/lerobot/processor/relative_action_processor.py

@@ -21,7 +21,7 @@ from torch import Tensor
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.types import EnvTransition, TransitionKey
-from lerobot.utils.constants import OBS_STATE
+from lerobot.utils.constants import ACTION, OBS_STATE
 
 from .delta_action_processor import MapDeltaActionToRobotActionStep, MapTensorToDeltaActionDictStep
 from .pipeline import ProcessorStep, ProcessorStepRegistry
@@ -34,57 +34,399 @@ __all__ = [
     "AbsoluteActionsProcessorStep",
     "to_relative_actions",
     "to_absolute_actions",
+    "to_relative_se3_pose",
+    "to_absolute_se3_pose",
+    "to_relative_se3_pose_6d",
+    "to_absolute_se3_pose_6d",
+    "rotation_6d_to_rotvec",
+    "rotvec_to_rotation_6d",
+    "relative_action_output_dim",
 ]
 
 
-def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
-    """Convert absolute actions to relative: relative = action - state (for masked dims).
+def _rotvec_to_quaternion(rotvec: Tensor) -> Tensor:
+    angle = torch.linalg.vector_norm(rotvec, dim=-1, keepdim=True)
+    angle_sq = angle.square()
+    small_scale = 0.5 - angle_sq / 48.0 + angle_sq.square() / 3840.0
+    scale = torch.where(angle > 1e-6, torch.sin(angle / 2.0) / angle.clamp_min(1e-12), small_scale)
+    return torch.cat((torch.cos(angle / 2.0), rotvec * scale), dim=-1)
+
+
+def _quaternion_to_rotvec(quaternion: Tensor) -> Tensor:
+    quaternion = quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
+    quaternion = quaternion * torch.where(quaternion[..., :1] < 0, -1.0, 1.0)
+    vector = quaternion[..., 1:]
+    sin_half_angle = torch.linalg.vector_norm(vector, dim=-1, keepdim=True)
+    angle = 2.0 * torch.atan2(sin_half_angle, quaternion[..., :1].clamp_min(0.0))
+    small_scale = 2.0 + sin_half_angle.square() / 3.0
+    scale = torch.where(
+        sin_half_angle > 1e-6,
+        angle / sin_half_angle.clamp_min(1e-12),
+        small_scale,
+    )
+    return vector * scale
+
+
+def _quaternion_multiply(left: Tensor, right: Tensor) -> Tensor:
+    left_w, left_xyz = left[..., :1], left[..., 1:]
+    right_w, right_xyz = right[..., :1], right[..., 1:]
+    return torch.cat(
+        (
+            left_w * right_w - (left_xyz * right_xyz).sum(dim=-1, keepdim=True),
+            left_w * right_xyz + right_w * left_xyz + torch.linalg.cross(left_xyz, right_xyz, dim=-1),
+        ),
+        dim=-1,
+    )
+
+
+def _quaternion_conjugate(quaternion: Tensor) -> Tensor:
+    return torch.cat((quaternion[..., :1], -quaternion[..., 1:]), dim=-1)
+
+
+def _quaternion_rotate(quaternion: Tensor, vector: Tensor) -> Tensor:
+    quaternion_xyz = quaternion[..., 1:]
+    uv = torch.linalg.cross(quaternion_xyz, vector, dim=-1)
+    uuv = torch.linalg.cross(quaternion_xyz, uv, dim=-1)
+    return vector + 2.0 * (quaternion[..., :1] * uv + uuv)
+
+
+def _quaternion_to_matrix(quaternion: Tensor) -> Tensor:
+    quaternion = quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
+    w, x, y, z = quaternion.unbind(-1)
+    two_s = 2.0
+    return torch.stack(
+        (
+            1.0 - two_s * (y * y + z * z),
+            two_s * (x * y - z * w),
+            two_s * (x * z + y * w),
+            two_s * (x * y + z * w),
+            1.0 - two_s * (x * x + z * z),
+            two_s * (y * z - x * w),
+            two_s * (x * z - y * w),
+            two_s * (y * z + x * w),
+            1.0 - two_s * (x * x + y * y),
+        ),
+        dim=-1,
+    ).reshape(quaternion.shape[:-1] + (3, 3))
+
+
+def _matrix_to_quaternion(matrix: Tensor) -> Tensor:
+    """Convert proper rotation matrices to normalized ``[w, x, y, z]`` quaternions."""
+    if matrix.shape[-2:] != (3, 3):
+        raise ValueError(f"Rotation matrices must have shape (..., 3, 3), got {matrix.shape}")
+
+    m00 = matrix[..., 0, 0]
+    m01 = matrix[..., 0, 1]
+    m02 = matrix[..., 0, 2]
+    m10 = matrix[..., 1, 0]
+    m11 = matrix[..., 1, 1]
+    m12 = matrix[..., 1, 2]
+    m20 = matrix[..., 2, 0]
+    m21 = matrix[..., 2, 1]
+    m22 = matrix[..., 2, 2]
+
+    # Each row is a quaternion candidate scaled by the magnitude of its
+    # best-conditioned component. Selecting the largest component avoids the
+    # trace singularity at rotations close to pi.
+    q_abs = torch.sqrt(
+        torch.clamp(
+            torch.stack(
+                (
+                    1.0 + m00 + m11 + m22,
+                    1.0 + m00 - m11 - m22,
+                    1.0 - m00 + m11 - m22,
+                    1.0 - m00 - m11 + m22,
+                ),
+                dim=-1,
+            ),
+            min=0.0,
+        )
+    )
+    quat_by_rijk = torch.stack(
+        (
+            torch.stack((q_abs[..., 0].square(), m21 - m12, m02 - m20, m10 - m01), dim=-1),
+            torch.stack((m21 - m12, q_abs[..., 1].square(), m10 + m01, m02 + m20), dim=-1),
+            torch.stack((m02 - m20, m10 + m01, q_abs[..., 2].square(), m12 + m21), dim=-1),
+            torch.stack((m10 - m01, m02 + m20, m12 + m21, q_abs[..., 3].square()), dim=-1),
+        ),
+        dim=-2,
+    )
+    candidates = quat_by_rijk / (2.0 * q_abs[..., :, None].clamp_min(0.1))
+    best = torch.nn.functional.one_hot(q_abs.argmax(dim=-1), num_classes=4).to(dtype=matrix.dtype)
+    quaternion = (candidates * best[..., :, None]).sum(dim=-2)
+    return quaternion / torch.linalg.vector_norm(quaternion, dim=-1, keepdim=True).clamp_min(1e-12)
+
+
+def rotvec_to_rotation_6d(rotvec: Tensor) -> Tensor:
+    """Encode an axis-angle rotation as the first two rotation-matrix rows."""
+    matrix = _quaternion_to_matrix(_rotvec_to_quaternion(rotvec))
+    return matrix[..., :2, :].reshape(matrix.shape[:-2] + (6,))
+
+
+def rotation_6d_to_rotvec(rotation_6d: Tensor) -> Tensor:
+    """Decode two predicted 3-D vectors into an axis-angle rotation.
+
+    Gram-Schmidt orthonormalization follows the continuous 6-D rotation
+    representation. Degenerate predictions fail closed instead of producing an
+    invalid physical rotation.
+    """
+    if rotation_6d.shape[-1] != 6:
+        raise ValueError(f"6-D rotations must have six values, got {rotation_6d.shape}")
+    first = rotation_6d[..., :3]
+    second = rotation_6d[..., 3:]
+    first_norm = torch.linalg.vector_norm(first, dim=-1, keepdim=True)
+    first_unit = first / first_norm.clamp_min(1e-12)
+    second_orthogonal = second - (first_unit * second).sum(dim=-1, keepdim=True) * first_unit
+    second_norm = torch.linalg.vector_norm(second_orthogonal, dim=-1, keepdim=True)
+    if bool(torch.any(first_norm <= 1e-8)) or bool(torch.any(second_norm <= 1e-8)):
+        raise ValueError("Cannot decode a degenerate 6-D rotation prediction")
+    second_unit = second_orthogonal / second_norm
+    third_unit = torch.linalg.cross(first_unit, second_unit, dim=-1)
+    matrix = torch.stack((first_unit, second_unit, third_unit), dim=-2)
+    return _quaternion_to_rotvec(_matrix_to_quaternion(matrix))
+
+
+def to_relative_se3_pose(target_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Encode a pose as ``inv(T_reference) @ T_target``.
+
+    Poses use ``[x, y, z, rx, ry, rz]`` with an axis-angle rotation vector.
+    The relative translation is therefore expressed in the reference EE frame.
+    """
+    if target_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
+        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
+    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
+    target_quaternion = _rotvec_to_quaternion(target_pose[..., 3:])
+    inverse_reference_quaternion = _quaternion_conjugate(reference_quaternion)
+    relative_translation = _quaternion_rotate(
+        inverse_reference_quaternion, target_pose[..., :3] - reference_pose[..., :3]
+    )
+    relative_quaternion = _quaternion_multiply(inverse_reference_quaternion, target_quaternion)
+    return torch.cat((relative_translation, _quaternion_to_rotvec(relative_quaternion)), dim=-1)
+
+
+def to_absolute_se3_pose(relative_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Decode a pose with ``T_target = T_reference @ T_relative``."""
+    if relative_pose.shape[-1] != 6 or reference_pose.shape[-1] != 6:
+        raise ValueError("SE(3) poses must have six values: xyz followed by a rotation vector")
+    reference_quaternion = _rotvec_to_quaternion(reference_pose[..., 3:])
+    relative_quaternion = _rotvec_to_quaternion(relative_pose[..., 3:])
+    target_translation = reference_pose[..., :3] + _quaternion_rotate(
+        reference_quaternion, relative_pose[..., :3]
+    )
+    target_quaternion = _quaternion_multiply(reference_quaternion, relative_quaternion)
+    return torch.cat((target_translation, _quaternion_to_rotvec(target_quaternion)), dim=-1)
+
+
+def to_relative_se3_pose_6d(target_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Encode ``inv(T_reference) @ T_target`` as xyz plus continuous 6-D rotation."""
+    relative_pose = to_relative_se3_pose(target_pose, reference_pose)
+    return torch.cat((relative_pose[..., :3], rotvec_to_rotation_6d(relative_pose[..., 3:])), dim=-1)
+
+
+def to_absolute_se3_pose_6d(relative_pose: Tensor, reference_pose: Tensor) -> Tensor:
+    """Decode xyz plus continuous 6-D rotation with ``T_target = T_reference @ T_relative``."""
+    if relative_pose.shape[-1] != 9:
+        raise ValueError("6-D encoded SE(3) poses must have nine values: xyz plus rotation-6D")
+    relative_rotvec_pose = torch.cat(
+        (relative_pose[..., :3], rotation_6d_to_rotvec(relative_pose[..., 3:])), dim=-1
+    )
+    return to_absolute_se3_pose(relative_rotvec_pose, reference_pose)
+
+
+def _broadcast_reference(actions: Tensor, state: Tensor) -> Tensor:
+    if state.device != actions.device or state.dtype != actions.dtype:
+        state = state.to(device=actions.device, dtype=actions.dtype)
+    if actions.ndim == state.ndim + 1:
+        state = state.unsqueeze(-2)
+    return state
+
+
+def _validate_se3_pose_groups(
+    pose_representation: str,
+    se3_pose_groups: Sequence[Sequence[int]] | None,
+    mask: Sequence[bool],
+    action_dim: int,
+) -> list[list[int]]:
+    if pose_representation not in {"componentwise", "se3", "se3_6d"}:
+        raise ValueError(
+            f"Unsupported pose_representation={pose_representation!r}; expected "
+            "'componentwise', 'se3', or 'se3_6d'"
+        )
+    if pose_representation == "componentwise":
+        return []
+    if not se3_pose_groups:
+        raise ValueError(
+            f"pose_representation={pose_representation!r} requires at least one six-index se3_pose_group"
+        )
+
+    normalized_groups: list[list[int]] = []
+    used_indices: set[int] = set()
+    for raw_group in se3_pose_groups:
+        group = [int(index) for index in raw_group]
+        if len(group) != 6:
+            raise ValueError(f"Each SE(3) pose group must contain six indices, got {group}")
+        if len(set(group)) != 6 or any(index < 0 or index >= action_dim for index in group):
+            raise ValueError(f"Invalid SE(3) pose group for action_dim={action_dim}: {group}")
+        if pose_representation == "se3_6d" and group != list(range(group[0], group[0] + 6)):
+            raise ValueError("se3_6d pose groups must contain six contiguous ascending indices")
+        if any(index >= len(mask) for index in group):
+            raise ValueError(f"SE(3) pose group lies outside the relative mask: {group}")
+        if used_indices.intersection(group):
+            raise ValueError(f"SE(3) pose groups must not overlap: {group}")
+        group_mask = [bool(mask[index]) for index in group]
+        if any(group_mask) and not all(group_mask):
+            raise ValueError(f"An SE(3) pose group must be wholly relative or wholly absolute: {group}")
+        used_indices.update(group)
+        if all(group_mask):
+            normalized_groups.append(group)
+    return normalized_groups
+
+
+def relative_action_output_dim(
+    source_dim: int,
+    pose_representation: str,
+    se3_pose_groups: Sequence[Sequence[int]] | None,
+) -> int:
+    """Return the model-space action width for a source action width."""
+    if pose_representation != "se3_6d":
+        return source_dim
+    groups = se3_pose_groups or []
+    return source_dim + 3 * len(groups)
+
+
+def _expand_se3_6d_actions(
+    actions: Tensor,
+    state: Tensor,
+    groups: Sequence[Sequence[int]],
+) -> Tensor:
+    group_by_start = {group[0]: list(group) for group in groups}
+    grouped_indices = {index for group in groups for index in group}
+    parts: list[Tensor] = []
+    for index in range(actions.shape[-1]):
+        group = group_by_start.get(index)
+        if group is not None:
+            parts.append(to_relative_se3_pose_6d(actions[..., group], state[..., group]))
+        elif index not in grouped_indices:
+            parts.append(actions[..., index : index + 1])
+    return torch.cat(parts, dim=-1)
+
+
+def _collapse_se3_6d_actions(
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    groups: Sequence[Sequence[int]],
+) -> Tensor:
+    source_dim = len(mask)
+    expected_dim = relative_action_output_dim(source_dim, "se3_6d", groups)
+    if actions.shape[-1] != expected_dim:
+        raise ValueError(
+            f"Expected se3_6d action width {expected_dim} for source width {source_dim}, "
+            f"got {actions.shape[-1]}"
+        )
+    group_by_start = {group[0]: list(group) for group in groups}
+    grouped_indices = {index for group in groups for index in group}
+    parts: list[Tensor] = []
+    cursor = 0
+    for index in range(source_dim):
+        group = group_by_start.get(index)
+        if group is not None:
+            parts.append(to_absolute_se3_pose_6d(actions[..., cursor : cursor + 9], state[..., group]))
+            cursor += 9
+        elif index not in grouped_indices:
+            value = actions[..., cursor : cursor + 1]
+            if mask[index]:
+                value = value + state[..., index : index + 1]
+            parts.append(value)
+            cursor += 1
+    if cursor != actions.shape[-1]:
+        raise RuntimeError(f"Consumed {cursor} action values from width {actions.shape[-1]}")
+    return torch.cat(parts, dim=-1)
+
+
+def to_relative_actions(
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    *,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: Sequence[Sequence[int]] | None = None,
+) -> Tensor:
+    """Convert absolute actions to a configured relative representation.
+
+    Component-wise mode computes ``action - state``. SE(3) modes compute
+    ``inv(T_state) @ T_action`` for each configured pose group. ``se3_6d``
+    replaces each three-value relative rotation vector with its continuous
+    six-value encoding, increasing the output width by three per pose group.
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
+    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, actions.shape[-1])
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    if state.device != actions.device or state.dtype != actions.dtype:
-        state = state.to(device=actions.device, dtype=actions.dtype)
-    state_offset = state[..., :dims] * mask_t
-    if actions.ndim == 3:
-        state_offset = state_offset.unsqueeze(-2)
+    state = _broadcast_reference(actions, state)
+    component_mask = mask_t.clone()
+    for group in groups:
+        component_mask[group] = 0
+    state_offset = state[..., :dims] * component_mask
     actions = actions.clone()
     actions[..., :dims] -= state_offset
+    if pose_representation == "se3_6d":
+        return _expand_se3_6d_actions(actions, state, groups)
+    for group in groups:
+        actions[..., group] = to_relative_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
-def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) -> Tensor:
-    """Convert relative actions back to absolute: absolute = relative + state (for masked dims).
+def to_absolute_actions(
+    actions: Tensor,
+    state: Tensor,
+    mask: Sequence[bool],
+    *,
+    pose_representation: str = "componentwise",
+    se3_pose_groups: Sequence[Sequence[int]] | None = None,
+) -> Tensor:
+    """Convert relative actions back to absolute actions.
+
+    Component-wise mode computes ``relative + state``. SE(3) mode computes
+    ``T_state @ T_relative`` for each configured pose group.
 
     Args:
         actions: (B, T, action_dim) or (B, action_dim).
         state: (B, state_dim). Broadcast across time dimension.
         mask: Which dims to convert. Can be shorter than action_dim.
     """
+    source_dim = len(mask)
+    groups = _validate_se3_pose_groups(pose_representation, se3_pose_groups, mask, source_dim)
+    state = _broadcast_reference(actions, state)
+    if pose_representation == "se3_6d":
+        return _collapse_se3_6d_actions(actions, state, mask, groups)
+
     mask_t = torch.tensor(mask, dtype=actions.dtype, device=actions.device)
     dims = mask_t.shape[0]
     # Align state to the same device/dtype as actions. _last_state is cached before
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
-    if state.device != actions.device or state.dtype != actions.dtype:
-        state = state.to(device=actions.device, dtype=actions.dtype)
-    state_offset = state[..., :dims] * mask_t
-    if actions.ndim == 3:
-        state_offset = state_offset.unsqueeze(-2)
+    state = _broadcast_reference(actions, state)
+    component_mask = mask_t.clone()
+    for group in groups:
+        component_mask[group] = 0
+    state_offset = state[..., :dims] * component_mask
     actions = actions.clone()
     actions[..., :dims] += state_offset
+    for group in groups:
+        actions[..., group] = to_absolute_se3_pose(actions[..., group], state[..., group])
     return actions
 
 
 @ProcessorStepRegistry.register("relative_actions_processor")
 @dataclass
 class RelativeActionsProcessorStep(ProcessorStep):
-    """Converts absolute actions to relative actions (action -= state) for masked dimensions.
+    """Converts absolute actions to the configured relative representation.
 
     Mirrors OpenPI's DeltaActions transform. Applied during preprocessing so the model
     trains on relative offsets instead of absolute positions.
@@ -101,7 +443,10 @@ class RelativeActionsProcessorStep(ProcessorStep):
     enabled: bool = False
     exclude_joints: list[str] = field(default_factory=list)
     action_names: list[str] | None = None
+    pose_representation: str = "componentwise"
+    se3_pose_groups: list[list[int]] = field(default_factory=list)
     _last_state: torch.Tensor | None = field(default=None, init=False, repr=False)
+    _last_mask: list[bool] | None = field(default=None, init=False, repr=False)
 
     def _build_mask(self, action_dim: int) -> list[bool]:
         if not self.exclude_joints or self.action_names is None:
@@ -126,42 +471,78 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
-        # Always cache state for the paired AbsoluteActionsProcessorStep.
-        if state is not None:
-            self._last_state = state
+        # State history has shape (B, H, D). Relative actions are referenced to
+        # the newest proprioceptive state, not the whole history tensor.
+        reference_state = state[:, -1] if state is not None and state.ndim == 3 else state
+
+        # Always cache state for the paired AbsoluteActionsProcessorStep
+        if reference_state is not None:
+            self._last_state = reference_state
+            self._last_mask = self._build_mask(reference_state.shape[-1])
 
         if not self.enabled:
             return transition
 
         new_transition = transition.copy()
         action = new_transition.get(TransitionKey.ACTION)
-        if action is None or state is None:
+        if action is None or reference_state is None:
             return new_transition
 
-        mask = self._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_relative_actions(action, state, mask)
+        mask = self._last_mask or self._build_mask(action.shape[-1])
+        new_transition[TransitionKey.ACTION] = to_relative_actions(
+            action,
+            reference_state,
+            mask,
+            pose_representation=self.pose_representation,
+            se3_pose_groups=self.se3_pose_groups,
+        )
         return new_transition
 
     def get_cached_state(self) -> torch.Tensor | None:
         """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
         return self._last_state
 
-    def set_cached_state(self, state: torch.Tensor | None) -> None:
-        """Override the cached anchor state, e.g. to re-pin a chunk's anchor after the
-        per-tick pipeline overwrote it (see ``SyncInferenceEngine``)."""
-        self._last_state = state
+    def get_cached_mask(self) -> list[bool] | None:
+        """Return the source-space mask cached with the latest state."""
+        return self._last_mask
+
+    def reset(self) -> None:
+        """Drop the inference reference so it cannot leak between sessions."""
+        self._last_state = None
+        self._last_mask = None
 
     def get_config(self) -> dict[str, Any]:
         return {
             "enabled": self.enabled,
             "exclude_joints": self.exclude_joints,
             "action_names": self.action_names,
+            "pose_representation": self.pose_representation,
+            "se3_pose_groups": self.se3_pose_groups,
         }
 
     def transform_features(
         self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
     ) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
-        return features
+        if not self.enabled or self.pose_representation != "se3_6d":
+            return features
+        transformed = {feature_type: dict(feature_group) for feature_type, feature_group in features.items()}
+        for feature_group in transformed.values():
+            action_feature = feature_group.get(ACTION)
+            if action_feature is None:
+                continue
+            source_dim = len(self.action_names) if self.action_names is not None else action_feature.shape[-1]
+            model_dim = relative_action_output_dim(source_dim, self.pose_representation, self.se3_pose_groups)
+            if action_feature.shape[-1] == source_dim:
+                feature_group[ACTION] = PolicyFeature(
+                    type=action_feature.type,
+                    shape=(model_dim,),
+                )
+            elif action_feature.shape[-1] != model_dim:
+                raise ValueError(
+                    f"Expected source/model action width {source_dim}/{model_dim}, "
+                    f"got {action_feature.shape[-1]}"
+                )
+        return transformed
 
 
 @ProcessorStepRegistry.register("absolute_actions_processor")
@@ -203,8 +584,16 @@ class AbsoluteActionsProcessorStep(ProcessorStep):
         if action is None:
             return new_transition
 
-        mask = self.relative_step._build_mask(action.shape[-1])
-        new_transition[TransitionKey.ACTION] = to_absolute_actions(action, cached_state, mask)
+        mask = self.relative_step.get_cached_mask()
+        if mask is None:
+            mask = self.relative_step._build_mask(cached_state.shape[-1])
+        new_transition[TransitionKey.ACTION] = to_absolute_actions(
+            action,
+            cached_state,
+            mask,
+            pose_representation=self.relative_step.pose_representation,
+            se3_pose_groups=self.relative_step.se3_pose_groups,
+        )
         return new_transition
 
     def get_config(self) -> dict[str, Any]:

src/lerobot/rewards/pretrained.py

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

src/lerobot/rollout/context.py

@@ -43,14 +43,17 @@ from lerobot.processor import (
     make_default_processors,
     rename_stats,
 )
+from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
 from lerobot.robots import make_robot_from_config
 from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
+from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
 
 from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
 from .inference import (
     InferenceEngine,
     RTCInferenceConfig,
+    SyncInferenceConfig,
     create_inference_engine,
 )
 from .robot_wrapper import ThreadSafeRobot
@@ -58,6 +61,35 @@ from .robot_wrapper import ThreadSafeRobot
 logger = logging.getLogger(__name__)
 
 
+def _validate_trained_rtc_rollout_config(policy_config, inference_config: RTCInferenceConfig) -> None:
+    """Fail fast when rollout cannot retain every trained RTC prefix."""
+    rtc = inference_config.rtc
+    if not rtc.enabled or rtc.mode != "trained":
+        return
+    if policy_config.type not in {"pi05", "pi052"}:
+        raise ValueError(
+            "--inference.rtc.mode=trained currently requires a PI05-compatible checkpoint; "
+            f"got policy type {policy_config.type!r}."
+        )
+
+    training_max_delay = int(getattr(policy_config, "rtc_training_max_delay", 0))
+    if training_max_delay <= 0:
+        raise ValueError(
+            "--inference.rtc.mode=trained requires a checkpoint trained with "
+            "--policy.rtc_training_max_delay > 0."
+        )
+    if rtc.execution_horizon < training_max_delay:
+        raise ValueError(
+            f"--inference.rtc.execution_horizon ({rtc.execution_horizon}) must be at least the "
+            f"checkpoint's rtc_training_max_delay ({training_max_delay})."
+        )
+    if inference_config.queue_threshold < training_max_delay:
+        raise ValueError(
+            f"--inference.queue_threshold ({inference_config.queue_threshold}) must be at least the "
+            f"checkpoint's rtc_training_max_delay ({training_max_delay})."
+        )
+
+
 def _resolve_action_key_order(
     policy_action_names: list[str] | None, dataset_action_names: list[str]
 ) -> list[str]:
@@ -78,6 +110,26 @@ def _resolve_action_key_order(
     return policy_action_names
 
 
+def _align_relative_state_feature_order(
+    hw_features: dict[str, dict], policy_action_names: list[str] | None
+) -> dict[str, dict]:
+    """Align policy-facing state with named relative-action dimensions."""
+    if not policy_action_names or OBS_STATE not in hw_features:
+        return hw_features
+
+    state_feature = hw_features[OBS_STATE]
+    state_names = state_feature.get("names")
+    if not state_names or len(state_names) != len(policy_action_names):
+        return hw_features
+    if set(state_names) != set(policy_action_names) or state_names == policy_action_names:
+        return hw_features
+
+    aligned = dict(hw_features)
+    aligned[OBS_STATE] = {**state_feature, "names": list(policy_action_names)}
+    logger.info("Aligned relative-action state order with checkpoint action names")
+    return aligned
+
+
 # ---------------------------------------------------------------------------
 # Sub-contexts
 # ---------------------------------------------------------------------------
@@ -176,6 +228,9 @@ def build_rollout_context(
     policy_config = cfg.policy
     policy_class = get_policy_class(policy_config.type)
 
+    if is_rtc:
+        _validate_trained_rtc_rollout_config(policy_config, cfg.inference)
+
     if hasattr(policy_config, "compile_model"):
         policy_config.compile_model = cfg.use_torch_compile
 
@@ -397,6 +452,27 @@ def build_rollout_context(
         },
     )
 
+    relative_action_step = next(
+        (
+            step
+            for step in getattr(preprocessor, "steps", ())
+            if isinstance(step, RelativeActionsProcessorStep) and step.enabled
+        ),
+        None,
+    )
+    if relative_action_step is not None:
+        relative_action_names = relative_action_step.action_names or policy_action_names
+        hw_features = _align_relative_state_feature_order(
+            hw_features,
+            list(relative_action_names) if relative_action_names else None,
+        )
+
+    if isinstance(cfg.inference, SyncInferenceConfig) and relative_action_step is not None:
+        raise NotImplementedError(
+            "SyncInferenceEngine does not support policies with relative actions for now."
+            "Use --inference.type=rtc or remove relative action processor steps from the policy pipeline."
+        )
+
     # --- 7. Inference strategy (needs policy + pre/post + hardware) --
     logger.info(
         "Creating inference engine (type=%s)...",

src/lerobot/rollout/inference/rtc.py

@@ -57,6 +57,18 @@ _RTC_MAX_CONSECUTIVE_ERRORS: int = 10
 _RTC_JOIN_TIMEOUT_S: float = 3.0
 
 
+class _FatalRTCInferenceError(RuntimeError):
+    """Base class for RTC errors that cannot become valid after a retry."""
+
+
+class _TrainedRTCDelayExceededError(_FatalRTCInferenceError):
+    """Raised when measured latency exceeds a trained RTC checkpoint's support."""
+
+
+class _TrainedRTCPrefixUnavailableError(_FatalRTCInferenceError):
+    """Raised when the queue cannot provide the prefix used for conditioning."""
+
+
 # ---------------------------------------------------------------------------
 # RTC helpers
 # ---------------------------------------------------------------------------
@@ -76,6 +88,50 @@ def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int
     return padded
 
 
+def _trained_rtc_chunk_can_merge(
+    *,
+    conditioned_delay: int,
+    measured_delay: int,
+    training_max_delay: int,
+    has_previous_actions: bool,
+) -> bool:
+    """Check that a trained RTC chunk covers the overlap observed during inference."""
+    if not has_previous_actions:
+        return True
+    if measured_delay > training_max_delay:
+        raise _TrainedRTCDelayExceededError(
+            f"Measured RTC inference delay ({measured_delay}) exceeds the checkpoint's "
+            f"rtc_training_max_delay ({training_max_delay})."
+        )
+    return measured_delay <= conditioned_delay
+
+
+def _estimate_rtc_delay(
+    *,
+    latency: float,
+    time_per_step: float,
+    mode: str,
+    training_max_delay: int,
+    has_previous_actions: bool,
+) -> int:
+    """Estimate overlap, using the trained capacity to bootstrap the first transition."""
+    if latency:
+        return math.ceil(latency / time_per_step)
+    if mode == "trained" and has_previous_actions:
+        return training_max_delay
+    return 0
+
+
+def _validate_trained_rtc_prefix_available(*, conditioned_delay: int, available_steps: int) -> None:
+    """Reject hard-prefix inference when the real queue is shorter than its delay."""
+    if conditioned_delay > available_steps:
+        raise _TrainedRTCPrefixUnavailableError(
+            f"Trained RTC needs {conditioned_delay} committed prefix actions, but the queue has "
+            f"only {available_steps}. Increase --inference.queue_threshold and "
+            "--inference.rtc.execution_horizon."
+        )
+
+
 # ---------------------------------------------------------------------------
 # RTCInferenceEngine
 # ---------------------------------------------------------------------------
@@ -272,9 +328,23 @@ class RTCInferenceEngine(InferenceEngine):
                         current_time = time.perf_counter()
                         idx_before = queue.get_action_index()
                         prev_actions = queue.get_left_over()
+                        has_previous_actions = prev_actions is not None and prev_actions.numel() > 0
 
+                        training_max_delay = int(getattr(self._policy.config, "rtc_training_max_delay", 0))
                         latency = latency_tracker.max()
-                        delay = math.ceil(latency / time_per_chunk) if latency else 0
+                        delay = _estimate_rtc_delay(
+                            latency=latency,
+                            time_per_step=time_per_chunk,
+                            mode=self._rtc_config.mode,
+                            training_max_delay=training_max_delay,
+                            has_previous_actions=has_previous_actions,
+                        )
+                        if self._rtc_config.mode == "trained" and delay > 0:
+                            available_steps = 0 if prev_actions is None else prev_actions.shape[0]
+                            _validate_trained_rtc_prefix_available(
+                                conditioned_delay=delay,
+                                available_steps=available_steps,
+                            )
 
                         obs_batch = build_dataset_frame(self._hw_features, obs, prefix="observation")
                         obs_batch = prepare_observation_for_inference(
@@ -316,11 +386,32 @@ class RTCInferenceEngine(InferenceEngine):
                         inference_count += 1
                         consecutive_errors = 0
                         is_warmup = self._use_torch_compile and inference_count <= warmup_required
-                        if is_warmup:
+                        is_initial_trained_chunk = (
+                            self._rtc_config.mode == "trained" and not has_previous_actions
+                        )
+                        if is_warmup or is_initial_trained_chunk:
                             latency_tracker.reset()
                         else:
                             latency_tracker.add(new_latency)
 
+                        if (
+                            not is_warmup
+                            and self._rtc_config.mode == "trained"
+                            and not _trained_rtc_chunk_can_merge(
+                                conditioned_delay=delay,
+                                measured_delay=new_delay,
+                                training_max_delay=training_max_delay,
+                                has_previous_actions=has_previous_actions,
+                            )
+                        ):
+                            logger.warning(
+                                "Discarding trained RTC chunk: measured delay %d exceeded "
+                                "conditioned delay %d; retrying with updated latency",
+                                new_delay,
+                                delay,
+                            )
+                            continue
+
                         queue.merge(original, processed, new_delay, idx_before)
 
                         if (
@@ -333,6 +424,8 @@ class RTCInferenceEngine(InferenceEngine):
 
                         logger.debug("RTC inference latency=%.2fs, queue=%d", new_latency, queue.qsize())
 
+                    except _FatalRTCInferenceError:
+                        raise
                     except Exception as e:
                         consecutive_errors += 1
                         logger.error(

src/lerobot/rollout/inference/sync.py

@@ -24,21 +24,26 @@ import torch
 
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
-from lerobot.processor import PolicyProcessorPipeline, RelativeActionsProcessorStep
+from lerobot.processor import PolicyProcessorPipeline
 
 from .base import InferenceEngine
 
 logger = logging.getLogger(__name__)
 
 
-# Relative-action support: a predicted chunk of offsets is anchored to the robot
-# state at prediction time, but the sync engine reruns the pre/post pipeline every
-# tick, so ``RelativeActionsProcessorStep`` would re-anchor cached actions to the
-# current (moved) state and drift through the chunk. We pin the anchor per chunk:
-# a probe on the policy's public ``predict_action_chunk`` flags the ticks that
-# predict a fresh chunk; on the others the engine restores the anchor the relative
-# step overwrote. ``select_action`` stays on the hot path, so per-tick side effects
-# (e.g. LingBot-VA keyframe feedback) are preserved.
+# TODO(Steven): support relative-action policies.  The per-tick flow refreshes
+# ``RelativeActionsProcessorStep._last_state`` every call, so cached chunk
+# actions popped on later ticks get reanchored to the *current* robot state and
+# absolute targets drift through the chunk.  Relative-action policies are
+# rejected at context-build time today; RTC postprocesses the whole chunk and
+# is unaffected.
+#
+# Candidate fix: drive the policy via ``predict_action_chunk`` and serve a
+# local FIFO of postprocessed actions.  Eliminates drift by construction and
+# saves per-tick pre/post work, but bypasses ``select_action`` — needs
+# fallbacks for SAC (raises), ACT temporal ensembling (ensembler lives in
+# ``select_action``), and Diffusion-family (obs-history queues populated as a
+# side effect of ``select_action``).
 
 
 class SyncInferenceEngine(InferenceEngine):
@@ -68,31 +73,6 @@ class SyncInferenceEngine(InferenceEngine):
         self._task = task
         self._device = torch.device(device or "cpu")
         self._robot_type = robot_type
-
-        # Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk
-        # (see module comment), mirroring the RTC engine.
-        self._relative_step = next(
-            (
-                s
-                for s in getattr(preprocessor, "steps", ())
-                if isinstance(s, RelativeActionsProcessorStep) and s.enabled
-            ),
-            None,
-        )
-        # Set by the probe for the current tick / ever, respectively.
-        self._chunk_predicted = False
-        self._ever_predicted_chunk = False
-        self._original_predict_action_chunk = None  # set while the probe is installed
-        if self._relative_step is not None:
-            # ``action_names`` is optional on the step; fill it lazily from the
-            # policy/dataset so the relative<->absolute mask is built correctly. This is
-            # a deliberate engine->step side effect (the step is configured by its consumer).
-            if self._relative_step.action_names is None:
-                cfg_names = getattr(policy.config, "action_feature_names", None)
-                self._relative_step.action_names = list(cfg_names) if cfg_names else list(ordered_action_keys)
-            self._install_chunk_probe()
-            logger.info("Relative actions enabled: chunk anchor pinned per predicted chunk")
-
         logger.info(
             "SyncInferenceEngine initialized (device=%s, action_keys=%d)",
             self._device,
@@ -105,11 +85,6 @@ class SyncInferenceEngine(InferenceEngine):
 
     def stop(self) -> None:
         """No background resources to stop."""
-        # Undo the probe so the policy object isn't left permanently patched
-        # (it may outlive this engine or be reused by another).
-        if self._original_predict_action_chunk is not None:
-            self._policy.predict_action_chunk = self._original_predict_action_chunk
-            self._original_predict_action_chunk = None
         logger.info("SyncInferenceEngine stopped")
 
     def reset(self) -> None:
@@ -118,27 +93,6 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
-        # New episode: the next tick predicts a fresh chunk and re-anchors.
-        self._chunk_predicted = False
-        self._ever_predicted_chunk = False
-
-    def _install_chunk_probe(self) -> None:
-        """Wrap the policy's public ``predict_action_chunk`` so we learn which ticks
-        predict a fresh chunk (when the anchor must advance) without introspecting any
-        private action queue. Chunking policies call it from ``select_action``.
-
-        Wraps whatever callable is currently bound (e.g. an already-``torch.compile``d
-        one, since ``build_rollout_context`` compiles before building the engine); undone
-        in ``stop()``."""
-        self._original_predict_action_chunk = self._policy.predict_action_chunk
-        inner = self._original_predict_action_chunk
-
-        def probe(*args, **kwargs):
-            self._chunk_predicted = True
-            self._ever_predicted_chunk = True
-            return inner(*args, **kwargs)
-
-        self._policy.predict_action_chunk = probe
 
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
         """Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
@@ -153,25 +107,12 @@ class SyncInferenceEngine(InferenceEngine):
             if self._device.type == "cuda" and self._policy.config.use_amp
             else nullcontext()
         )
-        # Snapshot the chunk anchor before the preprocessor overwrites it with this
-        # tick's state; restore it below if this tick only served a cached action.
-        # ``clone`` so the snapshot survives even if the cached tensor is ever mutated
-        # in place (today it is only rebound, but the copy is cheap for a state vector).
-        anchor_before = None
-        if self._relative_step is not None:
-            cached = self._relative_step.get_cached_state()
-            anchor_before = cached.clone() if cached is not None else None
-        self._chunk_predicted = False
         with torch.inference_mode(), autocast_ctx:
             observation = prepare_observation_for_inference(
                 observation, self._device, self._task, self._robot_type
             )
             observation = self._preprocessor(observation)
             action = self._policy.select_action(observation)
-            # Hold the anchor only for a chunking policy serving a cached action this
-            # tick; policies that never chunk or that recomputed keep refreshing.
-            if self._relative_step is not None and self._ever_predicted_chunk and not self._chunk_predicted:
-                self._relative_step.set_cached_state(anchor_before)
             action = self._postprocessor(action)
         action_tensor = action.squeeze(0).cpu()
 

src/lerobot/scripts/lerobot_annotate.py

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

src/lerobot/scripts/lerobot_edit_dataset.py

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

src/lerobot/scripts/lerobot_train.py

@@ -171,9 +171,6 @@ def update_policy(
     train_metrics.update_s = time.perf_counter() - start_time
     if torch.cuda.is_available():
         train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
-    # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
-    if output_dict:
-        train_metrics.update_metrics(output_dict)
     return train_metrics, output_dict
 
 
@@ -346,6 +343,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                 "enabled": True,
                 "exclude_joints": getattr(active_cfg, "relative_exclude_joints", []),
                 "action_names": getattr(active_cfg, "action_feature_names", None),
+                "pose_representation": getattr(active_cfg, "relative_pose_representation", "componentwise"),
+                "se3_pose_groups": getattr(active_cfg, "relative_se3_pose_groups", []),
             }
             postprocessor_overrides["absolute_actions_processor"] = {"enabled": True}
         processor_kwargs["preprocessor_overrides"] = preprocessor_overrides
@@ -575,7 +574,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        train_tracker, _ = update_policy(
+        train_tracker, output_dict = update_policy(
             train_tracker,
             policy,
             batch,
@@ -608,10 +607,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     train_tracker.samples_per_s = effective_batch_size / step_time
                 logging.info(train_tracker)
                 if wandb_logger:
-                    # Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
-                    # tracker by update_policy, so to_dict() already carries their windowed,
-                    # rank-reduced averages — no per-step output_dict passthrough needed.
                     wandb_log_dict = train_tracker.to_dict()
+                    if output_dict:
+                        wandb_log_dict.update(output_dict)
                     # Log sample weighting statistics if enabled
                     if sample_weighter is not None:
                         weighter_stats = sample_weighter.get_stats()

src/lerobot/utils/device_utils.py

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

src/lerobot/utils/logging_utils.py

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

tests/annotations/test_modules.py

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

tests/policies/pi0_pi05/test_pi05_training_time_rtc.py

@@ -0,0 +1,94 @@
+#!/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.
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.modeling_pi05 import (  # noqa: E402
+    _build_flow_matching_inputs,
+    _prepare_trained_rtc_prefix,
+    _reduce_training_rtc_loss,
+    create_sinusoidal_pos_embedding,
+)
+from lerobot.policies.pi_gemma import PiGemmaRMSNorm  # noqa: E402
+
+
+def test_pi05_training_rtc_uses_clean_prefix_and_per_token_time():
+    actions = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]])
+    noise = torch.tensor([[[10.0], [20.0], [30.0], [40.0]]])
+    time = torch.tensor([0.25])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
+
+    assert model_time.tolist() == [[0.0, 0.0, 0.25, 0.25]]
+    assert torch.equal(x_t[:, :2], actions[:, :2])
+    assert torch.equal(x_t[:, 2:], 0.25 * noise[:, 2:] + 0.75 * actions[:, 2:])
+
+
+def test_pi05_training_rtc_loss_excludes_clean_prefix():
+    losses = torch.tensor([[[100.0], [100.0], [2.0], [4.0]]])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean")
+
+    assert loss.item() == pytest.approx(3.0)
+
+
+def test_pi05_training_rtc_adaptive_norm_accepts_per_action_time_conditioning():
+    norm = PiGemmaRMSNorm(dim=4, cond_dim=3)
+    hidden = torch.randn(2, 5, 4)
+    per_action_condition = torch.randn(2, 5, 3)
+
+    output, gate = norm(hidden, per_action_condition)
+
+    assert output.shape == hidden.shape
+    assert gate.shape == hidden.shape
+
+
+def test_pi05_training_rtc_embeds_per_action_timesteps():
+    per_action_time = torch.tensor([[0.0, 0.0, 0.25, 0.25]])
+
+    embedding = create_sinusoidal_pos_embedding(
+        per_action_time,
+        dimension=8,
+        min_period=4e-3,
+        max_period=4.0,
+        device=per_action_time.device,
+    )
+
+    assert embedding.shape == (1, 4, 8)
+    torch.testing.assert_close(embedding[:, 0], embedding[:, 1])
+    torch.testing.assert_close(embedding[:, 2], embedding[:, 3])
+
+
+def test_pi05_trained_rtc_prefix_is_padded_to_model_width():
+    latent = torch.zeros(1, 5, 32)
+    previous = torch.arange(30, dtype=torch.float32).reshape(1, 3, 10)
+
+    prefix, mask = _prepare_trained_rtc_prefix(
+        latent,
+        previous,
+        inference_delay=2,
+        training_max_delay=4,
+    )
+
+    assert prefix.shape == latent.shape
+    assert mask.shape == latent.shape
+    torch.testing.assert_close(prefix[:, :2, :10], previous[:, :2])
+    assert torch.count_nonzero(prefix[:, :2, 10:]) == 0
+    assert mask[:, :2].all()
+    assert not mask[:, 2:].any()
+
+
+@pytest.mark.parametrize("max_delay", [-1, 5])
+def test_pi05_config_rejects_invalid_training_rtc_delay(max_delay):
+    with pytest.raises(ValueError, match="rtc_training_max_delay"):
+        PI05Config(chunk_size=5, n_action_steps=5, rtc_training_max_delay=max_delay)

tests/policies/rtc/test_configuration_rtc.py

@@ -16,6 +16,8 @@
 
 """Tests for RTC configuration module."""
 
+import pytest
+
 from lerobot.configs.types import RTCAttentionSchedule
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
 
@@ -27,6 +29,7 @@ def test_rtc_config_default_initialization():
     config = RTCConfig()
 
     assert config.enabled is True
+    assert config.mode == "guided"
     assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
     assert config.max_guidance_weight == 10.0
     assert config.execution_horizon == 10
@@ -34,10 +37,16 @@ def test_rtc_config_default_initialization():
     assert config.debug_maxlen == 100
 
 
+def test_rtc_config_rejects_unknown_mode():
+    with pytest.raises(ValueError, match="mode must be"):
+        RTCConfig(mode="unknown")
+
+
 def test_rtc_config_custom_initialization():
     """Test RTCConfig initializes with custom values."""
     config = RTCConfig(
         enabled=True,
+        mode="trained",
         prefix_attention_schedule=RTCAttentionSchedule.EXP,
         max_guidance_weight=5.0,
         execution_horizon=20,
@@ -46,6 +55,7 @@ def test_rtc_config_custom_initialization():
     )
 
     assert config.enabled is True
+    assert config.mode == "trained"
     assert config.prefix_attention_schedule == RTCAttentionSchedule.EXP
     assert config.max_guidance_weight == 5.0
     assert config.execution_horizon == 20

tests/policies/rtc/test_modeling_rtc.py

@@ -93,6 +93,19 @@ def test_rtc_processor_initialization_without_debug(rtc_config_debug_disabled):
     assert processor.tracker is None
 
 
+def test_rtc_processor_rejects_trained_mode_when_policy_does_not_support_it():
+    config = RTCConfig(mode="trained")
+
+    with pytest.raises(ValueError, match="requires a PI05-compatible checkpoint"):
+        RTCProcessor(config)
+
+    processor = RTCProcessor(config, trained_mode_supported=True)
+    assert processor.rtc_config.mode == "trained"
+
+    disabled = RTCProcessor(RTCConfig(enabled=False, mode="trained"))
+    assert disabled.rtc_config.enabled is False
+
+
 # ====================== Tracker Proxy Methods Tests ======================
 
 

tests/policies/rtc/test_rtc_relative_actions.py

@@ -505,6 +505,44 @@ class TestRTCReanchoringWithStateNormalizer:
         assert not torch.allclose(cached, post_normalize_state, atol=1e-3)
 
 
+def test_reanchor_se3_6d_prefix_uses_current_camera_frame_and_model_width():
+    """A leftover absolute EE chunk is recomposed relative to the latest camera-frame EE pose."""
+    names = ["pos_x", "pos_y", "pos_z", "rot_x", "rot_y", "rot_z", "gripper"]
+    relative_step = RelativeActionsProcessorStep(
+        enabled=True,
+        exclude_joints=["gripper"],
+        action_names=names,
+        pose_representation="se3_6d",
+        se3_pose_groups=[list(range(6))],
+    )
+    current_state = torch.tensor([[0.20, -0.10, 0.40, 0.31, -0.22, 0.17, 0.03]])
+    previous_absolute = torch.tensor(
+        [
+            [0.28, -0.03, 0.46, -0.18, 0.27, 0.41, 0.02],
+            [0.31, 0.02, 0.50, -0.11, 0.35, 0.52, 0.01],
+        ]
+    )
+
+    result = reanchor_relative_rtc_prefix(
+        prev_actions_absolute=previous_absolute,
+        current_state=current_state,
+        relative_step=relative_step,
+        normalizer_step=None,
+        policy_device="cpu",
+    )
+
+    expected = to_relative_actions(
+        previous_absolute,
+        current_state,
+        relative_step._build_mask(previous_absolute.shape[-1]),
+        pose_representation="se3_6d",
+        se3_pose_groups=[list(range(6))],
+    )
+    assert result.shape == (2, 10)
+    torch.testing.assert_close(result, expected, atol=1e-6, rtol=1e-6)
+    torch.testing.assert_close(result[:, -1], previous_absolute[:, -1])
+
+
 def _detect_relative_actions(preprocessor) -> bool:
     """Mirror of the helper in lerobot-rollout for testing without importing it."""
     return any(isinstance(step, RelativeActionsProcessorStep) and step.enabled for step in preprocessor.steps)

tests/policies/test_relative_actions.py

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

tests/processor/test_pi05_state_from_action.py

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

tests/processor/test_se3_relative_actions.py

@@ -0,0 +1,144 @@
+import math
+
+import pytest
+import torch
+
+from lerobot.processor.relative_action_processor import (
+    rotation_6d_to_rotvec,
+    rotvec_to_rotation_6d,
+    to_absolute_actions,
+    to_absolute_se3_pose,
+    to_absolute_se3_pose_6d,
+    to_relative_actions,
+    to_relative_se3_pose,
+    to_relative_se3_pose_6d,
+)
+
+POSE_GROUP = [list(range(6))]
+
+
+def test_se3_translation_is_expressed_in_reference_frame():
+    reference = torch.tensor([[1.0, 2.0, 3.0, 0.0, 0.0, math.pi / 2]])
+    target = torch.tensor([[1.0, 3.0, 3.0, 0.0, 0.0, math.pi / 2]])
+
+    relative = to_relative_se3_pose(target, reference)
+
+    torch.testing.assert_close(relative, torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0]]), atol=1e-6, rtol=1e-6)
+
+
+def test_se3_pose_roundtrip_for_batched_chunks():
+    torch.manual_seed(0)
+    reference = torch.randn(4, 6)
+    reference[:, 3:] *= 0.8
+    target = torch.randn(4, 11, 6)
+    target[..., 3:] *= 0.8
+
+    relative = to_relative_se3_pose(target, reference.unsqueeze(1))
+    recovered = to_absolute_se3_pose(relative, reference.unsqueeze(1))
+
+    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
+
+
+def test_mixed_se3_pose_and_absolute_gripper_roundtrip():
+    reference = torch.tensor([[0.2, -0.1, 0.4, 0.1, 0.2, -0.3, 0.06]])
+    target = torch.tensor([[[0.3, 0.2, 0.5, -0.2, 0.1, 0.4, 0.03], [0.1, -0.3, 0.2, 0.5, -0.1, 0.2, 0.05]]])
+    mask = [True, True, True, True, True, True, False]
+
+    relative = to_relative_actions(
+        target,
+        reference,
+        mask,
+        pose_representation="se3",
+        se3_pose_groups=POSE_GROUP,
+    )
+    recovered = to_absolute_actions(
+        relative,
+        reference,
+        mask,
+        pose_representation="se3",
+        se3_pose_groups=POSE_GROUP,
+    )
+
+    torch.testing.assert_close(relative[..., 6], target[..., 6])
+    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
+
+
+def test_se3_pose_group_cannot_be_partially_relative():
+    with pytest.raises(ValueError, match="wholly relative or wholly absolute"):
+        to_relative_actions(
+            torch.zeros(1, 7),
+            torch.zeros(1, 7),
+            [True, True, True, False, False, False, False],
+            pose_representation="se3",
+            se3_pose_groups=POSE_GROUP,
+        )
+
+
+@pytest.mark.parametrize(
+    "rotvec",
+    [
+        [0.0, 0.0, 0.0],
+        [0.2, -0.5, 0.8],
+        [math.pi - 1e-4, 0.0, 0.0],
+    ],
+)
+def test_rotation_6d_roundtrip(rotvec):
+    source = torch.tensor([rotvec], dtype=torch.float64)
+
+    recovered = rotation_6d_to_rotvec(rotvec_to_rotation_6d(source))
+
+    torch.testing.assert_close(recovered, source, atol=2e-6, rtol=2e-6)
+
+
+def test_rotation_6d_uses_umi_first_two_rows():
+    source = torch.tensor([[0.0, 0.0, math.pi / 2]], dtype=torch.float64)
+
+    encoded = rotvec_to_rotation_6d(source)
+
+    expected = torch.tensor([[0.0, -1.0, 0.0, 1.0, 0.0, 0.0]], dtype=torch.float64)
+    torch.testing.assert_close(encoded, expected, atol=1e-7, rtol=1e-7)
+    torch.testing.assert_close(rotation_6d_to_rotvec(expected), source, atol=1e-7, rtol=1e-7)
+
+
+def test_se3_6d_pose_roundtrip_for_batched_chunks():
+    torch.manual_seed(1)
+    reference = torch.randn(4, 6)
+    reference[:, 3:] *= 0.8
+    target = torch.randn(4, 11, 6)
+    target[..., 3:] *= 0.8
+
+    relative = to_relative_se3_pose_6d(target, reference.unsqueeze(1))
+    recovered = to_absolute_se3_pose_6d(relative, reference.unsqueeze(1))
+
+    assert relative.shape == (4, 11, 9)
+    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
+
+
+def test_mixed_se3_6d_pose_and_absolute_gripper_roundtrip():
+    reference = torch.tensor([[0.2, -0.1, 0.4, 0.1, 0.2, -0.3, 0.06]])
+    target = torch.tensor([[[0.3, 0.2, 0.5, -0.2, 0.1, 0.4, 0.03], [0.1, -0.3, 0.2, 0.5, -0.1, 0.2, 0.05]]])
+    mask = [True, True, True, True, True, True, False]
+
+    relative = to_relative_actions(
+        target,
+        reference,
+        mask,
+        pose_representation="se3_6d",
+        se3_pose_groups=POSE_GROUP,
+    )
+    recovered = to_absolute_actions(
+        relative,
+        reference,
+        mask,
+        pose_representation="se3_6d",
+        se3_pose_groups=POSE_GROUP,
+    )
+
+    assert relative.shape == (1, 2, 10)
+    torch.testing.assert_close(relative[..., 9], target[..., 6])
+    torch.testing.assert_close(recovered, target, atol=2e-5, rtol=2e-5)
+
+
+def test_rotation_6d_rejects_degenerate_prediction():
+    with pytest.raises(ValueError, match="degenerate"):
+        rotation_6d_to_rotvec(torch.zeros(1, 6))

tests/scripts/test_lerobot_annotate.py

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

tests/test_rollout.py

@@ -17,6 +17,7 @@
 from __future__ import annotations
 
 import dataclasses
+from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -98,6 +99,131 @@ def test_inference_config_types():
     assert rtc.rtc is not None
 
 
+def test_trained_rtc_retries_chunk_when_measured_delay_exceeds_conditioning():
+    from lerobot.rollout.inference.rtc import _trained_rtc_chunk_can_merge
+
+    assert not _trained_rtc_chunk_can_merge(
+        conditioned_delay=2,
+        measured_delay=3,
+        training_max_delay=4,
+        has_previous_actions=True,
+    )
+    assert _trained_rtc_chunk_can_merge(
+        conditioned_delay=2,
+        measured_delay=5,
+        training_max_delay=4,
+        has_previous_actions=False,
+    )
+
+
+def test_trained_rtc_bootstraps_first_overlap_with_checkpoint_capacity():
+    from lerobot.rollout.inference.rtc import _estimate_rtc_delay
+
+    assert (
+        _estimate_rtc_delay(
+            latency=0,
+            time_per_step=1 / 30,
+            mode="trained",
+            training_max_delay=10,
+            has_previous_actions=False,
+        )
+        == 0
+    )
+    assert (
+        _estimate_rtc_delay(
+            latency=0,
+            time_per_step=1 / 30,
+            mode="trained",
+            training_max_delay=10,
+            has_previous_actions=True,
+        )
+        == 10
+    )
+
+
+def test_trained_rtc_rejects_measured_delay_above_checkpoint_support():
+    from lerobot.rollout.inference.rtc import (
+        _trained_rtc_chunk_can_merge,
+        _TrainedRTCDelayExceededError,
+    )
+
+    with pytest.raises(_TrainedRTCDelayExceededError, match="rtc_training_max_delay"):
+        _trained_rtc_chunk_can_merge(
+            conditioned_delay=3,
+            measured_delay=5,
+            training_max_delay=4,
+            has_previous_actions=True,
+        )
+
+
+def test_trained_rtc_rejects_prefix_shorter_than_conditioned_delay():
+    from lerobot.rollout.inference.rtc import (
+        _TrainedRTCPrefixUnavailableError,
+        _validate_trained_rtc_prefix_available,
+    )
+
+    with pytest.raises(_TrainedRTCPrefixUnavailableError, match="only 2"):
+        _validate_trained_rtc_prefix_available(conditioned_delay=4, available_steps=2)
+
+
+@pytest.mark.parametrize(
+    ("execution_horizon", "queue_threshold", "match"),
+    [
+        (3, 4, "execution_horizon"),
+        (4, 3, "queue_threshold"),
+    ],
+)
+def test_trained_rtc_rollout_requires_capacity_for_max_delay(execution_horizon, queue_threshold, match):
+    from lerobot.policies.rtc.configuration_rtc import RTCConfig
+    from lerobot.rollout.context import _validate_trained_rtc_rollout_config
+    from lerobot.rollout.inference import RTCInferenceConfig
+
+    policy_config = SimpleNamespace(type="pi05", rtc_training_max_delay=4)
+    inference_config = RTCInferenceConfig(
+        rtc=RTCConfig(mode="trained", execution_horizon=execution_horizon),
+        queue_threshold=queue_threshold,
+    )
+
+    with pytest.raises(ValueError, match=match):
+        _validate_trained_rtc_rollout_config(policy_config, inference_config)
+
+
+def test_relative_state_order_follows_checkpoint_action_names():
+    from lerobot.rollout.context import _align_relative_state_feature_order
+    from lerobot.utils.constants import OBS_STATE
+    from lerobot.utils.feature_utils import build_dataset_frame
+
+    hw_features = {
+        OBS_STATE: {
+            "dtype": "float32",
+            "shape": (4,),
+            "names": ["left_joint.pos", "left_gripper.pos", "right_joint.pos", "right_gripper.pos"],
+        }
+    }
+    checkpoint_order = [
+        "right_joint.pos",
+        "right_gripper.pos",
+        "left_joint.pos",
+        "left_gripper.pos",
+    ]
+
+    aligned = _align_relative_state_feature_order(hw_features, checkpoint_order)
+    frame = build_dataset_frame(
+        aligned,
+        {
+            "left_joint.pos": 1.0,
+            "left_gripper.pos": 2.0,
+            "right_joint.pos": 3.0,
+            "right_gripper.pos": 4.0,
+        },
+        prefix="observation",
+    )
+
+    assert aligned[OBS_STATE]["names"] == checkpoint_order
+    assert frame[OBS_STATE].tolist() == [3.0, 4.0, 1.0, 2.0]
+    assert hw_features[OBS_STATE]["names"][0] == "left_joint.pos"
+
+
 def test_sentry_config_defaults():
     from lerobot.rollout import SentryStrategyConfig
 
@@ -348,225 +474,3 @@ def test_rollout_context_fields():
 
     field_names = {f.name for f in dataclasses.fields(RolloutContext)}
     assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
-
-
-# ---------------------------------------------------------------------------
-# Sync engine: relative-action anchoring (drift-free chunk execution)
-# ---------------------------------------------------------------------------
-
-_REL_ACTION_NAMES = ["j0.pos", "j1.pos", "j2.pos", "gripper.pos"]
-_REL_ACTION_DIM = len(_REL_ACTION_NAMES)
-
-
-def _relative_pre_post(exclude_joints=None):
-    """Pre/post processors wrapping the real relative (caches anchor) and absolute
-    (relative + cached state) steps, mirroring what the sync engine feeds them."""
-    from lerobot.processor import (
-        AbsoluteActionsProcessorStep,
-        RelativeActionsProcessorStep,
-        TransitionKey,
-        create_transition,
-    )
-    from lerobot.utils.constants import OBS_STATE
-
-    relative_step = RelativeActionsProcessorStep(
-        enabled=True, exclude_joints=exclude_joints or [], action_names=list(_REL_ACTION_NAMES)
-    )
-    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
-
-    class _Pre:
-        steps = [relative_step]
-
-        def __call__(self, observation):
-            # Run the relative step so it caches the anchor, then pass the batch through.
-            transition = create_transition(observation={OBS_STATE: observation[OBS_STATE]})
-            relative_step(transition)
-            return observation
-
-        def reset(self):
-            pass
-
-    class _Post:
-        def __call__(self, action):
-            transition = create_transition(action=action)
-            return absolute_step(transition)[TransitionKey.ACTION]
-
-        def reset(self):
-            pass
-
-    return _Pre(), _Post(), relative_step
-
-
-def _fake_relative_policy(chunk_rel, n_action_steps, chunking=True):
-    """Fake relative-action policy for the sync engine.
-
-    ``chunking=True`` buffers a chunk and serves it one action per tick, calling the
-    public ``predict_action_chunk`` only on refill (pi0/fastwam/lingbot). ``False``
-    returns an action directly and never calls it. The engine's anchor probe keys off
-    that public call, so the fake routes through it rather than any private queue.
-    """
-    from collections import deque
-
-    policy = MagicMock()
-    policy.config.use_amp = False
-    policy.config.action_feature_names = list(_REL_ACTION_NAMES)
-    state = {"predict_calls": 0}
-    queue = deque(maxlen=n_action_steps)
-
-    def predict_action_chunk(_batch=None, **_kwargs):
-        state["predict_calls"] += 1
-        return chunk_rel.unsqueeze(0)  # [B=1, n, dim]
-
-    def select_action(_observation):
-        if not chunking:
-            return chunk_rel[0].unsqueeze(0)
-        if len(queue) == 0:
-            actions = policy.predict_action_chunk(_observation)
-            queue.extend(actions.transpose(0, 1))  # [n, 1, dim]
-        return queue.popleft()
-
-    policy.predict_action_chunk.side_effect = predict_action_chunk
-    policy.select_action.side_effect = select_action
-    policy.reset.side_effect = queue.clear
-    policy._predict_state = state
-    return policy
-
-
-def _build_sync_engine(policy, pre, post):
-    from lerobot.rollout import SyncInferenceEngine
-
-    return SyncInferenceEngine(
-        policy=policy,
-        preprocessor=pre,
-        postprocessor=post,
-        dataset_features={"action": {"names": list(_REL_ACTION_NAMES)}},
-        ordered_action_keys=list(_REL_ACTION_NAMES),
-        task="test",
-        device="cpu",
-        robot_type="mock",
-    )
-
-
-def _obs_frame(state_values):
-    import numpy as np
-
-    return {"observation.state": np.asarray(state_values, dtype=np.float32)}
-
-
-def test_sync_relative_holds_anchor_across_chunk():
-    """Every action popped within a chunk must anchor to the tick-0 state (no drift)."""
-    n = 4
-    # A distinct relative offset per chunk step so a wrong anchor would be visible.
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
-    pre, post, relative_step = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
-    engine = _build_sync_engine(policy, pre, post)
-
-    assert engine._relative_step is relative_step  # introspection wired the step
-
-    s0 = [1.0, 2.0, 3.0, 4.0]
-    outputs = []
-    for tick in range(n):
-        # Feed a *different* state each tick; a drifting anchor would use it.
-        state = [v + tick for v in s0]
-        outputs.append(engine.get_action(_obs_frame(state)))
-
-    # Exactly one chunk was predicted across the n ticks.
-    assert policy._predict_state["predict_calls"] == 1
-    for tick in range(n):
-        expected = torch.tensor(s0) + chunk_rel[tick]
-        torch.testing.assert_close(outputs[tick], expected)
-
-    # Next tick empties the queue -> fresh chunk -> anchor advances to the new state.
-    s_next = [10.0, 20.0, 30.0, 40.0]
-    out = engine.get_action(_obs_frame(s_next))
-    assert policy._predict_state["predict_calls"] == 2
-    torch.testing.assert_close(out, torch.tensor(s_next) + chunk_rel[0])
-    # The anchor now reflects the fresh-chunk state, not the held one.
-    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next]))
-
-
-def test_sync_relative_reset_reanchors_new_episode():
-    """After ``reset()`` the first tick of the next episode anchors to the new state."""
-    n = 3
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
-    pre, post, relative_step = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
-    engine = _build_sync_engine(policy, pre, post)
-
-    # Episode 1: one tick anchors to s0 and leaves cached actions in the queue.
-    engine.get_action(_obs_frame([1.0, 1.0, 1.0, 1.0]))
-    assert policy._predict_state["predict_calls"] == 1
-
-    engine.reset()  # clears the queue and the per-episode chunk flags
-
-    # Episode 2: a fresh state must produce a fresh chunk anchored to that state,
-    # not carry over the previous episode's anchor.
-    s_new = [7.0, 8.0, 9.0, 10.0]
-    out = engine.get_action(_obs_frame(s_new))
-    assert policy._predict_state["predict_calls"] == 2
-    torch.testing.assert_close(out, torch.tensor(s_new) + chunk_rel[0])
-    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_new]))
-
-
-def test_sync_relative_non_chunking_policy_refreshes_every_tick():
-    """A policy that never calls ``predict_action_chunk`` must not freeze the anchor."""
-    n = 3
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.5) for _ in range(n)])
-    pre, post, _ = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n, chunking=False)
-    engine = _build_sync_engine(policy, pre, post)
-
-    s0 = [1.0, 1.0, 1.0, 1.0]
-    for tick in range(3):
-        state = [v + tick for v in s0]
-        out = engine.get_action(_obs_frame(state))
-        # Anchor must track the current state every tick (no chunk => no hold).
-        torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
-    assert policy._predict_state["predict_calls"] == 0
-
-
-def test_sync_engine_no_relative_step_is_none():
-    """Without an enabled relative step, the engine takes the plain select_action path."""
-    policy = MagicMock()
-    policy.config.use_amp = False
-    engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock())
-    assert engine._relative_step is None
-
-
-def test_sync_relative_exclude_joints_stay_absolute():
-    """With ``exclude_joints``, excluded dims pass through absolute while the relative
-    dims still hold the tick-0 anchor across the chunk."""
-    n = 4
-    # Distinct offset per step *and* per dim so a wrong anchor or a wrong mask shows up.
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
-    pre, post, relative_step = _relative_pre_post(exclude_joints=["gripper"])
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
-    engine = _build_sync_engine(policy, pre, post)
-
-    # gripper (last dim) is kept absolute; j0..j2 are relative.
-    mask = torch.tensor([1.0, 1.0, 1.0, 0.0])
-    s0 = [1.0, 2.0, 3.0, 4.0]
-    outputs = []
-    for tick in range(n):
-        state = [v + tick for v in s0]  # moving state; a drifting anchor would use it
-        outputs.append(engine.get_action(_obs_frame(state)))
-
-    assert policy._predict_state["predict_calls"] == 1  # one chunk held across n ticks
-    for tick in range(n):
-        # relative dims: anchor held at s0; excluded gripper dim: raw predicted value.
-        expected = chunk_rel[tick] + torch.tensor(s0) * mask
-        torch.testing.assert_close(outputs[tick], expected)
-
-
-def test_sync_relative_stop_restores_policy_method():
-    """``stop()`` un-patches the probe so the policy object isn't permanently modified."""
-    n = 3
-    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
-    pre, post, _ = _relative_pre_post()
-    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
-    original = policy.predict_action_chunk
-    engine = _build_sync_engine(policy, pre, post)
-    assert policy.predict_action_chunk is not original  # probe installed
-    engine.stop()
-    assert policy.predict_action_chunk is original  # restored

tests/utils/test_logging_utils.py

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