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 |
+| Flag                              | Description                                                                                                                                   | Default |
-| --------------------------------- | ----------------------------------------------------------------- | ------- |
+| --------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------- | ------- |
-| `--policy.path`                   | **Required.** HF Hub model ID or local checkpoint path            | --      |
+| `--policy.path`                   | **Required.** HF Hub model ID or local checkpoint path                                                                                        | --      |
-| `--robot.type`                    | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`) | --      |
+| `--robot.type`                    | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`)                                                                             | --      |
-| `--robot.port`                    | Serial port for the robot                                         | --      |
+| `--robot.port`                    | Serial port for the robot                                                                                                                     | --      |
-| `--robot.cameras`                 | Camera configuration (JSON dict)                                  | --      |
+| `--robot.cameras`                 | Camera configuration (JSON dict)                                                                                                              | --      |
-| `--fps`                           | Control loop frequency                                            | 30      |
+| `--fps`                           | Control loop frequency                                                                                                                        | 30      |
-| `--duration`                      | Run time in seconds (0 = infinite)                                | 0       |
+| `--duration`                      | Run time in seconds (0 = infinite)                                                                                                            | 0       |
-| `--device`                        | Torch device (`cpu`, `cuda`, `mps`)                               | auto    |
+| `--device`                        | Torch device (`cpu`, `cuda`, `mps`)                                                                                                           | auto    |
-| `--task`                          | Task description (used when no dataset is provided)               | --      |
+| `--task`                          | Task description (used when no dataset is provided)                                                                                           | --      |
-| `--display_data`                  | Stream telemetry to Rerun visualization                           | false   |
+| `--display_data`                  | Stream telemetry to Rerun visualization                                                                                                       | false   |
-| `--display_ip` / `--display_port` | Remote Rerun server address                                       | --      |
+| `--display_ip` / `--display_port` | Remote Rerun server address                                                                                                                   | --      |
-| `--interpolation_multiplier`      | Action interpolation factor                                       | 1       |
+| `--interpolation_multiplier`      | Action interpolation factor                                                                                                                   | 1       |
-| `--use_torch_compile`             | Enable `torch.compile` for inference                              | false   |
+| `--interactive`                   | Chat-style stdin session (see [Interactive Sessions](#interactive-sessions)); the robot stays idle until `/start`. Base and sentry strategies | false   |
-| `--resume`                        | Resume a previous recording session                               | false   |
+| `--use_torch_compile`             | Enable `torch.compile` for inference                                                                                                          | false   |
-| `--play_sounds`                   | Vocal synthesis for events                                        | true    |
+| `--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)
+With `torchrun`:
 You can specify all parameters directly in the command without running `accelerate config`:
 ```bash
-accelerate launch \
+torchrun --nproc-per-node=2 $(which lerobot-train) \
  --multi_gpu \
  --num_processes=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:**
+With `accelerate launch` (as a plain launcher):
 - `--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:
 ```bash
-accelerate config
+accelerate launch --num_processes=2 $(which lerobot-train) \
 ```
 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) \
  --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
+## Batch semantics, learning rate, and steps
 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
-## 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.
+```
-
+effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps
 ### 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
 ```
-**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
+torchrun --nproc-per-node=2 $(which lerobot-train) \
-# Original: batch_size=8, steps=100000
+  --batch_size=8 --accelerator.gradient_accumulation.steps=4 ...
 # 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
 ```
-## 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
+## Sharded training (FSDP)
 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.
-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
+### Wrap units
-compute_environment: LOCAL_MACHINE
+
-distributed_type: FSDP
+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:
-mixed_precision: bf16
+
-num_machines: 1
+```bash
-num_processes: 4
+--accelerator.fsdp.wrap_modules='["MyTransformerBlock"]'   # explicit class names
-fsdp_config:
+--accelerator.fsdp.min_num_params=1000000                  # or: wrap every submodule above 1M params
  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
 ```
-Set `fsdp_transformer_layer_cls_to_wrap` to your model's repeated transformer-block class so each
+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).
 block is sharded as its own unit. `fsdp_use_orig_params: true` is required because LeRobot builds the
 optimizer before `accelerator.prepare()`.
-### FSDP checkpoints
+Other sharding settings:
-LeRobot gathers the full state dict across all ranks and the main process writes it as a single
+- `--accelerator.fsdp.reshard_after_forward`: whether to keep each unit's parameters resident after forward.
-`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for:
+- `--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
+### HSDP
-  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
+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:
-  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
+```bash
-  first, or cast `model.safetensors` to the deployment dtype offline.
+# 16 GPUs = 2 nodes × 8: shard within each node, replicate across nodes
- The sharded optimizer state is gathered into a full (world-size-independent) state dict and saved
+torchrun --nnodes=2 --nproc-per-node=8 ... $(which lerobot-train) \
-  alongside the model in the same `optimizer_state.safetensors` / `optimizer_param_groups.json`
+  --parallelism.dp_replicate=2 --parallelism.dp_shard=8 ...
-  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
+## Checkpoints
-  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.
+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.
+- 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.
- 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.
+- Metrics are reduced across ranks before logging: losses are averaged, and `samples/s` reports cluster-wide throughput.
- 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 stepped once per training step regardless of the number of processes (`step_scheduler_with_optimizer=False` is baked in).
 - 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.
-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(
+def update_last_checkpoint(checkpoint_dir: Path) -> None:
-    step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None
+    """Point the `last` symlink in the checkpoints directory at the given checkpoint.
 ) -> 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)
    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:
+    Args:
-    training_step = load_json(save_dir / TRAINING_STEP)
+        checkpoint_dir (Path): The checkpoint step directory the `last` link should target.
-    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:
    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,
+    accelerator: "Accelerator | None" = None,
    batch_size: int | None = None,
    model_state_dict: dict | None = None,
    optim_state_dict: dict | 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)
+    │   ├── policy_preprocessor.json  # preprocessor config (if preprocessor provided)
-    │   └── step_*.safetensors  # processor state files (if any)
+    │   ├── 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_param_groups.json  # optimizer param groups (non-sharded runs)
-        ├── optimizer_state.safetensors  # optimizer state
+        ├── 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
+        ├── scheduler_state.json  # scheduler state (if scheduler provided)
-        └── training_step.json  # training step
+        └── 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.
+        preprocessor (PolicyProcessorPipeline | None, optional): The preprocessor/pipeline to save.
        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.
            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)
+    fmt = cfg.checkpoint_format
-    cfg.save_pretrained(pretrained_dir)
+    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
+        # PeftModel.save_pretrained is an external API with no internal rank gate, and the
-        # policy config which we need for loading the model. In this case we'll write it ourselves.
+        # adapters are replicated (PEFT x sharded is rejected at validation): main rank writes.
-        policy.config.save_pretrained(pretrained_dir)
+        if is_main_process():
-    if preprocessor is not None:
+            policy_to_save.save_pretrained(pretrained_dir)
-        preprocessor.save_pretrained(pretrained_dir)
+    elif fmt.wants_safetensors or not sharded:
-    if postprocessor is not None:
+        # Collective when sharded (full gather); writes happen on the main process only in all
-        postprocessor.save_pretrained(pretrained_dir)
+        # 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,
+        checkpoint_dir, step, cfg, optimizer, scheduler, accelerator, sharded=sharded, model=policy_to_save
        step,
        optimizer,
        scheduler,
        num_processes=num_processes,
        batch_size=batch_size,
        optim_state_dict=optim_state_dict,
    )
    if accelerator is not None:
        accelerator.wait_for_everyone()
 def save_training_state(
    checkpoint_dir: Path,
-    train_step: int,
+    step: int,
-    optimizer: Optimizer | None = None,
+    cfg: TrainPipelineConfig,
    optimizer: Optimizer | dict[str, Optimizer] | None = None,
    scheduler: LRScheduler | None = None,
-    num_processes: int | None = None,
+    accelerator: "Accelerator | None" = None,
-    batch_size: int | None = None,
+    *,
-    optim_state_dict: dict | None = None,
+    sharded: bool = False,
    model: PreTrainedPolicy | None = None,
 ) -> None:
-    """
+    """Write training_state/. Collective under sharding: call on every rank.
    Saves the training step, optimizer state, scheduler state, and rng state.
    Args:
-        save_dir (Path): The directory to save artifacts to.
+        checkpoint_dir (Path): The checkpoint step directory; `training_state/` is created inside it.
-        train_step (int): Current training step.
+        step (int): The training step at that checkpoint.
-        optimizer (Optimizer | None, optional): The optimizer from which to save the state_dict.
+        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.
+        accelerator (Accelerator | None, optional): Required when `sharded` is True — it owns
-            Defaults to None.
+            the DCP optimizer save channel. Defaults to None.
-        num_processes (int | None, optional): Distributed world size to record. Defaults to None.
+        sharded (bool): The model's sharding state, computed once in `save_checkpoint` and
-        batch_size (int | None, optional): Per-process batch size to record. Defaults to None.
+            threaded here so the two sites cannot disagree. Defaults to False.
-        optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of
+        model (PreTrainedPolicy | None, optional): Required only for the sharded optimizer
-            `optimizer.state_dict()` when provided. Defaults to None.
+            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 and sharded:
-    if optimizer is not None:
+        if accelerator is None or model is None:
-        save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict)
+            raise ValueError("Saving a sharded optimizer state requires the accelerator and model.")
-    if scheduler is not None:
+        # Collective — all ranks write their DCP shards into optimizer_0/.
-        save_scheduler_state(scheduler, save_dir)
+        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
+# Two-phase resume
-) -> tuple[int, Optimizer, LRScheduler | None]:
+# ---------------------------------------------------------------------------------------------
-    """
+
-    Loads the training step, optimizer state, scheduler state, and rng state.
+
-    This is used to resume a training run.
+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.
+        cfg (TrainPipelineConfig): The resumed training config; `cfg.checkpoint_path` locates
-        optimizer (Optimizer): The optimizer to load the state_dict to.
+            the checkpoint to restore from.
-        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
+    Returns:
-            True. Set to False under FSDP, where the sharded optimizer state must be loaded after
+        int: The training step recorded in the checkpoint (micro-batch counter).
            `accelerator.prepare()` via `load_fsdp_optimizer_state` (the optimizer is returned
            untouched here).
    Raises:
-        NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir
+        NotADirectoryError: If the checkpoint has no `training_state/` directory.
-
+        ValueError: If the resumed topology crosses the sharded/non-sharded boundary relative
-    Returns:
+            to the one recorded in the checkpoint.
        tuple[int, Optimizer, LRScheduler | None]: training step, optimizer and scheduler with their
            state_dict loaded.
    """
-    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)
+    return metadata["step"]
-    if load_optimizer:
+
-        optimizer = load_optimizer_state(optimizer, training_state_dir)
+
 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)
+        load_scheduler_state(scheduler, training_state_dir)
    return step, optimizer, scheduler
-def gather_fsdp_state_dicts(model, optimizer) -> tuple[dict, dict]:
+# ---------------------------------------------------------------------------------------------
-    """Gather the full (unsharded) model and optimizer state dicts under FSDP.
+# Hub: checkpoint push (resume artifact) and publishing (distribution artifact)
-
+# ---------------------------------------------------------------------------------------------
    `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)
 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
+                # Exact global quantiles cannot be recovered from quantile summaries.
-                aggregated[q_key] = weighted_quantiles.sum(axis=0) / total_count
+                # 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(
+def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None:
    optimizer: torch.optim.Optimizer, save_dir: Path, optim_state_dict: dict | None = None
 ) -> 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
+        cfg (PreTrainedConfig): The configuration for the policy to be created. If
-             set, the policy will be loaded with weights from that path.
+            `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
+        ds_meta (LeRobotDatasetMetadata | None): Dataset metadata used to infer feature shapes and
-                 statistics for normalization layers.
+            types. Also provides statistics for normalization layers.
-        env_cfg: Environment configuration used to infer feature shapes and types.
+        env_cfg (EnvConfig | None): Environment configuration used to infer feature shapes and
-                 One of `ds_meta` or `env_cfg` must be provided.
+            types. One of `ds_meta` or `env_cfg` must be provided.
-        rename_map: Optional mapping of dataset or environment feature keys to match
+        rename_map (dict[str, str] | None): Optional mapping of dataset or environment feature
-                 expected policy feature names (e.g., `"left"` → `"camera1"`).
+            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
+        NotImplementedError: If attempting to use an unsupported policy-backend combination
-                             combination (e.g., VQBeT with 'mps').
+            (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
+        if defer_weight_load:
-        # hyperparameters that we want to vary).
+            # Same construction path as from_pretrained (config already resolved from the
-        kwargs["pretrained_name_or_path"] = cfg.pretrained_path
+            # checkpoint by the caller; dataset_stats/dataset_meta kwargs identical), minus the
-        kwargs["revision"] = cfg.pretrained_revision
+            # weight load — parity by construction.
-        policy = policy_cls.from_pretrained(**kwargs)
+            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, Any, ClassVar, TypedDict, TypeVar, Unpack
 from typing import TYPE_CHECKING, 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")
-
+# Pinned far above any policy's total size so save_torch_state_dict always emits exactly one
-def _build_card_context(
+# `model.safetensors` (no shards, no index) — a constant, not a computed byte count.
-    cfg: TrainPipelineConfig | None,
+_SINGLE_FILE_SHARD_SIZE = "1TB"
    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
 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(
+    def _save_pretrained(self, save_directory: Path) -> None:
-        self,
+        """Serialize this policy's parameters (and config) into `save_directory`.
        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).
-        Overrides `HubMixin.save_pretrained` to add a `state_dict` argument (mirroring
+        Sharding is handled internally: under FSDP2 the full state dict is gathered through a
-        `transformers.PreTrainedModel.save_pretrained`). Under FSDP, `self.state_dict()` would
+        COLLECTIVE, so when the policy is sharded this method (via `save_pretrained`) must be
-        return sharded tensors, so the caller gathers the full state dict via a cross-rank
+        called on EVERY rank — a rank-0-gated call deadlocks. File writes happen on the main
-        collective and passes it here for `_save_pretrained` to write directly.
+        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)
+        # Lazy imports: the persistence layer pulls in lerobot.distributed only when saving.
-        save_directory.mkdir(parents=True, exist_ok=True)
+        from lerobot.distributed.checkpoint import full_model_state_dict, is_sharded_module
-        self._save_pretrained(save_directory, state_dict=state_dict)
+        from lerobot.distributed.utils import is_main_process
        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
    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:
+        if is_sharded_module(model_to_save):
-            save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE))
+            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.
+        self.config._save_pretrained(save_directory)
-        # `save_torch_state_dict` discards shared-tensor duplicates just like `save_model` does;
+        save_torch_state_dict(state_dict, str(save_directory), max_shard_size=_SINGLE_FILE_SHARD_SIZE)
        # 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))
    @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,
-    ):
+    ) -> None:
-        api = HfApi()
+        """Publish this policy to the Hub.
        repo_id = api.create_repo(
            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
        ).repo_id
-        # Push the files to the repo in a single commit
+        Deprecated: use :func:`lerobot.common.train_utils.publish_trained_model` instead, which
-        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
+        also publishes the pre/post-processors alongside the model.
            saved_path = Path(tmp) / repo_id
-            if peft_model is not None:
+        Args:
-                # Since PEFT just forwards calls to `push_model_to_hub`, `self` is not the PeftModel wrapper
+            cfg (TrainPipelineConfig): The training config; saved as `train_config.json` and
-                # but the actual policy which is why we need the PEFT model passed to us to save the adapter.
+                used to render the model card.
-                # That also means that we need to store the policy config ourselves since PEFT can't.
+            peft_model: The PEFT wrapper when training adapters, whose weights replace the full
-                peft_model.save_pretrained(saved_path)
+                model weights in the published repo. Defaults to None.
-                self.config.save_pretrained(saved_path)
+            state_dict (dict[str, Tensor] | None): Ignored; weights are now gathered internally
-            else:
+                when the policy is sharded. Defaults to None.
-                # Calls _save_pretrained and stores model tensors
+            dataset_meta (LeRobotDatasetMetadata | None): Dataset metadata for the model card,
-                self.save_pretrained(saved_path, state_dict=state_dict)
+                if available. Defaults to None.
        """
        from lerobot.common.train_utils import publish_trained_model
-            card = self.generate_model_card(
+        warnings.warn(
-                cfg.dataset.repo_id,
+            "PreTrainedPolicy.push_model_to_hub is deprecated and will be removed in a future "
-                self.config.type,
+            "version. Use lerobot.common.train_utils.publish_trained_model(cfg, model, "
-                self.config.license,
+            "preprocessor, postprocessor, dataset_meta) instead.",
-                self.config.tags,
+            FutureWarning,
-                cfg=cfg,
+            stacklevel=2,
-                dataset_meta=dataset_meta,
+        )
        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"))
+        publish_trained_model(cfg, self, None, None, dataset_meta, peft_model=peft_model)
            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
    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
+    # Generate model card through the free helper (PreTrainedPolicy.generate_model_card was
-    card = policy.generate_model_card(
+    # removed with the publisher redesign), then apply the metadata recovered above — the
-        dataset_repo_id=dataset_repo_id, model_type=policy_type, license=license, tags=tags
+    # 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"):
+    def push_model_to_hub(self, cfg: "TrainPipelineConfig") -> None:
-        api = HfApi()
+        """Publish this reward model to the Hub.
        repo_id = api.create_repo(
            repo_id=self.config.repo_id, private=self.config.private, exist_ok=True
        ).repo_id
-        # Push the files to the repo in a single commit
+        Deprecated: use :func:`lerobot.common.train_utils.publish_trained_model` instead.
        with TemporaryDirectory(ignore_cleanup_errors=True) as tmp:
            saved_path = Path(tmp) / repo_id
-            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(
+        warnings.warn(
-                cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags
+            "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, "
-            card.save(str(saved_path / "README.md"))
+            "preprocessor, postprocessor, dataset_meta) instead.",
-
+            FutureWarning,
-            cfg.save_pretrained(saved_path)  # Calls _save_pretrained and stores train config
+            stacklevel=2,
            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,
        )
-
+        publish_trained_model(cfg, self, None, None, None)
        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
@@ -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.
+        cfg (TrainRLServerPipelineConfig): Training configuration; `cfg.resume` gates the load and
-        optimizers: Optimizers to load state into.
+            `cfg.output_dir` locates the last checkpoint.
-        algorithm: Algorithm whose state dict should be restored.
+        optimizers (Optimizer | dict[str, Optimizer]): Optimizers to load state into.
-            Required for full main-equivalent resume;
+        algorithm (RLAlgorithm | None, optional): Algorithm whose state dict should be restored.
-            the policy itself is restored separately via ``make_policy``.
+            Required for full main-equivalent resume; the policy itself is restored separately via
-        device: Device on which to place loaded algorithm tensors.
+            `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(
+            if task_changed:
-                observation, self._device, self._task, self._robot_type
+                # 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):
+        # No dataset finalization here: run() must be restartable (interactive
-            try:
+        # segments), so the dataset stays open until teardown() finalizes it.
-                while not ctx.runtime.shutdown_event.is_set():
+        try:
-                    loop_start = time.perf_counter()
+            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:
+                if cfg.duration > 0 and (time.perf_counter() - start_time) >= cfg.duration:
-                        logger.info("Duration limit reached (%.0fs)", cfg.duration)
+                    logger.info("Duration limit reached (%.0fs)", cfg.duration)
-                        break
+                    break
-                    obs = robot.get_observation()
+                obs = robot.get_observation()
-                    obs_processed = self._process_observation_and_notify(ctx.processors, obs)
+                obs_processed = self._process_observation_and_notify(ctx.processors, obs)
-                    if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
+                if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval):
-                        continue
+                    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:
+                if action_dict is not None:
-                        self._log_telemetry(obs_processed, action_dict, ctx.runtime)
+                    self._log_telemetry(obs_processed, action_dict, ctx.runtime)
-                        obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
+                    obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR)
-                        action_frame = build_dataset_frame(features, action_dict, prefix=ACTION)
+                    action_frame = build_dataset_frame(features, action_dict, prefix=ACTION)
-                        frame = {**obs_frame, **action_frame, "task": task_str}
+                    # The task is read live from the engine (not snapshotted from
-                        # ``add_frame`` writes to the in-progress episode buffer; the
+                    # config) so an interactive /subtask relabels frames from the
-                        # background pusher only ever touches *finalised* episode
+                    # moment it re-instructs the policy; the writer stores a task
-                        # artifacts on disk.  The two operate on disjoint state, so
+                    # per frame.  At launch the engine holds the configured task.
-                        # ``add_frame`` does not need ``_episode_lock``.
+                    frame = {**obs_frame, **action_frame, "task": engine.task}
-                        dataset.add_frame(frame)
+                    # ``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.
+                # Episode rotation derived from video file-size target.
-                    # The duration is a conservative estimate so the actual
+                # The duration is a conservative estimate so the actual
-                    # video has crossed DEFAULT_VIDEO_FILE_SIZE_IN_MB by now,
+                # video has crossed DEFAULT_VIDEO_FILE_SIZE_IN_MB by now,
-                    # keeping push_to_hub efficient (uploads complete files).
+                # keeping push_to_hub efficient (uploads complete files).
-                    elapsed = time.perf_counter() - episode_start
+                elapsed = time.perf_counter() - episode_start
-                    if elapsed >= episode_duration_s:
+                if elapsed >= episode_duration_s:
-                        # ``save_episode`` finalises the in-progress episode and
+                    # ``save_episode`` finalises the in-progress episode and
-                        # flushes it to disk; ``_episode_lock`` serialises this with
+                    # flushes it to disk; ``_episode_lock`` serialises this with
-                        # ``push_to_hub`` (run in the background executor) so the
+                    # ``push_to_hub`` (run in the background executor) so the
-                        # pusher never reads a half-written episode.
+                    # 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):
                    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:
+def collect_episode_arrays(
-    """Process a single episode and return its statistics.
+    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
+        episode_idx: Index of the episode to read
-        use_sampling: If True, sub-sample image/video frames per episode to bound
+        use_sampling: If True, sub-sample image/video frames to bound memory.
-            memory. If False, use every frame (exact, higher 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
-
+            # (N, C, H, W) -> (N * H * W, C) so quantiles are computed per channel.
-            axes_to_reduce = (0, 2, 3)
+            channels = data.shape[1]
-            keepdims = True
+            values = data.transpose(0, 2, 3, 1).reshape(-1, channels)
        else:
-            axes_to_reduce = 0
+            values = data.reshape(-1, data.shape[-1]) if data.ndim > 1 else data.reshape(-1, 1)
-            keepdims = data.ndim == 1
+        episode_arrays[key] = (values, len(data_list))
-        ep_stats[key] = get_feature_stats(
+    return episode_arrays
            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
-def compute_quantile_stats_for_dataset(dataset: LeRobotDataset, use_sampling: bool = True) -> dict[str, dict]:
+def compute_quantile_stats_for_dataset(
-    """Compute quantile statistics for all episodes in the 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
+        Episodes are accumulated sequentially because the running accumulators are
-        when video keys are present. For datasets without videos, we use parallel processing
+        shared across all of them.
        with ThreadPoolExecutor for better performance.
    """
    logging.info(f"Computing quantile statistics for dataset with {dataset.num_episodes} episodes")
-    episode_stats_list = []
+    running_stats: dict[str, RunningQuantileStats] = {}
-    has_videos = len(dataset.meta.video_keys) > 0
+    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:
+    for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
-        logging.info("Dataset contains video keys - using sequential processing for thread safety")
+        episode_arrays = collect_episode_arrays(
-        for episode_idx in tqdm(range(dataset.num_episodes), desc="Processing episodes"):
+            dataset, episode_idx, use_sampling=use_sampling, skip_images=skip_images
-            ep_stats = process_single_episode(dataset, episode_idx, use_sampling)
+        )
-            episode_stats_list.append(ep_stats)
+        for key, (array, num_frames) in episode_arrays.items():
-    else:
+            running_stats.setdefault(key, RunningQuantileStats()).update(array)
-        logging.info("Dataset has no video keys - using parallel processing for better performance")
+            frame_counts[key] = frame_counts.get(key, 0) + num_frames
-        max_workers = min(dataset.num_episodes, int(os.environ.get("LEROBOT_STATS_MAX_WORKERS", 16)))
+            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:
+    if not running_stats:
            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:
        raise ValueError("No episode data found for computing statistics")
-    logging.info(f"Aggregating statistics from {len(episode_stats_list)} episodes")
+    aggregated_stats: dict[str, dict] = {}
-    return aggregate_stats(episode_stats_list)
+    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...")
+        if cfg.interactive:
-        strategy.run(ctx)
+            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)
+        # Lazy import: hub code must not import the distributed package at module load
-        repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id
+        # (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 * self._dp_world_size
        self.samples += self._batch_size * 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
+        logging. Outside distributed runs it is a no-op. Without it, metrics reported by the
-        reported by the main process only reflect rank 0; for bottleneck-style timings
+        main process only reflect rank 0; for bottleneck-style timings (``dataloading_s``,
-        (``dataloading_s``, ``update_s``, ...) that means the slowest worker's stall is invisible.
+        ``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)
+            dist.all_reduce(tensor, op=reduce_ops[reduction])
-            for name, value in zip(names, reduced.tolist(), strict=True):
+            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
+    # Lower quantiles use min; upper quantiles use max, regardless of counts.
-    assert "q01" in result
+    np.testing.assert_allclose(result["q01"], np.array([1.5]), atol=1e-6)
-    assert "q99" in result
+    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)
-    # Verify quantile aggregation (weighted average)
+    np.testing.assert_allclose(result["q90"], np.array([11.0]), atol=1e-6)
-    expected_q01 = (1.5 * 100 + 2.5 * 150) / 250  # ≈ 2.1
+    np.testing.assert_allclose(result["q99"], np.array([11.5]), atol=1e-6)
    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)
 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):
+def test_save_pretrained_single_file_artifact(dummy_dataset_metadata, tmp_path):
-    """Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict."""
+    """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"
+    save_dir = tmp_path / "single_file_artifact"
-    policy.save_pretrained(save_dir, state_dict=policy.state_dict())
+    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(
+    card = generate_model_card(model.config, cfg=_make_train_cfg("user/my_dataset"))
        dataset_repo_id="user/my_dataset",
        model_type=model.config.type,
        license=model.config.license,
        tags=model.config.tags,
    )
    # 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(
+    card = generate_model_card(model.config, cfg=_make_train_cfg("user/my_dataset"))
        dataset_repo_id="user/my_dataset",
        model_type=model.config.type,
        license=model.config.license,
        tags=None,
    )
    assert card.data.license == "apache-2.0"
-def test_reward_model_push_model_to_hub_uploads_expected_files(monkeypatch, _offline_model_card):
+def test_publish_trained_model_uploads_expected_reward_files(monkeypatch, _offline_model_card):
-    """``push_model_to_hub`` must:
+    """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,
+    2. push the model through ``HubMixin.push_to_hub`` (weights + config.json),
-    3. call ``api.upload_folder`` on that folder.
+    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 across the publish commits.
-    # Minimum required files that must be uploaded with a reward model.
+    assert CONFIG_NAME in all_files  # config.json (model commit)
-    assert CONFIG_NAME in uploaded["files"]  # config.json
+    assert TRAIN_CONFIG_NAME in all_files  # train_config.json (bundle commit)
-    assert TRAIN_CONFIG_NAME in uploaded["files"]  # train_config.json
+    assert "README.md" in all_files  # reward-specific card (bundle commit)
-    assert "README.md" in uploaded["files"]
+    assert any(name.endswith(".safetensors") for name in all_files)  # weights (model commit)
-    assert any(name.endswith(".safetensors") for name in uploaded["files"])
+
 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():
+def test_metrics_tracker_reduce_across_ranks_invokes_all_reduce(monkeypatch):
    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():
    captured = {}
-    def fake_reduce(tensor, reduction):
+    def fake_all_reduce(tensor, op):
-        captured["reduction"] = reduction
+        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(
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
        batch_size=32,
        num_frames=1000,
        num_episodes=50,
        metrics=metrics,
        accelerator=MockAccelerator(num_processes=4, reduce_fn=fake_reduce),
    )
    tracker.loss = 1.0
    tracker.update_s = 0.4
    tracker.reduce_across_ranks()
-    assert captured["reduction"] == "max"
+    assert captured["op"] == logging_utils.dist.ReduceOp.MAX
-    assert torch.allclose(captured["values"], torch.tensor([0.4]))
+    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_metadata,
    load_training_num_processes,
    load_training_state,
    load_training_step,
    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):
+def make_cfg(batch_size: int = 32) -> TrainPipelineConfig:
-    save_training_step(5000, tmp_path)
+    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):
+def test_save_training_state_records_topology(tmp_path, optimizer, scheduler):
-    step = 5000
+    save_training_state(tmp_path, 10, make_cfg(batch_size=32), optimizer, scheduler)
-    save_training_step(step, tmp_path)
+    metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR)
-    loaded_step = load_training_step(tmp_path)
+    assert metadata["step"] == 10
-    assert loaded_step == step
+    assert metadata["dp_world_size"] == 1
-
+    assert metadata["batch_size"] == 32
 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_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")
+# save_checkpoint round-trips (all formats, real policies) live in
-def test_save_checkpoint(mock_save_training_state, tmp_path, optimizer):
+# tests/common/test_checkpoint_save_resume.py.
    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()
-@patch("lerobot.common.train_utils.save_training_state")
+def test_save_training_state_layout(tmp_path, optimizer, scheduler):
-def test_save_checkpoint_peft(mock_save_training_state, tmp_path, optimizer):
+    save_training_state(tmp_path, 10, make_cfg(), optimizer, scheduler)
    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)
    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):
+# The two-phase resume (resume_before_prepare / resume_after_prepare) is covered in
-    save_training_state(tmp_path, 10, optimizer, scheduler)
+# tests/common/test_checkpoint_save_resume.py with real policies and optimizer state.
    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
 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