refactor(rollout): integrate feedback -> api, recording, log mut and keyboard

feat(rollout): add subtask command
feat(rollout): mute logs in interactive mode
2026-08-08 17:39:44 +00:00 · 2026-08-07 21:33:17 +02:00 · 2026-08-07 16:44:29 +02:00 · 2026-08-07 16:00:56 +02:00 · 2026-08-07 15:22:44 +02:00 · 2026-08-07 13:52:23 +02:00
74 changed files with 8763 additions and 1334 deletions
@@ -0,0 +1,457 @@
+# Interactive Rollout — Design Notes
+
+Branch: `feat/add_interactive_rollout` · Status: Phases 1–2 committed; Round 2
+(programmatic API, sentry support, muting v2, stdin move) implemented and tested,
+uncommitted.
+
+---
+
+## 1. Vision
+
+`lerobot-rollout` runs inference on a real robot: it connects hardware, loads the policy,
+builds the processor pipelines, optionally records a dataset, and spins the control loop.
+Today that is a **one-shot, fire-and-forget** program. You pass `--task="pick up the cube"`
+on the command line, the robot starts moving immediately, and the only interaction left is
+Ctrl-C. If you want a different instruction, you kill the process and pay the full startup
+cost again — reconnecting motors, re-homing, re-loading a multi-GB VLA onto the GPU.
+
+Since LeRobot gained subtask annotation and language conditioning, that model is the
+bottleneck. The **north star** is a chat-style CLI over stdin, where the operator issues
+commands *concurrently with the robot moving*:
+
+```
+/start                              begin (or resume) the policy control loop
+/subtask Grab the red cube          re-instruct the policy on the fly
+/ask what's the capital of France?  query an LLM while the robot keeps moving
+/reset                              stop movement, return home, clear the subtask —
+                                    but keep hardware and policy warm
+/stop                               graceful shutdown
+```
+
+The unifying idea: **the expensive things (hardware, policy weights, processors) stay warm
+across commands.** Only the cheap things — the instruction, the control loop — start and
+stop. That turns a rollout from a batch job into a session you can steer.
+
+## 2. Objective (scoped)
+
+Phased, so each phase lands as a reviewable unit:
+
+| Phase | Scope | Status |
+|---|---|---|
+| **1** | `--interactive` flag, non-blocking stdin listener, command parser, `/start` `/reset` `/stop` `/help` | ✅ done |
+| **1.5** | Mute system logs so they stop fighting the prompt for the terminal | ✅ done |
+| **2** | `/subtask <text>` — change the policy's instruction mid-run | ✅ done |
+| **2.5** | Round 2: `RolloutController` public API, sentry recording support, muting v2 (errors surface), stdin listener → `lerobot/utils` | ✅ done (see §5) |
+| **3** | `/ask` + hierarchical task-vs-subtask semantics (LLM in the loop) | not started |
+
+An explicit constraint through Phases 1–2: **do not couple this to the language runtime yet.**
+Build the mechanism; keep the door open.
+
+## 3. Inspiration — three reference PRs
+
+We read all three and deliberately implemented none of them verbatim.
+
+**PR #4108 — online subtask switching.** Introduces a `PromptBroker` + `PromptListenerBase`
+ `StdinPromptListener`, a `RuntimeContext.prompt_broker` field, `register_on_change`
+callbacks, an `--online_task_switching_flush` config flag, and `flush_action_queue()` /
+`_apply_pending_flush()` on `PreTrainedPolicy` — **with edits to 14 policy files** to call
+the flush at the top of `select_action`. Its architecture is designed for pluggable input
+sources (network, voice), which is the right long-term shape but more machinery than we
+need. *What we took:* the core insight that a mid-run instruction change must invalidate
+actions precomputed under the old instruction, and that the flush must happen on a thread
+that is safe to touch policy state from.
+
+**PR #4183 — experimental full-UX draft.** Achieves the whole north-star vision, but does
+so by adding a `lerobot.runtime` / `language_runtime.py` that **duplicates** `BaseStrategy`,
+`send_next_action`, and the rollout control loop. *What we took:* the UX target and the
+command vocabulary. *What we rejected:* the parallel runtime — a second control loop is a
+second thing to keep correct, and everything it does is already in `rollout/strategies/`.
+
+**PR #4234 — policy-side edits enabling #4183's runtime.** Read for context on where the
+language plumbing lands inside a policy. Relevant to Phase 3, not to what we built.
+
+## 4. What we built, and why
+
+Three commits on the branch:
+
+```
+072c697c0  feat(rollout): interactive v1
+d3ee0b820  feat(rollout): mute logs in interactive mode
+39c4e746f  feat(rollout): add subtask command
+```
+
+Cumulative footprint — one new module, one new test file, small surgical edits elsewhere:
+
+```
+ src/lerobot/rollout/interactive.py         | 580 +++++   (new)
+ tests/test_interactive_rollout.py          | 788 +++++   (new)
+ docs/source/inference.mdx                  |  87 +++
+ src/lerobot/rollout/inference/base.py      |  66 +++
+ src/lerobot/rollout/inference/rtc.py       |  61 +-
+ src/lerobot/rollout/inference/sync.py      |  21 +-
+ src/lerobot/scripts/lerobot_rollout.py     |  32 +-
+ src/lerobot/policies/pretrained.py         |  24 +
+ src/lerobot/rollout/strategies/core.py     |  21 +-
+ src/lerobot/rollout/configs.py             |  18 +
+ src/lerobot/rollout/__init__.py            |  16 +-
+ src/lerobot/rollout/strategies/episodic.py |   4 +-
+```
+
+The ratio matters: **~1400 of ~1680 added lines are the new module and its tests.** The
+existing rollout architecture was reused, not reshaped.
+
+### 4.1 Segments over a linked event — the load-bearing idea
+
+Every rollout strategy's control loop already polls `ctx.runtime.shutdown_event.is_set()`
+to know when to stop. So instead of teaching strategies about interactivity, we **swap in a
+smarter event**:
+
+```python
+class LinkedEvent(Event):
+    """is_set() reflects the local flag OR a parent event."""
+    def is_set(self) -> bool:
+        return super().is_set() or self.parent.is_set()
+```
+
+`lerobot-rollout` wraps the `ProcessSignalHandler`'s shutdown event in a `LinkedEvent` when
+`--interactive=true`. The session sets the **local** flag to end a run *segment*; SIGINT /
+SIGTERM still arrive through the **parent**, so Ctrl-C behaves exactly as before.
+
+`InteractiveSession.run()` then drives `strategy.run(ctx)` in restartable segments:
+
+```
+setup(ctx)  →  [idle]  →  /start → run(ctx) → /reset → [idle] → /start → run(ctx) → /stop  →  teardown(ctx)
+                                                  ↑ hardware + policy stay warm throughout
+```
+
+**Zero strategy code changed** to support this. The only additions to `strategies/core.py`
+were `reset_control_state()` (engine + interpolator + cached-observation reset, factored
+out of `_init_engine` so a segment can restart cleanly) and making
+`_return_to_initial_position` public.
+
+### 4.2 Threading model
+
+```
+ listener thread  ──publishes flags / strings──▶  main thread
+ (stdin reader)      never touches hardware        (session loop → strategy.run → control loop)
+                     never mutates policy state
+```
+
+The listener only ever writes `threading.Event` flags and a lock-guarded string. Everything
+that touches hardware or policy state happens on the thread that already owns it. This
+mirrors the existing DAgger events pattern rather than inventing a new concurrency idiom.
+
+### 4.3 stdin must be read with `os.read`, not `readline`
+
+Non-obvious and load-bearing. The first implementation used `select()` + `stream.readline()`
+and **two tests failed**: a buffered file object slurps *several* lines off the file
+descriptor in one syscall, after which `select` reports the drained fd as not-ready and the
+buffered lines are never delivered. Pasted or piped command batches got stuck. The reader
+now does `select()` + `os.read(fd, 4096)` + manual `\n` splitting, with a
+blocking-`readline` fallback for streams without a `fileno()` (non-POSIX, test doubles).
+
+Also: unlike `TerminalKeyListener`, this reader leaves the terminal in **canonical mode** —
+the operator is typing chat commands, not pressing hotkeys.
+
+### 4.4 EOF means stop
+
+A closed stdin means there is no way left to command the robot, so EOF (Ctrl-D, or an
+exhausted piped script) stops the session. An unexpected read error is treated the same way,
+for the same reason. Consequence, documented: piped scripts must hold stdin open —
+
+```bash
+(printf '/start\n'; sleep 60; printf '/stop\n') | lerobot-rollout ... --interactive=true
+```
+
+### 4.5 Commands are last-write-wins
+
+`/reset` and `/stop` cancel a still-pending `/start`, so the robot never starts moving after
+the operator's most recent command said not to. Handlers set their intent flag *first* and
+the segment-stop event *second*; `_run_segment` clears the segment-stop flag *before*
+re-checking the intent flags. A `/reset` racing a `/start` is therefore either seen before
+the segment begins or ends it on its first tick.
+
+### 4.6 Base strategy only (enforced by config validation)
+
+`--interactive=true` with a recording strategy raises a `ValueError`. Two reasons: recording
+strategies finalize their dataset inside `run()` (so `run()` is not restartable), and their
+keyboard listeners contend with the command reader for the same TTY. This is a deliberate,
+documented limitation — not an oversight.
+
+### 4.7 Log muting (Phase 1.5)
+
+Policy, robot and control-loop logs at every level interleave with the chat prompt and
+destroy the typing UX. Simplest workable answer, per explicit request: **mute console output
+for the duration of the session.**
+
+- Every logger's console `StreamHandler` is raised above `CRITICAL` — **not just root**,
+  because `transformers` and `datasets` attach their own stderr handlers with
+  `propagate=False`.
+- `warnings.simplefilter("ignore")`, with `warnings.filters` saved and restored.
+- **File handlers are untouched** — anyone wanting a persistent log can attach one.
+- Restored in `run()`'s `finally`, *before* the closing `log_say`, so teardown logs are visible.
+
+The obvious hazard: muting hides fatal errors. So `InferenceEngine` gained a
+`failure_traceback` property, RTC captures its traceback in the fatal handler, and the
+session prints it on failure. **Do not remove that when touching the failure path.**
+
+"See both logs and prompt" — a pinned input line, `prompt_toolkit`-style — was deliberately
+deferred: it needs a new dependency and a real TUI layer.
+
+### 4.8 `/subtask` — the engine *is* the broker
+
+The pivotal call on Phase 2: **skip PR #4108's `PromptBroker`.** After Phase 1, the session
+already owns the stdin thread and the parser, so a broker + listener base + on-change
+callbacks + a new `RuntimeContext` field would be duplicate machinery — and callbacks firing
+on the listener thread are exactly the cross-thread hazard we designed against.
+
+Instead, `InferenceEngine` (the ABC every backend already implements) became the thread-safe
+task holder:
+
+```python
+@property
+def task(self) -> str: ...              # lock-guarded read
+
+def set_task(self, task) -> bool:       # callable from ANY thread; True if it changed
+    ...
+
+def _take_task(self) -> tuple[str, bool]:   # consumed on the INFERENCE thread;
+    ...                                     # returns (task, changed) and clears the edge
+```
+
+`/subtask` is then three lines: read `engine.task`, call `engine.set_task(text)`, print the
+transition. No new module, no new context field, no callbacks.
+
+**The flush problem, and why it got small.** When the instruction changes, a chunking policy
+is still serving actions computed under the old one — up to `chunk_size` ticks of stale
+behavior. PR #4108 solved this by adding `flush_action_queue()` / `_apply_pending_flush()`
+to `PreTrainedPolicy` **and editing 14 policy files**, because its flush request arrived from
+a foreign thread and had to be deferred to a safe point inside `select_action`.
+
+Ours already runs *on* the thread that calls `select_action`. So: one concrete method on
+`PreTrainedPolicy` and **zero per-policy edits**.
+
+```python
+def drop_queued_actions(self) -> None:
+    queues = getattr(self, "_queues", None)
+    if isinstance(queues, dict) and ACTION in queues:
+        queues[ACTION].clear()
+    action_queue = getattr(self, "_action_queue", None)
+    if action_queue is not None:
+        action_queue.clear()
+```
+
+Two `getattr`s cover the repo's two queue idioms across all ~18 policies
+(`_queues[ACTION]`: diffusion, smolvla, tdmpc, vqbet, wall_x, xvla, multi_task_dit, vla_jepa;
+`_action_queue`: act, pi0, pi05, pi0_fast, eo1, evo1, groot, molmoact2, fastwam, lingbot_va).
+Policies with no queue inherit a no-op.
+
+**Why not `policy.reset()`?** That was the first implementation, and review caught it as too
+blunt. For Diffusion it wipes the observation history, so the next chunk is planned from a
+history of the current frame repeated — a visible discontinuity mid-motion. And ACT /
+Diffusion / VQBeT / TDMPC don't read `task` at all, so they'd pay that jerk for nothing.
+`drop_queued_actions` keeps episode state and drops only what is actually stale.
+
+**RTC deliberately does *not* flush.** Clearing its queue would leave the robot with no
+commands for a full inference latency (~1 s on a VLA). Instead the next chunk is generated
+under the new instruction and merged over the previous chunk's leftover prefix — the switch
+lands within one inference and the motion stays continuous. That is exactly what RTC's
+blending exists for. Documented per-backend in `inference.mdx`; no config flag, one sensible
+default per backend.
+
+**`/reset` restores the launch task on the listener thread.** Subtle and worth preserving:
+the restore lives in `_cmd_reset`, not in `_reset_robot` (which runs later, on the main
+thread). Otherwise `/reset` followed immediately by `/subtask` would be ordered by *service*
+time rather than *command* time, and the deferred restore would silently revert the new
+instruction — deterministically so, for pasted or piped input. Both writers now run on the
+same thread, so command order wins. There is a regression test driving this through a real pipe.
+
+## 5. Round 2 — the feature becomes a library API
+
+Four follow-up asks landed together (currently uncommitted on the branch):
+make the components programmatic-API friendly (the priority), extend interactive
+to recording where cheap, simplify muting / surface errors, and settle the
+ssh/headless + `keyboard_input` question.
+
+### 5.1 `RolloutController` — programmatic control
+
+`interactive.py` bisected cleanly, so the generic control logic moved to a new
+`rollout/controller.py`:
+
+```python
+controller = RolloutController(strategy, ctx, on_event=my_observer)
+controller.serve()      # blocking loop (run it on whatever thread you like)
+controller.start()      # -> bool: False when a segment is already running
+controller.set_task(t)  # -> bool: re-instruct mid-run, from any thread
+controller.reset()      # -> bool: True when the launch task was restored
+controller.stop()
+controller.task / .initial_task / .running / .failed / .failure_traceback
+```
+
+- **No I/O of its own** — no stdin, no prints, no log muting, no TTS. Every
+  state transition that used to be a `print` is now a `RolloutEvent`
+  (`SEGMENT_STARTED`, `SEGMENT_ENDED`, `RESET_STARTED/DONE/SKIPPED`,
+  `ENGINE_FAILED`, `STOPPED`) emitted on the serve thread.
+- **Thread-safe by lock, not by convention.** The old ordering guarantee
+  (`/subtask` right after `/reset` must win) relied on both writes running on
+  the single stdin thread. The controller serializes `start`/`reset`/`stop`/
+  `set_task` with an internal lock, so the guarantee now holds for arbitrary
+  caller threads — the prerequisite for network/voice front-ends.
+- `InteractiveSession` shrank to a thin adapter: stdin listener + parser +
+  rendering + muting; each command maps 1:1 onto a controller method, and the
+  controller is exposed as `session.controller`.
+- Exported from `lerobot.rollout`: `RolloutController`, `RolloutEvent`,
+  `LinkedEvent`. `docs/source/inference.mdx` gained a **Programmatic control**
+  section with a complete embedding example.
+
+### 5.2 Sentry + interactive — recording while you steer
+
+Decision, per the agreed criteria: the `/record` keyboard-handoff idea is
+**medium-to-large** (listeners have no suspend/resume API and start at
+creation, `esc` handlers are hardcoded and collide, pynput captures globally
+while you type, and each strategy carries per-run stale flags) → rejected.
+But the investigation showed **sentry has zero keyboard code** — the config
+comment lumping it with the keyboard strategies was simply wrong — and its
+only real blocker was one line: `with VideoEncodingManager(dataset)` inside
+`run()` finalizes the dataset the first time `run()` returns, after which a
+restarted segment would silently truncate the finalized parquet.
+
+So `--interactive=true` now supports `--strategy.type=sentry`:
+
+- **Finalization moved to `teardown()`** (which already called
+  `dataset.finalize()`); `run()` is segment-restartable. Each segment saves
+  complete episodes plus one tail partial episode; on a failed tail save the
+  in-flight streaming encode is cancelled *and* the half-mutated episode
+  buffer is discarded (see §6, round 2).
+- **Frames are labeled with the live `engine.task`** instead of a config
+  snapshot — the writer already stores a task per frame — so `/subtask`
+  changes the policy conditioning and the recorded label from the same frame
+  onward. This also resolved the "recorded frames ignore `/subtask`" open item
+  for sentry.
+- `episodes_since_push` hoisted to instance state so upload cadence survives
+  segments.
+- dagger / highlight / episodic stay excluded: keyboard conflicts plus per-run
+  recording state that does not survive a restart.
+
+### 5.3 Muting v2 — two lines, and errors surface
+
+The ~30-line per-handler walk became `logging.disable(logging.WARNING)` with
+the previous disable level restored afterwards. Strictly better coverage: the
+gate applies before handler dispatch, so it covers `propagate=False` library
+loggers *and* loggers created mid-session (the old snapshot missed those) —
+and **ERROR/CRITICAL now reach the console**, which the audit showed is safe:
+no ERROR-level emitter fires periodically in healthy operation (the periodic
+nuisances — slow-loop, camera hiccups — are WARNINGs and stay muted).
+Documented trade-off: the gate also withholds INFO/WARNING from file handlers
+during the session; acceptable because no default code path attaches one
+(only `rl/actor`, `rl/learner`, `async_inference` pass `log_file`). The
+`warnings` suppression stays (nothing calls `logging.captureWarnings`), and
+`failure_traceback` surfacing stays as the belt-and-suspenders for fatal
+engine errors.
+
+### 5.4 stdin listener → `lerobot/utils/stdin_input.py`
+
+The ssh/headless audit confirmed the listener was already the right design:
+`select`+`os.read` works over SSH (the session pty is a normal fd), from
+pipes, and headless — it's `keyboard_input`'s **pynput** backend that needs a
+display server. Nothing in `keyboard_input` overlaps enough to reuse
+(1-byte cbreak hotkey decoding vs canonical-mode line assembly), so
+`StdinCommandListener` moved to a **new** utils module — deliberately not
+into `keyboard_input.py`, which attempts a pynput import at module load.
+Canonical import only: `lerobot.utils.stdin_input` (removed from
+`lerobot.rollout`'s exports).
+
+The move fixed a real bug the audit found: with `sys.stdin is None`
+(daemonized processes), the blocking fallback died with an uncaught
+`AttributeError` without firing `on_eof` — leaving a session idling with no
+command channel. `start()` now treats a missing stream as immediate EOF.
+
+## 6. Bugs the adversarial reviews caught
+
+Four multi-agent review passes were run across the phases (28 / 5 / 27 / 12 agents;
+findings adversarially verified before acting). The ones that mattered:
+
+**Round 2 (2 confirmed, 0 refuted):**
+
+- **Controller `start()` race → phantom segment.** `start()` gated on `_running`, but the
+  serve loop cleared `_start_requested` *before* setting `_running` — a second `start()`
+  landing in that window (spanning `reset_control_state` and the SEGMENT_STARTED emission)
+  returned `True` and re-armed the flag, which nothing consumed during the segment; the
+  robot would start again, uncommanded, when the segment later ended on its own. Fixed:
+  the serve loop consumes the request and sets `_running` atomically under the control
+  lock, and `_running` spans the whole startup sequence.
+- **Sentry poisoned episode buffer.** `save_episode` mutates the buffer in place (pops
+  `size`/`task`) *before* the fallible writes; a failed tail save left a half-mutated dict
+  and the next segment's first `add_frame` crashed with `KeyError('size')`. Fixed: the
+  except branch discards the buffer so `add_frame` recreates it.
+
+**Rounds 1–3 (Phases 1–2):**
+
+- **RTC stale observation (critical).** `RTCInferenceEngine.reset()` never cleared
+  `_obs_holder["obs"]`. After `/reset` physically moved the robot home, the next `/start`
+  computed its first chunk from the **pre-reset pose** — a lurch back toward where the arm
+  used to be. Fixed by clearing the observation and adding a `_reset_epoch` counter so an
+  in-flight chunk computed across a reset is discarded rather than merged. This also fixes a
+  pre-existing DAgger staleness path.
+- **Muting hid fatal errors** → `failure_traceback` capture + session print (§4.7).
+- **Muting scope too narrow** → root-only missed `transformers` / `datasets`; `warnings`
+  output bypassed logging entirely.
+- **Command ordering** → `/reset` and `/stop` didn't cancel a pending `/start` (§4.5); the
+  `/reset`-then-`/subtask` clobber (§4.8).
+- **Flush too heavy** → `policy.reset()` → `drop_queued_actions()` (§4.8).
+- **Empty-task rendering** → `''` replaced with `(none — set one with /subtask <text>)`.
+- **Silent switch** → the confirmation now says "(applies from the next policy inference)",
+  since the explanatory logs are muted.
+
+## 7. Verification
+
+After Round 2:
+
+```
+uv run --extra dataset pytest tests/test_interactive_rollout.py \
+    tests/utils/test_stdin_input.py tests/test_rollout.py -q
+    → 81 passed
+
+pre-commit (all changed files)
+    → 0 failures
+```
+
+Phase 1–2 numbers (still green at the time): 64 rollout/interactive tests;
+223 passed / 5 skipped across `tests/policies/rtc`, factory, and common
+(confirming the shared `pretrained.py` change); pre-commit 0 failures.
+
+`tests/test_interactive_rollout.py` covers the parser, `LinkedEvent` semantics,
+`RolloutController` (start/reset/stop/set_task flows, events, startup-race
+rejection, failure surfacing, broken observers), session flows (start / reset /
+restart / stop, cancel-pending-start, engine failure with traceback, natural
+end, EOF, and a real `BaseStrategy` end-to-end), muting (INFO/WARNING blocked,
+ERROR surfaces, pre-existing disable level restored), `/subtask` semantics,
+sentry restartability + live labels + failed-tail-save recovery, the engine
+task holder, the sync flush, and `drop_queued_actions`.
+`tests/utils/test_stdin_input.py` covers the listener (select path, batched
+lines, blocking fallback, EOF, handler errors, None-stdin, broken streams).
+
+## 8. Extension points for Phase 3
+
+The design was built to make `/ask` an additive change:
+
+- **Command table.** `InteractiveSession._commands` is `name → (handler, arg hint, help)`.
+  `/help` and the startup banner render from it, so a new command is documented for free.
+- **Controller API.** New front-ends (network, voice, `/ask`'s LLM worker) call
+  `RolloutController.start/reset/stop/set_task` from their own threads — the internal lock
+  makes that safe — and observe `RolloutEvent`s instead of scraping terminal output.
+- **Thread discipline.** A command handler runs on the listener thread and must only call
+  controller methods. An LLM call belongs on its own worker thread so the robot keeps
+  moving — precisely the concurrency `/ask` is meant to demonstrate.
+- **Task holder.** `set_task` / `_take_task` already give any producer a safe way to
+  re-instruct the policy. Hierarchical task-vs-subtask semantics (per #4183 / #4234) layer
+  on top of it rather than replacing it.
+
+Open items, deliberately not addressed:
+
+- dagger / highlight / episodic remain non-interactive (keyboard conflicts + per-run
+  recording state); they also still snapshot the task label per run. Sentry is the
+  supported recording path for interactive sessions.
+- The "see logs and prompt simultaneously" TUI (pinned input line).
+- Non-stdin input sources (network, voice) — now unblocked by `RolloutController`; #4108's
+  pluggable-listener shape remains the reference for the transport layer.
@@ -161,6 +161,16 @@ The methods called by the train/eval loops:

 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.

+If your model is large enough to warrant [sharded multi-GPU training](./multi_gpu_training#sharded-training-fsdp), also declare its FSDP wrap units — the repeated block classes sharding operates on:
+
+```python
+class MyPolicy(PreTrainedPolicy):
+    ...
+    _fsdp_wrap_modules = ["MyTransformerBlock"]
+```
+
+With this one declaration, `--parallelism.dp_shard=N` works out of the box for your policy (users can still override it with `--accelerator.fsdp.wrap_modules`). Without any wrap source, sharded runs fail at startup by design.
+
 ### Processor functions

 LeRobot uses `PolicyProcessorPipeline`s to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see [`processor_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/processor_act.py) or [`processor_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/processor_diffusion.py).
@@ -300,7 +310,7 @@ The file names are load-bearing: the factory does lazy imports by name, and the
 Two 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.
+2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what the end-of-training publisher 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.

@@ -344,7 +354,7 @@ A new policy is much easier to review — and far more useful — when it ships

 **Pick at least one in-tree benchmark.** LeRobot ships sim benchmarks with per-benchmark Docker images (LIBERO, LIBERO-plus, Meta-World, RoboTwin 2.0, RoboCasa365, RoboCerebra, RoboMME, VLABench and more). Pick the one that matches your policy's modality — VLAs usually go to LIBERO or VLABench; image-only BC to LIBERO or Meta-World. The full list lives under [Benchmarks](./libero) in the docs sidebar.

-**Push the checkpoint & processors** to the Hub under `lerobot/<policy>_<benchmark>` (or your namespace if you don't have write access; a maintainer can mirror it). Use `PreTrainedPolicy.push_model_to_hub` so the repo gets `config.json`, `model.safetensors`, and a model card.
+**Push the checkpoint & processors** to the Hub under `lerobot/<policy>_<benchmark>` (or your namespace if you don't have write access; a maintainer can mirror it). The easiest way is training with `--policy.repo_id=<namespace>/<repo>` and `--policy.push_to_hub=true`: `lerobot-train` publishes the model, both processors, and a model card at the end of the run. To publish an existing checkpoint after the fact, upload its `pretrained_model/` directory (e.g. `huggingface-cli upload`), or use `lerobot-convert-dcp --push_to_hub=...` for sharded-format checkpoints.

 **Report results in your policy's MDX**, with the exact `lerobot-eval` command and hardware so anyone can re-run:

@@ -62,7 +62,10 @@ Reference data points on a 4×H100 80 GB cluster (`accelerate launch --num_proce
 | `smolvla`   | 27m 49s    |      0.312 |           0.011 | ~80%     | `--policy.path=lerobot/smolvla_base`, `freeze_vision_encoder=false`, `train_expert_only=false`                                 |
 | `pi05`      | 3h 41m     |      2.548 |           0.014 | ~95%     | `--policy.pretrained_path=lerobot/pi05_base`, `gradient_checkpointing=true`, `dtype=bfloat16`, vision encoder + expert trained |

-The `dataloading_s` vs. `update_s` ratio is the diagnostic that matters: when `dataloading_s` approaches `update_s`, more GPUs stop helping — your dataloader is the bottleneck and you should look at `--num_workers`, image resolution, and disk speed before adding compute.
+Training logs separate the full iteration into `dataloading_s` (`next(dl_iter)`), `preprocessing_s`
+(image conversion and the policy pipeline), and `update_s` (the optimizer update). `step_s` covers all
+three and drives `samples_per_s`. The benchmark above predates this split, so its `dataloading_s` includes
+preprocessing.

 ### Schedule and checkpoints

@@ -241,24 +241,129 @@ See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.

 ---

+## Interactive Sessions
+
+Add `--interactive=true` to drive the rollout from the terminal instead of starting immediately. Hardware connects and the policy loads as usual, but **the robot stays still until you type `/start`** — useful when you want to position the scene first, re-instruct the policy between attempts, or run several takes without paying the load time again.
+
+```bash
+lerobot-rollout \
+    --strategy.type=base \
+    --policy.path=${HF_USER}/my_smolvla_policy \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --robot.cameras="{ front: {type: opencv, index_or_path: 0, width: 640, height: 480, fps: 30}}" \
+    --task="pick up the cube" \
+    --interactive=true
+```
+
+| Command           | Action                                                                                                                                                                             |
+| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `/start`          | Start (or restart) the policy control loop                                                                                                                                         |
+| `/subtask <text>` | Change the instruction the policy follows, without stopping. No argument prints the current task. Only affects policies that condition on language (SmolVLA, π0/π0.5, and similar) |
+| `/reset`          | Stop movement, return the robot to its startup position, and restore the `--task` instruction                                                                                      |
+| `/stop`           | End the session and run the normal shutdown routines                                                                                                                               |
+| `/help`           | List the commands                                                                                                                                                                  |
+
+```text
+> /start
+Rollout running — task 'pick up the cube'. /subtask <text> to change it, ...
+> /subtask put the cube in the box
+Task: 'pick up the cube' → 'put the cube in the box' (applies from the next policy inference)
+> /reset
+Task restored to 'pick up the cube'
+Resetting — returning the robot to its initial position...
+Robot reset — holding at initial position. /start to run.
+> /stop
+```
+
+`Ctrl-C` still shuts down as usual, and closing stdin (`Ctrl-D`, or the end of a piped script) ends the session — so a piped script must keep stdin open for the intended duration:
+
+```bash
+(printf '/start\n'; sleep 60; printf '/stop\n') | lerobot-rollout ... --interactive=true
+```
+
+**How `/subtask` reaches the policy.** The stdin reader publishes the new instruction to the inference engine, which picks it up on its own inference thread, so nothing is mutated across threads while the robot is moving. How quickly the behavior changes depends on the backend:
+
+- **Sync** (`--inference.type=sync`) — precomputed chunk actions are dropped, so the new instruction applies on the very next control tick. Without this a chunking policy would keep executing up to `chunk_size` stale actions (seconds of the old behavior). Only the queued actions are discarded, so observation history and the rest of the episode state are preserved.
+- **RTC** (`--inference.type=rtc`) — the next chunk is generated under the new instruction and merged over the previous chunk's leftover prefix, so the switch lands within one inference and the motion stays continuous. The queue is deliberately not cleared: that would leave the robot without commands for a full inference latency. (With blending turned off via `--inference.rtc.enabled=false` the queued chunk drains first, so the switch lands up to one chunk later.)
+
+With `--use_torch_compile=true`, a switch whose instruction tokenizes to a different length can trigger a recompilation on the next forward pass, pausing inference for as long as the original warm-up took. Prefer leaving compilation off for sessions where you expect to re-instruct the policy often.
+
+**Logs below ERROR are muted while the session runs** so routine output doesn't interleave with what you're typing; errors and fatal inference failures still show, and normal logging resumes when the session ends. The gate is process-wide (it also withholds INFO/WARNING from any file handler you attached for the duration). Run without `--interactive` to watch the live log.
+
+Sessions work over SSH and on headless machines — the command reader uses the terminal (or pipe) directly and needs no display server.
+
+**Recording while interactive.** `--strategy.type=sentry` also supports `--interactive=true`: the session records continuously while you steer it. Each `/start`…`/reset` segment saves complete episodes plus one final partial episode, the dataset stays open until shutdown, and **frames are labeled with the live task** — a `/subtask` changes both the policy conditioning and the recorded label from the same frame onwards.
+
+```bash
+lerobot-rollout \
+    --strategy.type=sentry \
+    --policy.path=${HF_USER}/my_smolvla_policy \
+    --robot.type=so100_follower \
+    --robot.port=/dev/ttyACM0 \
+    --dataset.repo_id=${HF_USER}/rollout_cube_sessions \
+    --task="pick up the cube" \
+    --interactive=true
+```
+
+The other recording strategies (episodic, DAgger, highlight) are not supported: they bind their own keyboard controls, which would compete with the command prompt for the same terminal.
+
+### Programmatic control
+
+Everything the CLI session does is available as a library API: `RolloutController` exposes thread-safe `start()` / `reset()` / `stop()` / `set_task()` methods plus a `RolloutEvent` callback, with no stdin, printing, or log muting attached — embed it in your own application, network server, or notebook:
+
+```python
+from threading import Event, Thread
+
+from lerobot.rollout import (
+    LinkedEvent,
+    RolloutController,
+    RolloutEvent,
+    build_rollout_context,
+    create_strategy,
+)
+
+parent = Event()  # your application's shutdown signal
+ctx = build_rollout_context(cfg, LinkedEvent(parent))  # loads policy, connects robot
+strategy = create_strategy(cfg.strategy)
+strategy.setup(ctx)
+
+controller = RolloutController(strategy, ctx, on_event=print)  # or your own observer
+serve_thread = Thread(target=controller.serve)  # serve() blocks; run it where you like
+serve_thread.start()
+
+controller.start()                        # robot starts executing the policy
+controller.set_task("grab the red cube")  # re-instruct mid-run
+controller.reset()                        # stop movement, return home, stay warm
+controller.stop()                         # end serve()
+
+serve_thread.join()
+strategy.teardown(ctx)  # teardown stays with the caller
+```
+
+Set `play_sounds=False` in the config unless you want the vocal announcements, and note that `build_rollout_context` requires the shutdown event to be a `LinkedEvent` (the controller ends run segments through its local flag; your `parent` event still forces a full shutdown). `InteractiveSession` itself is a thin front-end over this controller — commands map 1:1 onto its methods.
+
+---
+
 ## Common Flags

-| Flag                              | Description                                                       | Default |
-| --------------------------------- | ----------------------------------------------------------------- | ------- |
-| `--policy.path`                   | **Required.** HF Hub model ID or local checkpoint path            | --      |
-| `--robot.type`                    | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`) | --      |
-| `--robot.port`                    | Serial port for the robot                                         | --      |
-| `--robot.cameras`                 | Camera configuration (JSON dict)                                  | --      |
-| `--fps`                           | Control loop frequency                                            | 30      |
-| `--duration`                      | Run time in seconds (0 = infinite)                                | 0       |
-| `--device`                        | Torch device (`cpu`, `cuda`, `mps`)                               | auto    |
-| `--task`                          | Task description (used when no dataset is provided)               | --      |
-| `--display_data`                  | Stream telemetry to Rerun visualization                           | false   |
-| `--display_ip` / `--display_port` | Remote Rerun server address                                       | --      |
-| `--interpolation_multiplier`      | Action interpolation factor                                       | 1       |
-| `--use_torch_compile`             | Enable `torch.compile` for inference                              | false   |
-| `--resume`                        | Resume a previous recording session                               | false   |
-| `--play_sounds`                   | Vocal synthesis for events                                        | true    |
+| Flag                              | Description                                                                                                                                   | Default |
+| --------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
+| `--policy.path`                   | **Required.** HF Hub model ID or local checkpoint path                                                                                        | --      |
+| `--robot.type`                    | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`)                                                                             | --      |
+| `--robot.port`                    | Serial port for the robot                                                                                                                     | --      |
+| `--robot.cameras`                 | Camera configuration (JSON dict)                                                                                                              | --      |
+| `--fps`                           | Control loop frequency                                                                                                                        | 30      |
+| `--duration`                      | Run time in seconds (0 = infinite)                                                                                                            | 0       |
+| `--device`                        | Torch device (`cpu`, `cuda`, `mps`)                                                                                                           | auto    |
+| `--task`                          | Task description (used when no dataset is provided)                                                                                           | --      |
+| `--display_data`                  | Stream telemetry to Rerun visualization                                                                                                       | false   |
+| `--display_ip` / `--display_port` | Remote Rerun server address                                                                                                                   | --      |
+| `--interpolation_multiplier`      | Action interpolation factor                                                                                                                   | 1       |
+| `--interactive`                   | Chat-style stdin session (see [Interactive Sessions](#interactive-sessions)); the robot stays idle until `/start`. Base and sentry strategies | false   |
+| `--use_torch_compile`             | Enable `torch.compile` for inference                                                                                                          | false   |
+| `--resume`                        | Resume a previous recording session                                                                                                           | false   |
+| `--play_sounds`                   | Vocal synthesis for events                                                                                                                    | true    |

 ---

@@ -242,6 +242,17 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
  --repo-id=your_dataset
 ```

+Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram:
+
+```bash
+python src/lerobot/scripts/augment_dataset_quantile_stats.py \
+  --repo-id=your_dataset \
+  --overwrite \
+  --skip-images
+```
+
+`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes MolmoAct2's normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected.
+
 Alternatively, train MolmoAct2 with mean/std normalization:

 ```bash
@@ -1,28 +1,29 @@
 # Multi-GPU Training

-This guide shows you how to train policies on multiple GPUs using [Hugging Face Accelerate](https://huggingface.co/docs/accelerate).
+LeRobot trains on multiple GPUs through [Hugging Face Accelerate](https://huggingface.co/docs/accelerate). Three data-parallel layouts are supported:
+
+| Layout   | What it does                                                  | Config                                                  |
+| -------- | ------------------------------------------------------------- | ------------------------------------------------------- |
+| **DDP**  | Replicates the full model on every GPU                        | default on any multi-GPU launch                         |
+| **FSDP** | Shards parameters, gradients, and optimizer state across GPUs | `--parallelism.dp_shard=N`                              |
+| **HSDP** | Shards within groups of GPUs, replicates across groups        | `--parallelism.dp_replicate=R --parallelism.dp_shard=S` |

 ## Installation

-`accelerate` is included in the `training` extra. Install it with:
+`accelerate` is included in the `training` extra:

 ```bash
 pip install 'lerobot[training]'
 ```

-## Training with Multiple GPUs
+## Launching

-You can launch training in two ways:
+Distributed training can be launched through both `torchrun` and `accelerate launch`. Accelerate is used as a plain launcher: it does not manage the training configuration, and every distributed training setting lives in LeRobot's own config system.

-### Option 1: Without config (specify parameters directly)
-
-You can specify all parameters directly in the command without running `accelerate config`:
+With `torchrun`:

 ```bash
-accelerate launch \
-  --multi_gpu \
-  --num_processes=2 \
-  $(which lerobot-train) \
+torchrun --nproc-per-node=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
@@ -31,32 +32,10 @@ accelerate launch \
  --wandb.enable=true
 ```

-**Key accelerate parameters:**
-
- `--multi_gpu`: Enable multi-GPU training
- `--num_processes=2`: Number of GPUs to use
- `--mixed_precision=fp16`: Use fp16 mixed precision (or `bf16` if supported)
-
-### Option 2: Using accelerate config
-
-If you prefer to save your configuration, you can optionally configure accelerate for your hardware setup by running:
+With `accelerate launch` (as a plain launcher):

 ```bash
-accelerate config
-```
-
-This interactive setup will ask you questions about your training environment (number of GPUs, mixed precision settings, etc.) and saves the configuration for future use. For a simple multi-GPU setup on a single machine, you can use these recommended settings:
-
- Compute environment: This machine
- Number of machines: 1
- Number of processes: (number of GPUs you want to use)
- GPU ids to use: (leave empty to use all)
- Mixed precision: fp16 or bf16 (recommended for faster training)
-
-Then launch training with:
-
-```bash
-accelerate launch $(which lerobot-train) \
+accelerate launch --num_processes=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
@@ -65,116 +44,133 @@ accelerate launch $(which lerobot-train) \
  --wandb.enable=true
 ```

-## How It Works
+With no `--parallelism.*` flags, a multi-process launch runs plain DDP. Multi-node runs use the standard `torchrun --nnodes/--node-rank/--rdzv-endpoint` flags (or `accelerate launch --num_machines/--machine_rank/--main_process_ip`).

-When you launch training with accelerate:
+> [!WARNING]
+> Accelerate's YAML config files (`accelerate launch --config_file some.yaml`, `accelerate config`) are not supported. They configure the engine through environment variables, bypassing LeRobot's configuration system, so `train_config.json` would no longer describe the settings a run actually used. `lerobot-train` therefore refuses to start when [accelerate environment variables](https://huggingface.co/docs/accelerate/usage_guides/fsdp) are set. Put the settings in `--parallelism.*` / `--accelerator.*` flags instead, or set `LEROBOT_ALLOW_ACCELERATE_ENV=1` to acknowledge the override and proceed anyway.

-1. **Automatic detection**: LeRobot automatically detects if it's running under accelerate
-2. **Data distribution**: Your batch is automatically split across GPUs
-3. **Gradient synchronization**: Gradients are synchronized across GPUs during backpropagation
-4. **Single process logging**: Only the main process logs to wandb and saves checkpoints
+## Batch semantics, learning rate, and steps

-## Learning Rate and Training Steps Scaling
+Each of the `dp_replicate × dp_shard` data-parallel workers loads its own `--batch_size` micro-batch every step, so one training step consumes `batch_size × dp_world_size` samples, and `× gradient_accumulation_steps` of those go into each optimizer update:

-**Important:** LeRobot does **NOT** automatically scale learning rates or training steps based on the number of GPUs. This gives you full control over your training hyperparameters.
-
-### Why No Automatic Scaling?
-
-Many distributed training frameworks automatically scale the learning rate by the number of GPUs (e.g., `lr = base_lr × num_gpus`).
-However, LeRobot keeps the learning rate exactly as you specify it.
-
-### When and How to Scale
-
-If you want to scale your hyperparameters when using multiple GPUs, you should do it manually:
-
-**Learning Rate Scaling:**
-
-```bash
-# Example: 2 GPUs with linear LR scaling
-# Base LR: 1e-4, with 2 GPUs -> 2e-4
-accelerate launch --num_processes=2 $(which lerobot-train) \
-  --optimizer.lr=2e-4 \
-  --dataset.repo_id=lerobot/pusht \
-  --policy.type=act
+```
+effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps
 ```

-**Training Steps Scaling:**
+The training banner prints this factorization at startup. `--steps` counts loop steps (micro-batches per worker), not optimizer updates.

-Since the effective batch size `bs` increases with multiple GPUs (batch_size × num_gpus), you may want to reduce the number of training steps proportionally:
+Gradient accumulation is a first-class flag:

 ```bash
-# Example: 2 GPUs with effective batch size 2x larger
-# Original: batch_size=8, steps=100000
-# With 2 GPUs: batch_size=8 (16 in total), steps=50000
-accelerate launch --num_processes=2 $(which lerobot-train) \
-  --batch_size=8 \
-  --steps=50000 \
-  --dataset.repo_id=lerobot/pusht \
-  --policy.type=act
+torchrun --nproc-per-node=2 $(which lerobot-train) \
+  --batch_size=8 --accelerator.gradient_accumulation.steps=4 ...
 ```

-## Training Large Models with FSDP
+**LeRobot does not auto-scale the learning rate or the number of steps** when the effective batch size grows. If you scale out and want equivalent training, please adjust manually, e.g. with 2 GPUs: double `--optimizer.lr` (linear scaling), or halve `--steps`.

-DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under
-DDP either. For large models, use **FSDP** (Fully Sharded Data Parallel), which shards parameters,
-gradients, and optimizer state across GPUs. See the [accelerate FSDP guide](https://huggingface.co/docs/accelerate/usage_guides/fsdp) for background.
+## Sharded training (FSDP)

-An example on how to launch LeRobot training with FSDP across 4 GPUs (1 machine):
+If a model is too large to train with DDP, shard it with FSDP2:

 ```bash
-accelerate launch --config_file fsdp.yaml --num_processes=4 $(which lerobot-train) \
+torchrun --nproc-per-node=4 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=<your_policy> \
+  --parallelism.dp_shard=4 \
+  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp
 ```

-A minimal `fsdp.yaml` (FSDP1; shards params/grads/optimizer — ZeRO-3-equivalent):
+`--parallelism.dp_shard=-1` shards over however many processes the launcher started.

-```yaml
-compute_environment: LOCAL_MACHINE
-distributed_type: FSDP
-mixed_precision: bf16
-num_machines: 1
-num_processes: 4
-fsdp_config:
-  fsdp_version: 1
-  fsdp_sharding_strategy: FULL_SHARD # params + grads + optimizer (ZeRO-3)
-  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  fsdp_transformer_layer_cls_to_wrap: <YourTransformerBlock> # repeated block class to shard
-  fsdp_use_orig_params: true # required: optimizer is built pre-prepare
-  fsdp_state_dict_type: FULL_STATE_DICT
+### Wrap units
+
+FSDP shards the model in units (typically the repeated transformer block) and gathers one unit at a time during forward/backward. Policies declare their wrap units via `_fsdp_wrap_modules` on the policy class. For example, ACT declares `["ACTEncoderLayer", "ACTDecoderLayer"]` and FastWAM declares `["MoTLayer"]`. For a policy without a `_fsdp_wrap_modules` declaration, pass one of the flags below. You can specify the module class name explicitly, or use a size-based policy instead:
+
+```bash
+--accelerator.fsdp.wrap_modules='["MyTransformerBlock"]'   # explicit class names
+--accelerator.fsdp.min_num_params=1000000                  # or: wrap every submodule above 1M params
 ```

-Set `fsdp_transformer_layer_cls_to_wrap` to your model's repeated transformer-block class so each
-block is sharded as its own unit. `fsdp_use_orig_params: true` is required because LeRobot builds the
-optimizer before `accelerator.prepare()`.
+If a policy doesn't declare `_fsdp_wrap_modules` and no `--accelerator.fsdp.wrap_modules` or `--accelerator.fsdp.min_num_params` is passed, the run fails at startup rather than silently wrapping only the root module (which would forfeit all sharding memory savings).

-### FSDP checkpoints
+Other sharding settings:

-LeRobot gathers the full state dict across all ranks and the main process writes it as a single
-`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for:
+- `--accelerator.fsdp.reshard_after_forward`: whether to keep each unit's parameters resident after forward.
+- `--accelerator.fsdp.cpu_offload`: keeps parameters, gradients and optimizer states on CPU.
+- `--accelerator.fsdp.ignored_modules`: a regex of module paths to keep unsharded.

- **Checkpoints store fp32 weights.** Under mixed precision (`bf16`/`fp16`) FSDP keeps an fp32 master
-  copy, and the checkpoint saves it (~2× the bf16 size on disk) so training can resume consistently
-  with the fp32 optimizer state; `from_pretrained` casts back to the policy dtype on load. FSDP-specific
-  caveat: an fp32 checkpoint is materialized in full precision on the target device _before_ casting,
-  so loading it for inference on a tight GPU can OOM even when the bf16 model would fit — load on CPU
-  first, or cast `model.safetensors` to the deployment dtype offline.
- The sharded optimizer state is gathered into a full (world-size-independent) state dict and saved
-  alongside the model in the same `optimizer_state.safetensors` / `optimizer_param_groups.json`
-  format as single-GPU training, so **resume-from-checkpoint is supported** with `--resume=true`.
-  Resume reshards both the model and the optimizer state to the _current_ FSDP topology, so you can
-  resume an FSDP checkpoint on a different number of GPUs. Note that the data sampler is only
-  sample-exact when the world size and batch size match the original run (a warning is logged
-  otherwise); the optimizer/model state itself is unaffected.
+### HSDP
+
+Hybrid Sharded Data Parallel: parameters, gradients and optimizer states are sharded across `dp_shard` ranks, and that sharding is replicated `dp_replicate` times. Parameter all-gathers and gradient reduce-scatters stay inside a shard group; only the all-reduce that synchronizes the replicas crosses between groups. The two degrees must multiply to the world size:
+
+```bash
+# 16 GPUs = 2 nodes × 8: shard within each node, replicate across nodes
+torchrun --nnodes=2 --nproc-per-node=8 ... $(which lerobot-train) \
+  --parallelism.dp_replicate=2 --parallelism.dp_shard=8 ...
+```
+
+## Checkpoints
+
+Every checkpoint contains a `pretrained_model/` directory and a `training_state/` directory:
+
+```text
+005000/  # the training step at that checkpoint
+├── pretrained_model/
+│   ├── config.json  # policy config
+│   ├── train_config.json  # the full training config
+│   ├── model.safetensors  # full weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run)
+│   ├── pytorch_model_fsdp_0/  # DCP weight shards (checkpoint_format ∈ {dcp, safetensors_dcp})
+│   ├── policy_preprocessor.json  # preprocessor config (when the run has a preprocessor)
+│   ├── policy_preprocessor_step_*.safetensors  # state of the stateful preprocessor steps
+│   ├── policy_postprocessor.json  # postprocessor config (when the run has a postprocessor)
+│   └── policy_postprocessor_step_*.safetensors  # state of the stateful postprocessor steps
+└── training_state/
+    ├── training_step.json  # step counter, topology, and batch semantics
+    ├── rng_state.safetensors  # rng states
+    ├── scheduler_state.json  # scheduler state (when the run has a scheduler)
+    ├── optimizer_state.safetensors  # full optimizer state (non-sharded runs)
+    ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
+    └── optimizer_0/  # DCP optimizer shards (sharded runs)
+```
+
+During single-GPU or DDP training, the pipeline serializes each state dict into a single file: `model.safetensors` for the model and `optimizer_state.safetensors` for the optimizer.
+
+During sharded training, the optimizer state is saved as DCP shards under `training_state/optimizer_0/`, and the layout of the model under `pretrained_model/` can be configured through `--checkpoint_format`:
+
+| `--checkpoint_format`     | Weights artifact                             | Use when                                                              |
+| ------------------------- | -------------------------------------------- | --------------------------------------------------------------------- |
+| `safetensors` _(default)_ | single `model.safetensors` only              | you want every checkpoint immediately loadable with `from_pretrained` |
+| `dcp`                     | `pytorch_model_fsdp_0/` shard directory only | gathering the full weights makes saves and resumes too slow           |
+| `safetensors_dcp`         | both                                         | you want fast resume _and_ immediately loadable checkpoints           |
+
+Two things to know about gathered (`safetensors`) checkpoints from sharded runs:
+
+- **They store fp32 weights.** Under mixed precision training, FSDP keeps an fp32 master copy, and the checkpoint saves the master copy to make sure training resumes consistently.
+- The gather is collective (all ranks participate) but only the main process writes.
+
+### Converting DCP checkpoints
+
+`lerobot-convert-dcp` merges a DCP shard directory into a regular `model.safetensors`, offline and without GPUs:
+
+```bash
+lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000
+lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=${HF_USER}/my_policy
+```
+
+`--push_to_hub` publishes the converted directory as a model repo.
+
+### Resuming
+
+Resume with `--resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json`. Resuming from a DCP checkpoint supports resharding the model and optimizer state to the _current_ topology, which means you can resume with a different `dp_replicate/dp_shard` split. The data sampler can always resume at the right epoch and offset, but is only _sample-exact_ when the world size and batch size match the original run (a warning is logged otherwise).
+
+> [!NOTE]
+> FSDP checkpoints written by LeRobot 0.6.x and earlier used a different on-disk layout (a gathered full optimizer state) and **cannot be resumed**.

 ## Notes

- The `--policy.use_amp` flag in `lerobot-train` is only used when **not** running with accelerate. When using accelerate, mixed precision is controlled by accelerate's configuration.
- Training logs, checkpoints, and hub uploads are only done by the main process to avoid conflicts. Non-main processes have console logging disabled to prevent duplicate output.
- The effective batch size is `batch_size × num_gpus`. If you use 4 GPUs with `--batch_size=8`, your effective batch size is 32.
- Learning rate scheduling is handled correctly across multiple processes—LeRobot sets `step_scheduler_with_optimizer=False` to prevent accelerate from adjusting scheduler steps based on the number of processes.
- When saving or pushing models, LeRobot automatically unwraps the model from accelerate's distributed wrapper to ensure compatibility.
- WandB integration automatically initializes only on the main process, preventing multiple runs from being created.
+- Checkpoint saves and end-of-training publishes are collective (every rank enters them). Gathered weights, sidecar files and Hub uploads are written by the main process alone.
+- Metrics are reduced across ranks before logging: losses are averaged, and `samples/s` reports cluster-wide throughput.
+- Learning-rate scheduling is stepped once per training step regardless of the number of processes (`step_scheduler_with_optimizer=False` is baked in).

-For more advanced configurations and troubleshooting, see the [Accelerate documentation](https://huggingface.co/docs/accelerate). If you want to learn more about how to train on a large number of GPUs, checkout this awesome guide: [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook).
+For background on the underlying machinery, see the [Accelerate FSDP guide](https://huggingface.co/docs/accelerate/usage_guides/fsdp). To go deeper on large-scale training, check out the [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook).
@@ -127,6 +127,17 @@ lerobot-edit-dataset \

 Or keep the dataset as-is and pass `--policy.normalization_mapping='{"ACTION": "MEAN_STD", "STATE": "MEAN_STD", "VISUAL": "IDENTITY"}'`.

+Recording, resuming, and merging aggregate quantiles from per-episode summaries, so `meta/stats.json` ends up holding a conservative envelope (`min` for `q <= 50`, `max` for `q > 50`) rather than whole-dataset quantiles. To estimate the latter, scan every episode with a running histogram:
+
+```bash
+python src/lerobot/scripts/augment_dataset_quantile_stats.py \
+    --repo-id=your_dataset \
+    --overwrite \
+    --skip-images
+```
+
+`--skip-images` keeps the existing image statistics and avoids video decoding when only `STATE`/`ACTION` need recomputing, and `--root` reads a local dataset instead of the Hub. These values are histogram estimates, subject to discretization and rebinning error, so they can differ from the conservative ones — which changes π₀.₅'s normalized targets and therefore its loss scale. Statistics already saved inside an existing checkpoint are not affected.
+
 ### Training Command Example

 The same finetune with the VLM frozen: less memory, at some cost in success rate. Swap `--dataset.repo_id` for your own dataset.
@@ -2,6 +2,25 @@

 https://diffusion-policy.cs.columbia.edu

+## Training
+
+The reference implementation maintains an exponential moving average (EMA) of the policy weights during training and evaluates the EMA weights. To reproduce this behavior, enable the trainer's EMA shadow:
+
+```bash
+lerobot-train \
+  --policy.type=diffusion \
+  --ema.enable=true \
+  ...
+```
+
+Checkpoints then contain a directly loadable copy of the EMA weights next to the live ones, e.g. for evaluation:
+
+```bash
+lerobot-eval --policy.path=outputs/train/.../checkpoints/last/pretrained_model_ema ...
+```
+
+The EMA decay schedule (`--ema.inv_gamma`, `--ema.power`, ...) defaults to the reference implementation's values. For a constant decay instead of the warmup schedule (e.g. to match openpi's pi0/pi05 training), set `--ema.decay=0.99`.
+
 ## Citation

 ```bibtex
@@ -59,6 +59,22 @@ When `use_relative_actions=true`, the training script automatically:

 ---

+## EMA of the policy weights
+
+OpenPI maintains an exponential moving average of the weights during training (`ema_decay=0.99` by default) and keeps the EMA copy for inference. To reproduce this with the LeRobot trainer, enable the EMA shadow with a constant decay:
+
+```bash
+python -m lerobot.scripts.lerobot_train \
+  --policy.type=pi05 \
+  --dataset.repo_id=your_org/your_dataset \
+  --ema.enable=true \
+  --ema.decay=0.99
+```
+
+Checkpoints then contain a directly loadable copy of the EMA weights in `pretrained_model_ema/` next to the live ones. Note that the shadow is a full extra copy of the parameters on the GPU. Like OpenPI (which disables EMA in its LoRA configs), EMA is not supported together with PEFT adapters.
+
+---
+
 ## Citation

 If you use this work, please cite both **OpenPI** and the π₀.₅ paper:
@@ -40,3 +40,15 @@ lerobot-eval \
 ```

 However, in most cases, presence of an accelerator is detected automatically and `policy.device` parameter can be omitted from CLI commands.
+
+## Mixed precision
+
+Training precision is owned by `--accelerator.mixed_precision`, which accepts `no` (default) and `bf16`:
+
+```bash
+lerobot-train \
+    --policy.type=act \
+    --accelerator.mixed_precision=bf16 ...
+```
+
+`bf16` requires an accelerator that supports it.
@@ -346,6 +346,7 @@ lerobot-record="lerobot.scripts.lerobot_record:main"
 lerobot-replay="lerobot.scripts.lerobot_replay:main"
 lerobot-setup-motors="lerobot.scripts.lerobot_setup_motors:main"
 lerobot-teleoperate="lerobot.scripts.lerobot_teleoperate:main"
+lerobot-convert-dcp="lerobot.scripts.lerobot_convert_dcp:main"
 lerobot-eval="lerobot.scripts.lerobot_eval:main"
 lerobot-train="lerobot.scripts.lerobot_train:main"
 lerobot-train-tokenizer="lerobot.scripts.lerobot_train_tokenizer:main"
@@ -475,6 +476,12 @@ default.extend-ignore-identifiers-re = [
 # TODO: Enable mypy gradually module by module across multiple PRs
 # Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations

+[tool.pytest.ini_options]
+markers = [
+    "multigpu: distributed tests needing 2-4 GPUs (CI: docker_publish.yml lane)",
+    "multigpu_heavy: 8-GPU sweeps and soak tests; never run in CI",
+]
+
 [tool.mypy]
 python_version = "3.12"
 ignore_missing_imports = true
@@ -521,6 +528,15 @@ disallow_untyped_defs = true
 disallow_incomplete_defs = true
 check_untyped_defs = true

+[[tool.mypy.overrides]]
+module = "lerobot.distributed.*"
+ignore_errors = false
+
+# extra strictness for the distributed engine
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+check_untyped_defs = true
+
 [[tool.mypy.overrides]]
 module = "lerobot.optim.*"
 ignore_errors = false
@@ -13,16 +13,41 @@
 # 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 pathlib import Path
+"""Training-output persistence: checkpoints, two-phase resume, and hub publishing.

-from huggingface_hub import HfApi, snapshot_download
+Rank discipline: every function here that can
+contain a collective is documented as such and must run on ALL ranks; rank-0-only file writes
+sit under one grouped ``is_main_process()`` gate per contiguous region, placed below all
+collectives. The leaf save/load helpers carry no rank gates of their own — the exception is
+``PreTrainedPolicy._save_pretrained``, whose gate is internal because its collective gather and
+its writes live in the same method.
+"""
+
+import logging
+from importlib.resources import files
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Any
+
+import torch.distributed as dist
+from huggingface_hub import HfApi, ModelCard, ModelCardData, snapshot_download
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler

+from lerobot.__version__ import __version__
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.configs.rewards import RewardModelConfig
 from lerobot.configs.train import TrainPipelineConfig
+from lerobot.distributed.checkpoint import (
+    is_sharded_module,
+    load_sharded_model,
+    load_sharded_optimizer,
+    save_sharded_model,
+    save_sharded_optimizer,
+)
+from lerobot.distributed.utils import is_main_process
 from lerobot.optim import (
    load_optimizer_state,
-    load_optimizer_state_dict,
    load_scheduler_state,
    save_optimizer_state,
    save_scheduler_state,
@@ -40,14 +65,39 @@ from lerobot.utils.hub import find_latest_hub_checkpoint
 from lerobot.utils.io_utils import load_json, write_json
 from lerobot.utils.random_utils import load_rng_state, save_rng_state

+if TYPE_CHECKING:
+    from accelerate import Accelerator
+
+    from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+    from lerobot.rewards.pretrained import PreTrainedRewardModel
+

 def get_step_identifier(step: int, total_steps: int) -> str:
+    """Format a step number as the zero-padded identifier used for checkpoint directory names.
+
+    Args:
+        step (int): The training step to format.
+        total_steps (int): The total number of training steps; sets the padding width
+            (minimum 6 digits).
+
+    Returns:
+        str: The zero-padded step identifier, e.g. `"005000"`.
+    """
    num_digits = max(6, len(str(total_steps)))
    return f"{step:0{num_digits}d}"


 def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Path:
-    """Returns the checkpoint sub-directory corresponding to the step number."""
+    """Returns the checkpoint sub-directory corresponding to the step number.
+
+    Args:
+        output_dir (Path): The training run's output directory.
+        total_steps (int): The total number of training steps; sets the identifier padding.
+        step (int): The training step of the checkpoint.
+
+    Returns:
+        Path: The checkpoint step directory, `output_dir/checkpoints/<step-identifier>`.
+    """
    step_identifier = get_step_identifier(step, total_steps)
    return output_dir / CHECKPOINTS_DIR / step_identifier

@@ -63,37 +113,15 @@ def should_save_checkpoint(step: int, save_freq: int, total_steps: int) -> bool:
    return (save_freq > 0 and step % save_freq == 0) or step == total_steps


-def save_training_step(
-    step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None
-) -> None:
-    state: dict = {"step": step}
-    # num_processes and batch_size are recorded so a resumed run can detect a changed world size or
-    # batch size: the sampler's resume offset is computed from the (num_processes, batch_size) that
-    # produced `step`, since both scale how many sampler positions a step consumes (see
-    # compute_sampler_state).
-    if num_processes is not None:
-        state["num_processes"] = num_processes
-    if batch_size is not None:
-        state["batch_size"] = batch_size
-    write_json(state, save_dir / TRAINING_STEP)
+def update_last_checkpoint(checkpoint_dir: Path) -> None:
+    """Point the `last` symlink in the checkpoints directory at the given checkpoint.

+    Any existing `last` symlink is replaced. The link target is relative to the checkpoints
+    directory, so the tree stays valid when the run directory is moved.

-def load_training_step(save_dir: Path) -> int:
-    training_step = load_json(save_dir / TRAINING_STEP)
-    return training_step["step"]
-
-
-def load_training_num_processes(checkpoint_dir: Path) -> int | None:
-    """World size recorded at checkpoint time, or None for checkpoints written before it was stored."""
-    return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("num_processes")
-
-
-def load_training_batch_size(checkpoint_dir: Path) -> int | None:
-    """Per-process batch size recorded at checkpoint time, or None for older checkpoints."""
-    return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("batch_size")
-
-
-def update_last_checkpoint(checkpoint_dir: Path) -> Path:
+    Args:
+        checkpoint_dir (Path): The checkpoint step directory the `last` link should target.
+    """
    last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK
    if last_checkpoint_dir.is_symlink():
        last_checkpoint_dir.unlink()
@@ -101,6 +129,68 @@ def update_last_checkpoint(checkpoint_dir: Path) -> Path:
    last_checkpoint_dir.symlink_to(relative_target)


+# ---------------------------------------------------------------------------------------------
+# training_step.json
+# ---------------------------------------------------------------------------------------------
+
+
+def save_training_metadata(step: int, save_dir: Path, cfg: TrainPipelineConfig) -> None:
+    """Record the step counter plus everything a resume needs to reason about topology changes.
+
+    `step` counts loop iterations (= micro-batches), so
+    the sampler resume offset is `step x batch_size x dp_world_size` with no grad-accum factor.
+    `grad_accum_steps` and the parallelism snapshot are recorded so a resume can warn precisely
+    when the optimizer-update cadence or the sharding topology changed.
+
+    Args:
+        step (int): The training step (micro-batch counter) to record.
+        save_dir (Path): The `training_state/` directory to write `training_step.json` into.
+        cfg (TrainPipelineConfig): The training config whose batch size, gradient-accumulation,
+            and parallelism settings are snapshotted alongside the step.
+    """
+    state: dict[str, Any] = {
+        "step": step,
+        "dp_world_size": cfg.parallelism.dp_world_size,
+        "batch_size": cfg.batch_size,
+        "grad_accum_steps": cfg.accelerator.gradient_accumulation.steps,
+        "parallelism": {
+            "dp_replicate": cfg.parallelism.dp_replicate,
+            "dp_shard": cfg.parallelism.dp_shard,
+            "ring_degree": cfg.parallelism.context_parallel.ring_degree,
+            "ulysses_degree": cfg.parallelism.context_parallel.ulysses_degree,
+        },
+    }
+    write_json(state, save_dir / TRAINING_STEP)
+
+
+def load_training_metadata(training_state_dir: Path) -> dict[str, Any]:
+    """Read everything `save_training_metadata` recorded, in a single pass.
+
+    Every key is always present: fields a checkpoint predates come back as None, so a caller
+    reading `metadata["batch_size"]` gets a KeyError on a typo rather than a silent None.
+
+    Args:
+        training_state_dir (Path): The checkpoint's `training_state/` directory.
+
+    Returns:
+        dict[str, Any]: `step` plus the `dp_world_size`, `batch_size`, `grad_accum_steps` and
+            `parallelism` snapshot recorded alongside it (None where not recorded).
+    """
+    state = load_json(training_state_dir / TRAINING_STEP)
+    return {
+        "step": int(state["step"]),
+        "dp_world_size": state.get("dp_world_size", state.get("num_processes")),
+        "batch_size": state.get("batch_size"),
+        "grad_accum_steps": state.get("grad_accum_steps"),
+        "parallelism": state.get("parallelism"),
+    }
+
+
+# ---------------------------------------------------------------------------------------------
+# Checkpoint save
+# ---------------------------------------------------------------------------------------------
+
+
 def save_checkpoint(
    checkpoint_dir: Path,
    step: int,
@@ -110,192 +200,301 @@ def save_checkpoint(
    scheduler: LRScheduler | None = None,
    preprocessor: PolicyProcessorPipeline | None = None,
    postprocessor: PolicyProcessorPipeline | None = None,
-    num_processes: int | None = None,
-    batch_size: int | None = None,
-    model_state_dict: dict | None = None,
-    optim_state_dict: dict | None = None,
+    accelerator: "Accelerator | None" = None,
 ) -> None:
    """This function creates the following directory structure:

    005000/  #  training step at checkpoint
    ├── pretrained_model/
    │   ├── config.json  # policy config
-    │   ├── model.safetensors  # policy weights
+    │   ├── model.safetensors  # policy weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run)
+    │   ├── pytorch_model_fsdp_0/  # DCP model shards (checkpoint_format ∈ {dcp, safetensors_dcp})
    │   ├── train_config.json  # train config
-    │   ├── processor.json  # processor config (if preprocessor provided)
-    │   └── step_*.safetensors  # processor state files (if any)
+    │   ├── policy_preprocessor.json  # preprocessor config (if preprocessor provided)
+    │   ├── policy_preprocessor_step_*.safetensors  # state of the stateful preprocessor steps
+    │   ├── policy_postprocessor.json  # postprocessor config (if postprocessor provided)
+    │   └── policy_postprocessor_step_*.safetensors  # state of the stateful postprocessor steps
    └── training_state/
-        ├── optimizer_param_groups.json  #  optimizer param groups
-        ├── optimizer_state.safetensors  # optimizer state
+        ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
+        ├── optimizer_state.safetensors  # optimizer state (non-sharded runs)
+        ├── optimizer_0/  # DCP optimizer shards (sharded runs)
        ├── rng_state.safetensors  # rng states
-        ├── scheduler_state.json  # scheduler state
-        └── training_step.json  # training step
+        ├── scheduler_state.json  # scheduler state (if scheduler provided)
+        └── training_step.json  # training step + dp_world_size/batch_size/grad_accum + topology
+
+    Collective: MUST be called on every rank. Rank-0-only writes are gated internally, so the
+    call site needs no rank branches.

    Args:
-        cfg (TrainPipelineConfig): The training config used for this run.
+        checkpoint_dir (Path): The checkpoint step directory to write (e.g. `.../checkpoints/005000`).
        step (int): The training step at that checkpoint.
+        cfg (TrainPipelineConfig): The training config used for this run.
        policy (PreTrainedPolicy): The policy to save.
-        optimizer (Optimizer | None, optional): The optimizer to save the state from. Defaults to None.
+        optimizer (Optimizer): The optimizer to save the state from.
        scheduler (LRScheduler | None, optional): The scheduler to save the state from. Defaults to None.
-        preprocessor: The preprocessor/pipeline to save. Defaults to None.
-        postprocessor: The postprocessor/pipeline to save. Defaults to None.
-        num_processes (int | None, optional): Distributed world size to record for sample-exact
-            resume. Defaults to None (not recorded).
-        batch_size (int | None, optional): Per-process batch size to record for sample-exact
-            resume. Defaults to None (not recorded).
-        model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP,
-            where `policy.state_dict()` would return sharded tensors; the caller gathers it via a
-            cross-rank collective and passes it here so rank 0 can write it directly. It holds
-            FSDP's fp32 master weights and is saved as-is (the loader casts to the policy dtype on
-            read). When None (DDP / single-GPU), the model is saved the normal way. Defaults to None.
-        optim_state_dict: Pre-gathered full (unsharded) optimizer state dict. Required under FSDP
-            (gathered alongside `model_state_dict` via `gather_fsdp_state_dicts`); saved in the same
-            safetensors format as the single-GPU path. When None, `optimizer.state_dict()` is used.
+        preprocessor (PolicyProcessorPipeline | None, optional): The preprocessor/pipeline to save.
            Defaults to None.
+        postprocessor (PolicyProcessorPipeline | None, optional): The postprocessor/pipeline to save.
+            Defaults to None.
+        accelerator (Accelerator | None, optional): The accelerator the policy was prepared with;
+            used to unwrap the model and required on sharded runs, where it owns the DCP save
+            channels. Defaults to None (plain single-process saves).
    """
    pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
-    policy.save_pretrained(pretrained_dir, state_dict=model_state_dict)
-    cfg.save_pretrained(pretrained_dir)
+    fmt = cfg.checkpoint_format
+    policy_to_save = accelerator.unwrap_model(policy) if accelerator is not None else policy
+    sharded = is_sharded_module(policy_to_save)
+
+    # -- model artifact(s): the two collective-capable calls ----------------------------------
    if cfg.peft is not None:
-        # When using PEFT, policy.save_pretrained will only write the adapter weights + config, not the
-        # policy config which we need for loading the model. In this case we'll write it ourselves.
-        policy.config.save_pretrained(pretrained_dir)
-    if preprocessor is not None:
-        preprocessor.save_pretrained(pretrained_dir)
-    if postprocessor is not None:
-        postprocessor.save_pretrained(pretrained_dir)
+        # PeftModel.save_pretrained is an external API with no internal rank gate, and the
+        # adapters are replicated (PEFT x sharded is rejected at validation): main rank writes.
+        if is_main_process():
+            policy_to_save.save_pretrained(pretrained_dir)
+    elif fmt.wants_safetensors or not sharded:
+        # Collective when sharded (full gather); writes happen on the main process only in all
+        # multi-rank layouts (the gate lives inside _save_pretrained, next to its collective gather).
+        policy_to_save.save_pretrained(pretrained_dir)
+    if fmt.wants_dcp and sharded:
+        save_sharded_model(accelerator, policy_to_save, pretrained_dir)
+
+    # -- sidecar configs: ONE gate for the whole contiguous rank-0-only region ----------------
+    if is_main_process():
+        if fmt.wants_dcp and not fmt.wants_safetensors:
+            # save_pretrained did not run: keep the DCP-only checkpoint self-describing.
+            policy_to_save.config.save_pretrained(pretrained_dir)
+        cfg.save_pretrained(pretrained_dir)
+        if cfg.peft is not None:
+            # PEFT's save_pretrained writes only adapter weights + config; the policy config
+            # needed to reload the base model is written explicitly.
+            policy_to_save.config.save_pretrained(pretrained_dir)
+        if preprocessor is not None:
+            preprocessor.save_pretrained(pretrained_dir)
+        if postprocessor is not None:
+            postprocessor.save_pretrained(pretrained_dir)
+
    save_training_state(
-        checkpoint_dir,
-        step,
-        optimizer,
-        scheduler,
-        num_processes=num_processes,
-        batch_size=batch_size,
-        optim_state_dict=optim_state_dict,
+        checkpoint_dir, step, cfg, optimizer, scheduler, accelerator, sharded=sharded, model=policy_to_save
    )
+    if accelerator is not None:
+        accelerator.wait_for_everyone()


 def save_training_state(
    checkpoint_dir: Path,
-    train_step: int,
-    optimizer: Optimizer | None = None,
+    step: int,
+    cfg: TrainPipelineConfig,
+    optimizer: Optimizer | dict[str, Optimizer] | None = None,
    scheduler: LRScheduler | None = None,
-    num_processes: int | None = None,
-    batch_size: int | None = None,
-    optim_state_dict: dict | None = None,
+    accelerator: "Accelerator | None" = None,
+    *,
+    sharded: bool = False,
+    model: PreTrainedPolicy | None = None,
 ) -> None:
-    """
-    Saves the training step, optimizer state, scheduler state, and rng state.
+    """Write training_state/. Collective under sharding: call on every rank.

    Args:
-        save_dir (Path): The directory to save artifacts to.
-        train_step (int): Current training step.
-        optimizer (Optimizer | None, optional): The optimizer from which to save the state_dict.
+        checkpoint_dir (Path): The checkpoint step directory; `training_state/` is created inside it.
+        step (int): The training step at that checkpoint.
+        cfg (TrainPipelineConfig): The training config used for this run (its topology and
+            accumulation settings are recorded in `training_step.json`).
+        optimizer (Optimizer | dict[str, Optimizer] | None, optional): The optimizer(s) to save
+            the state from. Defaults to None.
+        scheduler (LRScheduler | None, optional): The scheduler to save the state from.
            Defaults to None.
-        scheduler (LRScheduler | None, optional): The scheduler from which to save the state_dict.
-            Defaults to None.
-        num_processes (int | None, optional): Distributed world size to record. Defaults to None.
-        batch_size (int | None, optional): Per-process batch size to record. Defaults to None.
-        optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of
-            `optimizer.state_dict()` when provided. Defaults to None.
+        accelerator (Accelerator | None, optional): Required when `sharded` is True — it owns
+            the DCP optimizer save channel. Defaults to None.
+        sharded (bool): The model's sharding state, computed once in `save_checkpoint` and
+            threaded here so the two sites cannot disagree. Defaults to False.
+        model (PreTrainedPolicy | None, optional): Required only for the sharded optimizer
+            channel: torch's optimizer DCP APIs are model-coupled (the state dict is keyed by
+            model FQNs), so accelerate's `save_fsdp_optimizer` needs the sharded module
+            alongside the optimizer. Defaults to None.
    """
    save_dir = checkpoint_dir / TRAINING_STATE_DIR
+    # All ranks: the directory must exist before the DCP optimizer collective writes into it
+    # (exist_ok makes the concurrent mkdir race-free on shared filesystems).
    save_dir.mkdir(parents=True, exist_ok=True)
-    save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size)
-    save_rng_state(save_dir)
-    if optimizer is not None:
-        save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict)
-    if scheduler is not None:
-        save_scheduler_state(scheduler, save_dir)
+
+    if optimizer is not None and sharded:
+        if accelerator is None or model is None:
+            raise ValueError("Saving a sharded optimizer state requires the accelerator and model.")
+        # Collective — all ranks write their DCP shards into optimizer_0/.
+        save_sharded_optimizer(accelerator, optimizer, model, save_dir)
+
+    if is_main_process():  # ONE grouped gate for the whole rank-0-only region
+        save_training_metadata(step, save_dir, cfg)
+        save_rng_state(save_dir)
+        if scheduler is not None:
+            save_scheduler_state(scheduler, save_dir)
+        if optimizer is not None and not sharded:
+            save_optimizer_state(optimizer, save_dir)


-def load_training_state(
-    checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None, load_optimizer: bool = True
-) -> tuple[int, Optimizer, LRScheduler | None]:
-    """
-    Loads the training step, optimizer state, scheduler state, and rng state.
-    This is used to resume a training run.
+# ---------------------------------------------------------------------------------------------
+# Two-phase resume
+# ---------------------------------------------------------------------------------------------
+
+
+def resume_before_prepare(cfg: TrainPipelineConfig) -> int:
+    """Phase 1 — before `accelerator.prepare()`: restore RNG and return the step counter.
+
+    Pure loaders only. The sampler resume offset is *derived* from the returned step inside the
+    dataloader factory, and everything bound to sharded objects (model DCP shards, optimizer,
+    scheduler) loads in `resume_after_prepare`.

    Args:
-        checkpoint_dir (Path): The checkpoint directory. Should contain a 'training_state' dir.
-        optimizer (Optimizer): The optimizer to load the state_dict to.
-        scheduler (LRScheduler | None): The scheduler to load the state_dict to (can be None).
-        load_optimizer (bool, optional): Whether to load the optimizer state from disk. Defaults to
-            True. Set to False under FSDP, where the sharded optimizer state must be loaded after
-            `accelerator.prepare()` via `load_fsdp_optimizer_state` (the optimizer is returned
-            untouched here).
+        cfg (TrainPipelineConfig): The resumed training config; `cfg.checkpoint_path` locates
+            the checkpoint to restore from.
+
+    Returns:
+        int: The training step recorded in the checkpoint (micro-batch counter).

    Raises:
-        NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir
-
-    Returns:
-        tuple[int, Optimizer, LRScheduler | None]: training step, optimizer and scheduler with their
-            state_dict loaded.
+        NotADirectoryError: If the checkpoint has no `training_state/` directory.
+        ValueError: If the resumed topology crosses the sharded/non-sharded boundary relative
+            to the one recorded in the checkpoint.
    """
-    training_state_dir = checkpoint_dir / TRAINING_STATE_DIR
+    training_state_dir = cfg.checkpoint_path / TRAINING_STATE_DIR
    if not training_state_dir.is_dir():
        raise NotADirectoryError(training_state_dir)
-
+    metadata = load_training_metadata(training_state_dir)
+    _guard_resume_changes(cfg, metadata)
    load_rng_state(training_state_dir)
-    step = load_training_step(training_state_dir)
-    if load_optimizer:
-        optimizer = load_optimizer_state(optimizer, training_state_dir)
+    return metadata["step"]
+
+
+def _guard_resume_changes(cfg: TrainPipelineConfig, metadata: dict[str, Any]) -> None:
+    """Check the resumed run settings against the ones recorded in the checkpoint.
+
+    Two tiers, both driven by the checkpoint's recorded parallelism snapshot:
+
+    - **Hard error** when the resume crosses the sharded/non-sharded boundary in either
+      direction: the checkpoint's training-state artifacts only support resuming on the same
+      kind of topology (resharding works across sizes, not across kinds). Checkpoints without
+      a recorded snapshot skip this check.
+    - **One warning** naming every other recorded setting that differs — those changes are
+      legal (DCP reshards weights and optimizer state across topologies and the sampler offset
+      adapts), but a changed ``grad_accum_steps`` shifts the optimizer-update cadence, so the
+      resume says precisely what differs. The sampler-exactness warnings
+      (``dp_world_size``/``batch_size``) live with the sampler math in the dataloader factory.
+
+    Args:
+        cfg (TrainPipelineConfig): The resumed training config, compared against the settings
+            recorded in the checkpoint.
+        metadata (dict[str, Any]): The checkpoint's recorded training metadata, as returned by
+            `load_training_metadata`.
+
+    Raises:
+        ValueError: If the checkpoint records a sharded topology and the resumed run is
+            non-sharded, or vice versa.
+    """
+    snapshot = metadata["parallelism"]
+
+    if snapshot is not None:
+        recorded_sharded = (
+            snapshot.get("dp_shard", 1) != 1
+            or snapshot.get("ring_degree", 1) * snapshot.get("ulysses_degree", 1) > 1
+        )
+        if recorded_sharded != cfg.parallelism.is_sharded:
+            raise ValueError(
+                f"Cannot resume: the checkpoint was written with a "
+                f"{'sharded' if recorded_sharded else 'non-sharded'} topology "
+                f"(dp_replicate={snapshot.get('dp_replicate')}, dp_shard={snapshot.get('dp_shard')}) "
+                f"but this run is {'sharded' if cfg.parallelism.is_sharded else 'non-sharded'} "
+                f"(dp_replicate={cfg.parallelism.dp_replicate}, dp_shard={cfg.parallelism.dp_shard})."
+            )
+
+    recorded = {
+        "grad_accum_steps": (
+            metadata["grad_accum_steps"],
+            cfg.accelerator.gradient_accumulation.steps,
+        ),
+    }
+    if snapshot is not None:
+        recorded.update(
+            {
+                "dp_replicate": (snapshot.get("dp_replicate"), cfg.parallelism.dp_replicate),
+                "dp_shard": (snapshot.get("dp_shard"), cfg.parallelism.dp_shard),
+                "ring_degree": (
+                    snapshot.get("ring_degree"),
+                    cfg.parallelism.context_parallel.ring_degree,
+                ),
+                "ulysses_degree": (
+                    snapshot.get("ulysses_degree"),
+                    cfg.parallelism.context_parallel.ulysses_degree,
+                ),
+            }
+        )
+    changed = [f"{key}: {was} -> {now}" for key, (was, now) in recorded.items() if was not in (None, now)]
+    if changed and is_main_process():
+        logging.warning(
+            "Resuming with settings that differ from the checkpoint: " + "; ".join(changed) + ". "
+            "Topology changes reshard safely via DCP; a changed grad_accum_steps shifts the "
+            "optimizer-update cadence (the step counter keeps counting micro-batches)."
+        )
+
+
+def resume_after_prepare(
+    cfg: TrainPipelineConfig,
+    accelerator: "Accelerator",
+    policy: PreTrainedPolicy,
+    optimizer: Optimizer | dict[str, Optimizer],
+    scheduler: LRScheduler | None,
+) -> None:
+    """Phase 2 — after `accelerator.prepare()`: model (DCP) -> optimizer -> scheduler.
+
+    Collective under sharding: call on every rank. The model-weight source follows the
+    checkpoint's own recorded `checkpoint_format` (on resume, `cfg` was parsed from the
+    checkpoint's train_config.json): DCP-bearing formats load shards here into the prepared
+    model (whose construction skipped the safetensors load); the safetensors format was already
+    loaded by `from_pretrained` before sharding — no model step here.
+
+    Args:
+        cfg (TrainPipelineConfig): The resumed training config; `cfg.checkpoint_path` locates
+            the checkpoint and `cfg.checkpoint_format` selects the model-weight source.
+        accelerator (Accelerator): The accelerator the policy was prepared with; it unwraps the
+            model and owns the DCP load channels.
+        policy (PreTrainedPolicy): The prepared (possibly sharded) policy to load weights into.
+        optimizer (Optimizer | dict[str, Optimizer]): The prepared optimizer(s) to restore.
+        scheduler (LRScheduler | None): The scheduler to restore, or None if the run has none.
+
+    Raises:
+        FileNotFoundError: If the checkpoint format declares DCP model shards but the shard
+            directory is missing (e.g. it was pruned before upload).
+    """
+    checkpoint_dir = cfg.checkpoint_path
+    pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
+    training_state_dir = checkpoint_dir / TRAINING_STATE_DIR
+    unwrapped = accelerator.unwrap_model(policy)
+    sharded = is_sharded_module(unwrapped)
+
+    if cfg.checkpoint_format.wants_dcp:
+        from accelerate.utils.constants import FSDP_MODEL_NAME
+
+        dcp_dir = pretrained_dir / f"{FSDP_MODEL_NAME}_0"
+        if not dcp_dir.is_dir():
+            raise FileNotFoundError(
+                f"checkpoint_format={cfg.checkpoint_format.value} declares DCP model shards, "
+                f"but {dcp_dir} is missing. If the shards were pruned, convert what remains "
+                "with `lerobot-convert-dcp` or resume from a safetensors checkpoint."
+            )
+        load_sharded_model(accelerator, unwrapped, pretrained_dir)
+
+    if sharded:
+        # Requires the prepared optimizer: FSDP2's prepare rebinds param groups to DTensors but
+        # never migrates optimizer.state — DCP reshards it here (works across topology changes).
+        load_sharded_optimizer(accelerator, optimizer, unwrapped, training_state_dir)
+    else:
+        load_optimizer_state(optimizer, training_state_dir)
+
    if scheduler is not None:
-        scheduler = load_scheduler_state(scheduler, training_state_dir)
-
-    return step, optimizer, scheduler
+        load_scheduler_state(scheduler, training_state_dir)


-def gather_fsdp_state_dicts(model, optimizer) -> tuple[dict, dict]:
-    """Gather the full (unsharded) model and optimizer state dicts under FSDP.
-
-    `model.state_dict()` and `FSDP.optim_state_dict(...)` are cross-rank collectives, so this must be
-    called on *every* rank with the prepared (FSDP-wrapped) `model` and `optimizer`. With
-    `rank0_only=True` and `offload_to_cpu=True`, every rank runs the all-gather but only rank 0
-    materializes the full dicts (the others get empty dicts) and they are kept on CPU to bound GPU
-    memory. The returned optimizer state dict is keyed by parameter FQNs and is world-size
-    independent; `load_fsdp_optimizer_state` reshards it on resume.
-
-    Returns:
-        (model_state_dict, optim_state_dict): full dicts on rank 0, empty dicts on other ranks.
-    """
-    from torch.distributed.fsdp import (
-        FullOptimStateDictConfig,
-        FullStateDictConfig,
-        FullyShardedDataParallel as FSDP,  # noqa F401
-        StateDictType,
-    )
-
-    state_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
-    optim_cfg = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True)
-    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg):
-        model_state_dict = model.state_dict()
-        optim_state_dict = FSDP.optim_state_dict(model, optimizer)
-    return model_state_dict, optim_state_dict
-
-
-def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None:
-    """Load the FSDP optimizer state (saved as safetensors) and reshard it into the optimizer.
-
-    This is a cross-rank collective and must be called on every rank *after* `accelerator.prepare()`
-    with the prepared (FSDP-wrapped) `model` and `optimizer`. The saved state is the full,
-    world-size-independent optimizer state (keyed by parameter FQNs); `FSDP.optim_state_dict_to_load`
-    reshards it to the current FSDP topology, so resume on a different number of GPUs works.
-    """
-    from torch.distributed.fsdp import (
-        FullOptimStateDictConfig,
-        FullStateDictConfig,
-        FullyShardedDataParallel as FSDP,  # noqa F401
-        StateDictType,
-    )
-
-    # Every rank reads the same full state from the (shared) checkpoint dir, so rank0_only=False.
-    full_osd = load_optimizer_state_dict(checkpoint_dir / TRAINING_STATE_DIR)
-    state_cfg = FullStateDictConfig(rank0_only=False)
-    optim_cfg = FullOptimStateDictConfig(rank0_only=False)
-    with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg):
-        sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd)
-    optimizer.load_state_dict(sharded_osd)
+# ---------------------------------------------------------------------------------------------
+# Hub: checkpoint push (resume artifact) and publishing (distribution artifact)
+# ---------------------------------------------------------------------------------------------


 def push_checkpoint_to_hub(
@@ -311,6 +510,16 @@ def push_checkpoint_to_hub(
    The model repo is created idempotently, and the commit is tagged with the
    checkpoint step so a checkpoint can be recovered with
    --policy.pretrained_revision=<step> instead of a commit sha.
+
+    The directory is uploaded verbatim — including DCP shards under the DCP formats: this tree
+    exists for *resume*, not distribution, and `resolve_resume_checkpoint` downloads it back
+    symmetrically.
+
+    Args:
+        checkpoint_dir (Path): The local checkpoint step directory to upload.
+        repo_id (str): The Hub model repo to push to (created idempotently if missing).
+        private (bool | None): Whether a newly created repo should be private. Defaults to
+            None (public unless the organization's default is private).
    """
    api = HfApi()
    api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True)
@@ -338,6 +547,16 @@ def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path:
    into `output_dir/checkpoints/<step>/`, recreate the local `last` symlink, and return that local
    checkpoint dir. Used to resume training from the Hub on a machine (or HF Jobs pod) that does not
    have the original local run dir.
+
+    Args:
+        repo_id (str): The Hub model repo holding `checkpoints/<step>/` subtrees.
+        output_dir (Path): The local run directory to download the checkpoint into.
+
+    Returns:
+        Path: The local checkpoint step directory, `output_dir/checkpoints/<step>`.
+
+    Raises:
+        FileNotFoundError: If the repo contains no checkpoints under `checkpoints/`.
    """
    latest = find_latest_hub_checkpoint(repo_id)
    if latest is None:
@@ -354,3 +573,214 @@ def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path:
    checkpoint_dir = output_dir / latest
    update_last_checkpoint(checkpoint_dir)
    return checkpoint_dir
+
+
+def publish_trained_model(
+    cfg: TrainPipelineConfig,
+    model: "PreTrainedPolicy | PreTrainedRewardModel",
+    preprocessor: PolicyProcessorPipeline | None,
+    postprocessor: PolicyProcessorPipeline | None,
+    dataset_meta: "LeRobotDatasetMetadata | None",
+    *,
+    peft_model: Any | None = None,
+) -> None:
+    """Publish the complete training bundle as a distributable model repo.
+
+    Collective-safe: call on ALL ranks — the model commit gathers sharded weights through
+    `save_pretrained`; uploads happen on the main process only (gated inside
+    `HubMixin.push_to_hub` and here). Commits, in order: (1) the model (skipped for PEFT —
+    adapters replace full weights), (2) the preprocessor, (3) the postprocessor, (4) the bundle
+    sidecar: README.md model card + train_config.json (+ adapter weights and the wrapped
+    policy's config in the PEFT case). Every commit uploads a freshly assembled directory, so
+    a published repo carries only the distributable artifacts.
+
+    Args:
+        cfg (TrainPipelineConfig): The training config; saved as `train_config.json` and used
+            to render the model card.
+        model (PreTrainedPolicy | PreTrainedRewardModel): The trained model to publish; its
+            config supplies the target repo id, visibility, license, and tags.
+        preprocessor (PolicyProcessorPipeline | None): The preprocessor pipeline to publish
+            alongside the model, if any.
+        postprocessor (PolicyProcessorPipeline | None): The postprocessor pipeline to publish
+            alongside the model, if any.
+        dataset_meta (LeRobotDatasetMetadata | None): Dataset metadata for the model card, if
+            available.
+        peft_model (Any | None): The PEFT wrapper when training adapters; its adapter weights
+            replace the full model weights in the published repo. Defaults to None.
+
+    Raises:
+        ValueError: If the model config carries no repo id (`--policy.repo_id`).
+    """
+    model_cfg = model.config
+    repo_id = model_cfg.repo_id
+    if not repo_id:
+        raise ValueError("Publishing requires a repo id (--policy.repo_id).")
+    ignore = ["*.tmp", "*.log"]
+
+    if peft_model is None:
+        # Calls are made on the exact objects that own each method (never through PEFT's
+        # attribute forwarding), so the peft branch below never touches this path.
+        model.push_to_hub(repo_id, private=model_cfg.private, ignore_patterns=ignore)
+    if preprocessor is not None:
+        preprocessor.push_to_hub(repo_id, private=model_cfg.private)
+    if postprocessor is not None:
+        postprocessor.push_to_hub(repo_id, private=model_cfg.private)
+
+    if is_main_process():
+        api = HfApi()
+        repo_id = api.create_repo(repo_id=repo_id, private=model_cfg.private, exist_ok=True).repo_id
+        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
+            saved_path = Path(tmp) / repo_id
+            saved_path.mkdir(parents=True, exist_ok=True)
+            if peft_model is not None:
+                peft_model.save_pretrained(saved_path)  # adapter weights + adapter config
+                model.config.save_pretrained(saved_path)  # PEFT cannot write the policy config
+            card = generate_model_card(model_cfg, cfg=cfg, dataset_meta=dataset_meta)
+            card.save(str(saved_path / "README.md"))
+            cfg.save_pretrained(saved_path)  # train_config.json
+            commit_info = api.upload_folder(
+                repo_id=repo_id,
+                repo_type="model",
+                folder_path=saved_path,
+                commit_message="Upload model card and train config",
+                allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"],
+                ignore_patterns=ignore,
+            )
+        # Contract: lerobot.jobs.hf.submit_to_hf watches for this exact "Model pushed to <url>"
+        # line to end a remote run early. Keep the wording and URL format in sync.
+        logging.info(f"Model pushed to {commit_info.repo_url.url}")
+
+    if dist.is_initialized():
+        dist.barrier()
+
+
+# ---------------------------------------------------------------------------------------------
+# Model card
+# ---------------------------------------------------------------------------------------------
+
+_BASE_MODEL_MAPPING = {
+    "smolvla": "lerobot/smolvla_base",
+    "pi0": "lerobot/pi0_base",
+    "pi05": "lerobot/pi05_base",
+    "pi0_fast": "lerobot/pi0fast-base",
+    "xvla": "lerobot/xvla-base",
+}
+
+
+def build_card_context(
+    cfg: TrainPipelineConfig | None,
+    dataset_meta: "LeRobotDatasetMetadata | None",
+    input_features: dict | None,
+    output_features: dict | None,
+) -> dict:
+    """Collect optional data for the model-card template.
+
+    Returns plain values only (no Markdown) — the template in
+    ``lerobot/templates/lerobot_modelcard_template.md`` decides how and whether to show
+    each one. Everything is best-effort: anything unavailable is left empty/None and the
+    template simply skips that section, so this never breaks a Hub push.
+
+    Args:
+        cfg (TrainPipelineConfig | None): The training config supplying the training section,
+            if available.
+        dataset_meta (LeRobotDatasetMetadata | None): Dataset metadata supplying the dataset,
+            robot-type, and camera sections, if available.
+        input_features (dict | None): The policy's input feature declarations, if any.
+        output_features (dict | None): The policy's output feature declarations, if any.
+
+    Returns:
+        dict: Template context with `training`, `input_features`, `output_features`,
+            `dataset`, `robot_type`, and `cameras` entries; unavailable pieces stay
+            empty/None.
+    """
+    context = {
+        "training": None,
+        "input_features": input_features or {},
+        "output_features": output_features or {},
+        "dataset": None,
+        "robot_type": None,
+        "cameras": [],
+    }
+
+    if cfg is not None:
+        optimizer = getattr(cfg, "optimizer", None)
+        context["training"] = {
+            "steps": cfg.steps,
+            "batch_size": cfg.batch_size,
+            "seed": cfg.seed,
+            "optimizer": getattr(optimizer, "type", None) if optimizer else None,
+            "lr": getattr(optimizer, "lr", None) if optimizer else None,
+            "lerobot_version": __version__,
+        }
+
+    if dataset_meta is not None:
+        context["dataset"] = {
+            "repo_id": dataset_meta.repo_id,
+            "episodes": dataset_meta.total_episodes,
+            "frames": dataset_meta.total_frames,
+            "fps": dataset_meta.fps,
+            "tasks": [str(task) for task in dataset_meta.tasks.index],
+        }
+        context["robot_type"] = dataset_meta.robot_type
+        context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys]
+
+    return context
+
+
+def generate_model_card(
+    model_cfg: PreTrainedConfig | RewardModelConfig,
+    cfg: TrainPipelineConfig | None = None,
+    dataset_meta: "LeRobotDatasetMetadata | None" = None,
+) -> ModelCard:
+    """Render the LeRobot model card for a trained policy or reward model.
+
+    A free function on purpose: every template variable comes from arguments — the model
+    config, the training config, and the dataset metadata — none from a live model, so a card
+    can also be rendered from a checkpoint's `config.json` alone (see `lerobot-convert-dcp`).
+    The config type selects the template: reward models get the reward-model card, policies the
+    policy card with the training/dataset sections.
+
+    Args:
+        model_cfg (PreTrainedConfig | RewardModelConfig): The model config providing type,
+            license, tags, repo id, and — for policies — the feature declarations.
+        cfg (TrainPipelineConfig | None, optional): The training config for the training and
+            dataset card sections. Defaults to None.
+        dataset_meta (LeRobotDatasetMetadata | None, optional): Dataset metadata for the
+            dataset card sections. Defaults to None.
+
+    Returns:
+        ModelCard: The rendered and validated LeRobot model card.
+    """
+    model_type = model_cfg.type
+    base_model = _BASE_MODEL_MAPPING.get(model_type)
+
+    if isinstance(model_cfg, RewardModelConfig):
+        tags = {"robotics", "lerobot", "reward-model", model_type}
+        template_card = (
+            files("lerobot.templates")
+            .joinpath("lerobot_rewardmodel_modelcard_template.md")
+            .read_text("utf-8")
+        )
+        context: dict[str, Any] = {}  # the reward template renders from card_data alone
+    else:
+        tags = {"robotics", "lerobot", model_type}
+        template_card = (
+            files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text("utf-8")
+        )
+        context = build_card_context(cfg, dataset_meta, model_cfg.input_features, model_cfg.output_features)
+        # Used by the template to pre-fill commands and the "Fine-tuned from" line.
+        context["policy_repo_id"] = model_cfg.repo_id
+        context["base_model"] = base_model
+
+    card_data = ModelCardData(
+        license=model_cfg.license or "apache-2.0",
+        library_name="lerobot",
+        pipeline_tag="robotics",
+        tags=list(tags.union(model_cfg.tags or [])),
+        model_name=model_type,
+        datasets=cfg.dataset.repo_id if cfg is not None else None,
+        base_model=base_model,
+    )
+    card = ModelCard.from_template(card_data, template_str=template_card, **context)
+    card.validate()
+    return card
@@ -22,7 +22,7 @@ Import them directly: ``from lerobot.configs.train import TrainPipelineConfig``
 """

 from .dataset import DatasetRecordConfig
-from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EMAConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .recipe import MessageTurn, TrainingRecipe, load_recipe
 from .types import (
@@ -57,6 +57,7 @@ __all__ = [
    # Config classes
    "DatasetRecordConfig",
    "DatasetConfig",
+    "EMAConfig",
    "EvalConfig",
    "JobConfig",
    "MessageTurn",
@@ -0,0 +1,273 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Execution-runtime configuration: everything handed to (or applied by) the `Accelerator`.
+
+Each sub-config mirrors the plain-typed subset of the corresponding accelerate object and
+builds it at runtime (the way ``OptimizerConfig.build()`` constructs a ``torch.optim.Optimizer``),
+so the whole tree round-trips through the CLI and ``train_config.json`` and parsing a config
+never imports accelerate.
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from lerobot.configs.parallelism import ParallelismConfig
+
+if TYPE_CHECKING:
+    from accelerate import Accelerator
+    from accelerate.utils import (
+        DistributedDataParallelKwargs,
+        FullyShardedDataParallelPlugin,
+        GradientAccumulationPlugin,
+    )
+
+
+@dataclass
+class FSDPConfig:
+    """Mirror of the `FullyShardedDataParallelPlugin` subset LeRobot supports (FSDP2 only).
+
+    Exactly one wrap policy applies: `wrap_modules` (module *class names* forming the FSDP
+    units — and, later, the activation-checkpointing units) or `min_num_params` (size-based).
+    When both are None, the policy's own `_fsdp_wrap_modules` declaration is used; a run where
+    no wrap source exists at all fails loudly rather than silently wrapping only the root.
+    """
+
+    reshard_after_forward: bool = True
+    wrap_modules: list[str] | None = None
+    min_num_params: int | None = None
+    cpu_offload: bool = False
+    # Regex matched against module FQNs to exclude their parameters from sharding.
+    ignored_modules: str | None = None
+
+    def __post_init__(self) -> None:
+        """Validate the wrap-policy fields.
+
+        Raises:
+            ValueError: If both ``wrap_modules`` and ``min_num_params`` are set (they are
+                mutually exclusive wrap policies), or if ``min_num_params`` is < 1.
+        """
+        if self.wrap_modules is not None and self.min_num_params is not None:
+            raise ValueError(
+                "fsdp.wrap_modules and fsdp.min_num_params are mutually exclusive wrap policies."
+            )
+        if self.min_num_params is not None and self.min_num_params < 1:
+            raise ValueError(f"fsdp.min_num_params must be >= 1, got {self.min_num_params}.")
+
+    def build_plugin(self) -> "FullyShardedDataParallelPlugin":
+        """Build the FSDP2 plugin for `Accelerator(fsdp_plugin=...)`.
+
+        Returns:
+            FullyShardedDataParallelPlugin: FSDP2 (`fsdp_version=2`) plugin carrying the
+                mirrored wrap policy, resharding, CPU-offload, and ignored-modules settings.
+        """
+        from accelerate.utils import FullyShardedDataParallelPlugin
+
+        use_size_policy = self.min_num_params is not None
+        return FullyShardedDataParallelPlugin(
+            fsdp_version=2,
+            reshard_after_forward=self.reshard_after_forward,
+            auto_wrap_policy="size_based_wrap" if use_size_policy else "transformer_based_wrap",
+            # May legitimately still be None here: the policy-declared default is applied right
+            # before `accelerator.prepare()` (see lerobot.distributed.factory.set_fsdp_wrap_modules).
+            transformer_cls_names_to_wrap=list(self.wrap_modules) if self.wrap_modules else None,
+            min_num_params=self.min_num_params,
+            cpu_offload=self.cpu_offload,
+            ignored_modules=self.ignored_modules,
+            # state_dict_type stays at the FSDP2 default (SHARDED_STATE_DICT) and is never
+            # switched: full gathers go through torch's state-dict API, which does not consult
+            # the plugin. activation_checkpointing stays False: AC is LeRobot-owned.
+        )
+
+
+@dataclass
+class DDPConfig:
+    """Mirror of the `DistributedDataParallelKwargs` subset LeRobot exposes."""
+
+    # Today's in-script default, kept for models with conditional computation.
+    find_unused_parameters: bool = True
+    gradient_as_bucket_view: bool = False
+    static_graph: bool = False
+
+    def build_kwargs_handler(self) -> "DistributedDataParallelKwargs":
+        """Build the DDP kwargs handler for `Accelerator(kwargs_handlers=[...])`.
+
+        Returns:
+            DistributedDataParallelKwargs: Handler carrying the mirrored DDP fields, applied
+                by accelerate when it wraps the model in `DistributedDataParallel`.
+        """
+        from accelerate.utils import DistributedDataParallelKwargs
+
+        return DistributedDataParallelKwargs(
+            find_unused_parameters=self.find_unused_parameters,
+            gradient_as_bucket_view=self.gradient_as_bucket_view,
+            static_graph=self.static_graph,
+        )
+
+
+@dataclass
+class GradientAccumulationConfig:
+    """Mirror of the `GradientAccumulationPlugin` subset LeRobot supports.
+
+    Only the step count is a knob. ``sync_with_dataloader`` is pinned to False by
+    :meth:`build_plugin`: the training loop cycles a finite dataloader, so accelerate's default
+    of syncing at every dataloader end would force an optimizer step at every dataset epoch
+    boundary instead of every ``steps`` micro-batches.
+    """
+
+    steps: int = 1
+
+    def __post_init__(self) -> None:
+        """Validate the accumulation step count.
+
+        Raises:
+            ValueError: If ``steps`` is < 1.
+        """
+        if self.steps < 1:
+            raise ValueError(f"gradient_accumulation.steps must be >= 1, got {self.steps}.")
+
+    def build_plugin(self) -> "GradientAccumulationPlugin":
+        """Build the plugin for `Accelerator(gradient_accumulation_plugin=...)`.
+
+        A named plugin argument, not a `kwargs_handlers` entry: accelerate consumes this object
+        through its dedicated constructor parameter — the `KwargsHandler` base class only lends
+        it `to_kwargs()`, so the consumption site, not the inheritance, decides its role.
+
+        Returns:
+            GradientAccumulationPlugin: Carrying the mirrored step count, with
+                ``sync_with_dataloader=False`` pinned (see the class docstring).
+        """
+        from accelerate.utils import GradientAccumulationPlugin
+
+        return GradientAccumulationPlugin(num_steps=self.steps, sync_with_dataloader=False)
+
+
+@dataclass
+class CompileConfig:
+    """torch.compile knobs — a configured placeholder: wiring lands in a later round.
+
+    The setup-order contract it will follow is already fixed: compile applies
+    after CP dispatch install and activation checkpointing, before `fully_shard`, regionally
+    (per wrap unit) — the only combination proven with FSDP2.
+    """
+
+    enabled: bool = False
+    backend: str = "inductor"
+    mode: str | None = None
+    regional: bool = True
+
+
+class ActivationCheckpointingMode(str, Enum):
+    NONE = "none"
+    FULL = "full"
+
+
+@dataclass
+class ActivationCheckpointingConfig:
+    """Activation-checkpointing knobs — a configured placeholder: wiring lands in a later round.
+
+    AC units will coincide with the FSDP wrap units (one declaration drives both), applied
+    before torch.compile and `fully_shard` (the same ordering contract as CompileConfig).
+    """
+
+    mode: ActivationCheckpointingMode = ActivationCheckpointingMode.NONE
+
+
+@dataclass
+class AcceleratorConfig:
+    """Builds the `Accelerator` — the runtime counterpart of the `parallelism` topology.
+
+    `mixed_precision` selects accelerate-native AMP for DDP/single-GPU runs and the FSDP2
+    `MixedPrecisionPolicy` for sharded runs (accelerate derives it). Sharded runs support
+    "no" and "bf16" only; fp16's GradScaler-over-DTensor path is unverified and fails fast
+    at config validation.
+    """
+
+    mixed_precision: str = "no"
+    gradient_accumulation: GradientAccumulationConfig = field(default_factory=GradientAccumulationConfig)
+    fsdp: FSDPConfig = field(default_factory=FSDPConfig)
+    ddp: DDPConfig = field(default_factory=DDPConfig)
+    compile: CompileConfig = field(default_factory=CompileConfig)
+    activation_checkpointing: ActivationCheckpointingConfig = field(
+        default_factory=ActivationCheckpointingConfig
+    )
+
+    def __post_init__(self) -> None:
+        """Validate the accelerate-facing scalar fields.
+
+        Raises:
+            ValueError: If ``mixed_precision`` is not one of ``"no"``, ``"fp16"``, ``"bf16"``.
+        """
+        if self.mixed_precision not in ("no", "fp16", "bf16"):
+            raise ValueError(
+                f"mixed_precision must be one of 'no', 'fp16', 'bf16', got {self.mixed_precision!r}."
+            )
+
+    def build(self, parallelism: ParallelismConfig, *, cpu: bool = False) -> "Accelerator":
+        """Translate the mirrored fields into a ready `Accelerator` (call once per process).
+
+        `parallelism` must already be resolved against the world size. The degradation matrix
+        is encoded here and nowhere else: sharded -> FSDP2 (+HSDP via the accelerate
+        `ParallelismConfig` mesh), replicated-only -> DDP kwargs, single process -> plain.
+
+        Args:
+            parallelism (ParallelismConfig): The resolved process topology; selects which
+                accelerate path (FSDP2 mesh, DDP kwargs handler, or plain) is configured.
+            cpu (bool): Force CPU execution even when CUDA is available. Defaults to False.
+
+        Returns:
+            Accelerator: The configured accelerate entry point for this process.
+        """
+        from accelerate import Accelerator
+
+        kwargs: dict = {
+            # LeRobot steps its scheduler manually once per training step; accelerate must not
+            # rescale scheduler stepping by num_processes.
+            "step_scheduler_with_optimizer": False,
+            "gradient_accumulation_plugin": self.gradient_accumulation.build_plugin(),
+            "mixed_precision": self.mixed_precision,
+            "cpu": cpu,
+        }
+        if parallelism.is_sharded:
+            kwargs["fsdp_plugin"] = self.fsdp.build_plugin()
+            kwargs["parallelism_config"] = _accelerate_parallelism_config(parallelism)
+        elif parallelism.is_replicated_only:
+            kwargs["kwargs_handlers"] = [self.ddp.build_kwargs_handler()]
+        return Accelerator(**kwargs)
+
+
+def _accelerate_parallelism_config(parallelism: ParallelismConfig) -> object:
+    """LeRobot topology -> accelerate `ParallelismConfig`.
+
+    CP is declared honestly (`cp_size = ring x ulysses`) so accelerate builds the canonical
+    mesh, folds CP into the FSDP shard group (`dp_shard_cp`), and duplicates batches within CP
+    groups. The ring/ulysses sub-structure stays private to `lerobot.distributed.ParallelDims`.
+
+    Args:
+        parallelism (ParallelismConfig): The resolved LeRobot topology to translate.
+
+    Returns:
+        object: The accelerate `ParallelismConfig` mirroring `dp_replicate`, `dp_shard`, and
+            the collapsed `cp_size` (annotated as `object` so importing this module never
+            imports accelerate).
+    """
+    from accelerate.parallelism_config import ParallelismConfig as AccelerateParallelismConfig
+
+    return AccelerateParallelismConfig(
+        dp_replicate_size=parallelism.dp_replicate,
+        dp_shard_size=parallelism.dp_shard,
+        cp_size=parallelism.cp_size,
+    )
@@ -139,6 +139,59 @@ class EvalConfig:
        return min(by_cpu, self.n_episodes, 64)


+@dataclass
+class EMAConfig:
+    """Exponential moving average (EMA) of the policy weights.
+
+    Standard practice for diffusion-style policies (Chi et al. 2023, "Diffusion Policy", section V.D):
+    the reference implementation enables it in every config and evaluates the EMA weights. Off by
+    default here because it keeps a second full copy of the parameters in memory.
+
+    The decay follows the warmup schedule from diffusers' `EMAModel`:
+    `decay_t = 1 - (1 + t / inv_gamma) ** -power`, clamped to `[min_decay, max_decay]`.
+    The defaults mirror the reference implementation. Alternatively, set `decay` for a constant
+    decay at every step, as used by openpi for pi0/pi05 (`ema_decay=0.99`).
+    """
+
+    enable: bool = False
+    # Constant decay coefficient (openpi-style, e.g. 0.99 for pi0/pi05). When set, the warmup
+    # schedule below is bypassed and the shadow uses this decay at every step.
+    decay: float | None = None
+    # Number of optimizer steps during which the shadow stays a hard copy of the live weights.
+    update_after_step: int = 0
+    # Warmup schedule parameters (see class docstring).
+    inv_gamma: float = 1.0
+    power: float = 0.75
+    min_decay: float = 0.0
+    max_decay: float = 0.9999
+    # Evaluate the EMA weights (instead of the live ones) during periodic env eval.
+    # Offline eval-loss (--eval_steps) always uses the live weights: it runs on every rank
+    # while the EMA shadow only lives on the main process.
+    use_for_eval: bool = True
+
+    def __post_init__(self) -> None:
+        if not (0.0 <= self.min_decay <= self.max_decay <= 1.0):
+            raise ValueError(
+                "Expected 0 <= ema.min_decay <= ema.max_decay <= 1, got "
+                f"min_decay={self.min_decay} and max_decay={self.max_decay}."
+            )
+        if self.inv_gamma <= 0:
+            raise ValueError(f"ema.inv_gamma must be positive, got {self.inv_gamma}.")
+        if self.power <= 0:
+            raise ValueError(f"ema.power must be positive, got {self.power}.")
+        if self.update_after_step < 0:
+            raise ValueError(f"ema.update_after_step must be >= 0, got {self.update_after_step}.")
+        if self.decay is not None:
+            if not 0.0 <= self.decay <= 1.0:
+                raise ValueError(f"ema.decay must be in [0, 1], got {self.decay}.")
+            # Keep the literals in sync with the field defaults above.
+            if self.min_decay != 0.0 or self.max_decay != 0.9999:
+                raise ValueError(
+                    "ema.decay (constant decay) and ema.min_decay/ema.max_decay (schedule clamp) are "
+                    "mutually exclusive: set one or the other."
+                )
+
+
@dataclass
 class PeftConfig:
    # PEFT offers many fine-tuning methods, layer adapters being the most common and currently also the most
@@ -0,0 +1,190 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Declarative process topology for distributed training and inference.
+
+The mesh convention (canonical row-major rank layout, outermost first)::
+
+    (dp_replicate, dp_shard, ring, ulysses)
+
+- ``dp_replicate x dp_shard`` is the data-parallel world: HSDP replicates over
+  ``dp_replicate`` and shards parameters over ``dp_shard``. FSDP2's actual shard
+  group folds context parallelism in (``dp_shard x ring x ulysses``), matching
+  accelerate's ``dp_shard_cp`` flattening and torchtitan's ``fsdp`` axis.
+- ``ring`` is the outer and ``ulysses`` the inner context-parallel dim
+  (diffusers convention: ulysses all-to-all exchanges run over adjacent, typically
+  NVLink-connected ranks).
+- ``cfg_parallel`` (classifier-free-guidance parallelism) is a branch-parallel,
+  inference-only dim that sits between dp and the sequence dims. It never
+  affects weight sharding or checkpoints.
+
+This module is pure configuration: plain-typed dataclasses that draccus can
+round-trip through the CLI and ``train_config.json``. Runtime objects (device
+meshes, process groups) live in :mod:`lerobot.distributed`.
+"""
+
+import os
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ContextParallelConfig:
+    """Ring x Ulysses context parallelism (sequence parallelism for attention).
+
+    Both degrees are configured placeholders in this release: the CP engine is not implemented
+    yet, and enabling either degree > 1 fails fast at config validation. The fields exist now so
+    that the CLI surface, checkpoint metadata, and mesh math are stable when the engine lands.
+    """
+
+    ring_degree: int = 1
+    ulysses_degree: int = 1
+
+    def __post_init__(self) -> None:
+        """Validate the declared context-parallel degrees.
+
+        Raises:
+            ValueError: If ``ring_degree`` or ``ulysses_degree`` is < 1.
+        """
+        if self.ring_degree < 1 or self.ulysses_degree < 1:
+            raise ValueError(
+                f"Context-parallel degrees must be >= 1, got ring_degree={self.ring_degree}, "
+                f"ulysses_degree={self.ulysses_degree}."
+            )
+
+    @property
+    def size(self) -> int:
+        """Total number of ranks a full sequence is sharded across."""
+        return self.ring_degree * self.ulysses_degree
+
+
+@dataclass
+class ParallelismConfig:
+    """Degrees of every parallelism dim. Invariant: their product equals the world size.
+
+    Degradations are expressed purely through the degrees (no mode flags):
+
+    - single process: all degrees 1;
+    - DDP: ``dp_replicate == world_size`` (auto-filled when every sharding field is left at its
+      default — plain ``torchrun`` keeps today's out-of-the-box behavior);
+    - FSDP: ``dp_shard > 1`` (or ``-1`` to fill the remaining world into the shard dim);
+    - HSDP: ``dp_replicate > 1`` and ``dp_shard > 1``.
+
+    ``resolve()`` turns the declared degrees into concrete ones once the world size is known and
+    is the single place the world-size equation is enforced. It is called by
+    :func:`lerobot.distributed.factory.make_accelerator`; the config is inert until then.
+    """
+
+    dp_replicate: int = 1
+    # -1 is an explicit opt-in sentinel: shard over world_size // (dp_replicate * cp).
+    dp_shard: int = 1
+    context_parallel: ContextParallelConfig = field(default_factory=ContextParallelConfig)
+    # Classifier-free-guidance parallelism — inference-only (cosmos/vllm-omni precedent:
+    # cond/uncond branches on different ranks). Reserved for the serving round; training
+    # validates it to 1. Meaningful values are 1 or 2 (Cosmos3 has two CFG branches).
+    cfg_parallel: int = 1
+
+    def __post_init__(self) -> None:
+        """Validate the declared degrees (world-size-independent checks only).
+
+        Raises:
+            ValueError: If ``dp_replicate`` is < 1, ``dp_shard`` is neither >= 1 nor the
+                ``-1`` infer sentinel, or ``cfg_parallel`` is not 1 or 2.
+        """
+        if self.dp_replicate < 1:
+            raise ValueError(f"dp_replicate must be >= 1, got {self.dp_replicate}.")
+        if self.dp_shard < 1 and self.dp_shard != -1:
+            raise ValueError(f"dp_shard must be >= 1, or -1 to infer, got {self.dp_shard}.")
+        if self.cfg_parallel not in (1, 2):
+            raise ValueError(f"cfg_parallel must be 1 or 2, got {self.cfg_parallel}.")
+
+    @property
+    def cp_size(self) -> int:
+        """Total context-parallel size (``ring_degree * ulysses_degree``)."""
+        return self.context_parallel.size
+
+    @property
+    def is_sharded(self) -> bool:
+        """True when the run uses FSDP2 (parameters sharded); selects the sharded engine path."""
+        return self.dp_shard != 1 or self.cp_size > 1
+
+    @property
+    def is_replicated_only(self) -> bool:
+        """True for plain DDP (weights replicated, no sharding)."""
+        return not self.is_sharded and self.dp_replicate > 1
+
+    @property
+    def dp_world_size(self) -> int:
+        """Number of distinct data-parallel workers (batches are sharded this many ways).
+
+        Returns:
+            int: ``dp_replicate * dp_shard``.
+
+        Raises:
+            RuntimeError: If accessed while ``dp_shard`` is still the ``-1`` sentinel, i.e.
+                before :meth:`resolve` has bound the degrees to a world size.
+        """
+        if self.dp_shard == -1:
+            raise RuntimeError("dp_world_size is undefined before resolve() fills dp_shard=-1.")
+        return self.dp_replicate * self.dp_shard
+
+    def resolve(self, world_size: int) -> None:
+        """Bind the declared degrees to a concrete world size (idempotent).
+
+        Fills the ``dp_shard=-1`` sentinel, auto-fills ``dp_replicate`` for the DDP degradation,
+        and enforces ``dp_replicate * dp_shard * cp == world_size`` with every degree echoed on
+        failure.
+
+        Args:
+            world_size (int): Total number of launched processes (torchrun's ``WORLD_SIZE``).
+
+        Raises:
+            ValueError: If a context-parallel degree is > 1 (the CP engine is not implemented
+                yet), if ``dp_shard=-1`` cannot be inferred because ``world_size`` is not
+                divisible by ``dp_replicate * cp``, or if the resolved degrees do not multiply
+                to ``world_size``.
+        """
+        if self.cp_size > 1:
+            raise ValueError(
+                "Context parallelism is not implemented yet: ring_degree and ulysses_degree "
+                "must be 1. The fields are reserved for the CP engine round."
+            )
+        if self.is_sharded:
+            if self.dp_shard == -1:
+                self.dp_shard, remainder = divmod(world_size, self.dp_replicate * self.cp_size)
+                if remainder or self.dp_shard < 1:
+                    raise ValueError(
+                        f"Cannot infer dp_shard: world_size={world_size} is not divisible by "
+                        f"dp_replicate={self.dp_replicate} * cp={self.cp_size}."
+                    )
+        elif self.dp_replicate == 1:
+            # Untouched config on a multi-process launch: fill the DDP degradation.
+            self.dp_replicate = world_size
+        total = self.dp_replicate * self.dp_shard * self.cp_size
+        if total != world_size:
+            raise ValueError(
+                f"Parallelism degrees do not multiply to the world size: dp_replicate="
+                f"{self.dp_replicate} * dp_shard={self.dp_shard} * ring="
+                f"{self.context_parallel.ring_degree} * ulysses="
+                f"{self.context_parallel.ulysses_degree} = {total} != WORLD_SIZE={world_size}."
+            )
+
+
+def world_size_from_env() -> int:
+    """World size as set by torchrun (or 1 outside distributed launches).
+
+    Returns:
+        int: The ``WORLD_SIZE`` environment variable, or 1 when unset.
+    """
+    return int(os.environ.get("WORLD_SIZE", "1"))
@@ -18,6 +18,7 @@ import multiprocessing
 import os
 import tempfile
 from dataclasses import dataclass, field
+from enum import Enum
 from pathlib import Path
 from typing import Any

@@ -26,19 +27,49 @@ from huggingface_hub import hf_hub_download
 from huggingface_hub.errors import HfHubHTTPError

 from lerobot import envs
+from lerobot.configs.accelerator import AcceleratorConfig, ActivationCheckpointingMode
+from lerobot.configs.parallelism import ParallelismConfig
 from lerobot.optim import LRSchedulerConfig, OptimizerConfig
 from lerobot.utils.constants import PRETRAINED_MODEL_DIR
 from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint
 from lerobot.utils.sample_weighting import SampleWeightingConfig

 from . import parser
-from .default import DatasetConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
+from .default import DatasetConfig, EMAConfig, EvalConfig, JobConfig, PeftConfig, WandBConfig
 from .policies import PreTrainedConfig
 from .rewards import RewardModelConfig

 TRAIN_CONFIG_NAME = "train_config.json"


+class CheckpointFormat(str, Enum):
+    """Model-artifact format inside training checkpoints.
+
+    Selects only the *model* artifact; the training_state layout is format-independent (the
+    optimizer channel is always DCP under sharded runs, safetensors+json otherwise).
+
+    - SAFETENSORS (default): a full `model.safetensors` — maximum compatibility, one gather per
+      save under sharding.
+    - DCP: sharded `pytorch_model_fsdp_0/*.distcp` only — fastest save/resume; convert with
+      `lerobot-convert-dcp` before distributing.
+    - SAFETENSORS_AND_DCP: both artifacts, written independently.
+    """
+
+    SAFETENSORS = "safetensors"
+    DCP = "dcp"
+    SAFETENSORS_AND_DCP = "safetensors_dcp"
+
+    @property
+    def wants_safetensors(self) -> bool:
+        """True when a full `model.safetensors` artifact should be written."""
+        return self in (CheckpointFormat.SAFETENSORS, CheckpointFormat.SAFETENSORS_AND_DCP)
+
+    @property
+    def wants_dcp(self) -> bool:
+        """True when sharded DCP model shards (`pytorch_model_fsdp_0/`) should be written."""
+        return self in (CheckpointFormat.DCP, CheckpointFormat.SAFETENSORS_AND_DCP)
+
+
 def _migrate_legacy_rabc_fields(config: dict[str, Any]) -> dict[str, Any] | None:
    """Return migrated payload for legacy RA-BC fields, or None when no migration is needed."""
    legacy_fields = (
@@ -121,10 +152,19 @@ class TrainPipelineConfig(HubMixin):
    # Checkpoint is saved every `save_freq` training iterations and after the last training step.
    # A non-positive value disables periodic saving, keeping only the final checkpoint.
    save_freq: int = 20_000
+    # Model-artifact format inside checkpoints; non-default values require a sharded run.
+    checkpoint_format: CheckpointFormat = CheckpointFormat.SAFETENSORS
    use_policy_training_preset: bool = True
    optimizer: OptimizerConfig | None = None
    scheduler: LRSchedulerConfig | None = None
+    # Process topology: dp_replicate / dp_shard (HSDP) and context-parallel degree placeholders.
+    parallelism: ParallelismConfig = field(default_factory=ParallelismConfig)
+    # Execution runtime handed to the Accelerator: mixed precision, gradient accumulation,
+    # FSDP/DDP tuning knobs, compile & activation-checkpointing placeholders.
+    accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig)
    eval: EvalConfig = field(default_factory=EvalConfig)
+    # Maintain an EMA shadow of the policy weights during training (see EMAConfig).
+    ema: EMAConfig = field(default_factory=EMAConfig)
    wandb: WandBConfig = field(default_factory=WandBConfig)
    peft: PeftConfig | None = None

@@ -291,6 +331,60 @@ class TrainPipelineConfig(HubMixin):
        if self.save_checkpoint_to_hub and not (self.policy is not None and self.policy.repo_id):
            raise ValueError("save_checkpoint_to_hub requires --policy.repo_id.")

+        self._validate_distributed()
+
+    def _validate_distributed(self) -> None:
+        """Fail-fasts for the distributed-training scope.
+
+        Raises:
+            ValueError: If the config requests anything outside the verified scope: context
+                parallelism or CFG parallelism (reserved placeholders), the compile or
+                activation-checkpointing placeholders, a DCP checkpoint format on a
+                non-sharded run, or — under sharded training — fp16 mixed precision, PEFT,
+                reward-model training, in-training environment evaluation, or multi-optimizer
+                configs.
+        """
+        if self.parallelism.cp_size > 1:
+            raise ValueError(
+                "Context parallelism is not implemented yet: --parallelism.context_parallel.* "
+                "degrees must be 1 (reserved for the CP engine round)."
+            )
+        if self.parallelism.cfg_parallel != 1:
+            raise ValueError(
+                "CFG parallelism is inference-only and must be 1 for training "
+                "(cfg_parallel is reserved for the serving round)."
+            )
+        if self.accelerator.compile.enabled:
+            raise ValueError("--accelerator.compile is a placeholder and not wired yet.")
+        if self.accelerator.activation_checkpointing.mode is not ActivationCheckpointingMode.NONE:
+            raise ValueError("--accelerator.activation_checkpointing is a placeholder and not wired yet.")
+        if self.checkpoint_format is not CheckpointFormat.SAFETENSORS and not self.parallelism.is_sharded:
+            raise ValueError(
+                f"checkpoint_format={self.checkpoint_format.value} requires a sharded run "
+                "(--parallelism.dp_shard != 1); non-sharded checkpoints are always safetensors."
+            )
+        if self.parallelism.is_sharded:
+            if self.accelerator.mixed_precision == "fp16":
+                raise ValueError(
+                    "fp16 is not supported under sharded training (GradScaler over DTensor "
+                    "gradients is unverified); use bf16 or full precision."
+                )
+            if self.peft is not None:
+                raise ValueError("PEFT is not supported under sharded training yet.")
+            if self.is_reward_model_training:
+                raise ValueError(
+                    "Reward-model training is not supported under sharded training yet "
+                    "(reward models declare no FSDP wrap units and have no sharded save path)."
+                )
+            if self.env is not None and self.env_eval_freq > 0:
+                raise ValueError(
+                    "In-training environment evaluation is not supported under sharded training "
+                    "(a rank-0-only rollout of a sharded model deadlocks on collectives); set "
+                    "--env_eval_freq=0 and evaluate with lerobot-eval on saved checkpoints."
+                )
+            if self.optimizer is not None and self.optimizer.builds_multiple_optimizers:
+                raise ValueError("Multi-optimizer configs are not supported under sharded training.")
+
    @classmethod
    def __get_path_fields__(cls) -> list[str]:
        """Keys for draccus pretrained-path loading."""
@@ -613,8 +613,15 @@ def aggregate_feature_stats(stats_ft_list: list[dict[str, dict]]) -> dict[str, d
        for q_key in quantile_keys:
            if all(q_key in s for s in stats_ft_list):
                quantile_values = np.stack([s[q_key] for s in stats_ft_list])
-                weighted_quantiles = quantile_values * counts
-                aggregated[q_key] = weighted_quantiles.sum(axis=0) / total_count
+                # Exact global quantiles cannot be recovered from quantile summaries.
+                # Keep a conservative envelope of the available estimates: min
+                # for lower quantiles and max for upper quantiles. The resulting
+                # values are bounds across the inputs, not global quantile estimates.
+                q_percent = int(q_key[1:])
+                if q_percent <= 50:
+                    aggregated[q_key] = np.min(quantile_values, axis=0)
+                else:
+                    aggregated[q_key] = np.max(quantile_values, axis=0)

    return aggregated

@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Distributed-training runtime for LeRobot.
+
+This package owns everything that turns the declarative topology in
+:class:`lerobot.configs.parallelism.ParallelismConfig` into a running engine:
+mesh math (:class:`~lerobot.distributed.parallel_dims.ParallelDims`), the
+`Accelerator` factory (:func:`~lerobot.distributed.factory.make_accelerator`),
+sharding-aware checkpoint helpers, and small rank utilities.
+
+Setup-order contract (normative):
+CP dispatch install -> activation checkpointing -> torch.compile ->
+``fully_shard``/DDP (via ``accelerator.prepare``) -> optimizer rebind.
+Only the last two steps are active today; CP/AC/compile are configured
+placeholders wired in later rounds.
+"""
+
+from .factory import guard_against_env_interference, make_accelerator, set_fsdp_wrap_modules
+from .parallel_dims import ParallelDims
+from .utils import finalize_sharded_policy, is_main_process, strip_accelerate_cp_hooks
+
+__all__ = [
+    "ParallelDims",
+    "finalize_sharded_policy",
+    "guard_against_env_interference",
+    "is_main_process",
+    "make_accelerator",
+    "set_fsdp_wrap_modules",
+    "strip_accelerate_cp_hooks",
+]
@@ -0,0 +1,195 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Sharding-aware checkpoint primitives.
+
+Two artifact channels with distinct owners:
+
+- the **distributable** ``model.safetensors``: produced by ``PreTrainedPolicy.save_pretrained``
+  through :func:`full_model_state_dict` — a collective full gather when the model is sharded;
+- the **resume** channel (sharded runs): torch DCP directories written/read through accelerate's
+  ``save/load_fsdp_model`` and ``save/load_fsdp_optimizer`` (``pytorch_model_fsdp_0/`` and
+  ``optimizer_0/``, names imported from accelerate constants), which reshard on load across
+  topology changes.
+
+Every function that touches sharded state is a collective and must run on ALL ranks.
+"""
+
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+
+if TYPE_CHECKING:
+    from accelerate import Accelerator
+
+
+def is_sharded_module(module: nn.Module) -> bool:
+    """True when `fully_shard` owns this module's parameters (FSDP2's in-place class swap).
+
+    Args:
+        module (nn.Module): The module to inspect (a torch.compile wrapper is looked through
+            via `_orig_mod`).
+
+    Returns:
+        bool: True when the module (or its compiled `_orig_mod`) is an `FSDPModule`.
+    """
+    from torch.distributed.fsdp import FSDPModule
+
+    if isinstance(module, FSDPModule):
+        return True
+    # torch.compile wraps the sharded module; mirror accelerate's `_orig_mod` check.
+    orig_mod = getattr(module, "_orig_mod", None)
+    return orig_mod is not None and isinstance(orig_mod, FSDPModule)
+
+
+def full_model_state_dict(module: nn.Module) -> dict[str, torch.Tensor]:
+    """The module's full (unsharded) state dict, however its parameters are laid out.
+
+    Sharded modules gather through torch's DCP state-dict API: a COLLECTIVE that must run on
+    every rank; with ``cpu_offload=True`` the full dict materializes on the main rank only and
+    every other rank receives a literal ``{}`` (runtime-verified — a
+    rank-0-gated call deadlocks). Plain modules return ``module.state_dict()`` on every rank.
+
+    Args:
+        module (nn.Module): The (possibly sharded) module to read the state dict from.
+
+    Returns:
+        dict[str, torch.Tensor]: The full state dict — on the main rank only (``{}``
+            elsewhere) when the module is sharded, on every rank otherwise.
+    """
+    if not is_sharded_module(module):
+        return module.state_dict()
+
+    from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
+
+    return get_model_state_dict(module, options=StateDictOptions(full_state_dict=True, cpu_offload=True))
+
+
+def _fsdp_plugin(accelerator: "Accelerator") -> object:
+    """The accelerator's FSDP plugin, required by every DCP save/load helper below.
+
+    Args:
+        accelerator (Accelerator): The accelerator that prepared the sharded model.
+
+    Returns:
+        object: The FSDP plugin held by `accelerator.state`.
+
+    Raises:
+        RuntimeError: If the accelerator was not configured with an FSDP plugin.
+    """
+    plugin = getattr(accelerator.state, "fsdp_plugin", None)
+    if plugin is None:
+        raise RuntimeError("Sharded checkpointing requires an FSDP-prepared Accelerator.")
+    return plugin
+
+
+def save_sharded_model(accelerator: "Accelerator", model: nn.Module, output_dir: Path) -> None:
+    """Write the DCP model shards (`pytorch_model_fsdp_0/`). Collective: call on all ranks.
+
+    Args:
+        accelerator (Accelerator): The accelerator that prepared the sharded model.
+        model (nn.Module): The prepared (sharded) model to save.
+        output_dir (Path): The directory the shard subdirectory is created in.
+    """
+    from accelerate.utils import save_fsdp_model
+
+    # accelerate 1.14's DCP helpers do string containment checks on the path:
+    # always hand them str, never Path.
+    save_fsdp_model(_fsdp_plugin(accelerator), accelerator, model, str(output_dir))
+
+
+def load_sharded_model(accelerator: "Accelerator", model: nn.Module, input_dir: Path) -> None:
+    """Load DCP model shards into the prepared (sharded) model. Collective: call on all ranks.
+
+    Args:
+        accelerator (Accelerator): The accelerator that prepared the sharded model.
+        model (nn.Module): The prepared (sharded) model to load into.
+        input_dir (Path): The directory containing the `pytorch_model_fsdp_0/` shard
+            subdirectory.
+    """
+    from accelerate.utils import load_fsdp_model
+    from accelerate.utils.constants import FSDP_MODEL_NAME
+
+    # Pass the exact shard directory: accelerate's load resolves it with a substring check
+    # ("pytorch_model_fsdp" in the path -> use as-is), which misfires on run paths that happen
+    # to contain the marker; the exact dir makes the check deterministic.
+    load_fsdp_model(_fsdp_plugin(accelerator), accelerator, model, str(input_dir / f"{FSDP_MODEL_NAME}_0"))
+
+
+def save_sharded_optimizer(
+    accelerator: "Accelerator", optimizer: torch.optim.Optimizer, model: nn.Module, output_dir: Path
+) -> None:
+    """Write the DCP optimizer shards (`optimizer_0/`). Collective: call on all ranks.
+
+    Args:
+        accelerator (Accelerator): The accelerator that prepared the model and optimizer.
+        optimizer (torch.optim.Optimizer): The prepared optimizer to save the state from.
+        model (nn.Module): The prepared (sharded) model the optimizer state is keyed by.
+        output_dir (Path): The directory the shard subdirectory is created in.
+    """
+    from accelerate.utils import save_fsdp_optimizer
+
+    save_fsdp_optimizer(_fsdp_plugin(accelerator), accelerator, optimizer, model, str(output_dir))
+
+
+def load_sharded_optimizer(
+    accelerator: "Accelerator", optimizer: torch.optim.Optimizer, model: nn.Module, input_dir: Path
+) -> None:
+    """Load DCP optimizer shards into the prepared optimizer. Collective: call on all ranks.
+
+    Must run AFTER ``accelerator.prepare()``: FSDP2's prepare rebinds the optimizer's param
+    groups to sharded DTensors but never migrates ``optimizer.state`` — the resharding load is
+    the only correct way to restore it.
+
+    Args:
+        accelerator (Accelerator): The accelerator that prepared the model and optimizer.
+        optimizer (torch.optim.Optimizer): The prepared optimizer to restore the state into.
+        model (nn.Module): The prepared (sharded) model the optimizer state is keyed by.
+        input_dir (Path): The directory containing the `optimizer_0/` shard subdirectory.
+    """
+    from accelerate.utils import load_fsdp_optimizer
+    from accelerate.utils.constants import OPTIMIZER_NAME
+
+    # Exact shard directory for the same reason as load_sharded_model: accelerate's substring
+    # check ("optimizer" in the path) would misread e.g. --job_name=optimizer_sweep run paths.
+    load_fsdp_optimizer(
+        _fsdp_plugin(accelerator), accelerator, optimizer, model, str(input_dir / f"{OPTIMIZER_NAME}_0")
+    )
+
+
+def dcp_to_safetensors(dcp_dir: Path, output_dir: Path, *, delete_dcp: bool = False) -> Path:
+    """Merge a DCP shard directory into a single `model.safetensors` (offline, single process).
+
+    Thin wrapper over `accelerate.utils.merge_fsdp_weights`, which loads the shards without a
+    process group, writes safetensors directly, and — when asked — removes the merged shard
+    directory itself, only on the main process and only once the merge has succeeded.
+
+    Args:
+        dcp_dir (Path): The DCP shard directory to merge (e.g. `.../pytorch_model_fsdp_0`).
+        output_dir (Path): The directory the merged `model.safetensors` is written into.
+        delete_dcp (bool): Whether to remove the shard directory once it has been merged.
+            Defaults to False.
+
+    Returns:
+        Path: The written `model.safetensors` file's path.
+    """
+    from accelerate.utils import merge_fsdp_weights
+
+    merge_fsdp_weights(
+        str(dcp_dir), str(output_dir), safe_serialization=True, remove_checkpoint_dir=delete_dcp
+    )
+    return output_dir / "model.safetensors"
@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The `Accelerator` factory — the only place accelerate gets configured.
+
+`torchrun` is the launcher; every accelerate parameter comes from `TrainPipelineConfig`
+(`cfg.parallelism` + `cfg.accelerator`) so a run is reproducible from its `train_config.json`
+alone. `accelerate launch` without a `--config_file` remains equivalent (it only sets rendezvous
+env vars in that mode); the yaml flow is superseded.
+"""
+
+import os
+from typing import TYPE_CHECKING
+
+from lerobot.configs.parallelism import world_size_from_env
+from lerobot.configs.train import TrainPipelineConfig
+
+if TYPE_CHECKING:
+    from accelerate import Accelerator
+
+    from lerobot.policies.pretrained import PreTrainedPolicy
+
+# Env vars through which `accelerate launch --config_file` (or a stray shell) would configure
+# accelerate behind the config system's back, making train_config.json lie about what ran.
+_ACCELERATE_ENV_VARS = (
+    "ACCELERATE_USE_FSDP",
+    "ACCELERATE_USE_PARALLELISM_CONFIG",
+    "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+)
+_ENV_OVERRIDE = "LEROBOT_ALLOW_ACCELERATE_ENV"
+
+
+def guard_against_env_interference() -> None:
+    """Hard-error when accelerate-configuring env vars are set.
+
+    A silently env-overridden "reproducible" config is worse than a stop: users migrating from
+    the old `accelerate launch --config_file fsdp.yaml` flow get a precise error instead of a
+    config that lies. Set LEROBOT_ALLOW_ACCELERATE_ENV=1 to acknowledge and proceed.
+
+    Raises:
+        RuntimeError: If any accelerate-configuring environment variable is set and the
+            LEROBOT_ALLOW_ACCELERATE_ENV override is not.
+    """
+    if os.environ.get(_ENV_OVERRIDE):
+        return
+    offending = sorted(name for name in _ACCELERATE_ENV_VARS if name in os.environ)
+    if offending:
+        raise RuntimeError(
+            f"Accelerate-configuring environment variables are set: {', '.join(offending)}. "
+            "LeRobot manages accelerate exclusively through TrainPipelineConfig "
+            "(--parallelism.* / --accelerator.*); launch with plain torchrun and remove these "
+            "variables (the `accelerate launch --config_file` flow is superseded), or set "
+            f"{_ENV_OVERRIDE}=1 to acknowledge that they may override your config."
+        )
+
+
+def make_accelerator(cfg: TrainPipelineConfig) -> "Accelerator":
+    """Resolve the topology against the launched world and build the `Accelerator`.
+
+    Must run once per process, before any other component needs the device or the process
+    group (`Accelerator.__init__` initializes both and builds the device mesh).
+
+    Args:
+        cfg (TrainPipelineConfig): The full training config; `cfg.parallelism` is resolved in
+            place against the launched world size and `cfg.accelerator` builds the result.
+
+    Returns:
+        Accelerator: The configured accelerator, with device and process group initialized.
+
+    Raises:
+        ValueError: If `cfg.checkpoint_format` requires DCP but the topology resolved to a
+            non-sharded run.
+    """
+    guard_against_env_interference()
+    cfg.parallelism.resolve(world_size_from_env())
+    # The parse-time format check ran against the declared degrees, where the dp_shard=-1
+    # sentinel counts as sharded; it may resolve to an unsharded run (e.g. -1 at world size 1).
+    # Re-check against the concrete degrees so the recorded format never lies about the
+    # artifacts a checkpoint will actually contain.
+    if cfg.checkpoint_format.wants_dcp and not cfg.parallelism.is_sharded:
+        raise ValueError(
+            f"checkpoint_format={cfg.checkpoint_format.value} requires a sharded run, but the "
+            f"topology resolved to a non-sharded one (dp_replicate={cfg.parallelism.dp_replicate}, "
+            f"dp_shard={cfg.parallelism.dp_shard}); non-sharded checkpoints are always safetensors."
+        )
+    return cfg.accelerator.build(
+        cfg.parallelism,
+        cpu=cfg.trainable_config.device == "cpu",
+    )
+
+
+def set_fsdp_wrap_modules(accelerator: "Accelerator", policy: "PreTrainedPolicy") -> None:
+    """Resolve the FSDP wrap-unit class names onto the plugin before `accelerator.prepare()`.
+
+    Resolution order: user override (`--accelerator.fsdp.wrap_modules`, already on the plugin)
+    -> the policy's `_fsdp_wrap_modules` declaration -> hard error. Root-only wrapping — the
+    silent default when no wrap source exists — is never accepted: it quietly forfeits all
+    sharding memory savings.
+
+    No-op for the size-based policy (`--accelerator.fsdp.min_num_params`), which needs no class
+    names, and for non-sharded runs (no fsdp plugin).
+
+    Args:
+        accelerator (Accelerator): The accelerator whose FSDP plugin receives the wrap-unit
+            class names.
+        policy (PreTrainedPolicy): The trainable whose class may declare `_fsdp_wrap_modules`.
+
+    Raises:
+        ValueError: If sharded class-based wrapping is configured but neither a user override
+            nor a policy declaration supplies wrap-unit class names.
+    """
+    plugin = getattr(accelerator.state, "fsdp_plugin", None)
+    if plugin is None or plugin.min_num_params:
+        return
+    if plugin.transformer_cls_names_to_wrap:  # user override, set at build time
+        return
+    # getattr, not attribute access: non-policy trainables (no `_fsdp_wrap_modules` attribute)
+    # must reach the actionable error below, not an AttributeError.
+    declared = getattr(type(policy), "_fsdp_wrap_modules", None)
+    if not declared:
+        raise ValueError(
+            f"Policy '{type(policy).__name__}' declares no FSDP wrap units. Sharded training "
+            "requires wrap-unit class names: set --accelerator.fsdp.wrap_modules='[\"MyBlock\"]' "
+            "(or --accelerator.fsdp.min_num_params for a size-based policy), or declare "
+            "`_fsdp_wrap_modules` on the policy class."
+        )
+    plugin.transformer_cls_names_to_wrap = list(declared)
@@ -0,0 +1,112 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Runtime mesh math derived from the declarative :class:`ParallelismConfig`.
+
+`ParallelDims` is the training script's single source of truth for topology-derived numbers
+(data-parallel world size and rank, sample accounting inputs) and — once the CP engine lands —
+the owner of LeRobot's private ``(dp_replicate, dp_shard, ring, ulysses)`` mesh. It is a runtime
+object and is never serialized (the config it derives from is what lands in
+``train_config.json``).
+"""
+
+from dataclasses import dataclass
+
+import torch.distributed as dist
+
+from lerobot.configs.parallelism import ParallelismConfig
+
+
+@dataclass(frozen=True)
+class ParallelDims:
+    """Concrete parallelism degrees bound to a world size (canonical row-major rank layout)."""
+
+    dp_replicate: int
+    dp_shard: int
+    ring: int
+    ulysses: int
+    world_size: int
+    device_type: str
+
+    @classmethod
+    def from_config(cls, cfg: ParallelismConfig, world_size: int, device_type: str) -> "ParallelDims":
+        """Bind a *resolved* config to the actual runtime world size (cross-checked here).
+
+        Args:
+            cfg (ParallelismConfig): The declarative topology, already resolved via
+                `ParallelismConfig.resolve(world_size)`.
+            world_size (int): The launched world size the declared degrees must multiply to.
+            device_type (str): The accelerator device type backing the mesh (e.g. "cuda").
+
+        Returns:
+            ParallelDims: The concrete parallelism degrees bound to this world.
+
+        Raises:
+            ValueError: If the config is unresolved (`dp_shard == -1`) or its degrees do not
+                multiply to `world_size`.
+        """
+        total = cfg.dp_replicate * cfg.dp_shard * cfg.cp_size
+        if cfg.dp_shard == -1 or total != world_size:
+            raise ValueError(
+                f"ParallelismConfig is not resolved against this world: dp_replicate="
+                f"{cfg.dp_replicate} * dp_shard={cfg.dp_shard} * cp={cfg.cp_size} != "
+                f"world_size={world_size}. Call ParallelismConfig.resolve(world_size) first "
+                "(make_accelerator does this)."
+            )
+        return cls(
+            dp_replicate=cfg.dp_replicate,
+            dp_shard=cfg.dp_shard,
+            ring=cfg.context_parallel.ring_degree,
+            ulysses=cfg.context_parallel.ulysses_degree,
+            world_size=world_size,
+            device_type=device_type,
+        )
+
+    @property
+    def cp_size(self) -> int:
+        """Total context-parallel degree (`ring * ulysses`)."""
+        return self.ring * self.ulysses
+
+    @property
+    def is_sharded(self) -> bool:
+        """Whether parameters are sharded (`dp_shard > 1` or any context parallelism)."""
+        return self.dp_shard > 1 or self.cp_size > 1
+
+    @property
+    def dp_world_size(self) -> int:
+        """Number of distinct data-parallel workers — the divisor for all sample accounting."""
+        return self.dp_replicate * self.dp_shard
+
+    @property
+    def dp_rank(self) -> int:
+        """This process's data-parallel coordinate (CP peers share one dp_rank).
+
+        With the canonical row-major layout and (ring, ulysses) innermost, CP peers are
+        contiguous global ranks, so the dp coordinate is the integer quotient by cp_size —
+        the same arithmetic accelerate's mesh-aware dataloader applies.
+        """
+        global_rank = dist.get_rank() if dist.is_initialized() else 0
+        return global_rank // self.cp_size
+
+    def cp_mesh(self) -> None:
+        """Private (ring, ulysses) mesh for the CP engine — reserved for the CP round.
+
+        Raises:
+            NotImplementedError: Always — context parallelism is not implemented yet.
+        """
+        raise NotImplementedError(
+            "Context parallelism is not implemented yet; ParallelDims.cp_mesh is reserved for "
+            "the CP engine round (a private mesh aligned with accelerate's cp block)."
+        )
@@ -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.
+# 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.
+"""Rank utilities and post-`prepare()` sharding finalization."""
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch.distributed as dist
+from torch import nn
+
+if TYPE_CHECKING:
+    from lerobot.distributed.parallel_dims import ParallelDims
+
+
+def is_main_process() -> bool:
+    """True on the process that owns rank-0-only side effects (file writes, uploads, logging).
+
+    Torch-native on purpose: persistence code must not depend on an `Accelerator` handle —
+    `_save_pretrained` and the hub publishers run in contexts that have none. Outside
+    distributed runs every process is the main process.
+
+    Returns:
+        bool: True when this process is rank 0 or no process group is initialized.
+    """
+    return not dist.is_initialized() or dist.get_rank() == 0
+
+
+def strip_accelerate_cp_hooks(model: nn.Module) -> int:
+    """Remove accelerate's context-parallel forward-pre-hooks from every module.
+
+    When `cp_size > 1` is declared, `accelerator.prepare()` unconditionally attaches hooks that
+    silently replace any `attention_mask` kwarg of `*self_attn` modules with `is_causal=True`
+    (`accelerate.big_modeling._attach_context_parallel_hooks`) — mask corruption for policies
+    with non-causal attention. LeRobot implements CP itself and never enters accelerate's CP
+    context, so these hooks are pure hazard. Deterministically identified by their defining
+    module; a version canary pins that identity.
+
+    Args:
+        model (nn.Module): The prepared model to strip the hooks from (all submodules are
+            visited).
+
+    Returns:
+        int: The number of hooks removed.
+    """
+    removed = 0
+    for module in model.modules():
+        for hook_id, hook in list(module._forward_pre_hooks.items()):
+            if getattr(hook, "__module__", None) == "accelerate.big_modeling":
+                del module._forward_pre_hooks[hook_id]
+                module._forward_pre_hooks_with_kwargs.pop(hook_id, None)
+                removed += 1
+    return removed
+
+
+def finalize_sharded_policy(policy: nn.Module, parallel_dims: "ParallelDims") -> None:
+    """Sharding correctness protocol, applied once, immediately after `accelerator.prepare()`.
+
+    1. Strip accelerate's CP mask hooks (only attached when cp > 1 was declared).
+    2. Register the policy's non-`forward` entry points (`_fsdp_forward_methods`) so FSDP2
+       unshards parameters around `select_action` & co. — without this, any inference-style
+       call on a sharded policy crashes on mixed Tensor/DTensor.
+
+    No-op for DDP/single-process runs.
+
+    Args:
+        policy (nn.Module): The policy as returned by `accelerator.prepare()`.
+        parallel_dims (ParallelDims): The run's resolved topology; decides whether the protocol
+            applies.
+    """
+    if not parallel_dims.is_sharded:
+        return
+    if parallel_dims.cp_size > 1:
+        removed = strip_accelerate_cp_hooks(policy)
+        logging.info("Stripped %d accelerate context-parallel attention-mask hooks.", removed)
+
+    from torch.distributed.fsdp import FSDPModule, register_fsdp_forward_method
+
+    if isinstance(policy, FSDPModule):
+        for method_name in getattr(type(policy), "_fsdp_forward_methods", ()):
+            if callable(getattr(policy, method_name, None)):
+                register_fsdp_forward_method(policy, method_name)
@@ -432,7 +432,7 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None:

    # Finish as soon as the model is pushed, rather than waiting out the platform's
    # post-run finalization before the job stage flips to COMPLETED. This matches the
-    # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay
+    # exact log line emitted by lerobot.common.train_utils.publish_trained_model — the two must stay
    # in sync. If it ever stops matching we just fall back to stage-based completion
    # (~30s slower), so the contract is an optimization, not a correctness requirement.
    success_marker = f"Model pushed to https://huggingface.co/{repo_id}"
@@ -20,7 +20,6 @@ from .optimizers import (
    SGDConfig as SGDConfig,
    XVLAAdamWConfig as XVLAAdamWConfig,
    load_optimizer_state,
-    load_optimizer_state_dict,
    save_optimizer_state,
 )
 from .schedulers import (
@@ -51,7 +50,6 @@ __all__ = [
    "VQBeTSchedulerConfig",
    # State management
    "load_optimizer_state",
-    "load_optimizer_state_dict",
    "load_scheduler_state",
    "save_optimizer_state",
    "save_scheduler_state",
@@ -27,7 +27,7 @@ from lerobot.utils.constants import (
    OPTIMIZER_PARAM_GROUPS,
    OPTIMIZER_STATE,
 )
-from lerobot.utils.io_utils import deserialize_json_into_object, load_json, write_json
+from lerobot.utils.io_utils import deserialize_json_into_object, write_json
 from lerobot.utils.utils import flatten_dict, unflatten_dict

 # Type alias for parameters accepted by optimizer build() methods.
@@ -52,6 +52,11 @@ class OptimizerConfig(draccus.ChoiceRegistry, abc.ABC):
    def type(self) -> str:
        return self.get_choice_name(self.__class__)

+    @property
+    def builds_multiple_optimizers(self) -> bool:
+        """True when build() returns a dict of optimizers (unsupported under sharded training)."""
+        return False
+
    @classmethod
    def default_choice_name(cls) -> str | None:
        return "adam"
@@ -245,6 +250,10 @@ class MultiAdamConfig(OptimizerConfig):
    grad_clip_norm: float = 10.0
    optimizer_groups: dict[str, dict[str, Any]] = field(default_factory=dict)

+    @property
+    def builds_multiple_optimizers(self) -> bool:
+        return True
+
    def build(self, params: OptimizerParams) -> dict[str, torch.optim.Optimizer]:
        """Build multiple Adam optimizers.

@@ -283,35 +292,27 @@ class MultiAdamConfig(OptimizerConfig):
 def save_optimizer_state(
    optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer],
    save_dir: Path,
-    optim_state_dict: dict | None = None,
 ) -> None:
-    """Save optimizer state to disk.
+    """Save optimizer state to disk (non-sharded runs; sharded runs use the DCP channel).

    Args:
        optimizer: Either a single optimizer or a dictionary of optimizers.
        save_dir: Directory to save the optimizer state.
-        optim_state_dict: Pre-gathered optimizer state dict (for FSDP, where the sharded state must
-            be gathered across ranks first). If provided, it is saved directly instead of calling
-            ``optimizer.state_dict()``. Only supported for a single optimizer. Defaults to None.
    """
    if isinstance(optimizer, dict):
        # Handle dictionary of optimizers
-        if optim_state_dict is not None:
-            raise ValueError("optim_state_dict is not supported for a dict of optimizers")
        for name, opt in optimizer.items():
            optimizer_dir = save_dir / name
            optimizer_dir.mkdir(exist_ok=True, parents=True)
            _save_single_optimizer_state(opt, optimizer_dir)
    else:
        # Handle single optimizer
-        _save_single_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict)
+        _save_single_optimizer_state(optimizer, save_dir)


-def _save_single_optimizer_state(
-    optimizer: torch.optim.Optimizer, save_dir: Path, optim_state_dict: dict | None = None
-) -> None:
+def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None:
    """Save a single optimizer's state to disk."""
-    state = dict(optim_state_dict) if optim_state_dict is not None else optimizer.state_dict()
+    state = optimizer.state_dict()
    param_groups = state.pop("param_groups")
    flat_state = flatten_dict(state)
    save_file(flat_state, save_dir / OPTIMIZER_STATE)
@@ -365,19 +366,3 @@ def _load_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Pat

    optimizer.load_state_dict(loaded_state_dict)
    return optimizer
-
-
-def load_optimizer_state_dict(save_dir: Path) -> dict:
-    """Read a saved optimizer state dict (safetensors + json) back into a plain dict.
-
-    Unlike `load_optimizer_state`, this does not load into an optimizer and preserves the original
-    ``state`` keys verbatim (e.g. FSDP parameter FQNs, which are not integer-castable). It is used by
-    the FSDP resume path, where the full state must be resharded via `FSDP.optim_state_dict_to_load`
-    before being loaded into the (sharded) optimizer.
-    """
-    flat_state = load_file(save_dir / OPTIMIZER_STATE)
-    state = unflatten_dict(flat_state)
-    return {
-        "state": state.get("state", {}),
-        "param_groups": load_json(save_dir / OPTIMIZER_PARAM_GROUPS),
-    }
@@ -47,6 +47,8 @@ class ACTPolicy(PreTrainedPolicy):

    config_class = ACTConfig
    name = "act"
+    # FSDP2 wrap units: one unit per transformer layer of both stacks.
+    _fsdp_wrap_modules = ["ACTEncoderLayer", "ACTDecoderLayer"]

    def __init__(
        self,
@@ -242,6 +242,7 @@ def make_policy(
    ds_meta: LeRobotDatasetMetadata | None = None,
    env_cfg: EnvConfig | None = None,
    rename_map: dict[str, str] | None = None,
+    defer_weight_load: bool = False,
 ) -> PreTrainedPolicy:
    """
    Instantiate a policy model.
@@ -252,22 +253,27 @@ def make_policy(
    can either initialize a new policy from scratch or load a pretrained one.

    Args:
-        cfg: The configuration for the policy to be created. If `cfg.pretrained_path` is
-             set, the policy will be loaded with weights from that path.
-        ds_meta: Dataset metadata used to infer feature shapes and types. Also provides
-                 statistics for normalization layers.
-        env_cfg: Environment configuration used to infer feature shapes and types.
-                 One of `ds_meta` or `env_cfg` must be provided.
-        rename_map: Optional mapping of dataset or environment feature keys to match
-                 expected policy feature names (e.g., `"left"` → `"camera1"`).
+        cfg (PreTrainedConfig): The configuration for the policy to be created. If
+            `cfg.pretrained_path` is set, the policy will be loaded with weights from that path.
+        ds_meta (LeRobotDatasetMetadata | None): Dataset metadata used to infer feature shapes and
+            types. Also provides statistics for normalization layers.
+        env_cfg (EnvConfig | None): Environment configuration used to infer feature shapes and
+            types. One of `ds_meta` or `env_cfg` must be provided.
+        rename_map (dict[str, str] | None): Optional mapping of dataset or environment feature
+            keys to match expected policy feature names (e.g., `"left"` → `"camera1"`).
+        defer_weight_load (bool): Build the exact policy `from_pretrained` would build — same
+            config resolution, same stats-derived buffers, same device placement and eval mode —
+            but skip the safetensors weight load. Used when resuming from a DCP checkpoint, whose
+            sharded weights stream in after `accelerator.prepare()` (the distributed checkpoint
+            engine overwrites the random init).

    Returns:
-        An instantiated and device-placed policy model.
+        PreTrainedPolicy: An instantiated and device-placed policy model.

    Raises:
        ValueError: If both or neither of `ds_meta` and `env_cfg` are provided.
-        NotImplementedError: If attempting to use an unsupported policy-backend
-                             combination (e.g., VQBeT with 'mps').
+        NotImplementedError: If attempting to use an unsupported policy-backend combination
+            (e.g., VQBeT with 'mps').
    """
    if bool(ds_meta) == bool(env_cfg):
        raise ValueError("Either one of a dataset metadata or a sim env must be provided.")
@@ -332,11 +338,18 @@ def make_policy(
        )

    if cfg.pretrained_path and not cfg.use_peft:
-        # Load a pretrained policy and override the config if needed (for example, if there are inference-time
-        # hyperparameters that we want to vary).
-        kwargs["pretrained_name_or_path"] = cfg.pretrained_path
-        kwargs["revision"] = cfg.pretrained_revision
-        policy = policy_cls.from_pretrained(**kwargs)
+        if defer_weight_load:
+            # Same construction path as from_pretrained (config already resolved from the
+            # checkpoint by the caller; dataset_stats/dataset_meta kwargs identical), minus the
+            # weight load — parity by construction.
+            policy = policy_cls(**kwargs)
+            policy.eval()
+        else:
+            # Load a pretrained policy and override the config if needed (for example, if there
+            # are inference-time hyperparameters that we want to vary).
+            kwargs["pretrained_name_or_path"] = cfg.pretrained_path
+            kwargs["revision"] = cfg.pretrained_revision
+            policy = policy_cls.from_pretrained(**kwargs)
    elif cfg.pretrained_path and cfg.use_peft:
        # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
        # of the adapter and the adapter's config contains the path to the base policy. So we need the
@@ -54,6 +54,9 @@ class FastWAMPolicy(PreTrainedPolicy):

    config_class = FastWAMConfig
    name = "fastwam"
+    # FSDP2 wrap units: MoTLayer is the single FSDP owner of each layer's expert blocks
+    # (the blocks are re-parented onto it precisely so sharding has one boundary to hook).
+    _fsdp_wrap_modules = ["MoTLayer"]

    def __init__(
        self,
@@ -18,20 +18,18 @@ import builtins
 import dataclasses
 import logging
 import os
-from importlib.resources import files
+import warnings
 from pathlib import Path
-from tempfile import TemporaryDirectory
-from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
+from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, TypeVar, Unpack

-from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict
+from huggingface_hub import hf_hub_download, save_torch_state_dict
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
-from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
+from safetensors.torch import load_model as load_model_as_safetensor
 from torch import Tensor, nn

-from lerobot.__version__ import __version__
 from lerobot.configs import PreTrainedConfig
-from lerobot.configs.train import TrainPipelineConfig
+from lerobot.utils.constants import ACTION
 from lerobot.utils.device_utils import resolve_safetensors_device
 from lerobot.utils.hub import HubMixin
 from lerobot.utils.import_utils import _peft_available, require_package
@@ -46,56 +44,14 @@ else:
    get_peft_model = None

 if TYPE_CHECKING:
+    from lerobot.configs.train import TrainPipelineConfig
    from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata

 T = TypeVar("T", bound="PreTrainedPolicy")

-
-def _build_card_context(
-    cfg: TrainPipelineConfig | None,
-    dataset_meta: LeRobotDatasetMetadata | None,
-    input_features: dict | None,
-    output_features: dict | None,
-) -> dict:
-    """Collect optional data for the model-card template.
-
-    Returns plain values only (no Markdown) — the template in
-    ``lerobot/templates/lerobot_modelcard_template.md`` decides how and whether to show
-    each one. Everything is best-effort: anything unavailable is left empty/None and the
-    template simply skips that section, so this never breaks a Hub push.
-    """
-    context = {
-        "training": None,
-        "input_features": input_features or {},
-        "output_features": output_features or {},
-        "dataset": None,
-        "robot_type": None,
-        "cameras": [],
-    }
-
-    if cfg is not None:
-        optimizer = getattr(cfg, "optimizer", None)
-        context["training"] = {
-            "steps": cfg.steps,
-            "batch_size": cfg.batch_size,
-            "seed": cfg.seed,
-            "optimizer": getattr(optimizer, "type", None) if optimizer else None,
-            "lr": getattr(optimizer, "lr", None) if optimizer else None,
-            "lerobot_version": __version__,
-        }
-
-    if dataset_meta is not None:
-        context["dataset"] = {
-            "repo_id": dataset_meta.repo_id,
-            "episodes": dataset_meta.total_episodes,
-            "frames": dataset_meta.total_frames,
-            "fps": dataset_meta.fps,
-            "tasks": [str(task) for task in dataset_meta.tasks.index],
-        }
-        context["robot_type"] = dataset_meta.robot_type
-        context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys]
-
-    return context
+# Pinned far above any policy's total size so save_torch_state_dict always emits exactly one
+# `model.safetensors` (no shards, no index) — a constant, not a computed byte count.
+_SINGLE_FILE_SHARD_SIZE = "1TB"


 class ActionSelectKwargs(TypedDict, total=False):
@@ -110,6 +66,22 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
    config_class: None
    name: None

+    # --- declarative parallelism/acceleration surface ----------------------------------------
+    # Module CLASS names forming the FSDP2 wrap units (and, once wired, the activation-
+    # checkpointing units). Resolved onto the accelerate plugin right before
+    # `accelerator.prepare()` by `lerobot.distributed.set_fsdp_wrap_modules`; sharded training
+    # with no wrap source anywhere fails loudly instead of silently wrapping only the root.
+    _fsdp_wrap_modules: ClassVar[list[str] | None] = None
+    # Non-`forward` entry points that must trigger FSDP2 unshard/reshard hooks when called on a
+    # sharded policy (registered post-prepare via `torch.distributed.fsdp
+    # .register_fsdp_forward_method`); calling them unregistered crashes on mixed Tensor/DTensor.
+    _fsdp_forward_methods: ClassVar[tuple[str, ...]] = ("select_action", "predict_action_chunk")
+    # Capability gate for the (future) activation-checkpointing wiring.
+    supports_gradient_checkpointing: ClassVar[bool] = False
+    # Declarative context-parallel plan (diffusers `ContextParallelModelPlan` semantics:
+    # module FQN -> sequence split/gather spec). Reserved for the CP engine round.
+    _cp_plan: ClassVar[dict[str, Any] | None] = None
+
    def __init__(self, config: PreTrainedConfig, *inputs, **kwargs):
        super().__init__()
        if not isinstance(config, PreTrainedConfig):
@@ -127,43 +99,33 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        if not getattr(cls, "name", None):
            raise TypeError(f"Class {cls.__name__} must define 'name'")

-    def save_pretrained(
-        self,
-        save_directory: str | Path,
-        *,
-        state_dict: dict[str, Tensor] | None = None,
-        repo_id: str | None = None,
-        push_to_hub: bool = False,
-        card_kwargs: dict | None = None,
-        **push_to_hub_kwargs,
-    ) -> str | None:
-        """Save the policy to a directory (and optionally push to the Hub).
+    def _save_pretrained(self, save_directory: Path) -> None:
+        """Serialize this policy's parameters (and config) into `save_directory`.

-        Overrides `HubMixin.save_pretrained` to add a `state_dict` argument (mirroring
-        `transformers.PreTrainedModel.save_pretrained`). Under FSDP, `self.state_dict()` would
-        return sharded tensors, so the caller gathers the full state dict via a cross-rank
-        collective and passes it here for `_save_pretrained` to write directly.
+        Sharding is handled internally: under FSDP2 the full state dict is gathered through a
+        COLLECTIVE, so when the policy is sharded this method (via `save_pretrained`) must be
+        called on EVERY rank — a rank-0-gated call deadlocks. File writes happen on the main
+        process only, in all layouts (single, DDP, sharded).
+
+        Args:
+            save_directory (Path): Target directory for the policy config (`config.json`) and the
+                safetensors weight file(s).
        """
-        save_directory = Path(save_directory)
-        save_directory.mkdir(parents=True, exist_ok=True)
-        self._save_pretrained(save_directory, state_dict=state_dict)
-        if push_to_hub:
-            if repo_id is None:
-                repo_id = save_directory.name
-            return self.push_to_hub(repo_id=repo_id, card_kwargs=card_kwargs, **push_to_hub_kwargs)
-        return None
+        # Lazy imports: the persistence layer pulls in lerobot.distributed only when saving.
+        from lerobot.distributed.checkpoint import full_model_state_dict, is_sharded_module
+        from lerobot.distributed.utils import is_main_process

-    def _save_pretrained(self, save_directory: Path, state_dict: dict[str, Tensor] | None = None) -> None:
-        self.config._save_pretrained(save_directory)
        model_to_save = self.module if hasattr(self, "module") else self
-        if state_dict is None:
-            save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE))
+        if is_sharded_module(model_to_save):
+            logging.info("Gathering the full state dict from all ranks (sharded policy).")
+        state_dict = full_model_state_dict(model_to_save)  # collective when sharded; {} off-main
+        if not state_dict or not is_main_process():
+            # Sharded: the gather materializes on the main rank only (emptiness check).
+            # Non-sharded multi-rank (DDP): every rank holds a full dict — the explicit rank
+            # gate prevents N ranks racing on the same files. Single process: never taken.
            return
-        # A pre-gathered (e.g. FSDP full) state dict was supplied: write it directly.
-        # `save_torch_state_dict` discards shared-tensor duplicates just like `save_model` does;
-        # pin `max_shard_size` above the total size so the output stays a single `model.safetensors`
-        total_bytes = sum(t.numel() * t.element_size() for t in state_dict.values())
-        save_torch_state_dict(state_dict, str(save_directory), max_shard_size=max(total_bytes, 1))
+        self.config._save_pretrained(save_directory)
+        save_torch_state_dict(state_dict, str(save_directory), max_shard_size=_SINGLE_FILE_SHARD_SIZE)

    @classmethod
    def from_pretrained(
@@ -249,6 +211,29 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        """
        raise NotImplementedError

+    def drop_queued_actions(self) -> None:
+        """Discard actions precomputed by earlier ``select_action`` calls.
+
+        Chunking policies answer most control ticks from a queue filled by an
+        earlier forward pass, so a mid-episode change to the conditioning —
+        e.g. a new language instruction — would otherwise only take effect
+        once that queue drains (up to ``chunk_size`` ticks).  Dropping the
+        queue forces a fresh forward pass on the next ``select_action``.
+
+        Unlike :meth:`reset` this keeps the rest of the episode state (e.g.
+        observation history), so it does not perturb policies that condition
+        on it.  Call it from the thread that calls ``select_action``: it
+        mutates the same queues that thread pops from.
+
+        Policies that keep no action queue inherit a no-op.
+        """
+        queues = getattr(self, "_queues", None)
+        if isinstance(queues, dict) and ACTION in queues:
+            queues[ACTION].clear()
+        action_queue = getattr(self, "_action_queue", None)
+        if action_queue is not None:
+            action_queue.clear()
+
    def supports_rtc(self) -> bool:
        """Whether this policy implements Real-Time Chunking inference semantics."""
        return False
@@ -291,92 +276,39 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
        peft_model=None,
        state_dict: dict[str, Tensor] | None = None,
        dataset_meta: LeRobotDatasetMetadata | None = None,
-    ):
-        api = HfApi()
-        repo_id = api.create_repo(
-            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
-        ).repo_id
+    ) -> None:
+        """Publish this policy to the Hub.

-        # Push the files to the repo in a single commit
-        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
-            saved_path = Path(tmp) / repo_id
+        Deprecated: use :func:`lerobot.common.train_utils.publish_trained_model` instead, which
+        also publishes the pre/post-processors alongside the model.

-            if peft_model is not None:
-                # Since PEFT just forwards calls to `push_model_to_hub`, `self` is not the PeftModel wrapper
-                # but the actual policy which is why we need the PEFT model passed to us to save the adapter.
-                # That also means that we need to store the policy config ourselves since PEFT can't.
-                peft_model.save_pretrained(saved_path)
-                self.config.save_pretrained(saved_path)
-            else:
-                # Calls _save_pretrained and stores model tensors
-                self.save_pretrained(saved_path, state_dict=state_dict)
+        Args:
+            cfg (TrainPipelineConfig): The training config; saved as `train_config.json` and
+                used to render the model card.
+            peft_model: The PEFT wrapper when training adapters, whose weights replace the full
+                model weights in the published repo. Defaults to None.
+            state_dict (dict[str, Tensor] | None): Ignored; weights are now gathered internally
+                when the policy is sharded. Defaults to None.
+            dataset_meta (LeRobotDatasetMetadata | None): Dataset metadata for the model card,
+                if available. Defaults to None.
+        """
+        from lerobot.common.train_utils import publish_trained_model

-            card = self.generate_model_card(
-                cfg.dataset.repo_id,
-                self.config.type,
-                self.config.license,
-                self.config.tags,
-                cfg=cfg,
-                dataset_meta=dataset_meta,
+        warnings.warn(
+            "PreTrainedPolicy.push_model_to_hub is deprecated and will be removed in a future "
+            "version. Use lerobot.common.train_utils.publish_trained_model(cfg, model, "
+            "preprocessor, postprocessor, dataset_meta) instead.",
+            FutureWarning,
+            stacklevel=2,
+        )
+        if state_dict is not None:
+            warnings.warn(
+                "The `state_dict` argument is ignored: sharded weights are gathered internally "
+                "when the policy is saved.",
+                FutureWarning,
+                stacklevel=2,
            )
-            card.save(str(saved_path / "README.md"))
-
-            cfg.save_pretrained(saved_path)  # Calls _save_pretrained and stores train config
-
-            commit_info = api.upload_folder(
-                repo_id=repo_id,
-                repo_type="model",
-                folder_path=saved_path,
-                commit_message="Upload policy weights, train config and readme",
-                allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"],
-                ignore_patterns=["*.tmp", "*.log"],
-            )
-
-            # Contract: lerobot.jobs.hf.submit_to_hf watches for this exact
-            # "Model pushed to <url>" line to end a remote run early. Keep the wording
-            # and URL format in sync (it falls back to status polling if they drift).
-            logging.info(f"Model pushed to {commit_info.repo_url.url}")
-
-    def generate_model_card(
-        self,
-        dataset_repo_id: str,
-        model_type: str,
-        license: str | None,
-        tags: list[str] | None,
-        cfg: TrainPipelineConfig | None = None,
-        dataset_meta: LeRobotDatasetMetadata | None = None,
-    ) -> ModelCard:
-        base_model_mapping = {
-            "smolvla": "lerobot/smolvla_base",
-            "pi0": "lerobot/pi0_base",
-            "pi05": "lerobot/pi05_base",
-            "pi0_fast": "lerobot/pi0fast-base",
-            "xvla": "lerobot/xvla-base",
-        }
-
-        card_data = ModelCardData(
-            license=license or "apache-2.0",
-            library_name="lerobot",
-            pipeline_tag="robotics",
-            tags=list(set(tags or []).union({"robotics", "lerobot", model_type})),
-            model_name=model_type,
-            datasets=dataset_repo_id,
-            base_model=base_model_mapping.get(model_type),
-        )
-
-        context = _build_card_context(
-            cfg, dataset_meta, self.config.input_features, self.config.output_features
-        )
-        # Used by the template to pre-fill commands and the "Fine-tuned from" line.
-        context["policy_repo_id"] = getattr(self.config, "repo_id", None)
-        context["base_model"] = base_model_mapping.get(model_type)
-
-        template_card = (
-            files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8")
-        )
-        card = ModelCard.from_template(card_data, template_str=template_card, **context)
-        card.validate()
-        return card
+        publish_trained_model(cfg, self, None, None, dataset_meta, peft_model=peft_model)

    def wrap_with_peft(
        self,
@@ -647,10 +647,15 @@ def main():
    tags = set(tags).union({"robotics", "lerobot", policy_type})
    tags = list(tags)

-    # Generate model card
-    card = policy.generate_model_card(
-        dataset_repo_id=dataset_repo_id, model_type=policy_type, license=license, tags=tags
-    )
+    # Generate model card through the free helper (PreTrainedPolicy.generate_model_card was
+    # removed with the publisher redesign), then apply the metadata recovered above — the
+    # migrated policy config does not carry the original repo's card fields.
+    from lerobot.common.train_utils import generate_model_card
+
+    card = generate_model_card(policy.config)
+    card.data.datasets = dataset_repo_id
+    card.data.license = license
+    card.data.tags = sorted(tags)

    # Save model card locally
    card.save(str(output_dir / "README.md"))
@@ -16,12 +16,11 @@ import abc
 import builtins
 import logging
 import os
-from importlib.resources import files
+import warnings
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, TypeVar

-from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
+from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
@@ -61,6 +60,22 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
            raise TypeError(f"Class {cls.__name__} must define 'name'")

    def _save_pretrained(self, save_directory: Path) -> None:
+        """Serialize this reward model's parameters (and config) into `save_directory`.
+
+        Safe to call on every rank: replicas carry identical weights, so only the main process
+        writes (sharded reward models are rejected at config validation — no collective gather).
+
+        Args:
+            save_directory (Path): Target directory for the reward model config (`config.json`)
+                and `model.safetensors`.
+        """
+        from lerobot.distributed.utils import is_main_process
+
+        # save_checkpoint calls this on every rank; replicas carry identical
+        # weights, so the main process is the only writer. Sharded reward models are rejected
+        # at config validation, so no collective gather is needed here.
+        if not is_main_process():
+            return
        self.config._save_pretrained(save_directory)
        model_to_save = self.module if hasattr(self, "module") else self
        save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE))
@@ -175,53 +190,22 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
        """
        return type(self).forward is not PreTrainedRewardModel.forward

-    def push_model_to_hub(self, cfg: "TrainPipelineConfig"):
-        api = HfApi()
-        repo_id = api.create_repo(
-            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
-        ).repo_id
+    def push_model_to_hub(self, cfg: "TrainPipelineConfig") -> None:
+        """Publish this reward model to the Hub.

-        # Push the files to the repo in a single commit
-        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
-            saved_path = Path(tmp) / repo_id
+        Deprecated: use :func:`lerobot.common.train_utils.publish_trained_model` instead.

-            self.save_pretrained(saved_path)  # Calls _save_pretrained and stores model tensors
+        Args:
+            cfg (TrainPipelineConfig): The training config; saved as `train_config.json` and
+                used to render the model card.
+        """
+        from lerobot.common.train_utils import publish_trained_model

-            card = self.generate_model_card(
-                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
-            )
-            card.save(str(saved_path / "README.md"))
-
-            cfg.save_pretrained(saved_path)  # Calls _save_pretrained and stores train config
-
-            commit_info = api.upload_folder(
-                repo_id=repo_id,
-                repo_type="model",
-                folder_path=saved_path,
-                commit_message="Upload reward model weights, train config and readme",
-                allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"],
-                ignore_patterns=["*.tmp", "*.log"],
-            )
-
-            logging.info(f"Model pushed to {commit_info.repo_url.url}")
-
-    def generate_model_card(
-        self, dataset_repo_id: str, model_type: str, license: str | None, tags: list[str] | None
-    ) -> ModelCard:
-        card_data = ModelCardData(
-            license=license or "apache-2.0",
-            library_name="lerobot",
-            pipeline_tag="robotics",
-            tags=list(set(tags or []).union({"robotics", "lerobot", "reward-model", model_type})),
-            model_name=model_type,
-            datasets=dataset_repo_id,
+        warnings.warn(
+            "PreTrainedRewardModel.push_model_to_hub is deprecated and will be removed in a "
+            "future version. Use lerobot.common.train_utils.publish_trained_model(cfg, model, "
+            "preprocessor, postprocessor, dataset_meta) instead.",
+            FutureWarning,
+            stacklevel=2,
        )
-
-        template_card = (
-            files("lerobot.templates")
-            .joinpath("lerobot_rewardmodel_modelcard_template.md")
-            .read_text(encoding="utf-8")
-        )
-        card = ModelCard.from_template(card_data, template_str=template_card)
-        card.validate()
-        return card
+        publish_trained_model(cfg, self, None, None, None)
@@ -58,12 +58,11 @@ import builtins
 import logging
 import os
 from pathlib import Path
-from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, TypeVar

 import numpy as np
 import torch
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import hf_hub_download
 from huggingface_hub.constants import CONFIG_NAME
 from huggingface_hub.errors import HfHubHTTPError
 from torch import Tensor
@@ -75,9 +74,6 @@ from lerobot.rewards.topreward.configuration_topreward import TOPRewardConfig
 from lerobot.rewards.topreward.processor_topreward import TOPREWARD_FEATURE_PREFIX, TOPREWARD_INPUT_KEYS
 from lerobot.utils.import_utils import _transformers_available, require_package

-if TYPE_CHECKING:
-    from lerobot.configs.train import TrainPipelineConfig
-
 if TYPE_CHECKING or _transformers_available:
    from transformers import Qwen3VLForConditionalGeneration
 else:
@@ -205,34 +201,3 @@ class TOPRewardModel(PreTrainedRewardModel):
        instance.to(config.device)
        instance.eval()
        return instance
-
-    def push_model_to_hub(self, cfg: TrainPipelineConfig):
-        """Push the TOPReward ``config.json`` + model card to the Hub."""
-        api = HfApi()
-        repo_id = api.create_repo(
-            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
-        ).repo_id
-
-        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
-            saved_path = Path(tmp) / repo_id
-            saved_path.mkdir(parents=True, exist_ok=True)
-
-            self.config._save_pretrained(saved_path)
-
-            card = self.generate_model_card(
-                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
-            )
-            card.save(str(saved_path / "README.md"))
-
-            cfg.save_pretrained(saved_path)
-
-            commit_info = api.upload_folder(
-                repo_id=repo_id,
-                repo_type="model",
-                folder_path=saved_path,
-                commit_message="Upload TOPReward config and readme",
-                allow_patterns=["*.json", "*.yaml", "*.md"],
-                ignore_patterns=["*.tmp", "*.log", "*.safetensors"],
-            )
-
-            logger.info(f"Model pushed to {commit_info.repo_url.url}")
@@ -74,13 +74,14 @@ from torch.optim.optimizer import Optimizer
 from lerobot.cameras import opencv  # noqa: F401
 from lerobot.common.train_utils import (
    get_step_checkpoint_dir,
-    load_training_state as utils_load_training_state,
+    load_training_metadata,
    save_checkpoint,
    update_last_checkpoint,
 )
 from lerobot.common.wandb_utils import WandBLogger
 from lerobot.configs import parser
 from lerobot.datasets import LeRobotDataset, make_dataset
+from lerobot.optim import load_optimizer_state
 from lerobot.policies import make_policy, make_pre_post_processors
 from lerobot.robots import so_follower  # noqa: F401
 from lerobot.teleoperators import gamepad, so_leader  # noqa: F401
@@ -103,7 +104,7 @@ from lerobot.utils.constants import (
 from lerobot.utils.device_utils import get_safe_torch_device
 from lerobot.utils.io_utils import load_json, write_json
 from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method
-from lerobot.utils.random_utils import set_seed
+from lerobot.utils.random_utils import load_rng_state, set_seed
 from lerobot.utils.utils import (
    format_big_number,
    init_logging,
@@ -716,15 +717,18 @@ def load_training_state(
    algorithm-owned tensors) from the most recent checkpoint.

    Args:
-        cfg: Training configuration.
-        optimizers: Optimizers to load state into.
-        algorithm: Algorithm whose state dict should be restored.
-            Required for full main-equivalent resume;
-            the policy itself is restored separately via ``make_policy``.
-        device: Device on which to place loaded algorithm tensors.
+        cfg (TrainRLServerPipelineConfig): Training configuration; `cfg.resume` gates the load and
+            `cfg.output_dir` locates the last checkpoint.
+        optimizers (Optimizer | dict[str, Optimizer]): Optimizers to load state into.
+        algorithm (RLAlgorithm | None, optional): Algorithm whose state dict should be restored.
+            Required for full main-equivalent resume; the policy itself is restored separately via
+            `make_policy`. Defaults to None.
+        device (str | torch.device, optional): Device on which to place loaded algorithm tensors.
+            Defaults to "cpu".

    Returns:
-        tuple: (optimization_step, interaction_step) or (None, None) if not resuming
+        tuple[int | None, int | None]: `(optimization_step, interaction_step)`, or `(None, None)`
+        when not resuming or when loading the training state fails.
    """
    if not cfg.resume:
        return None, None
@@ -736,7 +740,10 @@ def load_training_state(

    try:
        # Restore optimizers + RNG + step from the standard `training_state/` folder
-        step, optimizers, _ = utils_load_training_state(checkpoint_dir, optimizers, None)
+        training_state_dir = checkpoint_dir / TRAINING_STATE_DIR
+        load_rng_state(training_state_dir)
+        step = load_training_metadata(training_state_dir)["step"]
+        optimizers = load_optimizer_state(optimizers, training_state_dir)

        # Restore algorithm-owned tensors
        if algorithm is not None:
@@ -38,6 +38,11 @@ from .context import (
    RuntimeContext,
    build_rollout_context,
 )
+from .controller import (
+    LinkedEvent,
+    RolloutController,
+    RolloutEvent,
+)
 from .inference import (
    InferenceEngine,
    InferenceEngineConfig,
@@ -47,6 +52,11 @@ from .inference import (
    SyncInferenceEngine,
    create_inference_engine,
 )
+from .interactive import (
+    InteractiveCommand,
+    InteractiveSession,
+    parse_command,
+)
 from .strategies import (
    BaseStrategy,
    DAggerStrategy,
@@ -65,19 +75,24 @@ __all__ = [
    "DAggerStrategy",
    "DAggerStrategyConfig",
    "DatasetContext",
+    "EpisodicStrategy",
+    "EpisodicStrategyConfig",
    "HardwareContext",
    "HighlightStrategy",
    "HighlightStrategyConfig",
-    "EpisodicStrategy",
-    "EpisodicStrategyConfig",
    "InferenceEngine",
    "InferenceEngineConfig",
+    "InteractiveCommand",
+    "InteractiveSession",
+    "LinkedEvent",
    "PolicyContext",
    "ProcessorContext",
    "RTCInferenceConfig",
    "RTCInferenceEngine",
    "RolloutConfig",
    "RolloutContext",
+    "RolloutController",
+    "RolloutEvent",
    "RolloutStrategy",
    "RolloutStrategyConfig",
    "RuntimeContext",
@@ -88,4 +103,5 @@ __all__ = [
    "build_rollout_context",
    "create_inference_engine",
    "create_strategy",
+    "parse_command",
 ]
@@ -239,6 +239,14 @@ class RolloutConfig:
    # Runtime
    fps: float = 30.0
    duration: float = 0.0  # 0 = infinite (24/7 mode)
+    # Interactive session: control the rollout from stdin with chat-style
+    # commands (/start, /subtask <text>, /reset, /stop) while hardware and
+    # policy stay warm.  The robot does not move until /start is received,
+    # `/subtask` re-instructs the policy mid-run, and logs below ERROR are
+    # muted while the session runs so they don't interleave with the prompt.
+    # Supported with --strategy.type=base (no recording) and sentry
+    # (continuous recording; frames are labeled with the live task).
+    interactive: bool = False
    interpolation_multiplier: int = 1
    device: str | None = None
    task: str = ""
@@ -294,6 +302,17 @@ class RolloutConfig:
                "Base strategy does not record data. Use sentry, highlight, or dagger for recording."
            )

+        # Interactive mode drives strategy.run() in restartable segments and reads
+        # commands from stdin.  Base and sentry qualify: their run() loops keep no
+        # per-run terminal or dataset-finalization state.  The other recording
+        # strategies are excluded for now: they bind their own keyboard controls
+        # (which fight the command prompt for the terminal) and their run() loops
+        # finalize the dataset on exit, so they cannot be restarted.
+        if self.interactive and not isinstance(self.strategy, (BaseStrategyConfig, SentryStrategyConfig)):
+            raise ValueError(
+                f"--interactive=true supports --strategy.type=base or sentry (got '{self.strategy.type}')."
+            )
+
        # Sentry MUST use streaming encoding to avoid disk I/O blocking the control loop
        if (
            isinstance(self.strategy, SentryStrategyConfig)
@@ -0,0 +1,382 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Programmatic control of a rollout: start, pause, re-instruct, and stop a
+policy while hardware and policy stay connected and warm.
+
+:class:`RolloutController` is the embedding-friendly core of interactive
+rollouts.  It has no I/O of its own — no stdin, no printing, no log
+manipulation — so it can be driven from any application code: a CLI
+(:class:`lerobot.rollout.interactive.InteractiveSession` is exactly that), a
+network server, a voice front-end, or a notebook.
+
+Typical embedding::
+
+    from threading import Event, Thread
+    from lerobot.rollout import (
+        LinkedEvent,
+        RolloutController,
+        build_rollout_context,
+        create_strategy,
+    )
+
+    parent = Event()  # your application's shutdown signal
+    ctx = build_rollout_context(cfg, LinkedEvent(parent))
+    strategy = create_strategy(cfg.strategy)
+    strategy.setup(ctx)
+
+    controller = RolloutController(strategy, ctx)
+    serve_thread = Thread(target=controller.serve)
+    serve_thread.start()  # or call serve() on your main thread
+
+    controller.start()  # robot starts executing the policy
+    controller.set_task("grab the red cube")  # re-instruct mid-run
+    controller.reset()  # stop movement, return home, stay warm
+    controller.stop()  # end serve()
+
+    serve_thread.join()
+    strategy.teardown(ctx)  # teardown stays with the caller
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable
+from enum import Enum
+from threading import Event, Lock
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .context import RolloutContext
+    from .strategies import RolloutStrategy
+
+logger = logging.getLogger(__name__)
+
+
+class LinkedEvent(Event):
+    """A ``threading.Event`` whose ``is_set`` also reflects a parent event.
+
+    ``set``/``clear`` act only on the local flag, so a controller can raise
+    and clear its own segment-stop requests without masking (or accidentally
+    re-arming) the process-wide shutdown event carried by ``parent``.  Every
+    rollout strategy control loop polls ``ctx.runtime.shutdown_event.is_set()``,
+    so installing a ``LinkedEvent`` there makes the loops react both to
+    controller commands and to real shutdown signals.
+    """
+
+    _WAIT_SLICE_S = 0.05
+
+    def __init__(self, parent: Event) -> None:
+        super().__init__()
+        self.parent = parent
+
+    def is_set(self) -> bool:
+        return super().is_set() or self.parent.is_set()
+
+    def wait(self, timeout: float | None = None) -> bool:
+        """Wait for either the local or the parent flag.
+
+        The base ``Event.wait`` only watches the local flag, so poll in short
+        slices to also observe the parent.  Strategy loops only call
+        ``is_set()``; this coarse wait exists for API completeness.
+        """
+        deadline = None if timeout is None else time.perf_counter() + timeout
+        while not self.is_set():
+            remaining = None if deadline is None else deadline - time.perf_counter()
+            if remaining is not None and remaining <= 0:
+                return False
+            wait_slice = self._WAIT_SLICE_S if remaining is None else min(self._WAIT_SLICE_S, remaining)
+            super().wait(wait_slice)
+        return True
+
+
+class RolloutEvent(Enum):
+    """Lifecycle notifications emitted by :class:`RolloutController`.
+
+    All events are emitted on the thread running :meth:`RolloutController.serve`;
+    callbacks must be quick and must not call back into the controller's
+    blocking methods.
+    """
+
+    SEGMENT_STARTED = "segment_started"
+    """A control-loop segment is about to run (control state freshly reset)."""
+
+    SEGMENT_ENDED = "segment_ended"
+    """The segment returned on its own (e.g. ``--duration`` elapsed); the
+    controller is idle again and the robot is holding position."""
+
+    RESET_STARTED = "reset_started"
+    """A reset is being executed: inference paused, robot about to move home."""
+
+    RESET_DONE = "reset_done"
+    """The robot is back at its initial position, holding."""
+
+    RESET_SKIPPED = "reset_skipped"
+    """No initial position was captured; the robot holds its current pose."""
+
+    ENGINE_FAILED = "engine_failed"
+    """The inference engine hit an unrecoverable error; ``serve()`` is about
+    to return.  Read :attr:`RolloutController.failure_traceback` for details."""
+
+    STOPPED = "stopped"
+    """``serve()`` is returning (after :meth:`RolloutController.stop`, EOF of
+    the driving front-end, an engine failure, or a parent shutdown signal)."""
+
+
+class RolloutController:
+    """Drive a rollout strategy through thread-safe start/reset/stop/set_task calls.
+
+    The controller owns the outer lifecycle between ``strategy.setup(ctx)``
+    and ``strategy.teardown(ctx)`` (both stay with the caller): the robot is
+    idle until :meth:`start`, each run *segment* executes ``strategy.run(ctx)``
+    on the thread that called :meth:`serve` until interrupted or until the
+    strategy returns on its own (e.g. ``--duration`` elapsed).  :meth:`reset`
+    pauses the inference engine, returns the robot to its initial position,
+    and restores the launch task, while hardware and policy stay warm.
+    :meth:`stop` ends :meth:`serve` so the caller can run
+    ``strategy.teardown(ctx)``.
+
+    Requires ``ctx.runtime.shutdown_event`` to be a :class:`LinkedEvent`: the
+    controller sets the local flag to end a segment, and process shutdown
+    signals still propagate through the parent.  Build the context with
+    ``build_rollout_context(cfg, LinkedEvent(shutdown_event))``.
+
+    Thread safety: the control methods (:meth:`start`, :meth:`reset`,
+    :meth:`stop`, :meth:`set_task`) may be called from any thread and are
+    serialized by an internal lock, so calls issued in order from one thread
+    keep that order — e.g. a ``set_task`` right after a ``reset`` is not
+    clobbered by the reset's task restore.  Commands are last-write-wins:
+    ``reset`` and ``stop`` cancel a still-pending ``start`` so the robot
+    never starts moving after the caller's most recent command asked it not
+    to.  Events are emitted on the :meth:`serve` thread via ``on_event``.
+    """
+
+    _POLL_INTERVAL_S = 0.2
+
+    def __init__(
+        self,
+        strategy: RolloutStrategy,
+        ctx: RolloutContext,
+        on_event: Callable[[RolloutEvent], None] | None = None,
+    ) -> None:
+        stop_event = ctx.runtime.shutdown_event
+        if not isinstance(stop_event, LinkedEvent):
+            raise TypeError(
+                "RolloutController requires ctx.runtime.shutdown_event to be a LinkedEvent so "
+                "reset() can end a run segment without triggering process shutdown. Build the "
+                "rollout context with build_rollout_context(cfg, LinkedEvent(shutdown_event))."
+            )
+        self._strategy = strategy
+        self._ctx = ctx
+        self._segment_stop = stop_event
+        self._global_shutdown = stop_event.parent
+        self._on_event = on_event
+        # The instruction the rollout was launched with; reset() restores it.
+        self._initial_task = ctx.policy.inference.task
+
+        # Serializes the control methods so multi-writer task updates (e.g.
+        # reset()'s restore followed by a set_task()) keep their call order.
+        self._control_lock = Lock()
+
+        # Written by control methods (any thread), consumed by the serve loop.
+        self._start_requested = Event()
+        self._reset_requested = Event()
+        self._stop_requested = Event()
+        self._wake = Event()
+        self._running = Event()
+
+    # ------------------------------------------------------------------
+    # Introspection
+    # ------------------------------------------------------------------
+
+    @property
+    def task(self) -> str:
+        """The language instruction currently conditioning inference."""
+        return self._ctx.policy.inference.task
+
+    @property
+    def initial_task(self) -> str:
+        """The instruction the rollout was launched with (restored by :meth:`reset`)."""
+        return self._initial_task
+
+    @property
+    def running(self) -> bool:
+        """True while a control-loop segment is executing."""
+        return self._running.is_set()
+
+    @property
+    def failed(self) -> bool:
+        """True if the inference engine hit an unrecoverable error."""
+        return self._ctx.policy.inference.failed
+
+    @property
+    def failure_traceback(self) -> str | None:
+        """Formatted traceback of the engine failure, when :attr:`failed` is True."""
+        return self._ctx.policy.inference.failure_traceback
+
+    # ------------------------------------------------------------------
+    # Control methods (callable from any thread)
+    # ------------------------------------------------------------------
+
+    def start(self) -> bool:
+        """Request a control-loop segment.
+
+        Returns ``False`` when a segment is already running (the request is
+        ignored); ``True`` when the segment was scheduled.  The segment itself
+        executes on the :meth:`serve` thread.
+        """
+        with self._control_lock:
+            if self._running.is_set():
+                return False
+            self._start_requested.set()
+            self._wake.set()
+            return True
+
+    def reset(self) -> bool:
+        """Stop movement, return the robot to its initial position, restore the launch task.
+
+        Hardware and policy stay warm; call :meth:`start` to run again.
+        Returns ``True`` when the task was restored to the launch task (i.e.
+        it had been changed), ``False`` when it was already the launch task.
+        """
+        with self._control_lock:
+            # Last command wins: a start() still waiting to be serviced is
+            # cancelled so the robot never starts moving after the caller
+            # asked it not to.  Flag first, segment-stop second (see the
+            # ordering note in _run_segment).
+            self._start_requested.clear()
+            # Restore the task here, under the control lock, rather than in
+            # _reset_robot (which runs later, on the serve thread) so that a
+            # set_task() issued right after this reset() is not silently
+            # reverted by a deferred restore.
+            restored = self._ctx.policy.inference.set_task(self._initial_task)
+            self._reset_requested.set()
+            self._segment_stop.set()
+            self._wake.set()
+            return restored
+
+    def stop(self) -> None:
+        """End :meth:`serve`; the caller then runs ``strategy.teardown(ctx)``."""
+        with self._control_lock:
+            self._start_requested.clear()  # last command wins, see reset()
+            self._stop_requested.set()
+            self._segment_stop.set()
+            self._wake.set()
+
+    def set_task(self, task: str) -> bool:
+        """Change the instruction the policy follows, effective from the next inference.
+
+        Returns ``True`` when the value actually changed.  Safe to call while
+        a segment is running: the engine applies the switch on its own
+        inference thread (sync backends also drop actions precomputed under
+        the previous instruction).
+        """
+        with self._control_lock:
+            return self._ctx.policy.inference.set_task(task)
+
+    # ------------------------------------------------------------------
+    # Serve loop (blocks the calling thread)
+    # ------------------------------------------------------------------
+
+    def serve(self) -> None:
+        """Service control requests until :meth:`stop`, engine failure, or parent shutdown.
+
+        Blocks the calling thread; run segments execute here.  Emits
+        :class:`RolloutEvent` notifications through ``on_event``.
+        """
+        try:
+            while not self._global_shutdown.is_set():
+                if self._ctx.policy.inference.failed:
+                    self._emit(RolloutEvent.ENGINE_FAILED)
+                    break
+                if self._stop_requested.is_set():
+                    break
+                if self._reset_requested.is_set():
+                    self._reset_requested.clear()
+                    self._reset_robot()
+                    continue
+                if self._start_requested.is_set():
+                    # Consume the request and mark the segment running in one
+                    # atomic step: start() gates on _running, so a concurrent
+                    # start() is rejected for the entire startup sequence
+                    # (reset_control_state, SEGMENT_STARTED emission), not just
+                    # once strategy.run() begins — otherwise it could re-arm
+                    # _start_requested behind the running segment and the robot
+                    # would start again, uncommanded, when the segment ends.
+                    with self._control_lock:
+                        starting = self._start_requested.is_set()
+                        if starting:
+                            self._start_requested.clear()
+                            self._running.set()
+                    if starting:
+                        self._run_segment()
+                    continue
+                self._wake.wait(timeout=self._POLL_INTERVAL_S)
+                self._wake.clear()
+        finally:
+            self._emit(RolloutEvent.STOPPED)
+
+    def _run_segment(self) -> None:
+        """Execute one ``strategy.run`` segment until interrupted or finished.
+
+        The serve loop has already set ``_running`` (under the control lock),
+        so this method must clear it on every exit path.
+        """
+        engine = self._ctx.policy.inference
+        try:
+            # Clear the local flag *before* checking the request flags: control
+            # methods set their flag first and the segment-stop event second, so
+            # a reset() or stop() racing with this start() is either seen here or
+            # ends the freshly started loop on its first tick.
+            self._segment_stop.clear()
+            if (
+                self._stop_requested.is_set()
+                or self._reset_requested.is_set()
+                or self._global_shutdown.is_set()
+            ):
+                return
+            self._strategy.reset_control_state()
+            self._emit(RolloutEvent.SEGMENT_STARTED)
+            try:
+                self._strategy.run(self._ctx)
+            finally:
+                engine.pause()
+        finally:
+            self._running.clear()
+        if engine.failed:
+            return  # the serve loop emits ENGINE_FAILED and shuts down
+        if not (
+            self._stop_requested.is_set() or self._reset_requested.is_set() or self._global_shutdown.is_set()
+        ):
+            self._emit(RolloutEvent.SEGMENT_ENDED)
+
+    def _reset_robot(self) -> None:
+        """Pause inference and return the robot home (the task was restored by :meth:`reset`)."""
+        self._emit(RolloutEvent.RESET_STARTED)
+        self._ctx.policy.inference.pause()
+        if self._ctx.hardware.initial_position:
+            self._strategy.return_to_initial_position(self._ctx.hardware)
+            self._emit(RolloutEvent.RESET_DONE)
+        else:
+            logger.warning("No initial position captured — skipping the return move")
+            self._emit(RolloutEvent.RESET_SKIPPED)
+
+    def _emit(self, event: RolloutEvent) -> None:
+        if self._on_event is None:
+            return
+        try:
+            self._on_event(event)
+        except Exception:  # a broken observer must not kill the serve loop
+            logger.exception("Error in RolloutController event callback for %s", event)
@@ -22,9 +22,13 @@ or asynchronously in a background thread (RTC).
 from __future__ import annotations

 import abc
+import logging
+from threading import Lock

 import torch

+logger = logging.getLogger(__name__)
+

 class InferenceEngine(abc.ABC):
    """Abstract backend for producing actions during rollout.
@@ -47,12 +51,69 @@ class InferenceEngine(abc.ABC):
    backends always compute from ``obs_frame``; async backends ignore
    it (they receive observations via ``notify_observation``).

+    Task
+    ----
+    ``task`` / ``set_task`` hold the language instruction the policy is
+    conditioned on.  ``set_task`` is safe to call from any thread (the
+    interactive session's ``/subtask`` command calls it from its stdin
+    reader); subclasses pick the new value up on their own inference
+    thread via :meth:`_take_task`, so no policy state is ever mutated
+    across threads.
+
    Optional hooks
    --------------
    ``notify_observation`` / ``pause`` / ``resume`` have a no-op default
    so rollout strategies can invoke them unconditionally.
+
+    Subclasses must call ``super().__init__(task=...)``; the task holder
+    is set up there.
    """

+    def __init__(self, task: str = "") -> None:
+        self._task = task
+        self._task_changed = False
+        self._task_lock = Lock()
+
+    # ------------------------------------------------------------------
+    # Task (language instruction)
+    # ------------------------------------------------------------------
+
+    @property
+    def task(self) -> str:
+        """The language instruction currently conditioning inference."""
+        with self._task_lock:
+            return self._task
+
+    def set_task(self, task: str) -> bool:
+        """Set the instruction used from the next inference onwards.
+
+        Callable from any thread.  Returns ``True`` when the value
+        actually changed, so callers can report no-op switches.
+        """
+        with self._task_lock:
+            if task == self._task:
+                return False
+            previous, self._task = self._task, task
+            self._task_changed = True
+        logger.info("Task changed: '%s' -> '%s'", previous, task)
+        return True
+
+    def _take_task(self) -> tuple[str, bool]:
+        """Read the task and whether it changed since the last read.
+
+        Call from the thread that runs inference: the "changed" edge is
+        consumed here so the backend can drop actions precomputed under
+        the previous instruction before using the new one.
+        """
+        with self._task_lock:
+            changed, self._task_changed = self._task_changed, False
+            return self._task, changed
+
+    def _discard_task_change(self) -> None:
+        """Drop a pending task-change edge, e.g. from ``reset`` (state is already cleared)."""
+        with self._task_lock:
+            self._task_changed = False
+
    @abc.abstractmethod
    def start(self) -> None:
        """Initialise the backend."""
@@ -87,3 +148,8 @@ class InferenceEngine(abc.ABC):
    def failed(self) -> bool:
        """True if an unrecoverable error occurred in the backend."""
        return False
+
+    @property
+    def failure_traceback(self) -> str | None:
+        """Formatted traceback of the unrecoverable error, when ``failed`` is True."""
+        return None
@@ -124,13 +124,13 @@ class RTCInferenceEngine(InferenceEngine):
        rtc_queue_threshold: int = 30,
        shutdown_event: Event | None = None,
    ) -> None:
+        super().__init__(task=task)
        self._policy = policy
        self._preprocessor = preprocessor
        self._postprocessor = postprocessor
        self._robot = robot_wrapper
        self._rtc_config = rtc_config
        self._hw_features = hw_features
-        self._task = task
        self._fps = fps
        self._device = device or "cpu"
        self._use_torch_compile = use_torch_compile
@@ -140,10 +140,14 @@ class RTCInferenceEngine(InferenceEngine):
        self._action_queue: ActionQueue | None = None
        self._obs_holder: dict[str, Any] = {}
        self._obs_lock = Lock()
+        # Bumped by reset() (under _obs_lock) so chunks whose inference started
+        # before a reset are discarded instead of merged into the fresh queue.
+        self._reset_epoch = 0
        self._policy_active = Event()
        self._compile_warmup_done = Event()
        self._shutdown_event = Event()
        self._rtc_error = Event()
+        self._failure_traceback: str | None = None
        self._global_shutdown_event = shutdown_event
        self._rtc_thread: Thread | None = None

@@ -190,6 +194,15 @@ class RTCInferenceEngine(InferenceEngine):
        """True if the RTC background thread exited due to an unrecoverable error."""
        return self._rtc_error.is_set()

+    @property
+    def failure_traceback(self) -> str | None:
+        """Traceback captured when the RTC thread died (see ``failed``).
+
+        Kept on the engine so consumers that mute console logging (the
+        interactive session) can still surface the fatal error.
+        """
+        return self._failure_traceback
+
    @property
    def action_queue(self) -> ActionQueue | None:
        """The shared action queue between the RTC thread and the main loop."""
@@ -235,13 +248,29 @@ class RTCInferenceEngine(InferenceEngine):
        self._policy_active.set()

    def reset(self) -> None:
-        """Reset the policy, processors, and action queue."""
+        """Reset the policy, processors, and action queue.
+
+        Call while the engine is paused (both DAgger transitions and the
+        interactive session do): the RTC thread may still be finishing an
+        inference started before the pause, so ``reset`` also drops the last
+        published observation — it can be arbitrarily stale by the time the
+        engine resumes (e.g. the robot was returned to its initial position
+        in the meantime), and a chunk computed from it would jerk the robot
+        toward the old pose — and bumps the reset epoch so any in-flight
+        chunk is discarded instead of merged into the cleared queue.
+        """
        logger.info("Resetting RTC inference state (policy + processors + queue)")
        self._policy.reset()
        self._preprocessor.reset()
        self._postprocessor.reset()
        if self._action_queue is not None:
            self._action_queue.clear()
+        with self._obs_lock:
+            self._obs_holder["obs"] = None
+            self._reset_epoch += 1
+        # The queue was just cleared, so a pending task change has nothing
+        # stale left to blend against.
+        self._discard_task_change()

    # ------------------------------------------------------------------
    # Action production (called from main thread)
@@ -281,6 +310,7 @@ class RTCInferenceEngine(InferenceEngine):
                queue = self._action_queue
                with self._obs_lock:
                    obs = self._obs_holder.get("obs")
+                    epoch_before = self._reset_epoch
                if queue is None or obs is None:
                    time.sleep(_RTC_IDLE_SLEEP_S)
                    continue
@@ -294,11 +324,24 @@ class RTCInferenceEngine(InferenceEngine):
                        latency = latency_tracker.max()
                        delay = math.ceil(latency / time_per_chunk) if latency else 0

+                        task, task_changed = self._take_task()
+                        if task_changed:
+                            # No queue flush on purpose: dropping the queued
+                            # actions would leave the robot without commands for
+                            # a full inference latency.  With RTC blending on
+                            # (the default) this chunk — already conditioned on
+                            # the new instruction — is merged over the previous
+                            # chunk's leftover prefix, so the switch lands within
+                            # one inference and the transition stays continuous.
+                            # With blending disabled the queue drains first, so
+                            # it lands up to one chunk later.
+                            logger.info("Task changed to '%s' — applied from this chunk on", task)
+
                        obs_batch = build_dataset_frame(self._hw_features, obs, prefix="observation")
                        obs_batch = prepare_observation_for_inference(
-                            obs_batch, policy_device, self._task, self._robot.robot_type
+                            obs_batch, policy_device, task, self._robot.robot_type
                        )
-                        obs_batch["task"] = [self._task]
+                        obs_batch["task"] = [task]

                        preprocessed = self._preprocessor(obs_batch)

@@ -339,7 +382,12 @@ class RTCInferenceEngine(InferenceEngine):
                        else:
                            latency_tracker.add(new_latency)

-                        queue.merge(original, processed, new_delay, idx_before)
+                        with self._obs_lock:
+                            epoch_unchanged = epoch_before == self._reset_epoch
+                        if epoch_unchanged:
+                            queue.merge(original, processed, new_delay, idx_before)
+                        else:
+                            logger.info("Discarding action chunk computed before an engine reset")

                        if (
                            is_warmup
@@ -368,8 +416,9 @@ class RTCInferenceEngine(InferenceEngine):
                    time.sleep(_RTC_IDLE_SLEEP_S)

        except Exception as e:
+            self._failure_traceback = traceback.format_exc()
            logger.error("Fatal error in RTC thread: %s", e)
-            logger.error(traceback.format_exc())
+            logger.error(self._failure_traceback)
            self._rtc_error.set()
            # Unblock any warmup waiters so the main loop doesn't spin forever
            self._compile_warmup_done.set()
@@ -65,12 +65,12 @@ class SyncInferenceEngine(InferenceEngine):
        device: str | None,
        robot_type: str,
    ) -> None:
+        super().__init__(task=task)
        self._policy = policy
        self._preprocessor = preprocessor
        self._postprocessor = postprocessor
        self._dataset_features = dataset_features
        self._ordered_action_keys = ordered_action_keys
-        self._task = task
        self._device = torch.device(device or "cpu")
        self._robot_type = robot_type
        logger.info(
@@ -93,6 +93,9 @@ class SyncInferenceEngine(InferenceEngine):
        self._policy.reset()
        self._preprocessor.reset()
        self._postprocessor.reset()
+        # The policy was just reset, so a pending task change has nothing
+        # stale left to flush.
+        self._discard_task_change()

    def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
        """Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
@@ -107,10 +110,20 @@ class SyncInferenceEngine(InferenceEngine):
            if self._device.type == "cuda" and self._policy.config.use_amp
            else nullcontext()
        )
+        task, task_changed = self._take_task()
        with torch.inference_mode(), autocast_ctx:
-            observation = prepare_observation_for_inference(
-                observation, self._device, self._task, self._robot_type
-            )
+            if task_changed:
+                # Chunking policies serve actions from an internal queue filled
+                # under the previous instruction (up to chunk_size ticks of stale
+                # behavior), so drop them and let the new instruction take effect
+                # on this very tick.  Deliberately narrower than ``policy.reset``:
+                # observation history and other episode state are kept, so a
+                # policy that conditions on them (and one that ignores the task
+                # entirely) sees no discontinuity.  Safe to mutate here — this is
+                # the thread that calls ``select_action``.
+                logger.info("Task changed to '%s' — dropping precomputed actions", task)
+                self._policy.drop_queued_actions()
+            observation = prepare_observation_for_inference(observation, self._device, task, self._robot_type)
            observation = self._preprocessor(observation)
            action = self._policy.select_action(observation)
            action = self._postprocessor(action)
@@ -0,0 +1,304 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Interactive rollout session: chat-style stdin commands for ``lerobot-rollout``.
+
+Enabled with ``--interactive=true``, this module lets the operator control a
+rollout from the terminal while hardware and policy stay connected and warm:
+
+    /start           start (or restart) the policy control loop
+    /subtask <text>  change the instruction the policy follows, mid-run
+    /reset           stop movement, return the robot to its initial position,
+                     and restore the instruction passed on the command line
+    /stop            end the session and run the normal shutdown routines
+    /help            show the available commands
+
+This module is only the CLI front-end: stdin reading
+(:class:`lerobot.utils.stdin_input.StdinCommandListener`), command parsing,
+terminal output, and log muting.  All control logic lives in
+:class:`lerobot.rollout.controller.RolloutController`, which is the public
+API for driving a rollout programmatically (from an application, a network
+server, a notebook, ...) without any of this module's terminal I/O.
+
+Threading model: a daemon stdin-listener thread parses lines and calls the
+controller's thread-safe methods (``start``/``reset``/``stop``/``set_task``);
+it never touches hardware or policy state.  ``RolloutController.serve()``
+runs on the main thread and executes ``strategy.run(ctx)`` in *segments*,
+ended through the session's :class:`LinkedEvent` (installed as
+``ctx.runtime.shutdown_event``).  Real shutdown signals (SIGINT/SIGTERM)
+propagate through the linked event's parent, so Ctrl-C behaves exactly as in
+non-interactive runs.
+
+While the session runs, log records below ERROR and Python warnings are
+suppressed process-wide (via ``logging.disable``) so routine system output
+does not interleave with the chat prompt; ERROR and CRITICAL records still
+reach the console, and a fatal inference-engine error is additionally
+reported with its captured traceback.  Normal logging resumes when the
+session ends (so teardown logs are visible).  Run without ``--interactive``
+to see the full live log output.
+
+The command table is intentionally a name → (handler, argument hint, help)
+mapping so further commands (``/ask`` and the rest of the language-runtime
+work in PR #4183/#4234) can be registered without restructuring the parser,
+the help output, or the session loop.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import warnings
+from collections.abc import Callable, Iterator
+from dataclasses import dataclass
+from typing import IO, TYPE_CHECKING
+
+from lerobot.utils.stdin_input import StdinCommandListener
+from lerobot.utils.utils import log_say
+
+from .controller import RolloutController, RolloutEvent
+
+if TYPE_CHECKING:
+    from .context import RolloutContext
+    from .strategies import RolloutStrategy
+
+logger = logging.getLogger(__name__)
+
+_BANNER_RULE = "─" * 60
+
+
+@contextlib.contextmanager
+def _mute_system_output() -> Iterator[None]:
+    """Suppress log records below ERROR and Python warnings, process-wide.
+
+    System logs (policy, robot, control loop) contend with the chat prompt
+    for the terminal.  ``logging.disable`` gates records before any handler
+    dispatch, which covers non-propagating library loggers (``transformers``,
+    ``datasets``) and loggers created mid-session alike; ERROR and CRITICAL
+    records still get through, so failures stay visible.  The gate applies to
+    every handler — including file handlers, which therefore also miss
+    INFO/WARNING records for the duration.  Python warnings bypass logging
+    entirely and are silenced separately.
+    """
+    previous_disable = logging.root.manager.disable
+    saved_warning_filters = warnings.filters[:]
+    logging.disable(logging.WARNING)
+    warnings.simplefilter("ignore")
+    try:
+        yield
+    finally:
+        logging.disable(previous_disable)
+        warnings.filters[:] = saved_warning_filters
+
+
+@dataclass(frozen=True)
+class InteractiveCommand:
+    """A parsed ``/name args`` line from the interactive prompt."""
+
+    name: str
+    args: str = ""
+
+
+def _format_task(task: str) -> str:
+    """Render a task string for the operator, naming the empty case explicitly."""
+    return repr(task) if task else "(none — set one with /subtask <text>)"
+
+
+def _strip_quotes(text: str) -> str:
+    """Drop one layer of matching surrounding quotes from a command argument."""
+    if len(text) >= 2 and text[0] == text[-1] and text[0] in ("'", '"'):
+        return text[1:-1]
+    return text
+
+
+def parse_command(line: str) -> InteractiveCommand | None:
+    """Parse an input line into an :class:`InteractiveCommand`.
+
+    Commands are ``/name`` optionally followed by free-text arguments (e.g.
+    ``/subtask grab the red cube``).  Returns ``None`` for lines that are not
+    commands (no leading ``/`` or a bare ``/``).
+    """
+    line = line.strip()
+    if not line.startswith("/"):
+        return None
+    head, *rest = line.split(maxsplit=1)
+    name = head[1:].lower()
+    if not name:
+        return None
+    return InteractiveCommand(name=name, args=rest[0].strip() if rest else "")
+
+
+class InteractiveSession:
+    """Drive a rollout from chat-style stdin commands.
+
+    A thin terminal front-end over :class:`RolloutController`: the stdin
+    listener parses lines into commands, each command calls one of the
+    controller's thread-safe methods, and controller events are rendered
+    back as terminal output.  The controller is exposed as
+    :attr:`controller` for tests and embedders.
+
+    Commands are last-write-wins: ``/reset`` and ``/stop`` cancel a pending
+    ``/start`` so the robot never starts moving after the operator's final
+    command asked it not to.  End-of-file on the command stream stops the
+    session (a closed stdin means there is no way left to command the
+    robot), so piped scripts must keep stdin open for the intended session
+    duration, e.g. ``(printf '/start\\n'; sleep 60; printf '/stop\\n') |
+    lerobot-rollout ... --interactive=true``.  The session works over SSH
+    and in headless setups — it reads the terminal (or pipe) directly and
+    needs no display server.
+    """
+
+    def __init__(
+        self,
+        strategy: RolloutStrategy,
+        ctx: RolloutContext,
+        input_stream: IO[str] | None = None,
+    ) -> None:
+        self.controller = RolloutController(strategy, ctx, on_event=self._on_event)
+        self._play_sounds = ctx.runtime.cfg.play_sounds
+        self._listener = StdinCommandListener(self._handle_line, on_eof=self._handle_eof, stream=input_stream)
+
+        # name -> (handler, argument hint, help line); /help and the banner
+        # render from this table, so future commands (e.g. /ask) stay
+        # documented for free.
+        self._commands: dict[str, tuple[Callable[[InteractiveCommand], None], str, str]] = {
+            "start": (self._cmd_start, "", "start (or restart) the policy control loop"),
+            "subtask": (self._cmd_subtask, " <text>", "set the instruction the policy follows"),
+            "reset": (self._cmd_reset, "", "stop movement, return to initial position, restore the task"),
+            "stop": (self._cmd_stop, "", "end the session and shut down"),
+            "help": (self._cmd_help, "", "show this help"),
+        }
+
+    def run(self) -> None:
+        """Run the session until ``/stop``, EOF, engine failure, or a shutdown signal."""
+        try:
+            with _mute_system_output():
+                self._print(self._render_banner())
+                self._listener.start()
+                try:
+                    self.controller.serve()
+                finally:
+                    self._listener.stop()
+        finally:
+            # Outside the muting context so the closing announcement and any
+            # teardown logs are visible again.
+            log_say("Interactive session ended", self._play_sounds)
+
+    # ------------------------------------------------------------------
+    # Controller events (fired on the serve thread) -> terminal output
+    # ------------------------------------------------------------------
+
+    def _on_event(self, event: RolloutEvent) -> None:
+        if event is RolloutEvent.SEGMENT_STARTED:
+            log_say("Starting rollout", self._play_sounds)
+            self._print(
+                f"Rollout running — task {_format_task(self.controller.task)}. "
+                "/subtask <text> to change it, /reset to return to initial position, /stop to shut down."
+            )
+        elif event is RolloutEvent.SEGMENT_ENDED:
+            self._print(
+                "Rollout run ended on its own (duration reached). Robot is holding position — "
+                "/start to run again, /reset to return to initial position, /stop to shut down."
+            )
+        elif event is RolloutEvent.RESET_STARTED:
+            log_say("Resetting robot to initial position", self._play_sounds)
+            self._print("Resetting — returning the robot to its initial position...")
+        elif event is RolloutEvent.RESET_DONE:
+            self._print("Robot reset — holding at initial position. /start to run.")
+        elif event is RolloutEvent.RESET_SKIPPED:
+            self._print("Robot paused — no initial position captured, holding current pose. /start to run.")
+        elif event is RolloutEvent.ENGINE_FAILED:
+            self._report_engine_failure()
+
+    def _report_engine_failure(self) -> None:
+        """Surface a fatal engine error despite the muted console logging."""
+        self._print("Inference engine failed — shutting down.")
+        failure_traceback = self.controller.failure_traceback
+        if failure_traceback:
+            self._print(failure_traceback)
+        else:
+            self._print("Re-run without --interactive=true to see the error output.")
+
+    # ------------------------------------------------------------------
+    # Command handlers (called from the listener thread; the controller's
+    # methods are thread-safe and only publish state)
+    # ------------------------------------------------------------------
+
+    def _handle_line(self, line: str) -> None:
+        cmd = parse_command(line)
+        if cmd is None:
+            self._print("Input not recognized — commands start with '/'. Type /help for the list.")
+            return
+        entry = self._commands.get(cmd.name)
+        if entry is None:
+            self._print(f"Unknown command '/{cmd.name}'. Type /help for the list.")
+            return
+        handler = entry[0]
+        handler(cmd)
+
+    def _handle_eof(self) -> None:
+        self._print("Input stream closed — stopping the session.")
+        self.controller.stop()
+
+    def _cmd_start(self, cmd: InteractiveCommand) -> None:
+        if not self.controller.start():
+            self._print("Already running — /reset to pause first, or /stop to shut down.")
+
+    def _cmd_subtask(self, cmd: InteractiveCommand) -> None:
+        if not cmd.args:
+            self._print(f"Current task: {_format_task(self.controller.task)}")
+            return
+        task = _strip_quotes(cmd.args)
+        previous = self.controller.task
+        if self.controller.set_task(task):
+            self._print(
+                f"Task: {_format_task(previous)} → {_format_task(task)} "
+                "(applies from the next policy inference)"
+            )
+        else:
+            self._print(f"Task unchanged: {_format_task(task)}")
+
+    def _cmd_reset(self, cmd: InteractiveCommand) -> None:
+        if self.controller.reset():
+            self._print(f"Task restored to {_format_task(self.controller.initial_task)}")
+
+    def _cmd_stop(self, cmd: InteractiveCommand) -> None:
+        self.controller.stop()
+
+    def _cmd_help(self, cmd: InteractiveCommand) -> None:
+        self._print(self._render_help())
+
+    # ------------------------------------------------------------------
+    # Rendering
+    # ------------------------------------------------------------------
+
+    def _render_help(self) -> str:
+        usages = {name: f"/{name}{entry[1]}" for name, entry in self._commands.items()}
+        width = max(len(usage) for usage in usages.values())
+        lines = [f"  {usages[name]:<{width}}   {entry[2]}" for name, entry in self._commands.items()]
+        return "Available commands:\n" + "\n".join(lines)
+
+    def _render_banner(self) -> str:
+        return (
+            f"{_BANNER_RULE}\n"
+            "Interactive rollout session — the robot will NOT move until you type /start.\n"
+            f"Task: {_format_task(self.controller.initial_task)}\n"
+            f"{self._render_help()}\n"
+            "System logs and warnings are muted during the session (errors still show).\n"
+            f"{_BANNER_RULE}"
+        )
+
+    @staticmethod
+    def _print(message: str) -> None:
+        """User-facing chat output; logging stays on stderr, replies on stdout."""
+        print(message, flush=True)
@@ -63,12 +63,25 @@ class RolloutStrategy(abc.ABC):
        self._interpolator = ActionInterpolator(multiplier=ctx.runtime.cfg.interpolation_multiplier)
        self._engine = ctx.policy.inference
        logger.info("Starting inference engine...")
-        self._engine.reset()
+        self.reset_control_state()
        self._engine.start()
        self._warmup_flushed = False
-        self._cached_obs_processed = None
        logger.info("Inference engine started")

+    def reset_control_state(self) -> None:
+        """Clear episode-scoped control state so a paused session can restart cleanly.
+
+        Resets the inference engine (policy hidden state, action queues), the
+        action interpolator, and the cached processed observation.  Used by the
+        interactive session between run segments; only call while the control
+        loop is not running.
+        """
+        if self._engine is not None:
+            self._engine.reset()
+        if self._interpolator is not None:
+            self._interpolator.reset()
+        self._cached_obs_processed = None
+
    def _process_observation_and_notify(self, processors: ProcessorContext, obs_raw: dict) -> dict:
        """Run the observation processor and notify the engine — throttled to policy ticks.

@@ -125,7 +138,7 @@ class RolloutStrategy(abc.ABC):
        if robot.is_connected:
            if return_to_initial_position and hw.initial_position:
                logger.info("Returning robot to initial position before shutdown...")
-                self._return_to_initial_position(hw)
+                self.return_to_initial_position(hw)
            elif not return_to_initial_position:
                logger.info(
                    "Skipping return-to-initial-position (disabled by config); leaving robot in final pose."
@@ -138,7 +151,7 @@ class RolloutStrategy(abc.ABC):
            teleop.disconnect()

    @staticmethod
-    def _return_to_initial_position(hw: HardwareContext, duration_s: float = 3.0, fps: int = 50) -> None:
+    def return_to_initial_position(hw: HardwareContext, duration_s: float = 3.0, fps: int = 50) -> None:
        """Smoothly interpolate the robot back to its initial position."""
        robot = hw.robot_wrapper
        target = hw.initial_position
@@ -165,7 +165,7 @@ class EpisodicStrategy(RolloutStrategy):

                        elif self.config.reset_to_initial_position:
                            # No teleop: return the robot to its startup position.
-                            self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
+                            self.return_to_initial_position(hw=ctx.hardware, duration_s=1)

                        self._reset_loop(
                            ctx=ctx,
@@ -187,7 +187,7 @@ class EpisodicStrategy(RolloutStrategy):

                        # returns to its initial joint positions captured at startup
                        if not teleop and self.config.reset_to_initial_position:
-                            self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
+                            self.return_to_initial_position(hw=ctx.hardware, duration_s=1)

                        continue

@@ -22,7 +22,6 @@ import time
 from concurrent.futures import Future, ThreadPoolExecutor
 from threading import Event, Lock

-from lerobot.datasets import VideoEncodingManager
 from lerobot.datasets.utils import DEFAULT_VIDEO_FILE_SIZE_IN_MB
 from lerobot.utils.constants import ACTION, OBS_STR
 from lerobot.utils.feature_utils import build_dataset_frame
@@ -55,6 +54,14 @@ class SentryStrategy(RolloutStrategy):

    Requires ``streaming_encoding=True`` (enforced in config validation)
    to prevent disk I/O from blocking the control loop.
+
+    ``run()`` is restartable: each call records complete episodes plus one
+    final partial episode, and the dataset is only finalized in
+    ``teardown()`` — this is what lets ``--interactive=true`` drive sentry
+    in start/reset/start segments while the dataset stays open.  Frames are
+    labeled with the inference engine's *live* task, so a mid-run
+    ``/subtask`` changes both the policy conditioning and the recorded
+    label from the same frame onwards.
    """

    config: SentryStrategyConfig
@@ -70,6 +77,9 @@ class SentryStrategy(RolloutStrategy):
        """Initialise the inference engine and background push executor."""
        self._init_engine(ctx)
        self._push_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sentry-push")
+        # Instance state (not run()-local) so the upload cadence survives
+        # interactive run segments.
+        self._episodes_since_push = 0
        target_mb = self.config.target_video_file_size_mb or DEFAULT_VIDEO_FILE_SIZE_IN_MB
        self._episode_duration_s = estimate_max_episode_seconds(
            ctx.data.dataset_features, ctx.runtime.cfg.fps, target_size_mb=target_mb
@@ -97,79 +107,95 @@ class SentryStrategy(RolloutStrategy):

        start_time = time.perf_counter()
        episode_start = time.perf_counter()
-        episodes_since_push = 0
-        task_str = cfg.dataset.single_task if cfg.dataset else cfg.task
        logger.info("Sentry recording started (episode_duration=%.0fs)", episode_duration_s)

-        with VideoEncodingManager(dataset):
-            try:
-                while not ctx.runtime.shutdown_event.is_set():
-                    loop_start = time.perf_counter()
+        # No dataset finalization here: run() must be restartable (interactive
+        # segments), so the dataset stays open until teardown() finalizes it.
+        try:
+            while not ctx.runtime.shutdown_event.is_set():
+                loop_start = time.perf_counter()

-                    if cfg.duration > 0 and (time.perf_counter() - start_time) >= cfg.duration:
-                        logger.info("Duration limit reached (%.0fs)", cfg.duration)
-                        break
+                if cfg.duration > 0 and (time.perf_counter() - start_time) >= cfg.duration:
+                    logger.info("Duration limit reached (%.0fs)", cfg.duration)
+                    break

-                    obs = robot.get_observation()
-                    obs_processed = self._process_observation_and_notify(ctx.processors, obs)
+                obs = robot.get_observation()
+                obs_processed = self._process_observation_and_notify(ctx.processors, obs)

-                    if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
-                        continue
+                if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
+                    continue

-                    action_dict = send_next_action(obs_processed, obs, ctx, interpolator)
+                action_dict = send_next_action(obs_processed, obs, ctx, interpolator)

-                    if action_dict is not None:
-                        self._log_telemetry(obs_processed, action_dict, ctx.runtime)
-                        obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
-                        action_frame = build_dataset_frame(features, action_dict, prefix=ACTION)
-                        frame = {**obs_frame, **action_frame, "task": task_str}
-                        # ``add_frame`` writes to the in-progress episode buffer; the
-                        # background pusher only ever touches *finalised* episode
-                        # artifacts on disk.  The two operate on disjoint state, so
-                        # ``add_frame`` does not need ``_episode_lock``.
-                        dataset.add_frame(frame)
+                if action_dict is not None:
+                    self._log_telemetry(obs_processed, action_dict, ctx.runtime)
+                    obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
+                    action_frame = build_dataset_frame(features, action_dict, prefix=ACTION)
+                    # The task is read live from the engine (not snapshotted from
+                    # config) so an interactive /subtask relabels frames from the
+                    # moment it re-instructs the policy; the writer stores a task
+                    # per frame.  At launch the engine holds the configured task.
+                    frame = {**obs_frame, **action_frame, "task": engine.task}
+                    # ``add_frame`` writes to the in-progress episode buffer; the
+                    # background pusher only ever touches *finalised* episode
+                    # artifacts on disk.  The two operate on disjoint state, so
+                    # ``add_frame`` does not need ``_episode_lock``.
+                    dataset.add_frame(frame)

-                    # Episode rotation derived from video file-size target.
-                    # The duration is a conservative estimate so the actual
-                    # video has crossed DEFAULT_VIDEO_FILE_SIZE_IN_MB by now,
-                    # keeping push_to_hub efficient (uploads complete files).
-                    elapsed = time.perf_counter() - episode_start
-                    if elapsed >= episode_duration_s:
-                        # ``save_episode`` finalises the in-progress episode and
-                        # flushes it to disk; ``_episode_lock`` serialises this with
-                        # ``push_to_hub`` (run in the background executor) so the
-                        # pusher never reads a half-written episode.
-                        with self._episode_lock:
-                            dataset.save_episode()
-                        episodes_since_push += 1
-                        self._needs_push.set()
-                        logger.info(
-                            "Episode saved (total: %d, elapsed: %.1fs)",
-                            dataset.num_episodes,
-                            elapsed,
-                        )
-                        log_say(f"Episode {dataset.num_episodes} saved", play_sounds)
-
-                        if episodes_since_push >= self.config.upload_every_n_episodes:
-                            self._background_push(dataset, cfg)
-                            episodes_since_push = 0
-
-                        episode_start = time.perf_counter()
-
-                    dt = time.perf_counter() - loop_start
-                    if (sleep_t := control_interval - dt) > 0:
-                        precise_sleep(sleep_t)
-                    else:
-                        logger.warning(
-                            f"Record loop is running slower ({1 / dt:.1f} Hz) than the target FPS ({cfg.fps} Hz). Dataset frames might be dropped and robot control might be unstable. Common causes are: 1) Camera FPS not keeping up 2) Policy inference taking too long 3) CPU starvation"
-                        )
-
-            finally:
-                logger.info("Sentry control loop ended — saving final episode")
-                with contextlib.suppress(Exception):
+                # Episode rotation derived from video file-size target.
+                # The duration is a conservative estimate so the actual
+                # video has crossed DEFAULT_VIDEO_FILE_SIZE_IN_MB by now,
+                # keeping push_to_hub efficient (uploads complete files).
+                elapsed = time.perf_counter() - episode_start
+                if elapsed >= episode_duration_s:
+                    # ``save_episode`` finalises the in-progress episode and
+                    # flushes it to disk; ``_episode_lock`` serialises this with
+                    # ``push_to_hub`` (run in the background executor) so the
+                    # pusher never reads a half-written episode.
                    with self._episode_lock:
                        dataset.save_episode()
+                    self._episodes_since_push += 1
                    self._needs_push.set()
+                    logger.info(
+                        "Episode saved (total: %d, elapsed: %.1fs)",
+                        dataset.num_episodes,
+                        elapsed,
+                    )
+                    log_say(f"Episode {dataset.num_episodes} saved", play_sounds)
+
+                    if self._episodes_since_push >= self.config.upload_every_n_episodes:
+                        self._background_push(dataset, cfg)
+                        self._episodes_since_push = 0
+
+                    episode_start = time.perf_counter()
+
+                dt = time.perf_counter() - loop_start
+                if (sleep_t := control_interval - dt) > 0:
+                    precise_sleep(sleep_t)
+                else:
+                    logger.warning(
+                        f"Record loop is running slower ({1 / dt:.1f} Hz) than the target FPS ({cfg.fps} Hz). Dataset frames might be dropped and robot control might be unstable. Common causes are: 1) Camera FPS not keeping up 2) Policy inference taking too long 3) CPU starvation"
+                    )
+
+        finally:
+            logger.info("Sentry control loop ended — saving final episode")
+            try:
+                with self._episode_lock:
+                    dataset.save_episode()
+                self._needs_push.set()
+            except Exception:
+                # The tail episode could not be committed (nothing was
+                # recorded, or the save failed mid-write).  Drop the in-flight
+                # streaming encode so teardown()'s finalize does not flush a
+                # half-written video, and discard the episode buffer: a failed
+                # save_episode leaves it half-mutated, which would crash the
+                # first add_frame of a restarted segment.  add_frame recreates
+                # a fresh buffer from None.
+                logger.warning("Tail episode was not saved — discarding it", exc_info=True)
+                if dataset.writer is not None:
+                    with contextlib.suppress(Exception):
+                        dataset.writer.cancel_pending_videos()
+                    dataset.writer.episode_buffer = None

    def teardown(self, ctx: RolloutContext) -> None:
        """Flush pending pushes, finalise the dataset, and disconnect hardware."""
@@ -25,6 +25,11 @@ quantile statistics (q01, q10, q50, q90, q99) in their metadata. This script:
 3. If missing, computes quantile statistics for all features
 4. Updates the dataset metadata with the new quantile statistics

+Statistics are accumulated into a single running histogram per feature across
+all episodes rather than aggregating per-episode quantile summaries. The
+resulting quantiles are histogram approximations, subject to discretization and
+range-rebinning error; image/video frames are sampled by default.
+
 Usage:

 ```bash
@@ -34,9 +39,7 @@ python src/lerobot/scripts/augment_dataset_quantile_stats.py \
 """

 import argparse
-import concurrent.futures
 import logging
-import os
 from pathlib import Path

 import numpy as np
@@ -49,11 +52,10 @@ from lerobot.datasets import (
    CODEBASE_VERSION,
    DEFAULT_QUANTILES,
    LeRobotDataset,
-    aggregate_stats,
    get_feature_stats,
    write_stats,
 )
-from lerobot.datasets.compute_stats import sample_indices
+from lerobot.datasets.compute_stats import RunningQuantileStats, sample_indices
 from lerobot.utils.utils import init_logging


@@ -79,20 +81,25 @@ def has_quantile_stats(stats: dict[str, dict] | None, quantile_list_keys: list[s
    return False


-def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampling: bool = True) -> dict:
-    """Process a single episode and return its statistics.
+def collect_episode_arrays(
+    dataset: LeRobotDataset,
+    episode_idx: int,
+    use_sampling: bool = True,
+    skip_images: bool = False,
+) -> dict[str, tuple[np.ndarray, int]]:
+    """Collect one episode's frames per feature, flattened to (num_samples, dim).

    Args:
        dataset: The LeRobot dataset
-        episode_idx: Index of the episode to process
-        use_sampling: If True, sub-sample image/video frames per episode to bound
-            memory. If False, use every frame (exact, higher memory).
+        episode_idx: Index of the episode to read
+        use_sampling: If True, sub-sample image/video frames to bound memory.
+            If False, use every frame (higher memory).
+        skip_images: If True, skip image/video features entirely.

    Returns:
-        Dictionary containing episode statistics
+        Mapping of feature name to that episode's values and the number of frames
+        they came from (which differs from the row count for image features).
    """
-    logging.info(f"Computing stats for episode {episode_idx}")
-
    start_idx = dataset.meta.episodes[episode_idx]["dataset_from_index"]
    end_idx = dataset.meta.episodes[episode_idx]["dataset_to_index"]

@@ -102,7 +109,9 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
    # numeric columns are cheap, so read them in full (exact).
    image_keys = [k for k in dataset.features if dataset.features[k]["dtype"] in ("image", "video")]
    numeric_keys = [
-        k for k in dataset.features if dataset.features[k]["dtype"] not in ("image", "video", "string")
+        k
+        for k in dataset.features
+        if dataset.features[k]["dtype"] not in ("image", "video", "string", "language")
    ]

    collected_data: dict[str, list] = {}
@@ -114,7 +123,7 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
            collected_data[key] = [torch.as_tensor(v) for v in numeric_cols[key]]

    # Image/video features: decode only a sampled subset of frames.
-    if image_keys:
+    if image_keys and not skip_images:
        sampled_offsets = sample_indices(episode_len) if use_sampling else list(range(episode_len))
        for offset in sampled_offsets:
            item = dataset[start_idx + offset]
@@ -122,87 +131,82 @@ def process_single_episode(dataset: LeRobotDataset, episode_idx: int, use_sampli
                if key in item:
                    collected_data.setdefault(key, []).append(item[key])

-    ep_stats = {}
+    episode_arrays: dict[str, tuple[np.ndarray, int]] = {}
    for key, data_list in collected_data.items():
-        if dataset.features[key]["dtype"] == "string":
-            continue
-
        data = torch.stack(data_list).cpu().numpy()
        if dataset.features[key]["dtype"] in ["image", "video"]:
            if data.dtype == np.uint8:
                data = data.astype(np.float32) / 255.0
-
-            axes_to_reduce = (0, 2, 3)
-            keepdims = True
+            # (N, C, H, W) -> (N * H * W, C) so quantiles are computed per channel.
+            channels = data.shape[1]
+            values = data.transpose(0, 2, 3, 1).reshape(-1, channels)
        else:
-            axes_to_reduce = 0
-            keepdims = data.ndim == 1
+            values = data.reshape(-1, data.shape[-1]) if data.ndim > 1 else data.reshape(-1, 1)
+        episode_arrays[key] = (values, len(data_list))

-        ep_stats[key] = get_feature_stats(
-            data, axis=axes_to_reduce, keepdims=keepdims, quantile_list=DEFAULT_QUANTILES
-        )
-
-        if dataset.features[key]["dtype"] in ["image", "video"]:
-            ep_stats[key] = {
-                k: v if k == "count" else np.squeeze(v, axis=0) for k, v in ep_stats[key].items()
-            }
-
-    return ep_stats
+    return episode_arrays


-def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
-    """Compute quantile statistics for all episodes in the dataset.
+def compute_quantile_stats_for_dataset(
+    dataset: LeRobotDataset,
+    use_sampling: bool = True,
+    skip_images: bool = False,
+) -> dict[str, dict]:
+    """Compute whole-dataset statistics with one running histogram per feature.

    Args:
        dataset: The LeRobot dataset to compute statistics for
        use_sampling: If True, sub-sample image/video frames per episode to bound
-            memory. If False, use every frame (exact, higher memory).
+            memory. If False, use every frame (higher memory).
+        skip_images: If True, skip image/video features and leave their stats untouched.

    Returns:
-        Dictionary containing aggregated statistics with quantiles
+        Dictionary containing statistics with histogram-based global quantile estimates

    Note:
-        Video decoding operations are not thread-safe, so we process episodes sequentially
-        when video keys are present. For datasets without videos, we use parallel processing
-        with ThreadPoolExecutor for better performance.
+        Episodes are accumulated sequentially because the running accumulators are
+        shared across all of them.
    """
    logging.info(f"Computing quantile statistics for dataset with {dataset.num_episodes} episodes")

-    episode_stats_list = []
-    has_videos = len(dataset.meta.video_keys) > 0
+    running_stats: dict[str, RunningQuantileStats] = {}
+    frame_counts: dict[str, int] = {}
+    row_counts: dict[str, int] = {}
+    # Kept only while a feature has a single row, so it can still be finalized.
+    single_row_arrays: dict[str, np.ndarray] = {}

-    if has_videos:
-        logging.info("Dataset contains video keys - using sequential processing for thread safety")
-        for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
-            ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
-            episode_stats_list.append(ep_stats)
-    else:
-        logging.info("Dataset has no video keys - using parallel processing for better performance")
-        max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
+    for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
+        episode_arrays = collect_episode_arrays(
+            dataset, episode_idx, use_sampling=use_sampling, skip_images=skip_images
+        )
+        for key, (array, num_frames) in episode_arrays.items():
+            running_stats.setdefault(key, RunningQuantileStats()).update(array)
+            frame_counts[key] = frame_counts.get(key, 0) + num_frames
+            row_counts[key] = row_counts.get(key, 0) + len(array)
+            if row_counts[key] < 2:
+                single_row_arrays[key] = array
+            else:
+                single_row_arrays.pop(key, None)

-        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
-            future_to_episode = {
-                executor.submit(process_single_episode, dataset, episode_idx, use_sampling): episode_idx
-                for episode_idx in range(dataset.num_episodes)
-            }
-
-            episode_results = {}
-            with tqdm(total=dataset.num_episodes, desc="Processing episodes") as pbar:
-                for future in concurrent.futures.as_completed(future_to_episode):
-                    episode_idx = future_to_episode[future]
-                    ep_stats = future.result()
-                    episode_results[episode_idx] = ep_stats
-                    pbar.update(1)
-
-        for episode_idx in range(dataset.num_episodes):
-            if episode_idx in episode_results:
-                episode_stats_list.append(episode_results[episode_idx])
-
-    if not episode_stats_list:
+    if not running_stats:
        raise ValueError("No episode data found for computing statistics")

-    logging.info(f"Aggregating statistics from {len(episode_stats_list)} episodes")
-    return aggregate_stats(episode_stats_list)
+    aggregated_stats: dict[str, dict] = {}
+    for key, accumulator in running_stats.items():
+        if row_counts[key] < 2:
+            # Histograms need at least two samples; mirror get_feature_stats' basic-stats path.
+            stats = get_feature_stats(single_row_arrays[key], axis=0, keepdims=False)
+        else:
+            stats = accumulator.get_statistics()
+        if dataset.features[key]["dtype"] in ["image", "video"]:
+            # Image stats are stored as (C, 1, 1) to broadcast over height and width.
+            stats = {k: v if k == "count" else v[:, np.newaxis, np.newaxis] for k, v in stats.items()}
+        # `get_feature_stats` counts frames, not the per-channel rows the accumulator sees.
+        stats["count"] = np.array([frame_counts[key]])
+        aggregated_stats[key] = stats
+
+    logging.info(f"Computed global histogram statistics for {len(aggregated_stats)} features")
+    return aggregated_stats


 def augment_dataset_with_quantile_stats(
@@ -210,6 +214,7 @@ def augment_dataset_with_quantile_stats(
    root: str | Path | None = None,
    overwrite: bool = False,
    use_sampling: bool = True,
+    skip_images: bool = False,
 ) -> None:
    """Augment a dataset with quantile statistics if they are missing.

@@ -218,7 +223,8 @@ def augment_dataset_with_quantile_stats(
        root: Local root directory for the dataset
        overwrite: Overwrite existing quantile statistics if they already exist
        use_sampling: If True, sub-sample image/video frames per episode to bound
-            memory. If False, use every frame (exact, higher memory).
+            memory. If False, use every frame (higher memory).
+        skip_images: If True, skip image/video features and keep their existing stats
    """
    logging.info(f"Loading dataset: {repo_id}")
    dataset = LeRobotDataset(
@@ -232,7 +238,13 @@ def augment_dataset_with_quantile_stats(

    logging.info("Dataset does not contain quantile statistics. Computing them now...")

-    new_stats = compute_quantile_stats_for_dataset(dataset, use_sampling=use_sampling)
+    new_stats = compute_quantile_stats_for_dataset(
+        dataset, use_sampling=use_sampling, skip_images=skip_images
+    )
+
+    if skip_images and dataset.meta.stats:
+        for key, feature_stats in dataset.meta.stats.items():
+            new_stats.setdefault(key, feature_stats)

    logging.info("Updating dataset metadata with new quantile statistics")
    dataset.meta.stats = new_stats
@@ -276,10 +288,15 @@ def main():
        "--no-sampling",
        action="store_true",
        help=(
-            "Compute stats over every frame (exact, higher memory). By default, "
+            "Compute stats over every frame (higher memory). By default, "
            "image/video frames are sub-sampled per episode to bound memory."
        ),
    )
+    parser.add_argument(
+        "--skip-images",
+        action="store_true",
+        help="Skip image/video features and preserve their existing stats",
+    )

    args = parser.parse_args()
    root = Path(args.root) if args.root else None
@@ -291,6 +308,7 @@ def main():
        root=root,
        overwrite=args.overwrite,
        use_sampling=not args.no_sampling,
+        skip_images=args.skip_images,
    )


@@ -0,0 +1,174 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Convert a DCP-format checkpoint into a distributable safetensors model, offline.
+
+Runs single-process (no GPUs, no process group). Example:
+
+```bash
+lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000
+lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=user/my-policy
+```
+
+`--push_to_hub` publishes the converted directory as a model repo, degrading gracefully: the
+core artifacts (model.safetensors, config.json, processor files) always upload; the README
+model card is enriched with training/dataset metadata only when `train_config.json` (and the
+dataset it names) are reachable, with a WARNING naming exactly what was skipped otherwise.
+DCP shard artifacts are never uploaded — published repos carry safetensors only.
+"""
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+
+from huggingface_hub import HfApi
+
+from lerobot.configs import parser
+from lerobot.distributed.checkpoint import dcp_to_safetensors
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR
+from lerobot.utils.utils import init_logging
+
+
+@dataclass
+class ConvertDcpConfig:
+    """CLI config for the offline DCP-to-safetensors checkpoint conversion."""
+
+    # A checkpoint step directory (containing pretrained_model/) or a pretrained_model
+    # directory itself.
+    checkpoint_dir: Path
+    # Remove the DCP shard directory after a successful conversion.
+    delete_dcp: bool = False
+    # Publish the converted directory to this Hub repo id (e.g. "user/my-policy").
+    push_to_hub: str | None = None
+    private: bool | None = None
+
+
+def _locate_pretrained_dir(checkpoint_dir: Path) -> Path:
+    """Resolve the pretrained_model/ directory from a user-supplied checkpoint path.
+
+    Args:
+        checkpoint_dir (Path): A checkpoint step directory (containing `pretrained_model/`) or a
+            `pretrained_model` directory itself.
+
+    Returns:
+        Path: The nested `pretrained_model/` directory when present, otherwise `checkpoint_dir`
+        unchanged.
+    """
+    nested = checkpoint_dir / PRETRAINED_MODEL_DIR
+    return nested if nested.is_dir() else checkpoint_dir
+
+
+def _publish_converted(pretrained_dir: Path, repo_id: str, private: bool | None) -> None:
+    """Best-effort publish of a converted checkpoint dir, degrading gracefully.
+
+    The core artifacts (model.safetensors, config.json, processor files) always upload; the README
+    model card gains training/dataset metadata only when `train_config.json` (and the dataset it
+    names) are reachable, with a WARNING naming what was skipped otherwise. DCP shard artifacts are
+    excluded from the upload.
+
+    Args:
+        pretrained_dir (Path): The converted `pretrained_model/` directory to upload.
+        repo_id (str): Target Hub model repo id (e.g. "user/my-policy"); created if missing.
+        private (bool | None): Repo visibility passed to `create_repo`; None keeps the Hub (or
+            existing repo's) default.
+    """
+    from lerobot.common.train_utils import generate_model_card
+    from lerobot.configs.policies import PreTrainedConfig
+    from lerobot.configs.train import TRAIN_CONFIG_NAME, TrainPipelineConfig
+
+    train_cfg = None
+    dataset_meta = None
+    if (pretrained_dir / TRAIN_CONFIG_NAME).is_file():
+        try:
+            train_cfg = TrainPipelineConfig.from_pretrained(pretrained_dir)
+        except Exception as e:  # noqa: BLE001 — degrade, never block the upload
+            logging.warning(f"Could not parse {TRAIN_CONFIG_NAME} ({e}); README will lack training metadata.")
+    else:
+        logging.warning(f"{TRAIN_CONFIG_NAME} missing; README will lack training metadata.")
+    if train_cfg is not None:
+        try:
+            from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
+
+            dataset_meta = LeRobotDatasetMetadata(
+                repo_id=train_cfg.dataset.repo_id,
+                root=train_cfg.dataset.root,
+                revision=train_cfg.dataset.revision,
+            )
+        except Exception as e:  # noqa: BLE001
+            logging.warning(
+                f"Dataset '{train_cfg.dataset.repo_id}' unreachable ({e}); README will lack dataset metadata."
+            )
+    try:
+        model_cfg = PreTrainedConfig.from_pretrained(pretrained_dir)
+        card = generate_model_card(model_cfg, cfg=train_cfg, dataset_meta=dataset_meta)
+        card.save(str(pretrained_dir / "README.md"))
+    except Exception as e:  # noqa: BLE001
+        logging.warning(f"Could not build the model card ({e}); publishing without README.")
+
+    api = HfApi()
+    repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
+    commit_info = api.upload_folder(
+        repo_id=repo_id,
+        repo_type="model",
+        folder_path=str(pretrained_dir),
+        commit_message="Upload converted policy (DCP -> safetensors)",
+        allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"],
+        # The checkpoint keeps its DCP shard directory unless --delete_dcp was passed; the
+        # allow list above admits neither `.distcp` shards nor their `.metadata` sidecar.
+        ignore_patterns=["*.tmp", "*.log"],
+    )
+    logging.info(f"Model pushed to {commit_info.repo_url.url}")
+
+
+@parser.wrap()
+def convert_checkpoint(cfg: ConvertDcpConfig) -> Path:
+    """Merge a checkpoint's DCP shards into `model.safetensors`, then optionally publish it.
+
+    Args:
+        cfg (ConvertDcpConfig): Conversion options — the checkpoint directory to convert, whether
+            to delete the DCP shards after a successful merge, and the optional Hub repo id (and
+            visibility) to publish the converted directory to.
+
+    Returns:
+        Path: The path to the merged `model.safetensors` file.
+
+    Raises:
+        FileNotFoundError: If the checkpoint has no DCP shard directory, i.e. it was not saved
+            with `checkpoint_format=dcp` (or `safetensors_dcp`).
+    """
+    from accelerate.utils.constants import FSDP_MODEL_NAME
+
+    pretrained_dir = _locate_pretrained_dir(cfg.checkpoint_dir)
+    dcp_dir = pretrained_dir / f"{FSDP_MODEL_NAME}_0"
+    if not dcp_dir.is_dir():
+        raise FileNotFoundError(
+            f"No DCP shard directory at {dcp_dir}. Point --checkpoint_dir at a checkpoint "
+            "saved with checkpoint_format=dcp (or safetensors_dcp)."
+        )
+    logging.info(f"Merging {dcp_dir} -> {pretrained_dir / 'model.safetensors'}")
+    safetensors_path = dcp_to_safetensors(dcp_dir, pretrained_dir, delete_dcp=cfg.delete_dcp)
+    if cfg.push_to_hub:
+        _publish_converted(pretrained_dir, cfg.push_to_hub, cfg.private)
+    return safetensors_path
+
+
+def main() -> None:
+    """`lerobot-convert-dcp` console entry point: set up logging and run the conversion."""
+    init_logging()
+    convert_checkpoint()
+
+
+if __name__ == "__main__":
+    main()
@@ -44,6 +44,19 @@ Usage examples
        --robot.port=/dev/ttyACM0 \\
        --task="pick up cube" --duration=30

+    # Interactive session (base or sentry strategy): the robot stays idle
+    # until /start is typed; /subtask <text> re-instructs the policy mid-run;
+    # /reset returns it to the initial position (hardware and policy stay
+    # warm); /stop shuts down gracefully.  With --strategy.type=sentry the
+    # session also records continuously, labeling frames with the live task.
+    lerobot-rollout \\
+        --strategy.type=base \\
+        --policy.path=lerobot/act_koch_real \\
+        --robot.type=koch_follower \\
+        --robot.port=/dev/ttyACM0 \\
+        --task="pick up cube" \\
+        --interactive=true
+
    # Base mode — RTC inference for slow VLAs (Pi0, Pi0.5, SmolVLA)
    lerobot-rollout \\
        --strategy.type=base \\
@@ -173,7 +186,13 @@ from lerobot.robots import (  # noqa: F401
    so_follower,
    unitree_g1 as unitree_g1_robot,
 )
-from lerobot.rollout import RolloutConfig, build_rollout_context, create_strategy
+from lerobot.rollout import (
+    InteractiveSession,
+    LinkedEvent,
+    RolloutConfig,
+    build_rollout_context,
+    create_strategy,
+)
 from lerobot.teleoperators import (  # noqa: F401
    Teleoperator,
    TeleoperatorConfig,
@@ -215,6 +234,10 @@ def rollout(cfg: RolloutConfig):

    signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
    shutdown_event = signal_handler.shutdown_event
+    if cfg.interactive:
+        # Session commands (/reset, /stop) end the running control loop by setting
+        # the local flag; process signals still propagate through the parent event.
+        shutdown_event = LinkedEvent(shutdown_event)

    logger.info("Building rollout context...")
    ctx = build_rollout_context(cfg, shutdown_event)
@@ -230,8 +253,12 @@ def rollout(cfg: RolloutConfig):

    try:
        strategy.setup(ctx)
-        logger.info("Rollout setup complete, starting rollout...")
-        strategy.run(ctx)
+        if cfg.interactive:
+            logger.info("Rollout setup complete — starting interactive session (robot idle until /start)")
+            InteractiveSession(strategy, ctx).run()
+        else:
+            logger.info("Rollout setup complete, starting rollout...")
+            strategy.run(ctx)
    except KeyboardInterrupt:
        logger.info("Interrupted by user")
    finally:
@@ -25,6 +25,9 @@ from .constants import CHECKPOINTS_DIR
 T = TypeVar("T", bound="HubMixin")


+# Sharded-training resume artifacts (torch DCP shard dirs + shard files). Published model repos
+# carry safetensors only, so publishing uploads exclude these — checkpoint pushes (which exist
+# for resume, not distribution) deliberately do not.
 def find_latest_hub_checkpoint(
    repo_id: str,
    *,
@@ -36,6 +39,16 @@ def find_latest_hub_checkpoint(
    Training runs push checkpoints to ``checkpoints/<step>/`` (see
    ``push_checkpoint_to_hub``). This lists those step dirs and returns
    ``checkpoints/<highest-step>``, or ``None`` if the repo has no checkpoints.
+
+    Args:
+        repo_id (str): The Hub model repo to inspect.
+        token (str | bool | None): Hub authentication token. Defaults to None (the token
+            cached by `huggingface-cli login`).
+        revision (str | None): Repo revision to list. Defaults to None (the default branch).
+
+    Returns:
+        str | None: The repo-relative path `checkpoints/<highest-step>`, or None if the repo
+            has no checkpoints.
    """
    files = HfApi().list_repo_files(repo_id=repo_id, repo_type="model", revision=revision, token=token)
    prefix = f"{CHECKPOINTS_DIR}/"
@@ -164,7 +177,7 @@ class HubMixin:
        ignore_patterns: list[str] | str | None = None,
        delete_patterns: list[str] | str | None = None,
        card_kwargs: dict[str, Any] | None = None,
-    ) -> str:
+    ) -> str | None:
        """
        Upload model checkpoint to the Hub.

@@ -172,6 +185,10 @@ class HubMixin:
        `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more
        details.

+        Distributed contract: call on EVERY rank. `save_pretrained` runs on all ranks — for
+        sharded objects it can contain a collective gather (rank-gating it would deadlock) —
+        while repo creation and the upload happen on the main process only.
+
        Args:
            repo_id (`str`):
                ID of the repository to push to (example: `"username/my-model"`).
@@ -197,11 +214,17 @@ class HubMixin:
                Additional arguments passed to the card template to customize the card.

        Returns:
-            The url of the commit of your object in the given repository.
+            `str` or `None`: The url of the commit of your object in the given repository, or
+            `None` on non-main ranks of a distributed run (only the main process uploads).
        """
-        api = HfApi(token=token)
-        repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
+        # Lazy import: hub code must not import the distributed package at module load
+        # (configs -> hub is on the import path of lerobot.distributed itself).
+        from lerobot.distributed.utils import is_main_process

+        # Distributed contract: `save_pretrained` runs on EVERY rank — for sharded policies it
+        # contains a collective gather (rank-gating it would deadlock) and it writes into this
+        # rank's private tmpdir only on the main process. Repo creation and upload are then
+        # main-process-only.
        if commit_message is None:
            if "Policy" in self.__class__.__name__:
                commit_message = "Upload policy"
@@ -210,10 +233,13 @@ class HubMixin:
            else:
                commit_message = f"Upload {self.__class__.__name__}"

-        # Push the files to the repo in a single commit
        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
            saved_path = Path(tmp) / repo_id
            self.save_pretrained(saved_path, card_kwargs=card_kwargs)
+            if not is_main_process():
+                return None
+            api = HfApi(token=token)
+            repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
            return api.upload_folder(
                repo_id=repo_id,
                repo_type="model",
@@ -14,10 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections import defaultdict
-from collections.abc import Callable
 from typing import Any

 import torch
+import torch.distributed as dist

 from .utils import format_big_number

@@ -69,12 +69,31 @@ class MetricsTracker:
    """
    A helper class to track and log metrics over time.

+    Args:
+        batch_size (int): Per-process batch size (samples per micro-batch on each
+            data-parallel worker).
+        num_frames (int): Total number of frames in the training dataset.
+        num_episodes (int): Total number of episodes in the training dataset.
+        metrics (dict[str, AverageMeter]): The meters to track, keyed by metric name.
+        initial_step (int): Step counter to start from (non-zero when resuming a run).
+            Defaults to 0.
+        dp_world_size (int): Number of distinct data-parallel workers
+            (`dp_replicate * dp_shard`), used to scale sample accounting; context-parallel
+            peers consume the same batch and must not be double counted. Defaults to 1.
+
    Usage pattern:

    ```python
    # initialize, potentially with non-zero initial step (e.g. if resuming run)
    metrics = {"loss": AverageMeter("loss", ":.3f")}
-    train_metrics = MetricsTracker(cfg, dataset, metrics, initial_step=step)
+    train_metrics = MetricsTracker(
+        batch_size,
+        dataset.num_frames,
+        dataset.num_episodes,
+        metrics,
+        initial_step=step,
+        dp_world_size=dp_world,
+    )

    # update metrics derived from step (samples, episodes, epochs) at each training step
    train_metrics.step()
@@ -98,12 +117,12 @@ class MetricsTracker:
        "_batch_size",
        "_num_frames",
        "_avg_samples_per_ep",
+        "_dp_world_size",
        "metrics",
        "steps",
        "samples",
        "episodes",
        "epochs",
-        "accelerator",
        "_caller_metrics",
    ]

@@ -114,22 +133,25 @@ class MetricsTracker:
        num_episodes: int,
        metrics: dict[str, AverageMeter],
        initial_step: int = 0,
-        accelerator: Callable | None = None,
+        dp_world_size: int = 1,
    ):
        self.__dict__.update(dict.fromkeys(self.__keys__))
        self._batch_size = batch_size
        self._num_frames = num_frames
        self._avg_samples_per_ep = num_frames / num_episodes
+        # Sample accounting scales by the number of DISTINCT data-parallel workers, which is
+        # dp_replicate * dp_shard — not the world size: context-parallel peers consume the same
+        # batch and must not be double counted. `step` counts micro-batches, so no
+        # grad-accumulation factor belongs here either.
+        self._dp_world_size = dp_world_size
        self.metrics = metrics

        self.steps = initial_step
-        world_size = accelerator.num_processes if accelerator else 1
        # A sample is an (observation,action) pair, where observation and action
        # can be on multiple timestamps. In a batch, we have `batch_size` number of samples.
-        self.samples = self.steps * self._batch_size * world_size
+        self.samples = self.steps * self._batch_size * self._dp_world_size
        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)
@@ -155,8 +177,7 @@ class MetricsTracker:
        Updates metrics that depend on 'step' for one step.
        """
        self.steps += 1
-        world_size = self.accelerator.num_processes if self.accelerator else 1
-        self.samples += self._batch_size * world_size
+        self.samples += self._batch_size * self._dp_world_size
        self.episodes = self.samples / self._avg_samples_per_ep
        self.epochs = self.samples / self._num_frames

@@ -181,11 +202,16 @@ class MetricsTracker:
        across all distributed processes (in-place).

        This is a collective operation and MUST be invoked on every rank — typically just before
-        logging. With no accelerator or in single-process runs it is a no-op. Without it, metrics
-        reported by the main process only reflect rank 0; for bottleneck-style timings
-        (``dataloading_s``, ``update_s``, ...) that means the slowest worker's stall is invisible.
+        logging. Outside distributed runs it is a no-op. Without it, metrics reported by the
+        main process only reflect rank 0; for bottleneck-style timings (``dataloading_s``,
+        ``update_s``, ...) that means the slowest worker's stall is invisible.
+
+        Torch-native on purpose: metrics code carries no Accelerator dependency.
+        Note the reduction spans the WORLD group — correct for count-free averages (loss values
+        are identical within a context-parallel group, so including CP peers is a weighted
+        no-op).
        """
-        if self.accelerator is None or self.accelerator.num_processes <= 1:
+        if not dist.is_initialized() or dist.get_world_size() <= 1:
            return

        buckets: dict[str, list[str]] = defaultdict(list)
@@ -195,11 +221,20 @@ class MetricsTracker:
        if not buckets:
            return

-        device = self.accelerator.device
+        device = (
+            torch.device("cuda", torch.cuda.current_device())
+            if torch.cuda.is_available()
+            else torch.device("cpu")
+        )
+        reduce_ops = {
+            "mean": dist.ReduceOp.AVG,
+            "sum": dist.ReduceOp.SUM,
+            "max": dist.ReduceOp.MAX,
+        }
        for reduction, names in buckets.items():
            tensor = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device)
-            reduced = self.accelerator.reduce(tensor, reduction=reduction)
-            for name, value in zip(names, reduced.tolist(), strict=True):
+            dist.all_reduce(tensor, op=reduce_ops[reduction])
+            for name, value in zip(names, tensor.tolist(), strict=True):
                meter = self.metrics[name]
                # Preserve avg == sum / count so a later .update() on this meter accumulates
                # against the cluster view, not the stale per-rank history.
@@ -0,0 +1,187 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Non-blocking, line-oriented stdin reading.
+
+This complements :mod:`lerobot.utils.keyboard_input`, which serves discrete
+hotkeys: :class:`TerminalKeyListener` reads single raw bytes in cbreak mode,
+whereas :class:`StdinCommandListener` here assembles whole typed lines and
+leaves the terminal in canonical (line-buffered, echoing) mode — the operator
+is typing chat-style commands, not pressing hotkeys.  The two cannot share
+stdin at the same time.
+
+Environment support: reading works over SSH (the session's pty is a regular
+TTY file descriptor), in headless setups (no display server is involved,
+unlike the ``pynput`` keyboard backend), and from piped stdin.  End-of-file
+means "no more commands": an interactive Ctrl-D, an exhausted piped script,
+or ``stdin`` redirected from ``/dev/null`` all trigger ``on_eof``, as does a
+missing ``sys.stdin`` (e.g. a daemonized process).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import select
+import sys
+from collections.abc import Callable
+from threading import Thread
+from typing import IO
+
+logger = logging.getLogger(__name__)
+
+
+class StdinCommandListener:
+    """Daemon thread that reads input lines and forwards them to a callback.
+
+    On POSIX the reader polls the stream with ``select`` so ``stop()`` can
+    end the thread promptly; elsewhere (or for file-like objects without a
+    file descriptor) it falls back to a blocking ``readline`` daemon thread
+    that dies with the process.  Blank lines are skipped; end-of-file and
+    unexpected read errors trigger ``on_eof`` — a dead command channel must
+    never leave the consumer waiting for input that can no longer arrive.
+    """
+
+    def __init__(
+        self,
+        on_line: Callable[[str], None],
+        on_eof: Callable[[], None] | None = None,
+        stream: IO[str] | None = None,
+        poll_interval_s: float = 0.2,
+    ) -> None:
+        self._on_line = on_line
+        self._on_eof = on_eof
+        # sys.stdin can itself be None (pythonw, daemonized processes);
+        # start() treats that as an immediately-closed stream.
+        self._stream = stream if stream is not None else sys.stdin
+        self._poll_interval_s = poll_interval_s
+        self._running = False
+        self._thread: Thread | None = None
+        self._use_select = False
+        if os.name == "posix":
+            try:
+                self._stream.fileno()
+                self._use_select = True
+            except (OSError, ValueError, AttributeError):
+                pass
+
+    def start(self) -> None:
+        """Start the reader thread (idempotent)."""
+        if self._thread is not None:
+            return
+        if self._stream is None:
+            logger.warning("No stdin available for command input — treating as EOF")
+            self._emit_eof()
+            return
+        self._running = True
+        self._thread = Thread(target=self._run, daemon=True, name="StdinCommandListener")
+        self._thread.start()
+        if not self._use_select:
+            logger.info("stdin listener running in blocking mode (select unavailable for this stream)")
+
+    def stop(self) -> None:
+        """Stop the reader thread.
+
+        Blocking-mode threads may be stuck inside ``readline`` and cannot be
+        joined; they are daemons and die with the process. Late lines are
+        ignored via the ``_running`` flag either way.
+        """
+        self._running = False
+        thread = self._thread
+        self._thread = None
+        if thread is not None and thread.is_alive() and self._use_select:
+            thread.join(timeout=1.0)
+
+    def _run(self) -> None:
+        if self._use_select:
+            self._run_select()
+        else:
+            self._run_blocking()
+
+    def _run_select(self) -> None:
+        """Poll the file descriptor and split lines from raw bytes.
+
+        Reading raw bytes (instead of ``stream.readline()``) matters: a
+        buffered file object can slurp several lines off the descriptor at
+        once, after which ``select`` reports the drained fd as not-ready and
+        the buffered lines would never be delivered — breaking pasted or
+        piped command sequences.
+        """
+        fd = self._stream.fileno()
+        buffer = b""
+        while self._running:
+            try:
+                ready, _, _ = select.select([fd], [], [], self._poll_interval_s)
+            except (OSError, ValueError):  # stream closed underneath us
+                self._emit_read_error()
+                return
+            if not ready:
+                continue
+            try:
+                chunk = os.read(fd, 4096)
+            except OSError:
+                self._emit_read_error()
+                return
+            if not self._running:
+                return
+            if chunk == b"":  # EOF: Ctrl-D or the piped input ended
+                self._emit_line(buffer)  # a final command without trailing newline still counts
+                self._emit_eof()
+                return
+            buffer += chunk
+            while b"\n" in buffer:
+                raw, buffer = buffer.split(b"\n", 1)
+                self._emit_line(raw)
+
+    def _run_blocking(self) -> None:
+        while self._running:
+            try:
+                line = self._stream.readline()
+            except (OSError, ValueError, AttributeError):
+                self._emit_read_error()
+                return
+            if not self._running:
+                return
+            if line == "":  # EOF
+                self._emit_eof()
+                return
+            self._emit_line(line.encode() if isinstance(line, str) else line)
+
+    def _emit_line(self, raw: bytes) -> None:
+        line = raw.decode(errors="replace").strip()
+        if not line:
+            return
+        try:
+            self._on_line(line)
+        except Exception:  # never let a handler error kill the reader thread
+            logger.exception("Error while handling input line %r", line)
+
+    def _emit_eof(self) -> None:
+        logger.info("Input stream closed (EOF)")
+        if self._on_eof is not None:
+            try:
+                self._on_eof()
+            except Exception:
+                logger.exception("Error while handling input EOF")
+
+    def _emit_read_error(self) -> None:
+        """Treat an unexpected read failure like EOF so consumers shut down.
+
+        A dead command channel must not leave the consumer running with no
+        way to reach it.  Deliberate ``stop()`` calls clear ``_running``
+        first and do not reach this path.
+        """
+        if self._running:
+            logger.warning("Input stream failed — treating as EOF")
+            self._emit_eof()
@@ -0,0 +1,221 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Legacy-checkpoint contracts.
+
+Two contracts are pinned here so they are documented behavior, not accidents:
+
+- **The v0.6.0 hard break.** The v0.6.0 #3810 FSDP checkpoint layout
+  (full gathered ``model.safetensors`` + full ``optimizer_state.safetensors``, no DCP dirs,
+  no ``checkpoint_format`` in ``train_config.json``) is a hard break with ZERO v0.6.0-aware
+  runtime code — not even layout detection. A sharded resume pointed at such a checkpoint
+  must fail through the ORDINARY missing-artifact path (torch DCP erroring on the absent
+  ``training_state/optimizer_0/``), while the model weights remain loadable forever via
+  ``from_pretrained`` and the old ``num_processes`` key keeps feeding the topology reader.
+- **Converter equivalence.** ``dcp_to_safetensors`` (real ``merge_fsdp_weights``, no mocks)
+  on accelerate's ``save_fsdp_model`` DCP layout reproduces exactly the tensors that the
+  direct-gather ``save_pretrained`` artifact contains.
+"""
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+import torch
+import torch.distributed.checkpoint as dist_cp
+from accelerate.utils.constants import FSDP_MODEL_NAME, OPTIMIZER_NAME
+from safetensors.torch import load_file
+from torch.distributed.checkpoint.api import CheckpointException
+from torch.distributed.fsdp import FSDPModule
+
+from lerobot.common.train_utils import (
+    load_training_metadata,
+    resume_after_prepare,
+    resume_before_prepare,
+)
+from lerobot.configs.accelerator import FSDPConfig
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TRAIN_CONFIG_NAME, CheckpointFormat, TrainPipelineConfig
+from lerobot.distributed.checkpoint import dcp_to_safetensors, is_sharded_module
+from lerobot.optim.optimizers import save_optimizer_state
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP
+from lerobot.utils.io_utils import write_json
+from lerobot.utils.random_utils import save_rng_state
+from tests.fixtures.dummy_checkpoint_policy import DummyCheckpointPolicy, make_dummy_policy
+
+
+@pytest.fixture
+def accelerate_state():
+    """accelerate's process state, as the trainer's `Accelerator()` would have initialized it.
+
+    `load_fsdp_optimizer` and `merge_fsdp_weights` both consult `PartialState` internals
+    (logging and main-process gating). Single-process CPU state; reset on teardown so no
+    global accelerate state leaks into other tests.
+    """
+    from accelerate.state import AcceleratorState, PartialState
+
+    PartialState()
+    yield
+    AcceleratorState._reset_state(reset_partial_state=True)
+
+
+def make_v060_fsdp_checkpoint(checkpoint_dir: Path) -> dict[str, torch.Tensor]:
+    """Reproduce the v0.6.0 #3810 FSDP checkpoint layout with real artifacts.
+
+    - ``pretrained_model/``: ``config.json`` + full gathered ``model.safetensors`` (real
+      ``save_pretrained`` outputs) and a ``train_config.json`` predating the v0.7 fields
+      (``checkpoint_format``/``parallelism``/``accelerator`` stripped from the draccus dump);
+    - ``training_state/``: old-style ``training_step.json`` (``{"step", "num_processes"}``,
+      no ``dp_world_size``), ``rng_state.safetensors``, and the gathered full optimizer
+      channel (``optimizer_state.safetensors`` + ``optimizer_param_groups.json``) — and,
+      crucially, NO ``optimizer_0/`` DCP directory.
+
+    Returns the saved model weights for later comparison.
+    """
+    policy = make_dummy_policy()
+    optimizer = torch.optim.Adam(policy.parameters())
+    policy.forward({"observation.state": torch.randn(2, 4)})[0].backward()
+    optimizer.step()  # real optimizer state, applied before the weights are saved
+
+    pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
+    policy.save_pretrained(pretrained_dir)
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=3)
+    cfg._save_pretrained(pretrained_dir)
+    config_path = pretrained_dir / TRAIN_CONFIG_NAME
+    raw = json.loads(config_path.read_text())
+    assert "checkpoint_format" in raw  # draccus dumps defaults; a v0.6.0 config predates the key
+    for key in ("checkpoint_format", "parallelism", "accelerator"):
+        raw.pop(key, None)
+    config_path.write_text(json.dumps(raw, indent=4))
+
+    training_state_dir = checkpoint_dir / TRAINING_STATE_DIR
+    training_state_dir.mkdir()
+    write_json({"step": 5000, "num_processes": 4}, training_state_dir / TRAINING_STEP)
+    save_rng_state(training_state_dir)
+    save_optimizer_state(optimizer, training_state_dir)
+    return {key: tensor.clone() for key, tensor in policy.state_dict().items()}
+
+
+def as_fsdp2_module(policy: DummyCheckpointPolicy) -> DummyCheckpointPolicy:
+    """Give the policy FSDP2's runtime identity via the in-place class swap `fully_shard` performs.
+
+    torch's `fully_shard` swaps ``module.__class__`` to a ``(FSDPModule, type(module))``
+    subclass; mirroring that swap is what makes `is_sharded_module` (and thus the sharded
+    branch of `resume_after_prepare`) see a sharded model on a CPU-only single process. The
+    parameters stay plain tensors — sufficient here, because the resume must fail at the DCP
+    read before any sharded state is touched.
+    """
+    policy.__class__ = type(f"FSDP{type(policy).__name__}", (FSDPModule, type(policy)), {})
+    assert is_sharded_module(policy)
+    return policy
+
+
+def sharded_passthrough_accelerator() -> SimpleNamespace:
+    """The accelerator surface the sharded resume touches, carrying the trainer's real plugin.
+
+    `FSDPConfig.build_plugin()` is the exact FSDP2 plugin construction `make_accelerator`
+    hands to accelerate (state_dict_type stays at the FSDP2 default, SHARDED_STATE_DICT).
+    """
+    return SimpleNamespace(
+        unwrap_model=lambda m: m,
+        wait_for_everyone=lambda: None,
+        state=SimpleNamespace(fsdp_plugin=FSDPConfig().build_plugin()),
+    )
+
+
+class TestV060HardBreak:
+    """Pin the v0.6.0 hard break as a contract.
+
+    Zero v0.6.0-aware code ships — not even layout detection — so every assertion here must
+    hold through ORDINARY code paths only: the recorded config parses with plain defaults,
+    phase-1 resume and the weights stay loadable, and the sharded phase-2 resume fails with
+    torch DCP's own missing-artifact error, never a bespoke v0.6.0 message.
+    """
+
+    def test_sharded_resume_fails_with_ordinary_missing_artifact_error(self, tmp_path, accelerate_state):
+        make_v060_fsdp_checkpoint(tmp_path)
+
+        # No checkpoint_format recorded -> plain draccus default, no layout detection anywhere.
+        cfg = TrainPipelineConfig.from_pretrained(tmp_path / PRETRAINED_MODEL_DIR / TRAIN_CONFIG_NAME)
+        assert cfg.checkpoint_format is CheckpointFormat.SAFETENSORS
+        cfg.checkpoint_path = tmp_path
+
+        # Phase 1 (RNG + step counter) is format-independent and still succeeds.
+        assert resume_before_prepare(cfg) == 5000
+
+        # Phase 2 under sharding: the recorded format skips the DCP model preflight (the
+        # weights were already loaded by from_pretrained), then the sharded optimizer load
+        # hits the absent optimizer_0/ and fails inside torch DCP — the ordinary error path.
+        assert not (tmp_path / TRAINING_STATE_DIR / f"{OPTIMIZER_NAME}_0").exists()
+        policy = as_fsdp2_module(make_dummy_policy())
+        optimizer = torch.optim.Adam(policy.parameters())
+        with pytest.raises(CheckpointException) as excinfo:
+            resume_after_prepare(cfg, sharded_passthrough_accelerator(), policy, optimizer, None)
+        message = str(excinfo.value)
+        assert "lerobot-convert-dcp" not in message  # the converter hint belongs to recorded-format=DCP
+        assert "v0.6" not in message  # no bespoke wording: the explanation lives in the migration docs
+
+    def test_weights_remain_loadable_via_from_pretrained(self, tmp_path):
+        saved_weights = make_v060_fsdp_checkpoint(tmp_path)
+        policy = DummyCheckpointPolicy.from_pretrained(tmp_path / PRETRAINED_MODEL_DIR)
+        for key, tensor in policy.state_dict().items():
+            assert torch.equal(tensor, saved_weights[key]), key
+
+    def test_topology_reader_falls_back_to_legacy_num_processes(self, tmp_path):
+        make_v060_fsdp_checkpoint(tmp_path)
+        assert load_training_metadata(tmp_path / TRAINING_STATE_DIR)["dp_world_size"] == 4
+
+
+class TestConverterEquivalence:
+    def test_dcp_to_safetensors_output_equals_direct_gather(self, tmp_path, accelerate_state):
+        """DCP -> safetensors conversion is exactly the direct-gather artifact.
+
+        The DCP checkpoint is written with torch's real `dist_cp.save` (single process, no
+        process group), replicating accelerate's `save_fsdp_model` SHARDED_STATE_DICT branch
+        byte for byte: the ``{"model": state_dict}`` nesting and the ``pytorch_model_fsdp_0``
+        directory name. The conversion runs the real `merge_fsdp_weights` — no mocks.
+        """
+        policy = make_dummy_policy()
+        with torch.no_grad():
+            for param in policy.parameters():
+                param.add_(torch.randn_like(param))  # make every tensor distinct from init
+        reference = {key: tensor.clone() for key, tensor in policy.state_dict().items()}
+
+        # The direct-gather artifact (on a single process the gather is state_dict itself).
+        direct_dir = tmp_path / "direct"
+        policy.save_pretrained(direct_dir)
+
+        # The DCP artifact, laid out exactly as accelerate's save_fsdp_model writes it.
+        pretrained_dir = tmp_path / "checkpoint" / PRETRAINED_MODEL_DIR
+        dcp_dir = pretrained_dir / f"{FSDP_MODEL_NAME}_0"
+        dcp_dir.mkdir(parents=True)
+        dist_cp.save(
+            state_dict={"model": policy.state_dict()},
+            storage_writer=dist_cp.FileSystemWriter(str(dcp_dir)),
+        )
+
+        merged_file = dcp_to_safetensors(dcp_dir, pretrained_dir)
+        assert merged_file == pretrained_dir / "model.safetensors"
+        merged = load_file(merged_file)
+        direct = load_file(direct_dir / "model.safetensors")
+        assert set(merged) == set(direct) == set(reference)
+        for key, tensor in reference.items():
+            assert torch.equal(merged[key], tensor), key
+            assert torch.equal(direct[key], tensor), key
+            assert merged[key].dtype == tensor.dtype, key
@@ -0,0 +1,186 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Checkpoint save/resume round-trips on the non-sharded paths."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+from lerobot.common.train_utils import (
+    load_training_metadata,
+    resume_after_prepare,
+    resume_before_prepare,
+    save_checkpoint,
+)
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP
+from lerobot.utils.io_utils import load_json, write_json
+from tests.fixtures.dummy_checkpoint_policy import make_dummy_policy
+
+
+def make_cfg(**overrides) -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=3)
+    cfg.parallelism.resolve(1)
+    for name, value in overrides.items():
+        setattr(cfg, name, value)
+    return cfg
+
+
+def passthrough_accelerator() -> SimpleNamespace:
+    """The accelerator surface save/resume touches on non-sharded runs."""
+    return SimpleNamespace(unwrap_model=lambda m: m, wait_for_everyone=lambda: None)
+
+
+class TestSaveCheckpoint:
+    def test_non_sharded_layout(self, tmp_path):
+        policy = make_dummy_policy()
+        optimizer = torch.optim.Adam(policy.parameters())
+        save_checkpoint(
+            tmp_path,
+            step=7,
+            cfg=make_cfg(),
+            policy=policy,
+            optimizer=optimizer,
+            accelerator=passthrough_accelerator(),
+        )
+        pretrained = tmp_path / PRETRAINED_MODEL_DIR
+        state = tmp_path / TRAINING_STATE_DIR
+        assert (pretrained / "model.safetensors").is_file()
+        assert (pretrained / "config.json").is_file()
+        assert (pretrained / "train_config.json").is_file()
+        assert (state / TRAINING_STEP).is_file()
+        assert (state / "rng_state.safetensors").is_file()
+        assert (state / "optimizer_state.safetensors").is_file()
+        # single-file artifact, no index, weights intact
+        weights = load_file(pretrained / "model.safetensors")
+        assert torch.allclose(weights["net.weight"], torch.full_like(weights["net.weight"], 0.5))
+        assert not list(pretrained.glob("*.index.json"))
+
+    def test_training_step_records_topology(self, tmp_path):
+        cfg = make_cfg()
+        cfg.accelerator.gradient_accumulation.steps = 4
+        policy = make_dummy_policy()
+        save_checkpoint(
+            tmp_path,
+            step=11,
+            cfg=cfg,
+            policy=policy,
+            optimizer=torch.optim.Adam(policy.parameters()),
+            accelerator=passthrough_accelerator(),
+        )
+        metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR)
+        assert metadata["dp_world_size"] == 1
+        assert metadata["batch_size"] == 3
+        assert metadata["grad_accum_steps"] == 4
+
+    def test_dp_world_size_legacy_fallback(self, tmp_path):
+        """Pre-v0.7 checkpoints recorded num_processes; the reader falls back to it."""
+        state_dir = tmp_path / TRAINING_STATE_DIR
+        state_dir.mkdir(parents=True)
+        write_json({"step": 5, "num_processes": 4}, state_dir / TRAINING_STEP)
+        metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR)
+        assert metadata["dp_world_size"] == 4
+        assert metadata["batch_size"] is None
+
+
+class TestResume:
+    def _checkpointed_run(self, tmp_path):
+        policy = make_dummy_policy()
+        optimizer = torch.optim.Adam(policy.parameters(), lr=0.123)
+        # give the optimizer real state
+        policy.forward({"observation.state": torch.randn(2, 4)})[0].backward()
+        optimizer.step()
+        cfg = make_cfg()
+        save_checkpoint(
+            tmp_path,
+            step=42,
+            cfg=cfg,
+            policy=policy,
+            optimizer=optimizer,
+            accelerator=passthrough_accelerator(),
+        )
+        cfg.checkpoint_path = tmp_path
+        return cfg, policy, optimizer
+
+    def test_two_phase_resume_round_trip(self, tmp_path):
+        cfg, _, optimizer = self._checkpointed_run(tmp_path)
+        assert resume_before_prepare(cfg) == 42
+
+        fresh_policy = make_dummy_policy()
+        fresh_optimizer = torch.optim.Adam(fresh_policy.parameters(), lr=0.999)
+        resume_after_prepare(cfg, passthrough_accelerator(), fresh_policy, fresh_optimizer, None)
+        restored = fresh_optimizer.state_dict()
+        original = optimizer.state_dict()
+        assert restored["param_groups"][0]["lr"] == original["param_groups"][0]["lr"]
+        for key, tensor in original["state"][0].items():
+            assert torch.equal(restored["state"][0][key], tensor), key
+
+    def test_resume_warns_on_changed_cadence_and_topology(self, tmp_path, caplog):
+        """The recorded grad-accum factor and parallelism snapshot must be compared on
+        resume, with one warning naming the diff."""
+        import logging
+
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        cfg.accelerator.gradient_accumulation.steps = 4
+        cfg.parallelism.dp_replicate = 2  # same dp_world_size story is irrelevant here
+        with caplog.at_level(logging.WARNING):
+            assert resume_before_prepare(cfg) == 42
+        warning = next(m for m in caplog.messages if "differ from the checkpoint" in m)
+        assert "grad_accum_steps: 1 -> 4" in warning
+        assert "dp_replicate: 1 -> 2" in warning
+
+    def test_resume_unchanged_settings_stay_silent(self, tmp_path, caplog):
+        import logging
+
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        with caplog.at_level(logging.WARNING):
+            resume_before_prepare(cfg)
+        assert not [m for m in caplog.messages if "differ from the checkpoint" in m]
+
+    def test_resume_rejects_non_sharded_checkpoint_on_sharded_run(self, tmp_path):
+        """Resharding works across sizes, not across kinds: non-sharded -> sharded is rejected."""
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        cfg.parallelism.dp_shard = 2
+        with pytest.raises(ValueError, match="Cannot resume"):
+            resume_before_prepare(cfg)
+
+    def test_resume_rejects_sharded_checkpoint_on_non_sharded_run(self, tmp_path):
+        """The symmetric direction: a checkpoint recorded sharded cannot resume non-sharded."""
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        state_file = tmp_path / TRAINING_STATE_DIR / TRAINING_STEP
+        state = load_json(state_file)
+        state["parallelism"]["dp_shard"] = 2
+        write_json(state, state_file)
+        with pytest.raises(ValueError, match="Cannot resume"):
+            resume_before_prepare(cfg)
+
+    def test_resume_before_prepare_requires_training_state(self, tmp_path):
+        cfg = make_cfg()
+        cfg.checkpoint_path = tmp_path
+        with pytest.raises(NotADirectoryError):
+            resume_before_prepare(cfg)
+
+    def test_dcp_format_integrity_preflight(self, tmp_path):
+        """A checkpoint declaring DCP shards without the shard dir fails with the converter hint."""
+        pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+        cfg, policy, optimizer = self._checkpointed_run(tmp_path)
+        cfg.parallelism.dp_shard = 2  # pretend the recorded run was sharded
+        cfg.checkpoint_format = CheckpointFormat.DCP
+        with pytest.raises(FileNotFoundError, match="lerobot-convert-dcp"):
+            resume_after_prepare(cfg, passthrough_accelerator(), policy, optimizer, None)
@@ -0,0 +1,148 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""publish_trained_model: commit set, card, log-line contract, PEFT branch (hub fully mocked)."""
+
+import logging
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+import lerobot.common.train_utils as train_utils
+import lerobot.utils.hub as hub
+from lerobot.common.train_utils import generate_model_card, publish_trained_model
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TrainPipelineConfig
+from tests.fixtures.dummy_checkpoint_policy import make_dummy_policy
+
+
+class FakeHfApi:
+    """Records every repo/upload interaction; shared across both HfApi import sites."""
+
+    calls: list[dict] = []
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def create_repo(self, repo_id, private=None, exist_ok=False, **kwargs):
+        return SimpleNamespace(repo_id=repo_id)
+
+    def upload_folder(self, *, repo_id, folder_path, commit_message, **kwargs):
+        FakeHfApi.calls.append(
+            {
+                "repo_id": repo_id,
+                "commit_message": commit_message,
+                "files": sorted(p.name for p in Path(folder_path).iterdir()),
+                "ignore_patterns": kwargs.get("ignore_patterns"),
+            }
+        )
+        return SimpleNamespace(repo_url=SimpleNamespace(url=f"https://huggingface.co/{repo_id}"))
+
+
+@pytest.fixture
+def mocked_hub(monkeypatch):
+    FakeHfApi.calls = []
+    monkeypatch.setattr(train_utils, "HfApi", FakeHfApi)
+    monkeypatch.setattr(hub, "HfApi", FakeHfApi)
+    # card.validate() hits the Hub; publishing must work offline in tests
+    monkeypatch.setattr(train_utils.ModelCard, "validate", lambda self: None)
+    return FakeHfApi
+
+
+def make_cfg() -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="user/dataset"))
+    cfg.parallelism.resolve(1)
+    return cfg
+
+
+class RecordingProcessor:
+    def __init__(self):
+        self.pushed_to = None
+
+    def push_to_hub(self, repo_id, **kwargs):
+        self.pushed_to = repo_id
+
+
+class TestPublishTrainedModel:
+    def test_commit_set_and_log_contract(self, mocked_hub, caplog):
+        policy = make_dummy_policy(repo_id="user/policy")
+        pre, post = RecordingProcessor(), RecordingProcessor()
+        with caplog.at_level(logging.INFO):
+            publish_trained_model(make_cfg(), policy, pre, post, dataset_meta=None)
+
+        # commit 1: the model through HubMixin (config.json + model.safetensors in a tmpdir)
+        model_commit = mocked_hub.calls[0]
+        assert {"config.json", "model.safetensors"} <= set(model_commit["files"])
+        # commits 2-3: processors
+        assert pre.pushed_to == "user/policy" and post.pushed_to == "user/policy"
+        # commit 4: the bundle sidecar
+        bundle = mocked_hub.calls[-1]
+        assert {"README.md", "train_config.json"} <= set(bundle["files"])
+        # the exact line lerobot.jobs.hf watches to end remote runs early
+        assert any(
+            m.startswith("Model pushed to https://huggingface.co/user/policy") for m in caplog.messages
+        )
+
+    def test_peft_branch_skips_model_commit(self, mocked_hub):
+        policy = make_dummy_policy(repo_id="user/policy")
+
+        class FakePeftModel:
+            def save_pretrained(self, path):
+                (Path(path) / "adapter_model.safetensors").write_bytes(b"x")
+
+        publish_trained_model(make_cfg(), policy, None, None, dataset_meta=None, peft_model=FakePeftModel())
+        assert len(mocked_hub.calls) == 1  # only the bundle commit
+        bundle = mocked_hub.calls[0]
+        # adapter weights + the wrapped policy's config + card + train config, no full weights
+        assert {"README.md", "adapter_model.safetensors", "config.json", "train_config.json"} <= set(
+            bundle["files"]
+        )
+        assert "model.safetensors" not in bundle["files"]
+
+    def test_missing_repo_id_fails_loudly(self, mocked_hub):
+        policy = make_dummy_policy(repo_id=None)
+        with pytest.raises(ValueError, match="repo id"):
+            publish_trained_model(make_cfg(), policy, None, None, dataset_meta=None)
+
+
+class TestGenerateModelCard:
+    def test_free_function_renders_from_arguments(self, monkeypatch):
+        monkeypatch.setattr(train_utils.ModelCard, "validate", lambda self: None)
+        policy = make_dummy_policy(repo_id="user/policy")
+        card = generate_model_card(policy.config, cfg=make_cfg(), dataset_meta=None)
+        assert card.data.library_name == "lerobot"
+        assert card.data.datasets == "user/dataset"
+        assert "lerobot" in card.data.tags
+
+
+class TestDeprecatedPushModelToHub:
+    """`push_model_to_hub` stays callable for external scripts, delegating to the publisher."""
+
+    def test_policy_shim_warns_and_publishes(self, mocked_hub):
+        policy = make_dummy_policy(repo_id="user/policy")
+        with pytest.warns(FutureWarning, match="push_model_to_hub is deprecated"):
+            policy.push_model_to_hub(make_cfg())
+
+        # Same artifacts the method produced before: weights + config, then card + train config.
+        model_commit = mocked_hub.calls[0]
+        assert {"config.json", "model.safetensors"} <= set(model_commit["files"])
+        bundle = mocked_hub.calls[-1]
+        assert {"README.md", "train_config.json"} <= set(bundle["files"])
+
+    def test_policy_shim_warns_that_state_dict_is_ignored(self, mocked_hub):
+        policy = make_dummy_policy(repo_id="user/policy")
+        with pytest.warns(FutureWarning, match="`state_dict` argument is ignored"):
+            policy.push_model_to_hub(make_cfg(), state_dict=policy.state_dict())
@@ -0,0 +1,135 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import draccus
+import pytest
+
+from lerobot.configs.accelerator import (
+    AcceleratorConfig,
+    ActivationCheckpointingConfig,
+    ActivationCheckpointingMode,
+    CompileConfig,
+    DDPConfig,
+    FSDPConfig,
+    GradientAccumulationConfig,
+)
+from lerobot.configs.parallelism import ParallelismConfig
+
+
+class TestFieldValidation:
+    def test_wrap_policies_mutually_exclusive(self):
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            FSDPConfig(wrap_modules=["Block"], min_num_params=1000)
+
+    def test_min_num_params_positive(self):
+        with pytest.raises(ValueError, match="min_num_params"):
+            FSDPConfig(min_num_params=0)
+
+    def test_mixed_precision_choices(self):
+        with pytest.raises(ValueError, match="mixed_precision"):
+            AcceleratorConfig(mixed_precision="tf32")
+
+    def test_gradient_accumulation_positive(self):
+        with pytest.raises(ValueError, match="gradient_accumulation.steps"):
+            GradientAccumulationConfig(steps=0)
+
+
+class TestDraccusRoundTrip:
+    @pytest.mark.parametrize(
+        "cfg",
+        [
+            AcceleratorConfig(),
+            AcceleratorConfig(
+                mixed_precision="bf16",
+                gradient_accumulation=GradientAccumulationConfig(steps=4),
+                fsdp=FSDPConfig(
+                    reshard_after_forward=False,
+                    wrap_modules=["ACTEncoderLayer", "ACTDecoderLayer"],
+                    cpu_offload=True,
+                    ignored_modules=r".*pos_embed.*",
+                ),
+                ddp=DDPConfig(find_unused_parameters=False, static_graph=True),
+                compile=CompileConfig(enabled=True, mode="max-autotune", regional=False),
+                activation_checkpointing=ActivationCheckpointingConfig(mode=ActivationCheckpointingMode.FULL),
+            ),
+            AcceleratorConfig(fsdp=FSDPConfig(min_num_params=1_000_000)),
+        ],
+    )
+    def test_encode_json_decode_identity(self, cfg):
+        payload = json.loads(json.dumps(draccus.encode(cfg)))
+        assert draccus.decode(AcceleratorConfig, payload) == cfg
+
+    def test_pre_existing_config_without_fields_gets_defaults(self):
+        assert draccus.decode(AcceleratorConfig, {}) == AcceleratorConfig()
+
+
+class TestRuntimeBuilders:
+    """The mirrors must translate into real accelerate objects (plugins built lazily)."""
+
+    @pytest.fixture(autouse=True)
+    def _requires_accelerate(self):
+        pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+    def test_fsdp_plugin_translation(self):
+        plugin = FSDPConfig(
+            reshard_after_forward=False, wrap_modules=["MyBlock"], cpu_offload=True
+        ).build_plugin()
+        assert plugin.fsdp_version == 2
+        assert plugin.reshard_after_forward is False
+        assert plugin.transformer_cls_names_to_wrap == ["MyBlock"]
+        # bools are normalized into torch offload policies by the plugin itself
+        assert type(plugin.cpu_offload).__name__ == "CPUOffloadPolicy"
+        # LeRobot never switches state_dict_type: FSDP2's SHARDED default must hold
+        assert plugin.state_dict_type.name == "SHARDED_STATE_DICT"
+        assert not plugin.activation_checkpointing
+
+    def test_fsdp_plugin_size_based_policy(self):
+        plugin = FSDPConfig(min_num_params=1024).build_plugin()
+        assert plugin.min_num_params == 1024
+        assert plugin.transformer_cls_names_to_wrap is None
+
+    def test_ddp_kwargs_translation(self):
+        handler = DDPConfig(find_unused_parameters=False, gradient_as_bucket_view=True).build_kwargs_handler()
+        assert handler.find_unused_parameters is False
+        assert handler.gradient_as_bucket_view is True
+
+    def test_gradient_accumulation_plugin_translation(self):
+        plugin = GradientAccumulationConfig(steps=4).build_plugin()
+        assert plugin.num_steps == 4
+        assert plugin.sync_with_dataloader is False
+
+    def test_gradient_accumulation_never_syncs_with_dataloader(self, monkeypatch):
+        """The loop cycles a finite dataloader, so accelerate's default
+        sync_with_dataloader=True would force an optimizer step at every dataset epoch
+        boundary instead of every num_steps micro-batches."""
+        captured = {}
+
+        class FakeAccelerator:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        monkeypatch.setattr("accelerate.Accelerator", FakeAccelerator)
+        parallelism = ParallelismConfig()
+        parallelism.resolve(1)
+        AcceleratorConfig(gradient_accumulation=GradientAccumulationConfig(steps=4)).build(
+            parallelism, cpu=True
+        )
+        ga_plugin = captured["gradient_accumulation_plugin"]
+        assert ga_plugin.num_steps == 4
+        assert ga_plugin.sync_with_dataloader is False
+        assert "gradient_accumulation_steps" not in captured
@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import draccus
+import pytest
+
+from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig
+
+
+class TestResolve:
+    def test_single_process_defaults(self):
+        cfg = ParallelismConfig()
+        cfg.resolve(1)
+        assert (cfg.dp_replicate, cfg.dp_shard) == (1, 1)
+        assert not cfg.is_sharded and not cfg.is_replicated_only
+        assert cfg.dp_world_size == 1
+
+    def test_untouched_config_fills_ddp(self):
+        """Plain `torchrun --nproc-per-node=8` with a default config resolves to DDP."""
+        cfg = ParallelismConfig()
+        cfg.resolve(8)
+        assert cfg.dp_replicate == 8
+        assert cfg.is_replicated_only and not cfg.is_sharded
+        assert cfg.dp_world_size == 8
+
+    def test_full_shard_sentinel(self):
+        cfg = ParallelismConfig(dp_shard=-1)
+        assert cfg.is_sharded  # sharded even before resolve: -1 is an explicit opt-in
+        cfg.resolve(8)
+        assert cfg.dp_shard == 8 and cfg.dp_replicate == 1
+
+    def test_hsdp_sentinel_infers_shard(self):
+        cfg = ParallelismConfig(dp_replicate=2, dp_shard=-1)
+        cfg.resolve(8)
+        assert (cfg.dp_replicate, cfg.dp_shard) == (2, 4)
+        assert cfg.dp_world_size == 8
+
+    def test_explicit_hsdp(self):
+        cfg = ParallelismConfig(dp_replicate=2, dp_shard=4)
+        cfg.resolve(8)
+        assert cfg.is_sharded and not cfg.is_replicated_only
+
+    def test_product_mismatch_lists_all_degrees(self):
+        cfg = ParallelismConfig(dp_replicate=2, dp_shard=2)
+        with pytest.raises(ValueError, match=r"dp_replicate=2 \* dp_shard=2.*WORLD_SIZE=8"):
+            cfg.resolve(8)
+
+    def test_explicit_replicate_must_match_world(self):
+        cfg = ParallelismConfig(dp_replicate=4)
+        with pytest.raises(ValueError, match="WORLD_SIZE=8"):
+            cfg.resolve(8)
+
+    def test_sentinel_indivisible_world(self):
+        cfg = ParallelismConfig(dp_replicate=3, dp_shard=-1)
+        with pytest.raises(ValueError, match="not divisible"):
+            cfg.resolve(8)
+
+    def test_cp_fails_fast(self):
+        cfg = ParallelismConfig(dp_shard=-1, context_parallel=ContextParallelConfig(ulysses_degree=2))
+        with pytest.raises(ValueError, match="not implemented"):
+            cfg.resolve(8)
+
+
+class TestFieldValidation:
+    @pytest.mark.parametrize("kwargs", [{"dp_replicate": 0}, {"dp_shard": 0}, {"dp_shard": -2}])
+    def test_bad_dp_degrees(self, kwargs):
+        with pytest.raises(ValueError):
+            ParallelismConfig(**kwargs)
+
+    def test_cfg_parallel_capped_at_two(self):
+        ParallelismConfig(cfg_parallel=2)  # reserved but representable
+        with pytest.raises(ValueError, match="cfg_parallel"):
+            ParallelismConfig(cfg_parallel=3)
+
+    @pytest.mark.parametrize("kwargs", [{"ring_degree": 0}, {"ulysses_degree": -1}])
+    def test_bad_cp_degrees(self, kwargs):
+        with pytest.raises(ValueError):
+            ContextParallelConfig(**kwargs)
+
+    def test_dp_world_size_undefined_before_resolve(self):
+        with pytest.raises(RuntimeError, match="resolve"):
+            _ = ParallelismConfig(dp_shard=-1).dp_world_size
+
+
+class TestDraccusRoundTrip:
+    @pytest.mark.parametrize(
+        "cfg",
+        [
+            ParallelismConfig(),
+            ParallelismConfig(dp_replicate=2, dp_shard=4, cfg_parallel=2),
+            ParallelismConfig(
+                dp_shard=-1,
+                context_parallel=ContextParallelConfig(ring_degree=2, ulysses_degree=4),
+            ),
+        ],
+    )
+    def test_encode_json_decode_identity(self, cfg):
+        payload = json.loads(json.dumps(draccus.encode(cfg)))
+        assert draccus.decode(ParallelismConfig, payload) == cfg
+
+    def test_pre_existing_config_without_fields_gets_defaults(self):
+        """Checkpoints written before this feature parse with default topology."""
+        assert draccus.decode(ParallelismConfig, {}) == ParallelismConfig()
@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TrainPipelineConfig integration for the distributed fields: fail-fasts + config compat."""
+
+import draccus
+import pytest
+
+from lerobot.configs.accelerator import ActivationCheckpointingMode
+from lerobot.configs.default import DatasetConfig, PeftConfig
+from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig
+from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig
+from lerobot.optim.optimizers import AdamConfig, MultiAdamConfig
+
+
+def make_cfg(**overrides) -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"))
+    for name, value in overrides.items():
+        setattr(cfg, name, value)
+    return cfg
+
+
+def sharded() -> ParallelismConfig:
+    return ParallelismConfig(dp_shard=-1)
+
+
+class TestDistributedFailFasts:
+    def test_defaults_pass(self):
+        make_cfg()._validate_distributed()
+
+    def test_cp_reserved(self):
+        cfg = make_cfg(parallelism=ParallelismConfig(context_parallel=ContextParallelConfig(ring_degree=2)))
+        with pytest.raises(ValueError, match="not implemented"):
+            cfg._validate_distributed()
+
+    def test_cfg_parallel_training_rejected(self):
+        cfg = make_cfg(parallelism=ParallelismConfig(cfg_parallel=2))
+        with pytest.raises(ValueError, match="inference-only"):
+            cfg._validate_distributed()
+
+    def test_compile_placeholder(self):
+        cfg = make_cfg()
+        cfg.accelerator.compile.enabled = True
+        with pytest.raises(ValueError, match="compile"):
+            cfg._validate_distributed()
+
+    def test_activation_checkpointing_placeholder(self):
+        cfg = make_cfg()
+        cfg.accelerator.activation_checkpointing.mode = ActivationCheckpointingMode.FULL
+        with pytest.raises(ValueError, match="activation_checkpointing"):
+            cfg._validate_distributed()
+
+    def test_dcp_format_requires_sharding(self):
+        cfg = make_cfg(checkpoint_format=CheckpointFormat.DCP)
+        with pytest.raises(ValueError, match="sharded"):
+            cfg._validate_distributed()
+        cfg.parallelism = sharded()
+        cfg._validate_distributed()
+
+    def test_fp16_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded())
+        cfg.accelerator.mixed_precision = "fp16"
+        with pytest.raises(ValueError, match="fp16"):
+            cfg._validate_distributed()
+        cfg.accelerator.mixed_precision = "bf16"
+        cfg._validate_distributed()
+
+    def test_peft_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded(), peft=PeftConfig())
+        with pytest.raises(ValueError, match="PEFT"):
+            cfg._validate_distributed()
+
+    def test_env_eval_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded(), env_eval_freq=1000)
+        cfg.env = object()  # any configured env triggers the check
+        with pytest.raises(ValueError, match="environment evaluation"):
+            cfg._validate_distributed()
+
+    def test_multi_optimizer_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded(), optimizer=MultiAdamConfig())
+        with pytest.raises(ValueError, match="Multi-optimizer"):
+            cfg._validate_distributed()
+        cfg.optimizer = AdamConfig()
+        cfg._validate_distributed()
+
+
+class TestConfigCompat:
+    def test_checkpoint_format_round_trip(self):
+        for fmt in CheckpointFormat:
+            assert draccus.decode(CheckpointFormat, draccus.encode(fmt)) is fmt
+
+    def test_wants_predicates(self):
+        assert CheckpointFormat.SAFETENSORS.wants_safetensors
+        assert not CheckpointFormat.SAFETENSORS.wants_dcp
+        assert CheckpointFormat.DCP.wants_dcp and not CheckpointFormat.DCP.wants_safetensors
+        both = CheckpointFormat.SAFETENSORS_AND_DCP
+        assert both.wants_safetensors and both.wants_dcp
+
+
+def test_reward_model_rejected_when_sharded():
+    """Sharded reward runs previously failed late (missing wrap
+    units, DTensor serialization at the first checkpoint) instead of at validation."""
+    cfg = make_cfg(parallelism=sharded())
+    cfg.reward_model = object()  # any configured reward model triggers the check
+    with pytest.raises(ValueError, match="Reward-model"):
+        cfg._validate_distributed()
@@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.

+from types import SimpleNamespace
+
 import numpy as np
 import pytest
+import torch

 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")

@@ -24,7 +27,9 @@ from lerobot.scripts.augment_dataset_quantile_stats import (


 def _numeric_keys(dataset):
-    return [k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string")]
+    return [
+        k for k, v in dataset.features.items() if v["dtype"] not in ("image", "video", "string", "language")
+    ]


 def _image_keys(dataset):
@@ -102,3 +107,112 @@ def test_quantile_stats_present_after_compute(tmp_path, lerobot_dataset_factory)
    )
    stats = compute_quantile_stats_for_dataset(dataset, use_sampling=True)
    assert has_quantile_stats(stats)
+
+
+class FakeHFDataset:
+    """Minimal stand-in exposing the column slicing used by the augment script."""
+
+    def __init__(self, columns: dict[str, list]):
+        self._columns = columns
+
+    def select_columns(self, keys):
+        return FakeHFDataset({key: self._columns[key] for key in keys})
+
+    def __getitem__(self, index):
+        return {key: values[index] for key, values in self._columns.items()}
+
+
+def test_compute_quantile_stats_skips_language_features():
+    class FakeDataset:
+        num_episodes = 1
+        features = {
+            "action": {"dtype": "float32"},
+            "observation.language": {"dtype": "language"},
+        }
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
+        hf_dataset = FakeHFDataset(
+            {
+                "action": [[0.0], [1.0]],
+                "observation.language": [
+                    [{"role": "user", "content": "pick"}],
+                    [{"role": "assistant", "content": "done"}],
+                ],
+            }
+        )
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset())
+
+    assert set(stats) == {"action"}
+
+
+def test_compute_quantile_stats_skip_images_avoids_decoding():
+    class FakeDataset:
+        num_episodes = 1
+        features = {
+            "action": {"dtype": "float32"},
+            "observation.images.cam": {"dtype": "video"},
+        }
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
+        hf_dataset = FakeHFDataset({"action": [[0.0], [1.0]]})
+
+        def __getitem__(self, index):
+            raise AssertionError(f"video frame {index} was decoded despite skip_images=True")
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset(), skip_images=True)
+
+    assert set(stats) == {"action"}
+
+
+def test_compute_quantile_stats_handles_single_frame():
+    class FakeDataset:
+        num_episodes = 1
+        features = {"action": {"dtype": "float32"}}
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 1}])
+        hf_dataset = FakeHFDataset({"action": [[5.0, 7.0]]})
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset())
+
+    np.testing.assert_array_equal(stats["action"]["count"], np.array([1]))
+    for key in ("min", "max", "mean", "q01", "q10", "q50", "q90", "q99"):
+        np.testing.assert_allclose(stats["action"][key], np.array([5.0, 7.0]))
+
+
+def test_compute_quantile_stats_image_count_uses_frames():
+    frames = [torch.zeros(3, 2, 2), torch.ones(3, 2, 2)]
+
+    class FakeDataset:
+        num_episodes = 1
+        features = {"observation.images.cam": {"dtype": "video"}}
+        meta = SimpleNamespace(episodes=[{"dataset_from_index": 0, "dataset_to_index": 2}])
+        hf_dataset = FakeHFDataset({})
+
+        def __getitem__(self, index):
+            return {"observation.images.cam": frames[index]}
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset(), use_sampling=False)
+    image_stats = stats["observation.images.cam"]
+
+    np.testing.assert_array_equal(image_stats["count"], np.array([2]))
+    assert image_stats["mean"].shape == (3, 1, 1)
+    np.testing.assert_allclose(image_stats["mean"], np.full((3, 1, 1), 0.5))
+
+
+def test_compute_quantile_stats_accumulates_across_episodes():
+    values = [[float(value)] for value in range(100)] + [[float(value)] for value in range(1000, 1010)]
+
+    class FakeDataset:
+        num_episodes = 2
+        features = {"action": {"dtype": "float32"}}
+        meta = SimpleNamespace(
+            episodes=[
+                {"dataset_from_index": 0, "dataset_to_index": 100},
+                {"dataset_from_index": 100, "dataset_to_index": 110},
+            ]
+        )
+        hf_dataset = FakeHFDataset({"action": values})
+
+    stats = compute_quantile_stats_for_dataset(FakeDataset())
+
+    np.testing.assert_array_equal(stats["action"]["count"], np.array([110]))
+    expected_q90 = np.percentile(np.asarray(values), 90, axis=0)
+    np.testing.assert_allclose(stats["action"]["q90"], expected_q90, atol=0.1)
@@ -688,7 +688,7 @@ def test_compute_episode_stats_string_features_skipped():


 def test_aggregate_feature_stats_with_quantiles():
-    """Test aggregating feature stats that include quantiles."""
+    """Test aggregating feature stats that include quantiles uses conservative bounds."""
    stats_ft_list = [
        {
            "min": np.array([1.0]),
@@ -697,6 +697,9 @@ def test_aggregate_feature_stats_with_quantiles():
            "std": np.array([2.0]),
            "count": np.array([100]),
            "q01": np.array([1.5]),
+            "q10": np.array([2.0]),
+            "q50": np.array([5.0]),
+            "q90": np.array([9.0]),
            "q99": np.array([9.5]),
        },
        {
@@ -706,22 +709,21 @@ def test_aggregate_feature_stats_with_quantiles():
            "std": np.array([2.5]),
            "count": np.array([150]),
            "q01": np.array([2.5]),
+            "q10": np.array([3.0]),
+            "q50": np.array([6.0]),
+            "q90": np.array([11.0]),
            "q99": np.array([11.5]),
        },
    ]

    result = aggregate_feature_stats(stats_ft_list)

-    # Should preserve quantiles
-    assert "q01" in result
-    assert "q99" in result
-
-    # Verify quantile aggregation (weighted average)
-    expected_q01 = (1.5 * 100 + 2.5 * 150) / 250  # ≈ 2.1
-    expected_q99 = (9.5 * 100 + 11.5 * 150) / 250  # ≈ 10.7
-
-    np.testing.assert_allclose(result["q01"], np.array([expected_q01]), atol=1e-6)
-    np.testing.assert_allclose(result["q99"], np.array([expected_q99]), atol=1e-6)
+    # Lower quantiles use min; upper quantiles use max, regardless of counts.
+    np.testing.assert_allclose(result["q01"], np.array([1.5]), atol=1e-6)
+    np.testing.assert_allclose(result["q10"], np.array([2.0]), atol=1e-6)
+    np.testing.assert_allclose(result["q50"], np.array([5.0]), atol=1e-6)
+    np.testing.assert_allclose(result["q90"], np.array([11.0]), atol=1e-6)
+    np.testing.assert_allclose(result["q99"], np.array([11.5]), atol=1e-6)


 def test_aggregate_stats_mixed_quantiles():
@@ -878,3 +880,60 @@ def test_fixed_quantiles_always_computed():
        for q_key in expected_quantiles:
            assert q_key in episode_stats[key]
            assert episode_stats[key][q_key].shape == (features[key]["shape"][0],)
+
+
+def test_aggregate_stats_incremental_resume():
+    """Verify conservative bounds remain associative across incremental additions."""
+    # Start with episode 1 stats (narrow distribution)
+    ep1_stats = {
+        "action": {
+            "min": np.array([-10.0, -5.0]),
+            "max": np.array([10.0, 5.0]),
+            "mean": np.array([0.0, 0.0]),
+            "std": np.array([3.0, 1.5]),
+            "count": np.array([500]),
+            "q01": np.array([-9.0, -4.5]),
+            "q99": np.array([9.0, 4.5]),
+        },
+    }
+
+    # Episode 2: wider distribution on dim 0
+    ep2_stats = {
+        "action": {
+            "min": np.array([-30.0, -5.0]),
+            "max": np.array([40.0, 6.0]),
+            "mean": np.array([5.0, 0.5]),
+            "std": np.array([15.0, 2.0]),
+            "count": np.array([100]),
+            "q01": np.array([-25.0, -4.0]),
+            "q99": np.array([35.0, 5.5]),
+        },
+    }
+
+    # First aggregation: ep1 + ep2 (simulates save_episode for ep2)
+    cumulative = aggregate_stats([ep1_stats, ep2_stats])
+
+    # q01 should take min (conservative lower bound)
+    np.testing.assert_allclose(cumulative["action"]["q01"], np.array([-25.0, -4.5]))
+    # q99 should take max (conservative upper bound)
+    np.testing.assert_allclose(cumulative["action"]["q99"], np.array([35.0, 5.5]))
+
+    # Episode 3: even wider on dim 1
+    ep3_stats = {
+        "action": {
+            "min": np.array([-8.0, -20.0]),
+            "max": np.array([8.0, 25.0]),
+            "mean": np.array([0.0, 3.0]),
+            "std": np.array([2.0, 8.0]),
+            "count": np.array([50]),
+            "q01": np.array([-7.0, -18.0]),
+            "q99": np.array([7.0, 22.0]),
+        },
+    }
+
+    # Second aggregation: cumulative + ep3 (simulates save_episode for ep3)
+    cumulative2 = aggregate_stats([cumulative, ep3_stats])
+
+    # Bounds should widen monotonically
+    np.testing.assert_allclose(cumulative2["action"]["q01"], np.array([-25.0, -18.0]))
+    np.testing.assert_allclose(cumulative2["action"]["q99"], np.array([35.0, 22.0]))
@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Version canaries for the accelerate/torch seams LeRobot's distributed engine relies on.
+
+LeRobot deliberately builds on a few accelerate internals that are not covered by a public
+stability promise. These tests exist to fail LOUDLY on a
+dependency upgrade — on a CPU runner, before any distributed job can be corrupted — whenever one
+of those seams moves. If a canary fails, re-audit the corresponding integration seam before bumping
+the pin; do not simply update the assertion.
+"""
+
+import inspect
+
+import pytest
+
+pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+
+def test_fsdp_checkpoint_name_constants():
+    """Checkpoint dir names are imported from accelerate; the on-disk layout depends on them."""
+    from accelerate.utils.constants import FSDP_MODEL_NAME, OPTIMIZER_NAME
+
+    assert FSDP_MODEL_NAME == "pytorch_model_fsdp"
+    assert OPTIMIZER_NAME == "optimizer"
+
+
+def test_parallelism_config_mesh_dim_contract():
+    """FSDP2 shards over the flattened dp_shard_cp dim; the dataloader keys on exact root names."""
+    from accelerate.parallelism_config import ParallelismConfig
+
+    pc = ParallelismConfig(dp_replicate_size=2, dp_shard_size=2, cp_size=2)
+    assert pc.fsdp_dim_names == ["dp_replicate", "dp_shard_cp"]
+    assert pc.dp_shard_cp_dim_names == ["dp_shard", "cp"]
+    assert pc.dp_cp_dim_names == ["dp_replicate", "dp_shard", "cp"]
+    # Degenerate FSDP-only case still shards over the flattened name.
+    pc_fsdp = ParallelismConfig(dp_replicate_size=1, dp_shard_size=4)
+    assert pc_fsdp.fsdp_dim_names == ["dp_shard_cp"]
+
+
+def test_accelerator_accepts_parallelism_config():
+    from accelerate import Accelerator
+
+    params = inspect.signature(Accelerator.__init__).parameters
+    assert "parallelism_config" in params
+    assert "fsdp_plugin" in params
+    assert "gradient_accumulation_plugin" in params
+
+
+def test_dataloader_is_mesh_aware():
+    """prepare_data_loader must accept the device mesh that makes CP peers share batches."""
+    from accelerate.data_loader import prepare_data_loader
+
+    assert "torch_device_mesh" in inspect.signature(prepare_data_loader).parameters
+
+
+def test_cp_mask_stripping_hook_seam():
+    """finalize_sharded_policy strips this exact hook.
+
+    If accelerate renames or moves it, the strip becomes a silent no-op and CP training would
+    inherit mask-corrupting hooks — hence a canary rather than a runtime hasattr.
+    """
+    from accelerate.big_modeling import _attach_context_parallel_hooks
+
+    assert callable(_attach_context_parallel_hooks)
+    assert _attach_context_parallel_hooks.__module__ == "accelerate.big_modeling"
+
+
+def test_fsdp_plugin_mirrored_fields_exist():
+    """AcceleratorConfig mirrors a plain-typed subset of the plugin; the fields must survive."""
+    from accelerate.utils import FullyShardedDataParallelPlugin
+
+    fields = {f.name for f in FullyShardedDataParallelPlugin.__dataclass_fields__.values()}
+    assert {
+        "fsdp_version",
+        "reshard_after_forward",
+        "auto_wrap_policy",
+        "transformer_cls_names_to_wrap",
+        "min_num_params",
+        "cpu_offload",
+        "ignored_modules",
+        "activation_checkpointing",
+        "state_dict_type",
+    } <= fields
+
+
+def test_merge_fsdp_weights_signature():
+    """The DCP->safetensors converter is a thin wrapper over this accelerate utility."""
+    from accelerate.utils import merge_fsdp_weights
+
+    params = inspect.signature(merge_fsdp_weights).parameters
+    assert {"checkpoint_dir", "output_path", "safe_serialization"} <= set(params)
+
+
+def test_fsdp_save_load_helpers_exist():
+    from accelerate.utils import (
+        load_fsdp_model,
+        load_fsdp_optimizer,
+        save_fsdp_model,
+        save_fsdp_optimizer,
+    )
+
+    for fn in (save_fsdp_model, load_fsdp_model, save_fsdp_optimizer, load_fsdp_optimizer):
+        assert callable(fn)
+
+
+def test_torch_fsdp2_seams():
+    """isinstance(FSDPModule) discrimination + non-forward entry registration + full gather."""
+    from torch.distributed.checkpoint.state_dict import (
+        StateDictOptions,
+        get_model_state_dict,  # noqa: F401
+    )
+    from torch.distributed.fsdp import FSDPModule, register_fsdp_forward_method  # noqa: F401
+
+    options = inspect.signature(StateDictOptions).parameters
+    assert {"full_state_dict", "cpu_offload"} <= set(options)
+
+
+def test_accelerate_version_floor():
+    import accelerate
+    from packaging import version
+
+    if version.parse(accelerate.__version__) < version.parse("1.14.0"):
+        pytest.fail(
+            f"accelerate {accelerate.__version__} < 1.14.0: the FSDP2 auto-wrap fallback fix "
+            "(#3999) and the bf16->fp32 master-weight upcast this design relies on are absent."
+        )
@@ -0,0 +1,70 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The DCP wrappers must hand accelerate exact shard directories.
+
+accelerate 1.14 resolves the load directory with a substring check ("optimizer" /
+"pytorch_model_fsdp" in the path -> use as-is) while the save side joins the shard name
+unconditionally, so a run path like `--job_name=optimizer_sweep` would save to
+`training_state/optimizer_0/` but load from `training_state/` itself. Passing the exact
+shard dir makes the containment check deterministically a no-op.
+"""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+
+def fake_accelerator() -> SimpleNamespace:
+    return SimpleNamespace(state=SimpleNamespace(fsdp_plugin=object()))
+
+
+# A parent path that trips both of accelerate's substring checks at once.
+POISONED_PARENT = Path("/outputs/train/optimizer_sweep_pytorch_model_fsdp_repro/training_state")
+
+
+def test_load_sharded_optimizer_passes_exact_shard_dir(monkeypatch):
+    import accelerate.utils
+
+    from lerobot.distributed.checkpoint import load_sharded_optimizer
+
+    seen = {}
+    monkeypatch.setattr(
+        accelerate.utils,
+        "load_fsdp_optimizer",
+        lambda plugin, accelerator, optimizer, model, input_dir: seen.update(path=input_dir),
+    )
+    load_sharded_optimizer(fake_accelerator(), optimizer=object(), model=object(), input_dir=POISONED_PARENT)
+    assert seen["path"] == str(POISONED_PARENT / "optimizer_0")
+    assert isinstance(seen["path"], str)  # str, never Path (accelerate does string checks)
+
+
+def test_load_sharded_model_passes_exact_shard_dir(monkeypatch):
+    import accelerate.utils
+
+    from lerobot.distributed.checkpoint import load_sharded_model
+
+    seen = {}
+    monkeypatch.setattr(
+        accelerate.utils,
+        "load_fsdp_model",
+        lambda plugin, accelerator, model, input_dir: seen.update(path=input_dir),
+    )
+    load_sharded_model(fake_accelerator(), model=object(), input_dir=POISONED_PARENT)
+    assert seen["path"] == str(POISONED_PARENT / "pytorch_model_fsdp_0")
+    assert isinstance(seen["path"], str)
@@ -0,0 +1,486 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""End-to-end multi-GPU tests for the distributed core.
+
+Sized for a 4-GPU CI lane, these tests execute the sharded code paths nothing else in the tree can
+reach — ``fully_shard`` via ``accelerator.prepare``, the DCP branches of ``save_checkpoint`` /
+``save_training_state`` / ``resume_after_prepare``, the collective gather inside
+``save_pretrained``, and HSDP/DDP gradient reduction — against the tiny
+``DummyCheckpointPolicy`` fixture on synthetic data (no datasets, no network, no site paths).
+
+Run on a node with at least 4 GPUs::
+
+    pytest -m multigpu tests/distributed/test_multigpu_training.py -v
+
+Mechanics:
+
+- Plain pytest, no ``torchrun``: each test launches its own ranks with
+  ``torch.multiprocessing.spawn`` (spawn start method) and a per-test free TCP port; workers set
+  the torchrun-equivalent env (``RANK``/``LOCAL_RANK``/``WORLD_SIZE``/``MASTER_*``) that
+  accelerate's ``env://`` initialization consumes.
+- Deadlock watchdog (:func:`_spawn`): the spawn context is polled with a deadline instead of a
+  blocking join, so a hung collective — the exact failure mode the all-ranks contracts guard
+  against — fails the test with ``TimeoutError`` (all workers SIGKILLed) rather than hanging CI.
+  A worker exception propagates through ``ProcessContext.join``, which tears down the survivors.
+- Workers configure accelerate exclusively through the LeRobot config mirrors
+  (``AcceleratorConfig.build(ParallelismConfig)`` after ``resolve(world_size)``) — the same
+  construction path ``make_accelerator`` takes; see :func:`_build_accelerator` for why the
+  factory itself is not called.
+- Without GPUs every test skips (``torch.cuda.device_count()`` gate), so the file is safe to
+  collect and run in the CPU lanes.
+"""
+
+import json
+import os
+import socket
+import time
+from pathlib import Path
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from safetensors.torch import load_file
+
+from lerobot.common.train_utils import resume_after_prepare, resume_before_prepare, save_checkpoint
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig
+from lerobot.distributed.checkpoint import full_model_state_dict, is_sharded_module
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR
+
+# The spawned children re-import this module by name, so this import must resolve there too:
+# torch.multiprocessing propagates the parent's sys.path through the spawn preparation data.
+from tests.fixtures.dummy_checkpoint_policy import DummyCheckpointConfig, DummyCheckpointPolicy
+
+SEED = 20260712
+HIDDEN = 8  # DummyCheckpointPolicy is one Linear(hidden, hidden): 4 ranks shard dim 0 evenly
+BATCH_SIZE = 2
+SAVE_STEP = 2  # optimizer steps run before saving in the round-trip workers
+PARITY_STEPS = 3
+GA_UPDATES = 3
+SAMPLES_PER_UPDATE = 4  # per rank per optimizer update — the fixed effective batch of test 5
+GRAD_CLIP_NORM = 100.0  # generous: exercises the clip call without perturbing parity
+# Generous headroom for cold NCCL init plus the lerobot re-import in 4 spawned children, while
+# still bounding a deadlocked collective to minutes instead of a hung CI job.
+WATCHDOG_TIMEOUT_S = 240.0
+_JOIN_POLL_S = 5.0
+
+
+# -------------------------------------------------------------------------------------------
+# Spawn infrastructure
+# -------------------------------------------------------------------------------------------
+
+
+def _find_free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        return sock.getsockname()[1]
+
+
+def _spawn(world_size: int, worker, *args, timeout_s: float = WATCHDOG_TIMEOUT_S) -> None:
+    """Run ``worker(rank, world_size, port, *args)`` on ``world_size`` fresh processes.
+
+    Watchdog approach: ``mp.spawn(join=False)`` returns a ``ProcessContext`` whose ``join`` is
+    polled under a deadline. On timeout every surviving worker is SIGKILLed and the test fails
+    with ``TimeoutError`` — a deadlock can never hang CI. When a worker raises, ``join`` itself
+    kills the remaining ranks and re-raises the worker's exception into the test.
+    """
+    port = _find_free_port()
+    context = mp.spawn(worker, args=(world_size, port, *args), nprocs=world_size, join=False)
+    deadline = time.monotonic() + timeout_s
+    while not context.join(timeout=_JOIN_POLL_S):
+        if time.monotonic() >= deadline:
+            for process in context.processes:
+                if process.is_alive():
+                    process.kill()
+            for process in context.processes:
+                process.join(timeout=10)
+            raise TimeoutError(
+                f"{getattr(worker, '__name__', worker)}: {world_size} workers still running "
+                f"after {timeout_s}s — presumed deadlock; all workers killed."
+            )
+
+
+def _init_worker_env(rank: int, world_size: int, port: int) -> None:
+    """Give the worker the torchrun-equivalent env accelerate's ``env://`` init consumes."""
+    # The tests configure accelerate through the config mirrors only; drop any accelerate env
+    # fallbacks inherited from the launching shell (what guard_against_env_interference would
+    # reject in production — here the env is simply owned by the test).
+    for name in list(os.environ):
+        if name.startswith(("FSDP_", "PARALLELISM_CONFIG_", "ACCELERATE_")):
+            del os.environ[name]
+    os.environ["MASTER_ADDR"] = "127.0.0.1"
+    os.environ["MASTER_PORT"] = str(port)
+    os.environ["RANK"] = str(rank)
+    os.environ["LOCAL_RANK"] = str(rank)
+    os.environ["WORLD_SIZE"] = str(world_size)
+    # The fp32 parity tolerances below assume true-fp32 matmuls.
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+
+# -------------------------------------------------------------------------------------------
+# Shared building blocks
+# -------------------------------------------------------------------------------------------
+
+
+def _make_cfg(
+    world_size: int,
+    *,
+    dp_replicate: int = 1,
+    dp_shard: int = 1,
+    checkpoint_format: CheckpointFormat = CheckpointFormat.SAFETENSORS,
+    grad_accum: int = 1,
+) -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=BATCH_SIZE)
+    cfg.checkpoint_format = checkpoint_format
+    cfg.parallelism.dp_replicate = dp_replicate
+    cfg.parallelism.dp_shard = dp_shard
+    cfg.accelerator.mixed_precision = "no"  # fp32 end to end: the parity tests depend on it
+    cfg.accelerator.gradient_accumulation.steps = grad_accum
+    # The dummy policy declares no _fsdp_wrap_modules; the size-based wrap policy shards its
+    # Linear without needing class names (the set_fsdp_wrap_modules no-op branch).
+    cfg.accelerator.fsdp.min_num_params = 1
+    cfg.parallelism.resolve(world_size)
+    return cfg
+
+
+def _build_accelerator(cfg: TrainPipelineConfig):
+    """``cfg.accelerator.build(cfg.parallelism)`` — make_accelerator's construction path.
+
+    Deliberately not ``make_accelerator`` itself: the factory additionally derives ``cpu=`` from
+    ``cfg.trainable_config`` (no policy config is attached to these synthetic cfgs) and re-runs
+    the env guard — both owned explicitly by the tests (see ``_init_worker_env``).
+    """
+    return cfg.accelerator.build(cfg.parallelism)
+
+
+def _make_policy(seed: int) -> DummyCheckpointPolicy:
+    """Identically seeded on every rank, so shard/replicate starts from one common init."""
+    torch.manual_seed(seed)
+    return DummyCheckpointPolicy(DummyCheckpointConfig(hidden=HIDDEN, device="cpu"))
+
+
+def _batch(step: int, rank: int, device: torch.device) -> dict[str, torch.Tensor]:
+    """Deterministic per-(step, rank) batch: every dp worker sees distinct, reproducible data."""
+    generator = torch.Generator().manual_seed(SEED + 1000 * step + rank)
+    return {"observation.state": torch.randn(BATCH_SIZE, HIDDEN, generator=generator).to(device)}
+
+
+def _gather_full(model, optimizer) -> tuple[dict, dict]:
+    """Full (unsharded) model + optimizer state via torch's DCP state-dict API — a COLLECTIVE.
+
+    With ``cpu_offload=True`` the dicts materialize on the main rank only; every other rank
+    receives a literal ``{}``.
+    """
+    from torch.distributed.checkpoint.state_dict import (
+        StateDictOptions,
+        get_model_state_dict,
+        get_optimizer_state_dict,
+    )
+
+    options = StateDictOptions(full_state_dict=True, cpu_offload=True)
+    return (
+        get_model_state_dict(model, options=options),
+        get_optimizer_state_dict(model, optimizer, options=options),
+    )
+
+
+def _assert_tree_equal(reference, actual, path: str) -> None:
+    """Exact (bitwise for tensors) equality of nested state dicts, with a failing path."""
+    if isinstance(reference, torch.Tensor):
+        assert isinstance(actual, torch.Tensor), f"{path}: {type(actual)} is not a tensor"
+        assert reference.dtype == actual.dtype, f"{path}: {reference.dtype} != {actual.dtype}"
+        assert reference.shape == actual.shape, f"{path}: {reference.shape} != {actual.shape}"
+        assert torch.equal(reference.cpu(), actual.cpu()), f"{path}: tensor values differ"
+    elif isinstance(reference, dict):
+        assert isinstance(actual, dict), f"{path}: {type(actual)} is not a dict"
+        assert set(reference) == set(actual), f"{path}: keys {set(reference) ^ set(actual)} differ"
+        for key in reference:
+            _assert_tree_equal(reference[key], actual[key], f"{path}.{key}")
+    elif isinstance(reference, list | tuple):
+        assert type(reference) is type(actual) and len(reference) == len(actual), path
+        for index, (ref_item, actual_item) in enumerate(zip(reference, actual, strict=True)):
+            _assert_tree_equal(ref_item, actual_item, f"{path}[{index}]")
+    else:
+        assert reference == actual, f"{path}: {reference!r} != {actual!r}"
+
+
+# -------------------------------------------------------------------------------------------
+# Workers (module-level: torch.multiprocessing.spawn pickles them by reference)
+# -------------------------------------------------------------------------------------------
+
+
+def _train_and_save_worker(rank: int, world_size: int, port: int, tmp_dir: str, fmt_value: str) -> None:
+    """FSDP2 (dp_shard=world_size): train SAVE_STEP steps, save_checkpoint, store the gathered
+    full model/optimizer state as the rank-0 reference for the resume workers."""
+    _init_worker_env(rank, world_size, port)
+    tmp = Path(tmp_dir)
+    fmt = CheckpointFormat(fmt_value)
+    cfg = _make_cfg(world_size, dp_shard=world_size, checkpoint_format=fmt)
+    accelerator = _build_accelerator(cfg)
+    policy = _make_policy(SEED)
+    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-2)
+    # FSDP2 requires model and optimizer in one prepare() call (accelerate rebinds param groups).
+    policy, optimizer = accelerator.prepare(policy, optimizer)
+    assert is_sharded_module(accelerator.unwrap_model(policy)), "prepare() did not shard the policy"
+
+    for step in range(SAVE_STEP):
+        loss, _ = policy(_batch(step, rank, accelerator.device))
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+
+    checkpoint_dir = tmp / "checkpoint"
+    save_checkpoint(
+        checkpoint_dir, step=SAVE_STEP, cfg=cfg, policy=policy, optimizer=optimizer, accelerator=accelerator
+    )
+
+    model_state, optimizer_state = _gather_full(policy, optimizer)
+    if accelerator.is_main_process:
+        from accelerate.utils.constants import FSDP_MODEL_NAME, OPTIMIZER_NAME
+
+        pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
+        assert (pretrained_dir / f"{FSDP_MODEL_NAME}_0").is_dir() == fmt.wants_dcp
+        assert (pretrained_dir / "model.safetensors").is_file() == fmt.wants_safetensors
+        assert (pretrained_dir / "config.json").is_file()
+        assert (pretrained_dir / "train_config.json").is_file()
+        # Sharded runs always use the DCP optimizer channel, never the safetensors one.
+        assert (checkpoint_dir / TRAINING_STATE_DIR / f"{OPTIMIZER_NAME}_0").is_dir()
+        assert not (checkpoint_dir / TRAINING_STATE_DIR / "optimizer_state.safetensors").exists()
+        torch.save({"model": model_state, "optimizer": optimizer_state}, tmp / "reference_state.pt")
+    accelerator.wait_for_everyone()
+    dist.destroy_process_group()
+
+
+def _resume_and_verify_worker(rank: int, world_size: int, port: int, tmp_dir: str, fmt_value: str) -> None:
+    """Two-phase resume at dp_shard=world_size; the gathered state must match the saved
+    reference exactly (DCP round-trips are bit-exact)."""
+    _init_worker_env(rank, world_size, port)
+    tmp = Path(tmp_dir)
+    cfg = _make_cfg(world_size, dp_shard=world_size, checkpoint_format=CheckpointFormat(fmt_value))
+    cfg.checkpoint_path = tmp / "checkpoint"
+    accelerator = _build_accelerator(cfg)
+
+    assert resume_before_prepare(cfg) == SAVE_STEP  # phase 1: RNG + step counter only
+
+    # Deliberately different init: the DCP load must overwrite every parameter.
+    policy = _make_policy(SEED + 1)
+    optimizer = torch.optim.Adam(policy.parameters(), lr=1e-2)
+    policy, optimizer = accelerator.prepare(policy, optimizer)
+    resume_after_prepare(cfg, accelerator, policy, optimizer, None)  # phase 2: DCP reshard-load
+
+    model_state, optimizer_state = _gather_full(policy, optimizer)
+    if accelerator.is_main_process:
+        reference = torch.load(tmp / "reference_state.pt", map_location="cpu", weights_only=True)
+        _assert_tree_equal(reference["model"], model_state, "model")
+        _assert_tree_equal(reference["optimizer"], optimizer_state, "optimizer")
+    accelerator.wait_for_everyone()
+    dist.destroy_process_group()
+
+
+def _loss_parity_worker(
+    rank: int, world_size: int, port: int, tmp_dir: str, dp_replicate: int, dp_shard: int, tag: str
+) -> None:
+    """Train PARITY_STEPS fp32 steps on per-rank deterministic data; rank 0 records the
+    dp-mean loss of every step. Gradient averaging spans the same rank set in any (R, S)
+    factorization of the world, so the loss trajectory is topology-invariant."""
+    _init_worker_env(rank, world_size, port)
+    cfg = _make_cfg(world_size, dp_replicate=dp_replicate, dp_shard=dp_shard)
+    accelerator = _build_accelerator(cfg)
+    policy = _make_policy(SEED)
+    optimizer = torch.optim.SGD(policy.parameters(), lr=0.05)
+    policy, optimizer = accelerator.prepare(policy, optimizer)
+    assert is_sharded_module(accelerator.unwrap_model(policy)) == (dp_shard > 1)
+
+    per_step_losses = []
+    for step in range(PARITY_STEPS):
+        loss, _ = policy(_batch(step, rank, accelerator.device))
+        per_step_losses.append(accelerator.gather(loss.detach().reshape(1)).double().mean().item())
+        accelerator.backward(loss)
+        optimizer.step()
+        optimizer.zero_grad()
+
+    if accelerator.is_main_process:
+        (Path(tmp_dir) / f"losses_{tag}.json").write_text(json.dumps(per_step_losses))
+    accelerator.wait_for_everyone()
+    dist.destroy_process_group()
+
+
+def _save_pretrained_all_ranks_worker(rank: int, world_size: int, port: int, tmp_dir: str) -> None:
+    """The all-ranks contract: every rank calls save_pretrained, the
+    collective gather completes (watchdog proves no deadlock), and only rank 0 writes files."""
+    _init_worker_env(rank, world_size, port)
+    cfg = _make_cfg(world_size, dp_shard=world_size)
+    accelerator = _build_accelerator(cfg)
+    policy = _make_policy(SEED)
+    # FSDP2 prepare requires an optimizer alongside the model even though this test never steps it.
+    optimizer = torch.optim.SGD(policy.parameters(), lr=0.1)
+    policy, optimizer = accelerator.prepare(policy, optimizer)
+    unwrapped = accelerator.unwrap_model(policy)
+    assert is_sharded_module(unwrapped)
+
+    # Gather semantics: the full dict materializes on the main rank; every other rank
+    # receives the literal empty dict.
+    reference = full_model_state_dict(unwrapped)
+    if accelerator.is_main_process:
+        assert set(reference) == {"net.weight", "net.bias"}
+    else:
+        assert reference == {}
+
+    # Every rank targets its own directory so writes are attributable per rank.
+    target = Path(tmp_dir) / f"rank_{rank}"
+    unwrapped.save_pretrained(target)
+    accelerator.wait_for_everyone()
+
+    if accelerator.is_main_process:
+        weights = load_file(target / "model.safetensors")
+        assert set(weights) == set(reference)
+        for key, tensor in reference.items():
+            assert torch.equal(weights[key], tensor), key
+        assert (target / "config.json").is_file()
+    else:
+        assert list(target.rglob("*")) == [], f"rank {rank} wrote files despite the rank-0 gate"
+    dist.destroy_process_group()
+
+
+def _grad_accum_worker(
+    rank: int, world_size: int, port: int, tmp_dir: str, micro_batch_size: int, grad_accum: int, tag: str
+) -> None:
+    """DDP fp32 with the exact accumulate/clip/step/zero_grad pattern of
+    ``lerobot_train.update_policy``; rank 0 records the final weights."""
+    _init_worker_env(rank, world_size, port)
+    assert micro_batch_size * grad_accum == SAMPLES_PER_UPDATE  # fixed effective batch
+    cfg = _make_cfg(world_size, dp_replicate=world_size, grad_accum=grad_accum)
+    accelerator = _build_accelerator(cfg)
+    # The GradientAccumulationPlugin wiring, un-overridden by any env fallback.
+    assert accelerator.gradient_accumulation_steps == grad_accum
+    policy = _make_policy(SEED)
+    optimizer = torch.optim.SGD(policy.parameters(), lr=0.05)
+    policy, optimizer = accelerator.prepare(policy, optimizer)
+
+    # One fixed per-rank sample stream, consumed in order by both variants: update k always
+    # covers rows [k * SAMPLES_PER_UPDATE, (k + 1) * SAMPLES_PER_UPDATE).
+    generator = torch.Generator().manual_seed(SEED + 7919 * rank)
+    stream = torch.randn(GA_UPDATES * SAMPLES_PER_UPDATE, HIDDEN, generator=generator)
+
+    updates_applied = 0
+    for micro_step in range(GA_UPDATES * grad_accum):
+        rows = stream[micro_step * micro_batch_size : (micro_step + 1) * micro_batch_size]
+        batch = {"observation.state": rows.to(accelerator.device)}
+        # update_policy's pattern: accumulate() suppresses grad sync and rescales the loss on
+        # non-final micro-batches, and AcceleratedOptimizer makes step()/zero_grad() no-ops
+        # until sync_gradients is True.
+        with accelerator.accumulate(policy):
+            loss, _ = policy(batch)
+            accelerator.backward(loss)
+            if accelerator.sync_gradients:
+                accelerator.clip_grad_norm_(policy.parameters(), GRAD_CLIP_NORM)
+                updates_applied += 1
+            optimizer.step()
+            optimizer.zero_grad()
+    assert updates_applied == GA_UPDATES  # exactly one optimizer update per accumulation window
+
+    if accelerator.is_main_process:
+        state = {key: value.cpu() for key, value in accelerator.unwrap_model(policy).state_dict().items()}
+        torch.save(state, Path(tmp_dir) / f"weights_{tag}.pt")
+    accelerator.wait_for_everyone()
+    dist.destroy_process_group()
+
+
+# -------------------------------------------------------------------------------------------
+# Tests
+# -------------------------------------------------------------------------------------------
+
+
+@pytest.mark.multigpu
+@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs")
+def test_fsdp2_train_save_resume_round_trip(tmp_path):
+    """FSDP2 dp_shard=4, checkpoint_format=safetensors_dcp: train -> save_checkpoint -> resume.
+
+    A second spawn resumes through the two-phase path and its gathered model weights and Adam
+    state tensors must match the pre-save gathered reference exactly (DCP round-trips are
+    bit-exact).
+    """
+    fmt = CheckpointFormat.SAFETENSORS_AND_DCP.value
+    _spawn(4, _train_and_save_worker, str(tmp_path), fmt)
+    _spawn(4, _resume_and_verify_worker, str(tmp_path), fmt)
+
+
+@pytest.mark.multigpu
+@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs")
+def test_hsdp_loss_parity_with_ddp(tmp_path):
+    """Same seed and per-rank data: DDP (dp_replicate=4) vs HSDP (2x2), fp32, no AMP.
+
+    Both topologies average gradients over the same four ranks, so per-step dp-mean losses must
+    match within tolerance. Exact parity is not expected: DDP all-reduces where HSDP
+    reduce-scatters within the shard group and all-reduces across replicas, and the different
+    reduction orders accumulate fp32 rounding — rtol=1e-4 leaves orders of magnitude of headroom
+    over that noise while still catching any real divergence (wrong averaging, wrong data).
+    """
+    _spawn(4, _loss_parity_worker, str(tmp_path), 4, 1, "ddp")
+    _spawn(4, _loss_parity_worker, str(tmp_path), 2, 2, "hsdp")
+    ddp_losses = json.loads((tmp_path / "losses_ddp.json").read_text())
+    hsdp_losses = json.loads((tmp_path / "losses_hsdp.json").read_text())
+    assert len(ddp_losses) == len(hsdp_losses) == PARITY_STEPS
+    for step, (ddp_loss, hsdp_loss) in enumerate(zip(ddp_losses, hsdp_losses, strict=True)):
+        assert hsdp_loss == pytest.approx(ddp_loss, rel=1e-4, abs=1e-6), f"step {step}"
+
+
+@pytest.mark.multigpu
+@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs")
+def test_changed_topology_resume(tmp_path):
+    """Save at dp_shard=4 (format=dcp), resume at dp_shard=2 on 2 ranks.
+
+    The DCP load reshards both the model weights and the optimizer state across the topology
+    change; the post-resume gathered state must equal the pre-save gathered reference exactly
+    (cross-topology resharding is runtime-verified).
+    """
+    fmt = CheckpointFormat.DCP.value
+    _spawn(4, _train_and_save_worker, str(tmp_path), fmt)
+    _spawn(2, _resume_and_verify_worker, str(tmp_path), fmt)
+
+
+@pytest.mark.multigpu
+@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs")
+def test_save_pretrained_all_ranks_no_deadlock(tmp_path):
+    """dp_shard=4: save_pretrained on ALL ranks completes under the watchdog.
+
+    Rank 0 writes model.safetensors (+ config.json) whose tensors equal the gathered full state;
+    ranks 1-3 write nothing. A rank-gated call would deadlock in the collective gather and be
+    killed by :func:`_spawn`'s timeout — completing at all is half of what this test asserts.
+    """
+    _spawn(4, _save_pretrained_all_ranks_worker, str(tmp_path))
+
+
+@pytest.mark.multigpu
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires 2 GPUs")
+def test_gradient_accumulation_equivalence(tmp_path):
+    """Fixed effective batch on 2-rank DDP fp32: (batch=4, GA=1) vs (batch=2, GA=2).
+
+    Both variants consume the identical per-rank sample stream in the same order for
+    GA_UPDATES optimizer updates, using update_policy's accumulate/clip/step pattern. The final
+    weights must agree: accumulate() rescales each micro-loss by 1/GA, so summed mean-of-2
+    gradients equal the mean-of-4 gradient up to fp32 summation order — hence allclose with
+    rtol=1e-5/atol=1e-6 (roughly 100x the observed associativity noise), not bitwise equality.
+    """
+    _spawn(2, _grad_accum_worker, str(tmp_path), 4, 1, "ga1")
+    _spawn(2, _grad_accum_worker, str(tmp_path), 2, 2, "ga2")
+    ga1 = torch.load(tmp_path / "weights_ga1.pt", weights_only=True)
+    ga2 = torch.load(tmp_path / "weights_ga2.pt", weights_only=True)
+    assert set(ga1) == set(ga2) == {"net.weight", "net.bias"}
+    for key in ga1:
+        assert torch.allclose(ga1[key], ga2[key], rtol=1e-5, atol=1e-6), key
@@ -0,0 +1,121 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig
+from lerobot.distributed import ParallelDims, guard_against_env_interference, is_main_process
+from lerobot.distributed.factory import _ENV_OVERRIDE
+
+
+class TestIsMainProcess:
+    def test_true_outside_distributed(self):
+        assert is_main_process() is True
+
+
+class TestParallelDims:
+    def _resolved(self, world_size: int = 8, **kwargs) -> ParallelismConfig:
+        cfg = ParallelismConfig(**kwargs)
+        cfg.resolve(world_size)
+        return cfg
+
+    def test_from_resolved_config(self):
+        dims = ParallelDims.from_config(self._resolved(dp_replicate=2, dp_shard=4), 8, "cpu")
+        assert dims.dp_world_size == 8
+        assert dims.is_sharded
+        assert dims.cp_size == 1
+        assert dims.dp_rank == 0  # no process group in unit tests
+
+    def test_rejects_unresolved_config(self):
+        with pytest.raises(ValueError, match="resolve"):
+            ParallelDims.from_config(ParallelismConfig(dp_shard=-1), 8, "cpu")
+
+    def test_rejects_world_mismatch(self):
+        with pytest.raises(ValueError, match="world_size=4"):
+            ParallelDims.from_config(self._resolved(8), 4, "cpu")
+
+    def test_cp_mesh_reserved(self):
+        dims = ParallelDims(dp_replicate=1, dp_shard=2, ring=2, ulysses=2, world_size=8, device_type="cpu")
+        assert dims.dp_rank == 0 and dims.dp_world_size == 2
+        with pytest.raises(NotImplementedError):
+            dims.cp_mesh()
+
+    def test_cp_peers_share_dp_rank_arithmetic(self):
+        """Row-major layout: cp is innermost, so dp_rank = global_rank // cp_size."""
+        dims = ParallelDims(dp_replicate=1, dp_shard=2, ring=1, ulysses=2, world_size=4, device_type="cpu")
+        # Without a process group the global rank is 0; the arithmetic contract is what matters.
+        assert dims.cp_size == 2
+        assert dims.dp_rank == 0 // dims.cp_size
+
+    def test_config_placeholder_degrees_flow_through(self):
+        cfg = ParallelismConfig(
+            dp_replicate=1,
+            dp_shard=2,
+            context_parallel=ContextParallelConfig(ring_degree=2, ulysses_degree=2),
+        )
+        # resolve() rejects cp>1 this round; ParallelDims math itself is already cp-aware.
+        dims = ParallelDims(
+            dp_replicate=cfg.dp_replicate,
+            dp_shard=cfg.dp_shard,
+            ring=cfg.context_parallel.ring_degree,
+            ulysses=cfg.context_parallel.ulysses_degree,
+            world_size=8,
+            device_type="cpu",
+        )
+        assert dims.dp_world_size == 2 and dims.cp_size == 4
+
+
+class TestEnvGuard:
+    # Silent config overrides inside accelerate itself — the guard must catch them.
+    _POISON = (
+        "ACCELERATE_USE_FSDP",
+        "ACCELERATE_USE_PARALLELISM_CONFIG",
+        "ACCELERATE_GRADIENT_ACCUMULATION_STEPS",
+    )
+
+    def test_clean_env_passes(self, monkeypatch):
+        for name in self._POISON + (_ENV_OVERRIDE,):
+            monkeypatch.delenv(name, raising=False)
+        guard_against_env_interference()
+
+    @pytest.mark.parametrize("name", _POISON)
+    def test_accelerate_env_rejected_with_actionable_error(self, name, monkeypatch):
+        monkeypatch.delenv(_ENV_OVERRIDE, raising=False)
+        monkeypatch.setenv(name, "true")
+        with pytest.raises(RuntimeError, match=name):
+            guard_against_env_interference()
+
+    def test_override_acknowledges(self, monkeypatch):
+        monkeypatch.setenv("ACCELERATE_USE_FSDP", "true")
+        monkeypatch.setenv(_ENV_OVERRIDE, "1")
+        guard_against_env_interference()
+
+
+def test_make_accelerator_rejects_format_after_sentinel_resolution(monkeypatch):
+    """dp_shard=-1 counts as sharded at parse time but can resolve
+    to an unsharded run (world size 1), which would write a safetensors-only checkpoint whose
+    recorded checkpoint_format=dcp fails its own validation on resume."""
+    from lerobot.configs.default import DatasetConfig
+    from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig
+    from lerobot.distributed.factory import make_accelerator
+
+    monkeypatch.delenv("WORLD_SIZE", raising=False)
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"))
+    cfg.parallelism.dp_shard = -1
+    cfg.checkpoint_format = CheckpointFormat.DCP
+    cfg._validate_distributed()  # passes: the sentinel is declared as sharded
+    with pytest.raises(ValueError, match="resolved to a non-sharded"):
+        make_accelerator(cfg)
@@ -0,0 +1,126 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The declarative policy surface and its distributed-side consumers."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from lerobot.configs.accelerator import FSDPConfig
+from lerobot.distributed import set_fsdp_wrap_modules, strip_accelerate_cp_hooks
+from lerobot.policies.pretrained import PreTrainedPolicy
+
+
+class TestDeclarativeAttributes:
+    def test_base_defaults(self):
+        assert PreTrainedPolicy._fsdp_wrap_modules is None
+        assert PreTrainedPolicy._fsdp_forward_methods == ("select_action", "predict_action_chunk")
+        assert PreTrainedPolicy.supports_gradient_checkpointing is False
+        assert PreTrainedPolicy._cp_plan is None
+
+    def test_act_wrap_units_name_real_classes(self):
+        """The declared class names must track the modeling code — this test pins the drift."""
+        from lerobot.policies.act import modeling_act
+
+        for name in modeling_act.ACTPolicy._fsdp_wrap_modules:
+            assert isinstance(getattr(modeling_act, name), type), name
+
+    def test_fastwam_wrap_units_name_real_classes(self):
+        from lerobot.policies.fastwam import modeling_fastwam
+        from lerobot.policies.fastwam.wan import modular
+
+        for name in modeling_fastwam.FastWAMPolicy._fsdp_wrap_modules:
+            assert isinstance(getattr(modular, name), type), name
+
+
+class _SelfAttn(nn.Module):
+    def forward(self, x, attention_mask=None, is_causal=False):
+        return x, attention_mask, is_causal
+
+
+class _TinyModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.self_attn = _SelfAttn()
+
+
+class TestStripAccelerateCpHooks:
+    def test_strips_the_real_accelerate_hook_and_restores_mask_semantics(self):
+        """Attach accelerate's actual mask-stripping hook, strip it, verify masks survive."""
+        pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+        from accelerate.big_modeling import _attach_context_parallel_hooks
+
+        model = _TinyModel()
+        mask = torch.ones(2, 2)
+
+        _attach_context_parallel_hooks(model)
+        _, hooked_mask, hooked_causal = model.self_attn(torch.zeros(1), attention_mask=mask)
+        assert hooked_mask is None and hooked_causal is True  # the hazard is real
+
+        assert strip_accelerate_cp_hooks(model) == 1
+        _, clean_mask, clean_causal = model.self_attn(torch.zeros(1), attention_mask=mask)
+        assert clean_mask is mask and clean_causal is False
+        assert not model.self_attn._forward_pre_hooks
+        assert not model.self_attn._forward_pre_hooks_with_kwargs
+
+    def test_user_hooks_survive(self):
+        model = _TinyModel()
+        model.self_attn.register_forward_pre_hook(lambda m, args: None)
+        assert strip_accelerate_cp_hooks(model) == 0
+        assert len(model.self_attn._forward_pre_hooks) == 1
+
+
+class _DeclaredPolicy:
+    _fsdp_wrap_modules = ["DeclaredBlock"]
+
+
+class _UndeclaredPolicy:
+    _fsdp_wrap_modules = None
+
+
+def _accelerator_with(plugin) -> SimpleNamespace:
+    return SimpleNamespace(state=SimpleNamespace(fsdp_plugin=plugin))
+
+
+class TestSetFsdpWrapModules:
+    @pytest.fixture(autouse=True)
+    def _requires_accelerate(self):
+        pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+    def test_policy_declaration_fills_plugin(self):
+        plugin = FSDPConfig().build_plugin()
+        set_fsdp_wrap_modules(_accelerator_with(plugin), _DeclaredPolicy())
+        assert plugin.transformer_cls_names_to_wrap == ["DeclaredBlock"]
+
+    def test_user_override_wins(self):
+        plugin = FSDPConfig(wrap_modules=["UserBlock"]).build_plugin()
+        set_fsdp_wrap_modules(_accelerator_with(plugin), _DeclaredPolicy())
+        assert plugin.transformer_cls_names_to_wrap == ["UserBlock"]
+
+    def test_no_wrap_source_fails_loudly(self):
+        plugin = FSDPConfig().build_plugin()
+        with pytest.raises(ValueError, match="_fsdp_wrap_modules"):
+            set_fsdp_wrap_modules(_accelerator_with(plugin), _UndeclaredPolicy())
+
+    def test_size_based_policy_needs_no_names(self):
+        plugin = FSDPConfig(min_num_params=1024).build_plugin()
+        set_fsdp_wrap_modules(_accelerator_with(plugin), _UndeclaredPolicy())
+        assert plugin.transformer_cls_names_to_wrap is None
+
+    def test_non_sharded_run_is_noop(self):
+        set_fsdp_wrap_modules(_accelerator_with(None), _UndeclaredPolicy())
@@ -0,0 +1,87 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A minimal real PreTrainedPolicy for checkpoint/publish unit tests (CPU, tiny)."""
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor, nn
+
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.optim.optimizers import AdamConfig, OptimizerConfig
+from lerobot.policies.pretrained import PreTrainedPolicy
+
+
+@PreTrainedConfig.register_subclass("dummy_checkpoint")
+@dataclass
+class DummyCheckpointConfig(PreTrainedConfig):
+    hidden: int = 4
+
+    @property
+    def observation_delta_indices(self) -> list | None:
+        return None
+
+    @property
+    def action_delta_indices(self) -> list | None:
+        return None
+
+    @property
+    def reward_delta_indices(self) -> list | None:
+        return None
+
+    def get_optimizer_preset(self) -> OptimizerConfig:
+        return AdamConfig(lr=1e-3)
+
+    def get_scheduler_preset(self) -> None:
+        return None
+
+    def validate_features(self) -> None:
+        pass
+
+
+class DummyCheckpointPolicy(PreTrainedPolicy):
+    config_class = DummyCheckpointConfig
+    name = "dummy_checkpoint"
+
+    def __init__(self, config: DummyCheckpointConfig, **kwargs):
+        super().__init__(config)
+        self.net = nn.Linear(config.hidden, config.hidden)
+
+    def get_optim_params(self) -> dict:
+        return self.parameters()
+
+    def reset(self) -> None:
+        pass
+
+    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
+        out = self.net(batch["observation.state"])
+        return out.mean(), None
+
+    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+        return self.net(batch["observation.state"])
+
+    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+        return self.net(batch["observation.state"])
+
+
+def make_dummy_policy(repo_id: str | None = None) -> DummyCheckpointPolicy:
+    config = DummyCheckpointConfig(device="cpu")
+    if repo_id is not None:
+        config.repo_id = repo_id
+    policy = DummyCheckpointPolicy(config)
+    with torch.no_grad():
+        policy.net.weight.fill_(0.5)
+    return policy
@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import (
    MultiAdamConfig,
    SGDConfig,
    load_optimizer_state,
-    load_optimizer_state_dict,
    save_optimizer_state,
 )
 from lerobot.utils.constants import (
@@ -66,44 +65,6 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path):
    torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict())


-def test_save_and_load_fsdp_optimizer_state_dict_roundtrip(tmp_path):
-    """The FSDP full optimizer state dict is keyed by parameter FQNs (dotted strings), not the
-    integer indices of the single-GPU path. Verify it survives the safetensors save -> read
-    round-trip used by the FSDP save/resume path (save_optimizer_state(optim_state_dict=...) then
-    load_optimizer_state_dict), which the flatten/unflatten "/" separator must not corrupt."""
-    full_osd = {
-        "state": {
-            "model.layers.0.weight": {
-                "step": torch.tensor(3.0),
-                "exp_avg": torch.randn(4, 4),
-                "exp_avg_sq": torch.randn(4, 4),
-            },
-            "model.layers.0.bias": {
-                "step": torch.tensor(3.0),
-                "exp_avg": torch.randn(4),
-                "exp_avg_sq": torch.randn(4),
-            },
-        },
-        "param_groups": [
-            {"lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0, "params": [0, 1]}
-        ],
-    }
-
-    save_optimizer_state(
-        torch.optim.Adam([torch.nn.Parameter(torch.randn(1))]), tmp_path, optim_state_dict=full_osd
-    )
-    assert (tmp_path / OPTIMIZER_STATE).is_file()
-    assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file()
-
-    loaded = load_optimizer_state_dict(tmp_path)
-    # FQN keys must be preserved verbatim (not int-cast, not split on their dots).
-    assert set(loaded["state"].keys()) == set(full_osd["state"].keys())
-    for fqn, sub in full_osd["state"].items():
-        for k, v in sub.items():
-            torch.testing.assert_close(loaded["state"][fqn][k], v)
-    assert loaded["param_groups"] == full_osd["param_groups"]
-
-
@pytest.fixture
 def base_params_dict():
    return {
@@ -301,8 +301,12 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name:
    torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0)


-def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path):
-    """Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict."""
+def test_save_pretrained_single_file_artifact(dummy_dataset_metadata, tmp_path):
+    """The distributable checkpoint is one unsharded safetensors file.
+
+    The former `state_dict=` variant of this test died with the #3810 save override: the
+    kwarg would now be silently swallowed by HubMixin's **push_to_hub_kwargs.
+    """
    policy_cls = get_policy_class("act")
    policy_cfg = make_policy_config("act")
    features = dataset_to_policy_features(dummy_dataset_metadata.features)
@@ -313,8 +317,8 @@ def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path):
    policy = policy_cls(policy_cfg)
    policy.to(policy_cfg.device)

-    save_dir = tmp_path / "fsdp_state_dict"
-    policy.save_pretrained(save_dir, state_dict=policy.state_dict())
+    save_dir = tmp_path / "single_file_artifact"
+    policy.save_pretrained(save_dir)

    # A single, unsharded safetensors file (no sharded set + index).
    assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file()
@@ -22,6 +22,7 @@ from types import SimpleNamespace
 import pytest
 import torch

+from lerobot.common.train_utils import generate_model_card
 from lerobot.configs.rewards import RewardModelConfig
 from lerobot.optim.optimizers import AdamWConfig
 from lerobot.rewards.pretrained import PreTrainedRewardModel
@@ -326,7 +327,7 @@ def test_train_pipeline_config_from_pretrained_strips_legacy_rabc_when_disabled(


 # ---------------------------------------------------------------------------
-# PreTrainedRewardModel hub upload: push_model_to_hub + generate_model_card.
+# PreTrainedRewardModel hub upload: publish_trained_model + generate_model_card.
 # We test the generation side (offline) fully, and the upload side with HfApi
 # mocked so nothing actually hits the network.
 # ---------------------------------------------------------------------------
@@ -336,6 +337,13 @@ def _make_dummy_reward_model(**config_kwargs):
    return _DummyHubReward(_DummyHubRewardConfig(**config_kwargs)), _DummyHubRewardConfig


+def _make_train_cfg(dataset_repo_id: str):
+    from lerobot.configs.default import DatasetConfig
+    from lerobot.configs.train import TrainPipelineConfig
+
+    return TrainPipelineConfig(dataset=DatasetConfig(repo_id=dataset_repo_id))
+
+
@pytest.fixture
 def _offline_model_card(monkeypatch):
    """``ModelCard.validate`` does a live ``POST`` to huggingface.co — bypass it
@@ -353,12 +361,7 @@ def test_reward_model_generate_model_card_renders_expected_fields(_offline_model
        tags=["robot", "sim"],
    )

-    card = model.generate_model_card(
-        dataset_repo_id="user/my_dataset",
-        model_type=model.config.type,
-        license=model.config.license,
-        tags=model.config.tags,
-    )
+    card = generate_model_card(model.config, cfg=_make_train_cfg("user/my_dataset"))

    # Metadata (YAML header) — ModelCardData fields.
    assert card.data.license == "mit"
@@ -380,21 +383,16 @@ def test_reward_model_generate_model_card_uses_default_license(_offline_model_ca
    """When config.license is None the card falls back to apache-2.0."""
    model, _ = _make_dummy_reward_model()

-    card = model.generate_model_card(
-        dataset_repo_id="user/my_dataset",
-        model_type=model.config.type,
-        license=model.config.license,
-        tags=None,
-    )
+    card = generate_model_card(model.config, cfg=_make_train_cfg("user/my_dataset"))

    assert card.data.license == "apache-2.0"


-def test_reward_model_push_model_to_hub_uploads_expected_files(monkeypatch, _offline_model_card):
-    """``push_model_to_hub`` must:
+def test_publish_trained_model_uploads_expected_reward_files(monkeypatch, _offline_model_card):
+    """Publishing a reward model through ``publish_trained_model`` must:
    1. create the repo,
-    2. assemble a temp folder with weights + config.json + train_config.json + README.md,
-    3. call ``api.upload_folder`` on that folder.
+    2. push the model through ``HubMixin.push_to_hub`` (weights + config.json),
+    3. upload a bundle sidecar with train_config.json + the reward-specific README.md.
    All network calls are mocked.
    """
    from huggingface_hub.constants import CONFIG_NAME
@@ -430,18 +428,80 @@ def test_reward_model_push_model_to_hub_uploads_expected_files(monkeypatch, _off
            uploaded["files"] = sorted(p.name for p in Path(folder_path).iterdir())
            return fake_commit_info

-    from lerobot.rewards import pretrained as reward_pretrained
+    import lerobot.common.train_utils as train_utils
+    import lerobot.utils.hub as hub_module
+    from lerobot.common.train_utils import publish_trained_model

-    monkeypatch.setattr(reward_pretrained, "HfApi", lambda *a, **kw: _FakeHfApi())
+    all_files: set[str] = set()

-    model.push_model_to_hub(train_cfg)
+    class _RecordingFakeHfApi(_FakeHfApi):
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def upload_folder(self, *, repo_id, repo_type, folder_path, commit_message, **_kwargs):
+            result = super().upload_folder(
+                repo_id=repo_id,
+                repo_type=repo_type,
+                folder_path=folder_path,
+                commit_message=commit_message,
+            )
+            all_files.update(uploaded["files"])
+            return result
+
+    monkeypatch.setattr(train_utils, "HfApi", _RecordingFakeHfApi)
+    monkeypatch.setattr(hub_module, "HfApi", _RecordingFakeHfApi)
+
+    publish_trained_model(train_cfg, model, None, None, dataset_meta=None)

    assert uploaded["create_repo_id"] == "user/my_reward"
    assert uploaded["upload_repo_id"] == "user/my_reward"
    assert uploaded["upload_repo_type"] == "model"
-    assert uploaded["commit_message"] == "Upload reward model weights, train config and readme"
-    # Minimum required files that must be uploaded with a reward model.
-    assert CONFIG_NAME in uploaded["files"]  # config.json
-    assert TRAIN_CONFIG_NAME in uploaded["files"]  # train_config.json
-    assert "README.md" in uploaded["files"]
-    assert any(name.endswith(".safetensors") for name in uploaded["files"])
+    # Minimum required files across the publish commits.
+    assert CONFIG_NAME in all_files  # config.json (model commit)
+    assert TRAIN_CONFIG_NAME in all_files  # train_config.json (bundle commit)
+    assert "README.md" in all_files  # reward-specific card (bundle commit)
+    assert any(name.endswith(".safetensors") for name in all_files)  # weights (model commit)
+
+
+def test_save_pretrained_writes_nothing_off_main_rank(tmp_path, monkeypatch):
+    """save_checkpoint calls save_pretrained on every rank; the
+    reward serializer must gate its writes so DDP replicas do not race on the same files."""
+    import lerobot.distributed.utils as dist_utils
+
+    model, _ = _make_dummy_reward_model()
+    monkeypatch.setattr(dist_utils, "is_main_process", lambda: False)
+    model.save_pretrained(tmp_path)
+    assert not any(tmp_path.iterdir())
+
+
+def test_reward_model_push_model_to_hub_shim_warns_and_publishes(monkeypatch, _offline_model_card):
+    """The deprecated ``push_model_to_hub`` stays callable, delegating to the publisher."""
+    from huggingface_hub.constants import CONFIG_NAME
+
+    import lerobot.common.train_utils as train_utils
+    import lerobot.utils.hub as hub_module
+    from lerobot.configs.train import TRAIN_CONFIG_NAME
+
+    all_files: set[str] = set()
+
+    class _FakeHfApi:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def create_repo(self, repo_id, private=None, exist_ok=False, **kwargs):
+            return SimpleNamespace(repo_id=repo_id)
+
+        def upload_folder(self, *, repo_id, folder_path, **_kwargs):
+            all_files.update(p.name for p in Path(folder_path).iterdir())
+            return SimpleNamespace(repo_url=SimpleNamespace(url=f"https://huggingface.co/{repo_id}"))
+
+    monkeypatch.setattr(train_utils, "HfApi", _FakeHfApi)
+    monkeypatch.setattr(hub_module, "HfApi", _FakeHfApi)
+
+    model, _ = _make_dummy_reward_model(repo_id="user/my_reward")
+    with pytest.warns(FutureWarning, match="push_model_to_hub is deprecated"):
+        model.push_model_to_hub(_make_train_cfg("user/my_dataset"))
+
+    assert CONFIG_NAME in all_files
+    assert TRAIN_CONFIG_NAME in all_files
+    assert "README.md" in all_files
@@ -0,0 +1,123 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""lerobot-convert-dcp: locating, converting, and graceful-degradation publishing."""
+
+import logging
+import shutil
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+import lerobot.distributed.checkpoint as dist_checkpoint
+from lerobot.scripts.lerobot_convert_dcp import (
+    ConvertDcpConfig,
+    _locate_pretrained_dir,
+    _publish_converted,
+    convert_checkpoint,
+)
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR
+
+
+@pytest.fixture
+def fake_merge(monkeypatch):
+    """Stand in for accelerate.utils.merge_fsdp_weights: writes a marker safetensors file."""
+    import accelerate.utils
+
+    def merge(checkpoint_dir, output_path, safe_serialization=True, remove_checkpoint_dir=False):
+        assert isinstance(checkpoint_dir, str) and isinstance(output_path, str)  # str, not Path
+        (Path(output_path) / "model.safetensors").write_bytes(b"merged")
+        # Mirror accelerate: the shard directory is removed by the merge itself, when asked.
+        if remove_checkpoint_dir:
+            shutil.rmtree(checkpoint_dir)
+
+    monkeypatch.setattr(accelerate.utils, "merge_fsdp_weights", merge)
+
+
+def make_dcp_checkpoint(tmp_path: Path) -> Path:
+    pretrained = tmp_path / PRETRAINED_MODEL_DIR
+    dcp_dir = pretrained / "pytorch_model_fsdp_0"
+    dcp_dir.mkdir(parents=True)
+    (dcp_dir / "__0_0.distcp").write_bytes(b"shard")
+    (pretrained / "config.json").write_text("{}")
+    return tmp_path
+
+
+class TestConvert:
+    def test_locate_accepts_step_dir_or_pretrained_dir(self, tmp_path):
+        step_dir = make_dcp_checkpoint(tmp_path)
+        pretrained = step_dir / PRETRAINED_MODEL_DIR
+        assert _locate_pretrained_dir(step_dir) == pretrained
+        assert _locate_pretrained_dir(pretrained) == pretrained
+
+    def test_convert_keeps_dcp_by_default(self, tmp_path, fake_merge):
+        step_dir = make_dcp_checkpoint(tmp_path)
+        out = convert_checkpoint(ConvertDcpConfig(checkpoint_dir=step_dir))
+        assert out.read_bytes() == b"merged"
+        assert (step_dir / PRETRAINED_MODEL_DIR / "pytorch_model_fsdp_0").is_dir()
+
+    def test_convert_delete_dcp(self, tmp_path, fake_merge):
+        step_dir = make_dcp_checkpoint(tmp_path)
+        convert_checkpoint(ConvertDcpConfig(checkpoint_dir=step_dir, delete_dcp=True))
+        assert not (step_dir / PRETRAINED_MODEL_DIR / "pytorch_model_fsdp_0").exists()
+
+    def test_missing_shards_error_names_the_format(self, tmp_path):
+        with pytest.raises(FileNotFoundError, match="checkpoint_format=dcp"):
+            convert_checkpoint(ConvertDcpConfig(checkpoint_dir=tmp_path))
+
+
+class TestPublishGracefulDegradation:
+    def _mock_api(self, monkeypatch):
+        calls = {}
+
+        class FakeApi:
+            def create_repo(self, repo_id, private=None, exist_ok=False):
+                return SimpleNamespace(repo_id=repo_id)
+
+            def upload_folder(self, *, repo_id, folder_path, allow_patterns, **kwargs):
+                calls["repo_id"] = repo_id
+                calls["files"] = sorted(p.name for p in Path(folder_path).iterdir())
+                calls["allow_patterns"] = allow_patterns
+                return SimpleNamespace(repo_url=SimpleNamespace(url=f"https://huggingface.co/{repo_id}"))
+
+        import lerobot.scripts.lerobot_convert_dcp as mod
+
+        monkeypatch.setattr(mod, "HfApi", FakeApi)
+        return calls
+
+    def test_missing_train_config_warns_and_uploads_core(self, tmp_path, monkeypatch, caplog):
+        calls = self._mock_api(monkeypatch)
+        pretrained = make_dcp_checkpoint(tmp_path) / PRETRAINED_MODEL_DIR
+        (pretrained / "model.safetensors").write_bytes(b"w")
+        with caplog.at_level(logging.WARNING):
+            _publish_converted(pretrained, "user/converted", private=None)
+        assert any("train_config.json missing" in m for m in caplog.messages)
+        assert "model.safetensors" in calls["files"]
+        # The DCP shard directory is still on disk (--delete_dcp defaults to False) but the
+        # allow list admits neither `.distcp` shards nor their `.metadata` sidecar.
+        assert set(calls["allow_patterns"]) == {"*.safetensors", "*.json", "*.yaml", "*.md"}
+        # config.json is not parseable as a policy config here -> card skipped with a warning
+        assert any("model card" in m for m in caplog.messages)
+
+    def test_dcp_to_safetensors_passes_str_paths(self, tmp_path, fake_merge):
+        """accelerate 1.14's DCP helpers do string containment checks."""
+        dcp_dir = tmp_path / "pytorch_model_fsdp_0"
+        dcp_dir.mkdir()
+        out = dist_checkpoint.dcp_to_safetensors(dcp_dir, tmp_path, delete_dcp=True)
+        assert out == tmp_path / "model.safetensors"
+        assert not dcp_dir.exists()
@@ -0,0 +1,314 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the opt-in EMA shadow maintained by the training pipeline (--ema.enable=true)."""
+
+import draccus
+import numpy as np
+import pytest
+import torch
+
+pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
+
+from lerobot.configs.default import EMAConfig
+from lerobot.configs.train import TrainPipelineConfig
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR
+
+DUMMY_REPO_ID = "dummy/repo"
+DUMMY_STATE_DIM = 6
+DUMMY_ACTION_DIM = 6
+IMAGE_SIZE = 32
+N_EPISODES = 2
+EPISODE_LENGTH = 12
+
+
+def test_ema_config_defaults_match_reference():
+    cfg = EMAConfig()
+    assert not cfg.enable
+    assert cfg.inv_gamma == 1.0
+    assert cfg.power == 0.75
+    assert cfg.update_after_step == 0
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"min_decay": 0.5, "max_decay": 0.1},
+        {"max_decay": 1.5},
+        {"min_decay": -0.1},
+        {"inv_gamma": 0.0},
+        {"power": -1.0},
+        {"update_after_step": -1},
+        {"decay": 1.5},
+        {"decay": -0.1},
+        {"decay": 0.99, "min_decay": 0.5},
+        {"decay": 0.99, "max_decay": 0.9},
+    ],
+)
+def test_ema_config_rejects_invalid_values(kwargs):
+    with pytest.raises(ValueError):
+        EMAConfig(**kwargs)
+
+
+def test_ema_config_cli_parsing():
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        None,
+        args=[
+            f"--dataset.repo_id={DUMMY_REPO_ID}",
+            "--ema.enable=true",
+            "--ema.power=0.8",
+            "--ema.update_after_step=10",
+        ],
+    )
+    assert cfg.ema.enable
+    assert cfg.ema.power == 0.8
+    assert cfg.ema.update_after_step == 10
+
+
+def test_ema_config_cli_parsing_constant_decay():
+    cfg = draccus.parse(
+        TrainPipelineConfig,
+        None,
+        args=[
+            f"--dataset.repo_id={DUMMY_REPO_ID}",
+            "--ema.enable=true",
+            "--ema.decay=0.99",
+        ],
+    )
+    assert cfg.ema.enable
+    assert cfg.ema.decay == 0.99
+
+
+def test_ema_constant_decay_pins_the_schedule():
+    """min_decay == max_decay clamps the warmup curve to a constant (how --ema.decay is implemented)."""
+    pytest.importorskip("diffusers")
+    from diffusers.training_utils import EMAModel
+
+    model = torch.nn.Linear(4, 4)
+    ema = EMAModel(
+        model.parameters(), decay=0.99, min_decay=0.99, use_ema_warmup=True, inv_gamma=1.0, power=0.75
+    )
+    # The first update is a hard copy (decay 0); every one after uses the constant decay.
+    for step in range(1, 6):
+        ema.step(model.parameters())
+        if step > 1:
+            assert ema.cur_decay_value == 0.99
+
+
+def test_ema_weights_context_swaps_and_restores():
+    pytest.importorskip("diffusers")
+    from diffusers.training_utils import EMAModel
+
+    from lerobot.scripts.lerobot_train import _ema_weights
+
+    torch.manual_seed(0)
+    model = torch.nn.Linear(4, 4)
+    ema = EMAModel(model.parameters(), decay=0.9999, use_ema_warmup=True, inv_gamma=1.0, power=0.75)
+
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
+    for _ in range(3):
+        model(torch.randn(2, 4)).sum().backward()
+        optimizer.step()
+        optimizer.zero_grad()
+        ema.step(model.parameters())
+
+    live = [p.detach().clone() for p in model.parameters()]
+    with _ema_weights(ema, model):
+        swapped = [p.detach().clone() for p in model.parameters()]
+    restored = list(model.parameters())
+
+    assert any(not torch.equal(a, b) for a, b in zip(live, swapped, strict=True))
+    assert all(torch.equal(a, b.detach()) for a, b in zip(live, restored, strict=True))
+
+
+def make_dummy_dataset(tmp_path):
+    features = {
+        "action": {"dtype": "float32", "shape": (DUMMY_ACTION_DIM,), "names": None},
+        "observation.state": {"dtype": "float32", "shape": (DUMMY_STATE_DIM,), "names": None},
+        "observation.images.top": {
+            "dtype": "image",
+            "shape": (IMAGE_SIZE, IMAGE_SIZE, 3),
+            "names": ["height", "width", "channel"],
+        },
+    }
+    root = tmp_path / "_dataset"
+    dataset = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=30, features=features, root=root)
+    rng = np.random.default_rng(0)
+    for ep_idx in range(N_EPISODES):
+        for _ in range(EPISODE_LENGTH):
+            dataset.add_frame(
+                {
+                    "action": rng.standard_normal(DUMMY_ACTION_DIM).astype(np.float32),
+                    "observation.state": rng.standard_normal(DUMMY_STATE_DIM).astype(np.float32),
+                    "observation.images.top": rng.integers(
+                        0, 255, size=(IMAGE_SIZE, IMAGE_SIZE, 3), dtype=np.uint8
+                    ),
+                    "task": f"task_{ep_idx}",
+                }
+            )
+        dataset.save_episode()
+    dataset.finalize()
+    return root
+
+
+def make_train_config(root, output_dir, steps, ema_enable, ema_decay=None):
+    from lerobot.configs.default import DatasetConfig
+    from lerobot.policies.factory import make_policy_config
+
+    policy_config = make_policy_config(
+        "diffusion",
+        device="cpu",
+        push_to_hub=False,
+        n_obs_steps=2,
+        horizon=8,
+        n_action_steps=4,
+        drop_n_last_frames=0,
+        down_dims=(32, 64),
+        diffusion_step_embed_dim=32,
+        spatial_softmax_num_keypoints=8,
+        num_inference_steps=2,
+        pretrained_backbone_weights=None,
+        use_group_norm=True,
+    )
+    cfg = TrainPipelineConfig(
+        dataset=DatasetConfig(repo_id=DUMMY_REPO_ID, root=str(root)),
+        policy=policy_config,
+        output_dir=output_dir,
+        steps=steps,
+        batch_size=2,
+        num_workers=0,
+        seed=42,
+        log_freq=0,
+        env_eval_freq=0,
+        save_freq=2,
+        ema=EMAConfig(enable=ema_enable, decay=ema_decay),
+    )
+    cfg.optimizer = policy_config.get_optimizer_preset()
+    cfg.scheduler = policy_config.get_scheduler_preset()
+    # The config is built in-process, so skip the CLI-oriented validation.
+    cfg.validate = lambda: None
+    return cfg
+
+
+def load_safetensors(path):
+    from safetensors.torch import load_file
+
+    return load_file(path)
+
+
+def test_train_diffusion_with_ema_checkpoint_and_resume(tmp_path):
+    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
+    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
+
+    root = make_dummy_dataset(tmp_path)
+    output_dir = tmp_path / "_output"
+
+    cfg = make_train_config(root, output_dir, steps=4, ema_enable=True)
+    train(cfg)
+
+    checkpoint_dir = output_dir / "checkpoints" / "000004"
+    ema_state_path = checkpoint_dir / TRAINING_STATE_DIR / EMA_STATE_FILENAME
+    ema_model_dir = checkpoint_dir / f"{PRETRAINED_MODEL_DIR}_ema"
+
+    # The shadow state is saved for resume and tracks every optimizer step.
+    assert ema_state_path.exists()
+    ema_state = torch.load(ema_state_path, weights_only=True)
+    assert ema_state["optimization_step"] == 4
+
+    # A directly loadable EMA model is saved next to the live one, with different weights.
+    live_weights = load_safetensors(checkpoint_dir / PRETRAINED_MODEL_DIR / "model.safetensors")
+    ema_weights = load_safetensors(ema_model_dir / "model.safetensors")
+    assert set(live_weights) == set(ema_weights)
+    assert any(not torch.equal(live_weights[k], ema_weights[k]) for k in live_weights)
+
+    from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
+
+    policy = DiffusionPolicy.from_pretrained(str(ema_model_dir))
+    assert isinstance(policy, DiffusionPolicy)
+
+    # Resuming picks the shadow up where it left off instead of restarting it.
+    resume_cfg = make_train_config(root, output_dir, steps=6, ema_enable=True)
+    resume_cfg.resume = True
+    resume_cfg.checkpoint_path = checkpoint_dir
+    train(resume_cfg)
+
+    resumed_state = torch.load(
+        output_dir / "checkpoints" / "000006" / TRAINING_STATE_DIR / EMA_STATE_FILENAME,
+        weights_only=True,
+    )
+    assert resumed_state["optimization_step"] == 6
+
+
+def test_train_with_constant_ema_decay(tmp_path):
+    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
+    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
+
+    root = make_dummy_dataset(tmp_path)
+    output_dir = tmp_path / "_output"
+
+    cfg = make_train_config(root, output_dir, steps=2, ema_enable=True, ema_decay=0.99)
+    train(cfg)
+
+    ema_state = torch.load(
+        output_dir / "checkpoints" / "000002" / TRAINING_STATE_DIR / EMA_STATE_FILENAME,
+        weights_only=True,
+    )
+    # The constant decay is implemented by pinning the schedule clamp to that value.
+    assert ema_state["decay"] == 0.99
+    assert ema_state["min_decay"] == 0.99
+    assert ema_state["optimization_step"] == 2
+
+
+def test_train_with_ema_and_gradient_accumulation(tmp_path):
+    """The shadow tracks optimizer steps, not micro-batches, under gradient accumulation."""
+    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
+    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
+
+    root = make_dummy_dataset(tmp_path)
+    output_dir = tmp_path / "_output"
+
+    cfg = make_train_config(root, output_dir, steps=4, ema_enable=True)
+    cfg.accelerator.gradient_accumulation.steps = 2
+    train(cfg)
+
+    ema_state = torch.load(
+        output_dir / "checkpoints" / "000004" / TRAINING_STATE_DIR / EMA_STATE_FILENAME,
+        weights_only=True,
+    )
+    # 4 micro-batches / 2 accumulation steps = 2 optimizer updates.
+    assert ema_state["optimization_step"] == 2
+
+
+def test_train_without_ema_writes_no_ema_files(tmp_path):
+    pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+    pytest.importorskip("diffusers", reason="diffusers is required (install lerobot[diffusion])")
+    from lerobot.scripts.lerobot_train import EMA_STATE_FILENAME, train
+
+    root = make_dummy_dataset(tmp_path)
+    output_dir = tmp_path / "_output"
+
+    cfg = make_train_config(root, output_dir, steps=2, ema_enable=False)
+    train(cfg)
+
+    checkpoint_dir = output_dir / "checkpoints" / "000002"
+    assert (checkpoint_dir / PRETRAINED_MODEL_DIR / "model.safetensors").exists()
+    assert not (checkpoint_dir / TRAINING_STATE_DIR / EMA_STATE_FILENAME).exists()
+    assert not (checkpoint_dir / f"{PRETRAINED_MODEL_DIR}_ema").exists()
@@ -17,6 +17,7 @@
 import pytest
 import torch

+import lerobot.utils.logging_utils as logging_utils
 from lerobot.utils.logging_utils import AverageMeter, MetricsTracker


@@ -25,19 +26,6 @@ def mock_metrics():
    return {"loss": AverageMeter("loss", ":.3f"), "accuracy": AverageMeter("accuracy", ":.2f")}


-class MockAccelerator:
-    def __init__(self, num_processes: int, reduce_fn=None):
-        self.num_processes = num_processes
-        self.device = torch.device("cpu")
-        self._reduce_fn = reduce_fn
-
-    def reduce(self, tensor, reduction="mean"):
-        # In single-process tests we just want a deterministic stand-in for accelerate's reduce.
-        if self._reduce_fn is not None:
-            return self._reduce_fn(tensor, reduction)
-        return tensor
-
-
 def test_average_meter_initialization():
    meter = AverageMeter("loss", ":.2f")
    assert meter.name == "loss"
@@ -96,14 +84,14 @@ def test_metrics_tracker_step(mock_metrics):
    assert tracker.epochs == tracker.samples / 1000


-def test_metrics_tracker_initialization_with_accelerator(mock_metrics):
+def test_metrics_tracker_initialization_with_dp_world(mock_metrics):
    tracker = MetricsTracker(
        batch_size=32,
        num_frames=1000,
        num_episodes=50,
        metrics=mock_metrics,
        initial_step=10,
-        accelerator=MockAccelerator(num_processes=2),
+        dp_world_size=2,
    )
    assert tracker.steps == 10
    assert tracker.samples == 10 * 32 * 2
@@ -111,14 +99,14 @@ def test_metrics_tracker_initialization_with_accelerator(mock_metrics):
    assert tracker.epochs == tracker.samples / 1000


-def test_metrics_tracker_step_with_accelerator(mock_metrics):
+def test_metrics_tracker_step_with_dp_world(mock_metrics):
    tracker = MetricsTracker(
        batch_size=32,
        num_frames=1000,
        num_episodes=50,
        metrics=mock_metrics,
        initial_step=5,
-        accelerator=MockAccelerator(num_processes=2),
+        dp_world_size=2,
    )
    tracker.step()
    assert tracker.steps == 6
@@ -178,54 +166,38 @@ def test_average_meter_reduction_stored():
    assert meter.reduction == "max"


-def test_metrics_tracker_reduce_across_ranks_no_accelerator():
+def test_metrics_tracker_reduce_across_ranks_outside_distributed():
    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
    tracker.update_s = 0.5
-    tracker.reduce_across_ranks()  # no-op without accelerator
+    tracker.reduce_across_ranks()  # no-op without an initialized process group
    assert tracker.update_s.avg == 0.5


-def test_metrics_tracker_reduce_across_ranks_single_process():
-    metrics = {"update_s": AverageMeter("update_s", reduction="max")}
-    tracker = MetricsTracker(
-        batch_size=32,
-        num_frames=1000,
-        num_episodes=50,
-        metrics=metrics,
-        accelerator=MockAccelerator(num_processes=1),
-    )
-    tracker.update_s = 0.5
-    tracker.reduce_across_ranks()  # no-op when world size is 1
-    assert tracker.update_s.avg == 0.5
-
-
-def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
+def test_metrics_tracker_reduce_across_ranks_invokes_all_reduce(monkeypatch):
    captured = {}

-    def fake_reduce(tensor, reduction):
-        captured["reduction"] = reduction
+    def fake_all_reduce(tensor, op):
+        captured["op"] = op
        captured["values"] = tensor.clone()
        # Pretend the slowest rank reported 0.9 instead of this rank's 0.4.
-        return torch.tensor([0.9], dtype=tensor.dtype, device=tensor.device)
+        tensor.fill_(0.9)
+
+    monkeypatch.setattr(logging_utils.dist, "is_initialized", lambda: True)
+    monkeypatch.setattr(logging_utils.dist, "get_world_size", lambda: 4)
+    monkeypatch.setattr(logging_utils.dist, "all_reduce", fake_all_reduce)

    metrics = {
        "loss": AverageMeter("loss"),  # reduction="none" -> not touched
        "update_s": AverageMeter("update_s", reduction="max"),
    }
-    tracker = MetricsTracker(
-        batch_size=32,
-        num_frames=1000,
-        num_episodes=50,
-        metrics=metrics,
-        accelerator=MockAccelerator(num_processes=4, reduce_fn=fake_reduce),
-    )
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
    tracker.loss = 1.0
    tracker.update_s = 0.4
    tracker.reduce_across_ranks()

-    assert captured["reduction"] == "max"
-    assert torch.allclose(captured["values"], torch.tensor([0.4]))
+    assert captured["op"] == logging_utils.dist.ReduceOp.MAX
+    assert torch.allclose(captured["values"], torch.tensor([0.4], device=captured["values"].device))
    assert tracker.update_s.avg == pytest.approx(0.9)
    # Metrics without a reduction stay untouched.
    assert tracker.loss.avg == 1.0
@@ -0,0 +1,136 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the non-blocking stdin line reader."""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import os
+import sys
+import time
+from threading import Event
+
+from lerobot.utils.stdin_input import StdinCommandListener
+
+
+def _wait_for(predicate, timeout: float = 2.0) -> bool:
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if predicate():
+            return True
+        time.sleep(0.005)
+    return predicate()
+
+
+@contextlib.contextmanager
+def _pipe_stream():
+    """A held-open pipe so the listener never sees EOF until we close it."""
+    read_fd, write_fd = os.pipe()
+    reader = os.fdopen(read_fd, "r")
+    writer = os.fdopen(write_fd, "w")
+    try:
+        yield reader, writer
+    finally:
+        with contextlib.suppress(OSError, ValueError):
+            writer.close()
+        with contextlib.suppress(OSError, ValueError):
+            reader.close()
+
+
+def test_stdin_listener_reads_lines_and_eof():
+    lines: list[str] = []
+    eof = Event()
+
+    with _pipe_stream() as (reader, writer):
+        listener = StdinCommandListener(lines.append, on_eof=eof.set, stream=reader)
+        listener.start()
+        writer.write("/start\n")
+        writer.write("   \n")  # blank lines are skipped
+        writer.write("/help\n")
+        writer.flush()
+        assert _wait_for(lambda: len(lines) == 2)
+        assert lines == ["/start", "/help"]
+
+        writer.close()
+        assert _wait_for(eof.is_set)
+        listener.stop()
+
+
+def test_stdin_listener_delivers_batched_lines():
+    """Several lines arriving in one chunk (paste, piped script) are all delivered."""
+    lines: list[str] = []
+
+    with _pipe_stream() as (reader, writer):
+        listener = StdinCommandListener(lines.append, stream=reader)
+        listener.start()
+        writer.write("/start\n/subtask grab the cube\n/stop\n")
+        writer.flush()
+        assert _wait_for(lambda: len(lines) == 3)
+        assert lines == ["/start", "/subtask grab the cube", "/stop"]
+        listener.stop()
+
+
+def test_stdin_listener_handler_errors_do_not_kill_reader():
+    lines: list[str] = []
+
+    def flaky(line: str) -> None:
+        if line == "/boom":
+            raise RuntimeError("boom")
+        lines.append(line)
+
+    with _pipe_stream() as (reader, writer):
+        listener = StdinCommandListener(flaky, stream=reader)
+        listener.start()
+        writer.write("/boom\n/start\n")
+        writer.flush()
+        assert _wait_for(lambda: lines == ["/start"])
+        listener.stop()
+
+
+def test_stdin_listener_blocking_fallback():
+    """Streams without a file descriptor (e.g. StringIO) use the blocking readline path."""
+    lines: list[str] = []
+    eof = Event()
+    listener = StdinCommandListener(lines.append, on_eof=eof.set, stream=io.StringIO("/start\n\n/help\n"))
+    assert not listener._use_select
+    listener.start()
+    assert _wait_for(eof.is_set)
+    assert lines == ["/start", "/help"]
+    listener.stop()
+
+
+def test_stdin_listener_none_stdin_treated_as_eof(monkeypatch):
+    """A missing sys.stdin (daemonized process) must fire on_eof, not hang silently."""
+    monkeypatch.setattr(sys, "stdin", None)
+    eof = Event()
+    listener = StdinCommandListener(lambda line: None, on_eof=eof.set)
+    listener.start()
+    assert eof.is_set()
+    listener.stop()
+
+
+def test_stdin_listener_broken_stream_treated_as_eof():
+    """A stream whose readline blows up mid-run must fire on_eof (dead command channel)."""
+
+    class _BrokenStream:
+        def readline(self):
+            raise AttributeError("broken")
+
+    eof = Event()
+    listener = StdinCommandListener(lambda line: None, on_eof=eof.set, stream=_BrokenStream())
+    listener.start()
+    assert _wait_for(eof.is_set)
+    listener.stop()
@@ -15,24 +15,22 @@
 # limitations under the License.

 from pathlib import Path
-from unittest.mock import MagicMock, Mock, patch
+from unittest.mock import MagicMock

 import pytest

 from lerobot.common.train_utils import (
    get_step_checkpoint_dir,
    get_step_identifier,
-    load_training_batch_size,
-    load_training_num_processes,
-    load_training_state,
-    load_training_step,
+    load_training_metadata,
    push_checkpoint_to_hub,
-    save_checkpoint,
+    save_training_metadata,
    save_training_state,
-    save_training_step,
    should_save_checkpoint,
    update_last_checkpoint,
 )
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TrainPipelineConfig
 from lerobot.utils.constants import (
    CHECKPOINTS_DIR,
    LAST_CHECKPOINT_LINK,
@@ -69,38 +67,23 @@ def test_get_step_checkpoint_dir():
    assert step_dir == output_dir / CHECKPOINTS_DIR / "000005"


-def test_save_load_training_step(tmp_path):
-    save_training_step(5000, tmp_path)
+def make_cfg(batch_size: int = 32) -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=batch_size)
+    cfg.parallelism.resolve(1)
+    return cfg
+
+
+def test_save_training_metadata_writes_the_step_file(tmp_path):
+    save_training_metadata(5000, tmp_path, make_cfg())
    assert (tmp_path / TRAINING_STEP).is_file()


-def test_load_training_step(tmp_path):
-    step = 5000
-    save_training_step(step, tmp_path)
-    loaded_step = load_training_step(tmp_path)
-    assert loaded_step == step
-
-
-def test_save_training_state_records_num_processes(tmp_path, optimizer, scheduler):
-    save_training_state(tmp_path, 10, optimizer, scheduler, num_processes=4)
-    assert load_training_num_processes(tmp_path) == 4
-
-
-def test_load_training_num_processes_absent_returns_none(tmp_path, optimizer, scheduler):
-    # Checkpoints written before the world size was recorded must still load (back-compat).
-    save_training_state(tmp_path, 10, optimizer, scheduler)
-    assert load_training_num_processes(tmp_path) is None
-
-
-def test_save_training_state_records_batch_size(tmp_path, optimizer, scheduler):
-    save_training_state(tmp_path, 10, optimizer, scheduler, batch_size=32)
-    assert load_training_batch_size(tmp_path) == 32
-
-
-def test_load_training_batch_size_absent_returns_none(tmp_path, optimizer, scheduler):
-    # Checkpoints written before the batch size was recorded must still load (back-compat).
-    save_training_state(tmp_path, 10, optimizer, scheduler)
-    assert load_training_batch_size(tmp_path) is None
+def test_save_training_state_records_topology(tmp_path, optimizer, scheduler):
+    save_training_state(tmp_path, 10, make_cfg(batch_size=32), optimizer, scheduler)
+    metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR)
+    assert metadata["step"] == 10
+    assert metadata["dp_world_size"] == 1
+    assert metadata["batch_size"] == 32


 def test_update_last_checkpoint(tmp_path):
@@ -112,32 +95,12 @@ def test_update_last_checkpoint(tmp_path):
    assert last_checkpoint.resolve() == checkpoint


-@patch("lerobot.common.train_utils.save_training_state")
-def test_save_checkpoint(mock_save_training_state, tmp_path, optimizer):
-    policy = Mock()
-    cfg = Mock()
-    save_checkpoint(tmp_path, 10, cfg, policy, optimizer)
-    policy.save_pretrained.assert_called_once()
-    cfg.save_pretrained.assert_called_once()
-    mock_save_training_state.assert_called_once()
+# save_checkpoint round-trips (all formats, real policies) live in
+# tests/common/test_checkpoint_save_resume.py.


-@patch("lerobot.common.train_utils.save_training_state")
-def test_save_checkpoint_peft(mock_save_training_state, tmp_path, optimizer):
-    policy = Mock()
-    policy.config = Mock()
-    policy.config.save_pretrained = Mock()
-    cfg = Mock()
-    cfg.use_peft = True
-    save_checkpoint(tmp_path, 10, cfg, policy, optimizer)
-    policy.save_pretrained.assert_called_once()
-    cfg.save_pretrained.assert_called_once()
-    policy.config.save_pretrained.assert_called_once()
-    mock_save_training_state.assert_called_once()
-
-
-def test_save_training_state(tmp_path, optimizer, scheduler):
-    save_training_state(tmp_path, 10, optimizer, scheduler)
+def test_save_training_state_layout(tmp_path, optimizer, scheduler):
+    save_training_state(tmp_path, 10, make_cfg(), optimizer, scheduler)
    assert (tmp_path / TRAINING_STATE_DIR).is_dir()
    assert (tmp_path / TRAINING_STATE_DIR / TRAINING_STEP).is_file()
    assert (tmp_path / TRAINING_STATE_DIR / RNG_STATE).is_file()
@@ -146,27 +109,8 @@ def test_save_training_state(tmp_path, optimizer, scheduler):
    assert (tmp_path / TRAINING_STATE_DIR / SCHEDULER_STATE).is_file()


-def test_save_load_training_state(tmp_path, optimizer, scheduler):
-    save_training_state(tmp_path, 10, optimizer, scheduler)
-    loaded_step, loaded_optimizer, loaded_scheduler = load_training_state(tmp_path, optimizer, scheduler)
-    assert loaded_step == 10
-    assert loaded_optimizer is optimizer
-    assert loaded_scheduler is scheduler
-
-
-def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler):
-    # FSDP loads optimizer separately (after accelerator.prepare)
-    # load_training_state(load_optimizer=False) must restore step + scheduler but leave the
-    # optimizer untouched and never touch the on-disk optimizer state.
-    save_training_state(tmp_path, 10, optimizer, scheduler)
-    with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state:
-        loaded_step, loaded_optimizer, loaded_scheduler = load_training_state(
-            tmp_path, optimizer, scheduler, load_optimizer=False
-        )
-    mock_load_optimizer_state.assert_not_called()
-    assert loaded_step == 10
-    assert loaded_optimizer is optimizer
-    assert loaded_scheduler is scheduler
+# The two-phase resume (resume_before_prepare / resume_after_prepare) is covered in
+# tests/common/test_checkpoint_save_resume.py with real policies and optimizer state.


 def test_push_checkpoint_to_hub_creates_repo_and_uploads(tmp_path, monkeypatch):
Author	SHA1	Message	Date
Steven Palma	a56fc0b174	refactor(rollout): integrate feedback -> api, recording, log mut and keyboard	2026-08-07 21:33:17 +02:00
Steven Palma	39c4e746f1	feat(rollout): add subtask command	2026-08-07 16:44:29 +02:00
Steven Palma	d3ee0b820c	feat(rollout): mute logs in interactive mode	2026-08-07 16:00:56 +02:00
Steven Palma	072c697c0e	feat(rollout): interactive v1	2026-08-07 15:22:44 +02:00
alejodosr	266be2bd17	feat(train): add opt-in EMA of the policy weights (--ema.enable=true) (#4323 ) * feat(train): add opt-in EMA of the policy weights (--ema.enable=true) Maintain an EMA shadow via diffusers' EMAModel (lazy import, no new dependency) with the reference Diffusion Policy schedule. Saves the shadow for exact resume plus a loadable pretrained_model_ema/ per checkpoint, evaluates the EMA weights during env eval, and pushes them to a sibling <repo_id>-ema repo. Fixes huggingface/lerobot#4259. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * docs(diffusion): document the --ema.enable training flag Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> * fix(tests): skip EMA training tests when accelerate/diffusers are missing * feat(train): support constant EMA decay (--ema.decay) for openpi-style policies * fix(train): gate EMA step on sync_gradients; use parallel_dims.is_sharded guard --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>	2026-08-07 13:52:23 +02:00
Haoming Song	ff7cc3de1d	fix(train): narrow the accelerate env guard, and fix a device-bound assert (#4347 ) Two post-merge CI failures on main, both from #4010. Benchmark Integration Tests (Libero) — `accelerate launch` exports whole groups of variables unconditionally (the five ACCELERATE_DYNAMO_* it writes default the backend to "no"), so matching on prefixes refused launches that configure nothing, contradicting the documented flow where accelerate is supported as a plain launcher. The guard now watches only the three switches that hand a subsystem to the environment. GPU Tests — `test_metrics_tracker_reduce_across_ranks_invokes_all_reduce` compared the captured reduction buffer against a CPU tensor, so the assert raised "Expected all tensors to be on the same device" wherever CUDA is available. The expected tensor is built on the buffer's device instead.	2026-08-06 18:39:24 +02:00
Hiroaki.Ishikawa	31fedfd9dd	fix(dataset): use conservative bounds for quantile aggregation instead of incorrect weighted mean (#3804 ) * fix(stats): use conservative bounds for quantile aggregation instead of incorrect weighted mean * docs: add --overwrite/--skip-images/--root options to augment_dataset_quantile_stats usage * fix(dataset): clarify quantile aggregation semantics * fix(augment): handle quantile stats edge cases	2026-08-06 18:39:01 +02:00
Pepijn	b1bf24f565	feat(train): split phase timing metrics (#4344 )	2026-08-06 17:20:35 +02:00
Haoming Song	ef88d4e52b	feat(train): parallel training framework — FSDP2, HSDP, gradient accumulation, and DCP checkpoints (#4010 ) * feat(train): parallel training engine with FSDP2, HSDP, and DCP checkpoints Replace the FSDP1 training path with a config-owned parallel-training engine: - Topology and runtime configs (--parallelism., --accelerator.): dp_replicate x dp_shard degrees select single-process, DDP (unchanged default), FSDP2, or HSDP; mixed precision, first-class gradient accumulation, and FSDP/DDP tuning knobs are mirrored as plain dataclasses that build the accelerate objects at runtime, so every run is reproducible from its train_config.json alone. Accelerate env vars are guarded against configuring the engine behind the config system's back. - Declarative policy surface: policies declare FSDP2 wrap units (_fsdp_wrap_modules) and non-forward entry points (_fsdp_forward_methods); a shared engine resolves them around accelerator.prepare(). Context-parallel fields are reserved and validated to 1. - Checkpoints: selectable --checkpoint_format (safetensors \| dcp \| safetensors_dcp); the sharded optimizer channel is always DCP; two-phase resume (step+RNG before prepare, DCP model/optimizer after) reshards across GPU-topology changes; lerobot-convert-dcp merges DCP shards into a distributable model.safetensors offline. - Publishing: PreTrainedPolicy.push_model_to_hub is replaced by the free publish_trained_model (model + processors + card + train config, all-ranks gather with main-rank writes); PreTrainedPolicy._save_pretrained gathers state dicts internally, removing the state_dict= threading from save_pretrained. - lerobot_train is restructured around the engine: optimizer built before the single prepare() call, deferred weight load on DCP resumes, collective save_checkpoint with no call-site rank branches, dp-world-size-based sample accounting. Breaking changes: FSDP checkpoints from lerobot <= 0.6.x are not resumable (weights stay loadable via from_pretrained; pin lerobot==0.6.x to finish old runs); the `accelerate launch --config_file` yaml flow is superseded by the config flags; training autocast is owned exclusively by --accelerator.mixed_precision (policy.dtype only casts parameters). Also fixes: reward-model hub publishing crash (TypeError on extra kwargs). Verified by ~200 new CPU tests (config round-trips, checkpoint round-trips per format, two-phase resume, publisher contracts, converter equivalence, accelerate canaries), a 5-test 4-GPU suite (FSDP2 save/resume bit-exactness, HSDP/DDP loss parity, changed-topology resume, all-ranks save_pretrained, grad-accum equivalence), and end-to-end ACT (1/4/8 GPUs) + FastWAM 6B (FSDP2 + HSDP) training runs.	2026-08-06 19:16:41 +08:00