diff --git a/INTERACTIVE_ROLLOUT.md b/INTERACTIVE_ROLLOUT.md new file mode 100644 index 000000000..65d155de7 --- /dev/null +++ b/INTERACTIVE_ROLLOUT.md @@ -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 ` — 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 )`. +- **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. diff --git a/docs/source/inference.mdx b/docs/source/inference.mdx index bee16a24a..bd511199f 100644 --- a/docs/source/inference.mdx +++ b/docs/source/inference.mdx @@ -289,31 +289,81 @@ Robot reset — holding at initial position. /start to run. 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. -**Console logs are muted while the session runs** so they don't interleave with what you're typing; they resume when it ends. A fatal inference error is still printed. Run without `--interactive` to watch the live log. +**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. -Interactive sessions currently require `--strategy.type=base`: the recording strategies finalize their dataset when their loop exits, so they cannot be restarted by `/start`, and their keyboard controls would compete for the same terminal. +Sessions work over SSH and on headless machines — the command reader uses the terminal (or pipe) directly and needs no display server. + +**Recording while interactive.** `--strategy.type=sentry` also supports `--interactive=true`: the session records continuously while you steer it. Each `/start`…`/reset` segment saves complete episodes plus one final partial episode, the dataset stays open until shutdown, and **frames are labeled with the live task** — a `/subtask` changes both the policy conditioning and the recorded label from the same frame onwards. + +```bash +lerobot-rollout \ + --strategy.type=sentry \ + --policy.path=${HF_USER}/my_smolvla_policy \ + --robot.type=so100_follower \ + --robot.port=/dev/ttyACM0 \ + --dataset.repo_id=${HF_USER}/rollout_cube_sessions \ + --task="pick up the cube" \ + --interactive=true +``` + +The other recording strategies (episodic, DAgger, highlight) are not supported: they bind their own keyboard controls, which would compete with the command prompt for the same terminal. + +### Programmatic control + +Everything the CLI session does is available as a library API: `RolloutController` exposes thread-safe `start()` / `reset()` / `stop()` / `set_task()` methods plus a `RolloutEvent` callback, with no stdin, printing, or log muting attached — embed it in your own application, network server, or notebook: + +```python +from threading import Event, Thread + +from lerobot.rollout import ( + LinkedEvent, + RolloutController, + RolloutEvent, + build_rollout_context, + create_strategy, +) + +parent = Event() # your application's shutdown signal +ctx = build_rollout_context(cfg, LinkedEvent(parent)) # loads policy, connects robot +strategy = create_strategy(cfg.strategy) +strategy.setup(ctx) + +controller = RolloutController(strategy, ctx, on_event=print) # or your own observer +serve_thread = Thread(target=controller.serve) # serve() blocks; run it where you like +serve_thread.start() + +controller.start() # robot starts executing the policy +controller.set_task("grab the red cube") # re-instruct mid-run +controller.reset() # stop movement, return home, stay warm +controller.stop() # end serve() + +serve_thread.join() +strategy.teardown(ctx) # teardown stays with the caller +``` + +Set `play_sounds=False` in the config unless you want the vocal announcements, and note that `build_rollout_context` requires the shutdown event to be a `LinkedEvent` (the controller ends run segments through its local flag; your `parent` event still forces a full shutdown). `InteractiveSession` itself is a thin front-end over this controller — commands map 1:1 onto its methods. --- ## Common Flags -| Flag | Description | Default | -| --------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------- | ------- | -| `--policy.path` | **Required.** HF Hub model ID or local checkpoint path | -- | -| `--robot.type` | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`) | -- | -| `--robot.port` | Serial port for the robot | -- | -| `--robot.cameras` | Camera configuration (JSON dict) | -- | -| `--fps` | Control loop frequency | 30 | -| `--duration` | Run time in seconds (0 = infinite) | 0 | -| `--device` | Torch device (`cpu`, `cuda`, `mps`) | auto | -| `--task` | Task description (used when no dataset is provided) | -- | -| `--display_data` | Stream telemetry to Rerun visualization | false | -| `--display_ip` / `--display_port` | Remote Rerun server address | -- | -| `--interpolation_multiplier` | Action interpolation factor | 1 | -| `--interactive` | Chat-style stdin session (see [Interactive Sessions](#interactive-sessions)); the robot stays idle until `/start`. Base strategy only | false | -| `--use_torch_compile` | Enable `torch.compile` for inference | false | -| `--resume` | Resume a previous recording session | false | -| `--play_sounds` | Vocal synthesis for events | true | +| Flag | Description | Default | +| --------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------- | ------- | +| `--policy.path` | **Required.** HF Hub model ID or local checkpoint path | -- | +| `--robot.type` | **Required.** Robot type (e.g. `so100_follower`, `koch_follower`) | -- | +| `--robot.port` | Serial port for the robot | -- | +| `--robot.cameras` | Camera configuration (JSON dict) | -- | +| `--fps` | Control loop frequency | 30 | +| `--duration` | Run time in seconds (0 = infinite) | 0 | +| `--device` | Torch device (`cpu`, `cuda`, `mps`) | auto | +| `--task` | Task description (used when no dataset is provided) | -- | +| `--display_data` | Stream telemetry to Rerun visualization | false | +| `--display_ip` / `--display_port` | Remote Rerun server address | -- | +| `--interpolation_multiplier` | Action interpolation factor | 1 | +| `--interactive` | Chat-style stdin session (see [Interactive Sessions](#interactive-sessions)); the robot stays idle until `/start`. Base and sentry strategies | false | +| `--use_torch_compile` | Enable `torch.compile` for inference | false | +| `--resume` | Resume a previous recording session | false | +| `--play_sounds` | Vocal synthesis for events | true | --- diff --git a/src/lerobot/rollout/__init__.py b/src/lerobot/rollout/__init__.py index 982b65dec..675a02b47 100644 --- a/src/lerobot/rollout/__init__.py +++ b/src/lerobot/rollout/__init__.py @@ -38,6 +38,11 @@ from .context import ( RuntimeContext, build_rollout_context, ) +from .controller import ( + LinkedEvent, + RolloutController, + RolloutEvent, +) from .inference import ( InferenceEngine, InferenceEngineConfig, @@ -50,8 +55,6 @@ from .inference import ( from .interactive import ( InteractiveCommand, InteractiveSession, - LinkedEvent, - StdinCommandListener, parse_command, ) from .strategies import ( @@ -88,12 +91,13 @@ __all__ = [ "RTCInferenceEngine", "RolloutConfig", "RolloutContext", + "RolloutController", + "RolloutEvent", "RolloutStrategy", "RolloutStrategyConfig", "RuntimeContext", "SentryStrategy", "SentryStrategyConfig", - "StdinCommandListener", "SyncInferenceConfig", "SyncInferenceEngine", "build_rollout_context", diff --git a/src/lerobot/rollout/configs.py b/src/lerobot/rollout/configs.py index 4686acea5..fd3e7fef0 100644 --- a/src/lerobot/rollout/configs.py +++ b/src/lerobot/rollout/configs.py @@ -242,9 +242,10 @@ class RolloutConfig: # Interactive session: control the rollout from stdin with chat-style # commands (/start, /subtask , /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 console logs are muted - # while the session runs so they don't interleave with the prompt. - # Currently limited to --strategy.type=base. + # `/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 @@ -302,14 +303,14 @@ class RolloutConfig: ) # Interactive mode drives strategy.run() in restartable segments and reads - # commands from stdin. Recording strategies are excluded for now: their - # run() loops finalize the dataset on exit (so they cannot be restarted) - # and their keyboard listeners read the same terminal as the command - # prompt. - if self.interactive and not isinstance(self.strategy, BaseStrategyConfig): + # 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 currently supports only --strategy.type=base " - f"(got '{self.strategy.type}')." + 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 diff --git a/src/lerobot/rollout/controller.py b/src/lerobot/rollout/controller.py new file mode 100644 index 000000000..687833965 --- /dev/null +++ b/src/lerobot/rollout/controller.py @@ -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) diff --git a/src/lerobot/rollout/interactive.py b/src/lerobot/rollout/interactive.py index dbde19e6d..899dd47af 100644 --- a/src/lerobot/rollout/interactive.py +++ b/src/lerobot/rollout/interactive.py @@ -24,28 +24,29 @@ rollout from the terminal while hardware and policy stay connected and warm: /stop end the session and run the normal shutdown routines /help show the available commands -Threading model (mirrors the DAgger events pattern): a daemon -:class:`StdinCommandListener` thread reads lines and only ever publishes -thread-safe state — flags for the session loop, and the instruction string -via :meth:`InferenceEngine.set_task`; it never touches hardware, and never -mutates policy state (the engine applies a task change on its own inference -thread). The -:class:`InteractiveSession` driver runs on the main thread and executes -``strategy.run(ctx)`` in *segments*: each ``/start`` begins a segment, and -``/reset`` / ``/stop`` end it by setting the session's :class:`LinkedEvent`, -which every strategy control loop already polls as -``ctx.runtime.shutdown_event``. Real shutdown signals (SIGINT/SIGTERM) +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, console log handlers are muted (including -non-propagating library loggers like ``transformers``) and Python warnings -are suppressed, so system output does not interleave with the chat prompt; -only the session's own output is shown. File log handlers are unaffected, -and console logging resumes when the session ends (so teardown logs are -visible). A fatal inference-engine error is still surfaced: the session -prints the engine's captured traceback. Run without ``--interactive`` to -see the full live log output. +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 @@ -55,19 +56,18 @@ the help output, or the session loop. from __future__ import annotations +import contextlib import logging -import os -import select -import sys -import time import warnings -from collections.abc import Callable +from collections.abc import Callable, Iterator from dataclasses import dataclass -from threading import Event, Thread 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 @@ -77,71 +77,28 @@ logger = logging.getLogger(__name__) _BANNER_RULE = "─" * 60 -def _mute_console_log_handlers() -> list[tuple[logging.Handler, int]]: - """Mute console log handlers for the interactive session. +@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, so raise every console handler above ``CRITICAL`` - while the session runs. All loggers are covered, not just the root: - libraries like ``transformers`` and ``datasets`` attach their own - stderr handlers with ``propagate=False``. File handlers are left - untouched — anyone who wants a persistent log can attach one — and the - previous levels are returned so :func:`_restore_log_handlers` can undo - the muting. + 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. """ - loggers = [logging.getLogger()] - loggers += [lg for lg in logging.Logger.manager.loggerDict.values() if isinstance(lg, logging.Logger)] - muted = [] - for lg in loggers: - for handler in lg.handlers: - if isinstance(handler, logging.StreamHandler) and not isinstance(handler, logging.FileHandler): - muted.append((handler, handler.level)) - handler.setLevel(logging.CRITICAL + 1) - return muted - - -def _restore_log_handlers(muted: list[tuple[logging.Handler, int]]) -> None: - """Restore handler levels changed by :func:`_mute_console_log_handlers`.""" - for handler, level in muted: - handler.setLevel(level) - - -class LinkedEvent(Event): - """A ``threading.Event`` whose ``is_set`` also reflects a parent event. - - ``set``/``clear`` act only on the local flag, so the interactive session - 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 session 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 + 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) @@ -167,10 +124,9 @@ def _strip_quotes(text: str) -> str: def parse_command(line: str) -> InteractiveCommand | None: """Parse an input line into an :class:`InteractiveCommand`. - Commands are ``/name`` optionally followed by free-text arguments - (unused by the built-in commands, but the grammar already supports - future ones like ``/subtask grab the red cube``). Returns ``None`` for - lines that are not commands (no leading ``/`` or a bare ``/``). + 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("/"): @@ -182,166 +138,14 @@ def parse_command(line: str) -> InteractiveCommand | None: return InteractiveCommand(name=name, args=rest[0].strip() if rest else "") -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`` (an interactive Ctrl-D or an - exhausted piped script both mean "no more commands" — the session must - not keep the robot running with no way to command it). - - Unlike :class:`lerobot.utils.keyboard_input.TerminalKeyListener`, this - reader leaves the terminal in canonical (line-buffered, echoing) mode — - the operator is typing chat-style commands, not pressing hotkeys. - """ - - 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 - 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 - self._running = True - self._thread = Thread(target=self._run, daemon=True, name="InteractiveStdin") - 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): - 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 interactive input %r", line) - - def _emit_eof(self) -> None: - logger.info("Interactive input stream closed (EOF)") - if self._on_eof is not None: - try: - self._on_eof() - except Exception: - logger.exception("Error while handling interactive input EOF") - - def _emit_read_error(self) -> None: - """Treat an unexpected read failure like EOF so the session shuts down. - - A dead command channel must not leave the robot running with no way - to stop it. Deliberate ``stop()`` calls clear ``_running`` first and - do not reach this path. - """ - if self._running: - logger.warning("Interactive input stream failed — treating as EOF") - self._emit_eof() - - class InteractiveSession: - """Drive a rollout strategy from chat-style stdin commands. + """Drive a rollout from chat-style stdin commands. - The session owns the outer lifecycle: after ``strategy.setup(ctx)`` the - robot stays idle until ``/start``. Each run *segment* executes - ``strategy.run(ctx)`` on the calling (main) thread until the operator - interrupts it or the strategy returns on its own (e.g. ``--duration`` - elapsed). ``/reset`` pauses the inference engine and returns the robot - to its initial position while hardware and policy stay warm; ``/stop`` - ends the session so the caller can run ``strategy.teardown(ctx)`` — the - same shutdown routine as non-interactive rollouts. - - Requires ``ctx.runtime.shutdown_event`` to be a :class:`LinkedEvent` - (installed by ``lerobot-rollout`` when ``--interactive=true``): the - session sets the local flag to end a segment, and process signals still - propagate through the parent. + 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 @@ -349,39 +153,21 @@ class InteractiveSession: 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``. + 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. """ - _POLL_INTERVAL_S = 0.2 - def __init__( self, strategy: RolloutStrategy, ctx: RolloutContext, input_stream: IO[str] | None = None, ) -> None: - stop_event = ctx.runtime.shutdown_event - if not isinstance(stop_event, LinkedEvent): - raise TypeError( - "InteractiveSession 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 - # The instruction the rollout was launched with; /reset restores it. - self._initial_task = ctx.policy.inference.task + 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) - # Written by the listener thread, consumed by the main loop. - self._start_requested = Event() - self._reset_requested = Event() - self._stop_requested = Event() - self._wake = Event() - self._running = Event() - # name -> (handler, argument hint, help line); /help and the banner # render from this table, so future commands (e.g. /ask) stay # documented for free. @@ -393,98 +179,59 @@ class InteractiveSession: "help": (self._cmd_help, "", "show this help"), } - # ------------------------------------------------------------------ - # Main-thread session loop - # ------------------------------------------------------------------ - def run(self) -> None: """Run the session until ``/stop``, EOF, engine failure, or a shutdown signal.""" - play_sounds = self._ctx.runtime.cfg.play_sounds - muted_handlers: list[tuple[logging.Handler, int]] = [] - saved_warning_filters = warnings.filters[:] try: - muted_handlers = _mute_console_log_handlers() - warnings.simplefilter("ignore") - self._print(self._render_banner()) - self._listener.start() - while not self._global_shutdown.is_set(): - if self._ctx.policy.inference.failed: - self._report_engine_failure() - 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(): - self._start_requested.clear() - self._run_segment() - continue - self._wake.wait(timeout=self._POLL_INTERVAL_S) - self._wake.clear() + with _mute_system_output(): + self._print(self._render_banner()) + self._listener.start() + try: + self.controller.serve() + finally: + self._listener.stop() finally: - self._listener.stop() - # Restore before log_say so teardown logs are visible again. - _restore_log_handlers(muted_handlers) - warnings.filters[:] = saved_warning_filters - log_say("Interactive session ended", play_sounds) + # 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 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._ctx.policy.inference.failure_traceback + 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.") - def _run_segment(self) -> None: - """Execute one ``strategy.run`` segment until interrupted or finished.""" - engine = self._ctx.policy.inference - # Clear the local flag *before* checking the request flags: command - # handlers 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() - log_say("Starting rollout", self._ctx.runtime.cfg.play_sounds) - self._print( - f"Rollout running — task {_format_task(engine.task)}. " - "/subtask to change it, /reset to return to initial position, /stop to shut down." - ) - self._running.set() - try: - self._strategy.run(self._ctx) - finally: - self._running.clear() - engine.pause() - if engine.failed: - return # the session loop reports the failure and shuts down - if not ( - self._stop_requested.is_set() or self._reset_requested.is_set() or self._global_shutdown.is_set() - ): - 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." - ) - - def _reset_robot(self) -> None: - """Pause inference and return the robot home (the task was restored by ``/reset``).""" - self._print("Resetting — returning the robot to its initial position...") - self._ctx.policy.inference.pause() - log_say("Resetting robot to initial position", self._ctx.runtime.cfg.play_sounds) - if self._ctx.hardware.initial_position: - self._strategy.return_to_initial_position(self._ctx.hardware) - self._print("Robot reset — holding at initial position. /start to run.") - else: - logger.warning("No initial position captured — skipping the return move") - self._print("Robot paused — no initial position captured, holding current pose. /start to run.") - # ------------------------------------------------------------------ - # Command handlers (called from the listener thread; only set flags) + # 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: @@ -501,25 +248,19 @@ class InteractiveSession: def _handle_eof(self) -> None: self._print("Input stream closed — stopping the session.") - self._request_stop() + self.controller.stop() def _cmd_start(self, cmd: InteractiveCommand) -> None: - if self._running.is_set(): + if not self.controller.start(): self._print("Already running — /reset to pause first, or /stop to shut down.") - return - self._start_requested.set() - self._wake.set() def _cmd_subtask(self, cmd: InteractiveCommand) -> None: - engine = self._ctx.policy.inference if not cmd.args: - self._print(f"Current task: {_format_task(engine.task)}") + self._print(f"Current task: {_format_task(self.controller.task)}") return task = _strip_quotes(cmd.args) - previous = engine.task - # Publishing the string is all this thread does: the engine applies the - # switch on its own inference thread. - if engine.set_task(task): + 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)" @@ -528,28 +269,11 @@ class InteractiveSession: self._print(f"Task unchanged: {_format_task(task)}") def _cmd_reset(self, cmd: InteractiveCommand) -> None: - # Last command wins: a /start still waiting to be serviced is cancelled - # so the robot never starts moving after the operator asked it not to. - # Flag first, segment-stop second (see the ordering note in _run_segment). - self._start_requested.clear() - # Restore the task here rather than in _reset_robot (which runs later, on - # the main thread) so that both task writers run on this thread and are - # ordered by command order — otherwise a /subtask typed right after - # /reset would be silently reverted by the deferred restore. - if self._ctx.policy.inference.set_task(self._initial_task): - self._print(f"Task restored to {_format_task(self._initial_task)}") - self._reset_requested.set() - self._segment_stop.set() - self._wake.set() + if self.controller.reset(): + self._print(f"Task restored to {_format_task(self.controller.initial_task)}") def _cmd_stop(self, cmd: InteractiveCommand) -> None: - self._request_stop() - - def _request_stop(self) -> None: - self._start_requested.clear() # last command wins, see _cmd_reset - self._stop_requested.set() - self._segment_stop.set() - self._wake.set() + self.controller.stop() def _cmd_help(self, cmd: InteractiveCommand) -> None: self._print(self._render_help()) @@ -568,9 +292,9 @@ class InteractiveSession: return ( f"{_BANNER_RULE}\n" "Interactive rollout session — the robot will NOT move until you type /start.\n" - f"Task: {_format_task(self._initial_task)}\n" + f"Task: {_format_task(self.controller.initial_task)}\n" f"{self._render_help()}\n" - "System logs and warnings are muted during the session; they resume when it ends.\n" + "System logs and warnings are muted during the session (errors still show).\n" f"{_BANNER_RULE}" ) diff --git a/src/lerobot/rollout/strategies/sentry.py b/src/lerobot/rollout/strategies/sentry.py index 61e38aa68..e9c8fef0a 100644 --- a/src/lerobot/rollout/strategies/sentry.py +++ b/src/lerobot/rollout/strategies/sentry.py @@ -22,7 +22,6 @@ import time from concurrent.futures import Future, ThreadPoolExecutor from threading import Event, Lock -from lerobot.datasets import VideoEncodingManager from lerobot.datasets.utils import DEFAULT_VIDEO_FILE_SIZE_IN_MB from lerobot.utils.constants import ACTION, OBS_STR from lerobot.utils.feature_utils import build_dataset_frame @@ -55,6 +54,14 @@ class SentryStrategy(RolloutStrategy): Requires ``streaming_encoding=True`` (enforced in config validation) to prevent disk I/O from blocking the control loop. + + ``run()`` is restartable: each call records complete episodes plus one + final partial episode, and the dataset is only finalized in + ``teardown()`` — this is what lets ``--interactive=true`` drive sentry + in start/reset/start segments while the dataset stays open. Frames are + labeled with the inference engine's *live* task, so a mid-run + ``/subtask`` changes both the policy conditioning and the recorded + label from the same frame onwards. """ config: SentryStrategyConfig @@ -70,6 +77,9 @@ class SentryStrategy(RolloutStrategy): """Initialise the inference engine and background push executor.""" self._init_engine(ctx) self._push_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sentry-push") + # Instance state (not run()-local) so the upload cadence survives + # interactive run segments. + self._episodes_since_push = 0 target_mb = self.config.target_video_file_size_mb or DEFAULT_VIDEO_FILE_SIZE_IN_MB self._episode_duration_s = estimate_max_episode_seconds( ctx.data.dataset_features, ctx.runtime.cfg.fps, target_size_mb=target_mb @@ -97,79 +107,95 @@ class SentryStrategy(RolloutStrategy): start_time = time.perf_counter() episode_start = time.perf_counter() - episodes_since_push = 0 - task_str = cfg.dataset.single_task if cfg.dataset else cfg.task logger.info("Sentry recording started (episode_duration=%.0fs)", episode_duration_s) - with VideoEncodingManager(dataset): - try: - while not ctx.runtime.shutdown_event.is_set(): - loop_start = time.perf_counter() + # No dataset finalization here: run() must be restartable (interactive + # segments), so the dataset stays open until teardown() finalizes it. + try: + while not ctx.runtime.shutdown_event.is_set(): + loop_start = time.perf_counter() - if cfg.duration > 0 and (time.perf_counter() - start_time) >= cfg.duration: - logger.info("Duration limit reached (%.0fs)", cfg.duration) - break + if cfg.duration > 0 and (time.perf_counter() - start_time) >= cfg.duration: + logger.info("Duration limit reached (%.0fs)", cfg.duration) + break - obs = robot.get_observation() - obs_processed = self._process_observation_and_notify(ctx.processors, obs) + obs = robot.get_observation() + obs_processed = self._process_observation_and_notify(ctx.processors, obs) - if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval): - continue + if self._handle_warmup(cfg.use_torch_compile, loop_start, control_interval): + continue - action_dict = send_next_action(obs_processed, obs, ctx, interpolator) + action_dict = send_next_action(obs_processed, obs, ctx, interpolator) - if action_dict is not None: - self._log_telemetry(obs_processed, action_dict, ctx.runtime) - obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR) - action_frame = build_dataset_frame(features, action_dict, prefix=ACTION) - frame = {**obs_frame, **action_frame, "task": task_str} - # ``add_frame`` writes to the in-progress episode buffer; the - # background pusher only ever touches *finalised* episode - # artifacts on disk. The two operate on disjoint state, so - # ``add_frame`` does not need ``_episode_lock``. - dataset.add_frame(frame) + if action_dict is not None: + self._log_telemetry(obs_processed, action_dict, ctx.runtime) + obs_frame = build_dataset_frame(features, obs_processed, prefix=OBS_STR) + action_frame = build_dataset_frame(features, action_dict, prefix=ACTION) + # The task is read live from the engine (not snapshotted from + # config) so an interactive /subtask relabels frames from the + # moment it re-instructs the policy; the writer stores a task + # per frame. At launch the engine holds the configured task. + frame = {**obs_frame, **action_frame, "task": engine.task} + # ``add_frame`` writes to the in-progress episode buffer; the + # background pusher only ever touches *finalised* episode + # artifacts on disk. The two operate on disjoint state, so + # ``add_frame`` does not need ``_episode_lock``. + dataset.add_frame(frame) - # Episode rotation derived from video file-size target. - # The duration is a conservative estimate so the actual - # video has crossed DEFAULT_VIDEO_FILE_SIZE_IN_MB by now, - # keeping push_to_hub efficient (uploads complete files). - elapsed = time.perf_counter() - episode_start - if elapsed >= episode_duration_s: - # ``save_episode`` finalises the in-progress episode and - # flushes it to disk; ``_episode_lock`` serialises this with - # ``push_to_hub`` (run in the background executor) so the - # pusher never reads a half-written episode. - with self._episode_lock: - dataset.save_episode() - episodes_since_push += 1 - self._needs_push.set() - logger.info( - "Episode saved (total: %d, elapsed: %.1fs)", - dataset.num_episodes, - elapsed, - ) - log_say(f"Episode {dataset.num_episodes} saved", play_sounds) - - if episodes_since_push >= self.config.upload_every_n_episodes: - self._background_push(dataset, cfg) - episodes_since_push = 0 - - episode_start = time.perf_counter() - - dt = time.perf_counter() - loop_start - if (sleep_t := control_interval - dt) > 0: - precise_sleep(sleep_t) - else: - logger.warning( - f"Record loop is running slower ({1 / dt:.1f} Hz) than the target FPS ({cfg.fps} Hz). Dataset frames might be dropped and robot control might be unstable. Common causes are: 1) Camera FPS not keeping up 2) Policy inference taking too long 3) CPU starvation" - ) - - finally: - logger.info("Sentry control loop ended — saving final episode") - with contextlib.suppress(Exception): + # Episode rotation derived from video file-size target. + # The duration is a conservative estimate so the actual + # video has crossed DEFAULT_VIDEO_FILE_SIZE_IN_MB by now, + # keeping push_to_hub efficient (uploads complete files). + elapsed = time.perf_counter() - episode_start + if elapsed >= episode_duration_s: + # ``save_episode`` finalises the in-progress episode and + # flushes it to disk; ``_episode_lock`` serialises this with + # ``push_to_hub`` (run in the background executor) so the + # pusher never reads a half-written episode. with self._episode_lock: dataset.save_episode() + self._episodes_since_push += 1 self._needs_push.set() + logger.info( + "Episode saved (total: %d, elapsed: %.1fs)", + dataset.num_episodes, + elapsed, + ) + log_say(f"Episode {dataset.num_episodes} saved", play_sounds) + + if self._episodes_since_push >= self.config.upload_every_n_episodes: + self._background_push(dataset, cfg) + self._episodes_since_push = 0 + + episode_start = time.perf_counter() + + dt = time.perf_counter() - loop_start + if (sleep_t := control_interval - dt) > 0: + precise_sleep(sleep_t) + else: + logger.warning( + f"Record loop is running slower ({1 / dt:.1f} Hz) than the target FPS ({cfg.fps} Hz). Dataset frames might be dropped and robot control might be unstable. Common causes are: 1) Camera FPS not keeping up 2) Policy inference taking too long 3) CPU starvation" + ) + + finally: + logger.info("Sentry control loop ended — saving final episode") + try: + with self._episode_lock: + dataset.save_episode() + self._needs_push.set() + except Exception: + # The tail episode could not be committed (nothing was + # recorded, or the save failed mid-write). Drop the in-flight + # streaming encode so teardown()'s finalize does not flush a + # half-written video, and discard the episode buffer: a failed + # save_episode leaves it half-mutated, which would crash the + # first add_frame of a restarted segment. add_frame recreates + # a fresh buffer from None. + logger.warning("Tail episode was not saved — discarding it", exc_info=True) + if dataset.writer is not None: + with contextlib.suppress(Exception): + dataset.writer.cancel_pending_videos() + dataset.writer.episode_buffer = None def teardown(self, ctx: RolloutContext) -> None: """Flush pending pushes, finalise the dataset, and disconnect hardware.""" diff --git a/src/lerobot/scripts/lerobot_rollout.py b/src/lerobot/scripts/lerobot_rollout.py index ef3e7445b..f57222397 100644 --- a/src/lerobot/scripts/lerobot_rollout.py +++ b/src/lerobot/scripts/lerobot_rollout.py @@ -44,10 +44,11 @@ Usage examples --robot.port=/dev/ttyACM0 \\ --task="pick up cube" --duration=30 - # Base mode — interactive session: the robot stays idle until /start is - # typed; /subtask re-instructs the policy mid-run; /reset returns - # it to the initial position (hardware and policy stay warm); /stop shuts - # down gracefully + # Interactive session (base or sentry strategy): the robot stays idle + # until /start is typed; /subtask 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 \\ diff --git a/src/lerobot/utils/stdin_input.py b/src/lerobot/utils/stdin_input.py new file mode 100644 index 000000000..bd2c316a1 --- /dev/null +++ b/src/lerobot/utils/stdin_input.py @@ -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() diff --git a/tests/test_interactive_rollout.py b/tests/test_interactive_rollout.py index ca6f018db..bc4a789ae 100644 --- a/tests/test_interactive_rollout.py +++ b/tests/test_interactive_rollout.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the interactive rollout session (--interactive=true).""" +"""Tests for interactive rollout control: the programmatic RolloutController +and the stdin-driven InteractiveSession (--interactive=true).""" from __future__ import annotations @@ -35,7 +36,8 @@ from lerobot.rollout import ( # noqa: E402 InteractiveCommand, InteractiveSession, LinkedEvent, - StdinCommandListener, + RolloutController, + RolloutEvent, parse_command, ) @@ -146,68 +148,15 @@ def test_linked_event_wait_wakes_on_parent_set(): # --------------------------------------------------------------------------- -# StdinCommandListener -# --------------------------------------------------------------------------- - - -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_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() - - -# --------------------------------------------------------------------------- -# InteractiveSession +# Shared fakes # --------------------------------------------------------------------------- class _FakeEngine(InferenceEngine): """Real task-holder semantics with mocked lifecycle methods. - Subclassing the ABC (instead of using a bare MagicMock) means the - session tests exercise the actual ``set_task``/``task`` plumbing. + Subclassing the ABC (instead of using a bare MagicMock) means these + tests exercise the actual ``set_task``/``task`` plumbing. ``failed``/``failure_traceback`` shadow the base properties as plain class attributes so tests can assign them. """ @@ -232,8 +181,8 @@ class _FakeEngine(InferenceEngine): self.get_action = MagicMock(return_value=None) -def _make_session(input_stream, run_behavior=None): - """Build a session around a mock strategy and a minimal fake context.""" +def _make_ctx(run_behavior=None): + """A mock strategy plus the minimal fake context the controller needs.""" parent = Event() stop_event = LinkedEvent(parent) engine = _FakeEngine() @@ -255,7 +204,224 @@ def _make_session(input_stream, run_behavior=None): time.sleep(0.005) strategy.run.side_effect = run_behavior or default_run + return ctx, strategy, engine, parent, run_started + +# --------------------------------------------------------------------------- +# RolloutController (the programmatic API) +# --------------------------------------------------------------------------- + + +def _make_controller(run_behavior=None): + ctx, strategy, engine, parent, run_started = _make_ctx(run_behavior) + events: list[RolloutEvent] = [] + controller = RolloutController(strategy, ctx, on_event=events.append) + return controller, events, strategy, engine, parent, run_started + + +def _serve_thread(controller) -> Thread: + thread = Thread(target=controller.serve, daemon=True) + thread.start() + return thread + + +def test_controller_requires_linked_event(): + ctx = SimpleNamespace(runtime=SimpleNamespace(shutdown_event=Event())) + with pytest.raises(TypeError, match="LinkedEvent"): + RolloutController(MagicMock(), ctx) + + +def test_controller_start_reset_stop_flow_and_events(): + controller, events, strategy, engine, _parent, run_started = _make_controller() + thread = _serve_thread(controller) + + # Idle until start(): the strategy loop must not run on its own. + time.sleep(0.05) + strategy.run.assert_not_called() + assert not controller.running + + assert controller.start() is True + assert _wait_for(run_started.is_set) + assert controller.running + assert strategy.reset_control_state.call_count == 1 + assert RolloutEvent.SEGMENT_STARTED in events + + # reset() ends the segment, pauses the engine, returns the robot home. + controller.reset() + assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1) + assert engine.pause.call_count >= 1 + assert thread.is_alive() + assert _wait_for(lambda: RolloutEvent.RESET_DONE in events) + assert RolloutEvent.RESET_STARTED in events + + # start() again runs a fresh segment with freshly reset control state. + run_started.clear() + assert controller.start() is True + assert _wait_for(run_started.is_set) + assert strategy.run.call_count == 2 + assert strategy.reset_control_state.call_count == 2 + + # stop() ends serve(); teardown stays with the caller. + controller.stop() + thread.join(timeout=2.0) + assert not thread.is_alive() + strategy.teardown.assert_not_called() + assert events[-1] is RolloutEvent.STOPPED + + +def test_controller_start_rejected_during_segment_startup(): + """A start() racing the segment startup must be rejected, not queued. + + Otherwise the re-armed request survives the whole segment and the robot + would start again, uncommanded, when the segment ends on its own. + """ + ctx, strategy, _engine, _parent, run_started = _make_ctx() + in_startup = Event() + startup_gate = Event() + + def slow_reset_control_state(): + in_startup.set() + startup_gate.wait(timeout=2.0) + + strategy.reset_control_state.side_effect = slow_reset_control_state + controller = RolloutController(strategy, ctx) + thread = _serve_thread(controller) + + assert controller.start() is True + assert _wait_for(in_startup.is_set) + # The serve thread is inside reset_control_state: the segment counts as + # running for concurrent callers even though strategy.run hasn't begun. + assert controller.start() is False + startup_gate.set() + assert _wait_for(run_started.is_set) + + # Ending the segment must not trigger a phantom second segment. + controller.reset() + assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1) + time.sleep(0.05) + assert strategy.run.call_count == 1 + + controller.stop() + thread.join(timeout=2.0) + + +def test_controller_start_returns_false_while_running(): + controller, _events, strategy, _engine, _parent, run_started = _make_controller() + thread = _serve_thread(controller) + + assert controller.start() is True + assert _wait_for(run_started.is_set) + assert controller.start() is False + time.sleep(0.05) + assert strategy.run.call_count == 1 + + controller.stop() + thread.join(timeout=2.0) + + +def test_controller_set_task_and_reset_restores_launch_task(): + controller, _events, strategy, engine, _parent, _run_started = _make_controller() + thread = _serve_thread(controller) + initial = controller.initial_task + + # reset() with the launch task still in place reports no restore. + assert controller.reset() is False + assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1) + + assert controller.set_task("fold the towel") is True + assert controller.task == "fold the towel" + assert engine.task == "fold the towel" + assert controller.set_task("fold the towel") is False + + assert controller.reset() is True # the task had been changed + assert controller.task == initial + + controller.stop() + thread.join(timeout=2.0) + + +def test_controller_engine_failure_emits_event(): + def failing_run(c): + c.policy.inference.failed = True + c.policy.inference.failure_traceback = "RuntimeError: boom-traceback" + c.runtime.shutdown_event.set() + + controller, events, strategy, _engine, _parent, _run_started = _make_controller(failing_run) + thread = _serve_thread(controller) + controller.start() + thread.join(timeout=2.0) + assert not thread.is_alive() + strategy.return_to_initial_position.assert_not_called() + assert RolloutEvent.ENGINE_FAILED in events + assert controller.failed + assert controller.failure_traceback == "RuntimeError: boom-traceback" + + +def test_controller_segment_ended_event_on_natural_end(): + def finite_run(c): + return None # e.g. --duration elapsed + + controller, events, strategy, _engine, _parent, _run_started = _make_controller(finite_run) + thread = _serve_thread(controller) + + controller.start() + assert _wait_for(lambda: RolloutEvent.SEGMENT_ENDED in events) + assert thread.is_alive() # back to idle, not shut down + assert not controller.running + + controller.stop() + thread.join(timeout=2.0) + + +def test_controller_reset_skipped_without_initial_position(): + ctx, strategy, _engine, _parent, _run_started = _make_ctx() + ctx.hardware.initial_position = {} + events: list[RolloutEvent] = [] + controller = RolloutController(strategy, ctx, on_event=events.append) + thread = _serve_thread(controller) + + controller.reset() + assert _wait_for(lambda: RolloutEvent.RESET_SKIPPED in events) + strategy.return_to_initial_position.assert_not_called() + + controller.stop() + thread.join(timeout=2.0) + + +def test_controller_callback_errors_do_not_kill_serve(): + ctx, strategy, _engine, _parent, run_started = _make_ctx() + + def broken_observer(event): + raise RuntimeError("observer boom") + + controller = RolloutController(strategy, ctx, on_event=broken_observer) + thread = _serve_thread(controller) + + controller.start() + assert _wait_for(run_started.is_set) + controller.stop() + thread.join(timeout=2.0) + assert not thread.is_alive() + + +def test_controller_works_without_event_callback(): + ctx, strategy, _engine, _parent, _run_started = _make_ctx() + controller = RolloutController(strategy, ctx) + thread = _serve_thread(controller) + controller.start() + controller.stop() + thread.join(timeout=2.0) + assert not thread.is_alive() + + +# --------------------------------------------------------------------------- +# InteractiveSession (the stdin CLI front-end) +# --------------------------------------------------------------------------- + + +def _make_session(input_stream, run_behavior=None): + """Build a session around a mock strategy and a minimal fake context.""" + ctx, strategy, engine, parent, run_started = _make_ctx(run_behavior) session = InteractiveSession(strategy, ctx, input_stream=input_stream) return session, strategy, engine, parent, run_started @@ -480,56 +646,59 @@ def test_session_commands_via_stream(): assert strategy.run.call_count == 1 -def test_session_mutes_console_logging_and_restores_on_exit(): +def test_session_mutes_logs_below_error_and_restores_on_exit(): import warnings - root = logging.getLogger() - console_handler = logging.StreamHandler(io.StringIO()) - console_handler.setLevel(logging.INFO) - root.addHandler(console_handler) # Libraries like transformers attach their own console handler with - # propagate=False; those must be muted too. + # propagate=False; the process-wide logging.disable gate covers those too. + lib_stream = io.StringIO() lib_logger = logging.getLogger("test_interactive_fake_lib") lib_logger.propagate = False - lib_handler = logging.StreamHandler(io.StringIO()) - lib_handler.setLevel(logging.WARNING) + lib_logger.setLevel(logging.DEBUG) + lib_handler = logging.StreamHandler(lib_stream) + lib_handler.setLevel(logging.INFO) lib_logger.addHandler(lib_handler) n_warning_filters = len(warnings.filters) try: with _pipe_stream() as (reader, _writer): session, _strategy, _engine, _parent, _run_started = _make_session(reader) thread = _start_session_thread(session) - assert _wait_for( - lambda: console_handler.level == logging.CRITICAL + 1 - and lib_handler.level == logging.CRITICAL + 1 - ) + assert _wait_for(lambda: logging.root.manager.disable == logging.WARNING) + + lib_logger.info("muted-info") + lib_logger.warning("muted-warning") + lib_logger.error("visible-error") # errors must surface mid-session + session._handle_line("/stop") thread.join(timeout=2.0) - assert console_handler.level == logging.INFO - assert lib_handler.level == logging.WARNING + + output = lib_stream.getvalue() + assert "muted-info" not in output + assert "muted-warning" not in output + assert "visible-error" in output + + # Everything is restored once the session ends. + assert logging.root.manager.disable == logging.NOTSET assert len(warnings.filters) == n_warning_filters + lib_logger.info("post-session-info") + assert "post-session-info" in lib_stream.getvalue() finally: - root.removeHandler(console_handler) lib_logger.removeHandler(lib_handler) -def test_session_does_not_mute_file_log_handlers(tmp_path): - root = logging.getLogger() - file_handler = logging.FileHandler(tmp_path / "session.log") - file_handler.setLevel(logging.INFO) - root.addHandler(file_handler) +def test_session_restores_preexisting_disable_level(): + """An embedding application's own logging.disable level survives the session.""" + logging.disable(logging.DEBUG) try: with _pipe_stream() as (reader, _writer): - session, _strategy, _engine, _parent, run_started = _make_session(reader) + session, _strategy, _engine, _parent, _run_started = _make_session(reader) thread = _start_session_thread(session) - session._handle_line("/start") - assert _wait_for(run_started.is_set) - assert file_handler.level == logging.INFO + assert _wait_for(lambda: logging.root.manager.disable == logging.WARNING) session._handle_line("/stop") thread.join(timeout=2.0) + assert logging.root.manager.disable == logging.DEBUG finally: - root.removeHandler(file_handler) - file_handler.close() + logging.disable(logging.NOTSET) def test_session_drives_real_base_strategy(): @@ -628,7 +797,7 @@ def test_session_subtask_strips_quotes_and_works_while_running(): session._handle_line('/subtask "fold the towel"') assert engine.task == "fold the towel" time.sleep(0.05) - assert session._running.is_set() + assert session.controller.running session._handle_line("/stop") thread.join(timeout=2.0) @@ -769,20 +938,164 @@ def test_drop_queued_actions_clears_both_queue_conventions(): flush(SimpleNamespace()) +# --------------------------------------------------------------------------- +# Sentry strategy: restartable run() segments + live task labels +# --------------------------------------------------------------------------- + + +def _make_sentry(monkeypatch): + from lerobot.rollout import SentryStrategy, SentryStrategyConfig + + # Keep setup independent of camera features and never rotate episodes + # mid-test; frame assembly is not under test here. + monkeypatch.setattr("lerobot.rollout.strategies.sentry.estimate_max_episode_seconds", lambda *a, **k: 1e9) + monkeypatch.setattr( + "lerobot.rollout.strategies.sentry.send_next_action", lambda *a, **k: {"joint.pos": 0.5} + ) + monkeypatch.setattr("lerobot.rollout.strategies.sentry.build_dataset_frame", lambda *a, **k: {}) + + engine = _FakeEngine() + dataset = MagicMock() + robot = MagicMock() + robot.get_observation.return_value = {"joint.pos": 0.0} + stop_event = Event() + + def identity(x): + return x + + ctx = SimpleNamespace( + runtime=SimpleNamespace( + cfg=SimpleNamespace( + play_sounds=False, + fps=200.0, + duration=0.0, + use_torch_compile=False, + interpolation_multiplier=1, + display_data=False, + return_to_initial_position=False, + task="pick up the cube", + dataset=SimpleNamespace(push_to_hub=False, tags=None, private=False), + ), + shutdown_event=stop_event, + ), + policy=SimpleNamespace(inference=engine), + hardware=SimpleNamespace(robot_wrapper=robot, teleop=None, initial_position={"joint.pos": 0.0}), + processors=SimpleNamespace( + teleop_action_processor=identity, + robot_action_processor=identity, + robot_observation_processor=identity, + ), + data=SimpleNamespace(dataset=dataset, dataset_features={}, hw_features={}, ordered_action_keys=[]), + ) + + strategy = SentryStrategy(SentryStrategyConfig()) + strategy.setup(ctx) + return strategy, ctx, dataset, engine, stop_event + + +def test_sentry_run_is_restartable_and_finalizes_only_in_teardown(monkeypatch): + strategy, ctx, dataset, _engine, stop_event = _make_sentry(monkeypatch) + + # Segment 1. + thread = Thread(target=strategy.run, args=(ctx,), daemon=True) + thread.start() + assert _wait_for(lambda: dataset.add_frame.call_count >= 2) + stop_event.set() + thread.join(timeout=2.0) + assert not thread.is_alive() + + # The segment saved its partial episode but left the dataset open. + assert dataset.save_episode.call_count == 1 + dataset.finalize.assert_not_called() + + # Segment 2: run() is restartable on the same instance. + stop_event.clear() + frames_before = dataset.add_frame.call_count + thread = Thread(target=strategy.run, args=(ctx,), daemon=True) + thread.start() + assert _wait_for(lambda: dataset.add_frame.call_count >= frames_before + 2) + stop_event.set() + thread.join(timeout=2.0) + assert dataset.save_episode.call_count == 2 + dataset.finalize.assert_not_called() + + # Only teardown finalizes. + strategy.teardown(ctx) + dataset.finalize.assert_called_once() + + +def test_sentry_labels_frames_with_live_engine_task(monkeypatch): + strategy, ctx, dataset, engine, stop_event = _make_sentry(monkeypatch) + + thread = Thread(target=strategy.run, args=(ctx,), daemon=True) + thread.start() + assert _wait_for(lambda: dataset.add_frame.call_count >= 2) + + # A live task change (e.g. /subtask through the controller) relabels + # recorded frames from that moment on. + engine.set_task("fold the towel") + assert _wait_for( + lambda: any(call.args[0]["task"] == "fold the towel" for call in dataset.add_frame.call_args_list) + ) + stop_event.set() + thread.join(timeout=2.0) + + tasks = [call.args[0]["task"] for call in dataset.add_frame.call_args_list] + assert tasks[0] == "pick up the cube" + assert tasks[-1] == "fold the towel" + + strategy.teardown(ctx) + + +def test_sentry_discards_tail_episode_when_final_save_fails(monkeypatch): + strategy, ctx, dataset, _engine, stop_event = _make_sentry(monkeypatch) + dataset.save_episode.side_effect = ValueError("disk full mid-write") + + thread = Thread(target=strategy.run, args=(ctx,), daemon=True) + thread.start() + assert _wait_for(lambda: dataset.add_frame.call_count >= 1) + stop_event.set() + thread.join(timeout=2.0) + assert not thread.is_alive() + + # The failed tail save must not leave a half-written streaming encode for + # teardown's finalize to flush, nor a half-mutated episode buffer that + # would crash the first add_frame of a restarted segment. + dataset.writer.cancel_pending_videos.assert_called_once() + assert dataset.writer.episode_buffer is None + dataset.finalize.assert_not_called() + + # --------------------------------------------------------------------------- # Config validation # --------------------------------------------------------------------------- -def test_interactive_requires_base_strategy(): +def test_interactive_rejects_keyboard_bound_strategies(): + from lerobot.configs.dataset import DatasetRecordConfig + from lerobot.rollout import HighlightStrategyConfig, RolloutConfig + from tests.mocks.mock_robot import MockRobotConfig + + with pytest.raises(ValueError, match="--interactive=true supports"): + RolloutConfig( + robot=MockRobotConfig(), + strategy=HighlightStrategyConfig(), + dataset=DatasetRecordConfig(repo_id="user/rollout_test", single_task="test"), + interactive=True, + ) + + +def test_interactive_allows_sentry(): from lerobot.configs.dataset import DatasetRecordConfig from lerobot.rollout import RolloutConfig, SentryStrategyConfig from tests.mocks.mock_robot import MockRobotConfig - with pytest.raises(ValueError, match="--interactive=true currently supports only"): - RolloutConfig( - robot=MockRobotConfig(), - strategy=SentryStrategyConfig(), - dataset=DatasetRecordConfig(repo_id="user/rollout_test", single_task="test"), - interactive=True, - ) + cfg = RolloutConfig( + robot=MockRobotConfig(), + strategy=SentryStrategyConfig(), + dataset=DatasetRecordConfig(repo_id="user/rollout_test", single_task="test"), + policy=SimpleNamespace(device="cpu"), # stands in for a PreTrainedConfig + interactive=True, + ) + assert cfg.interactive is True + assert cfg.dataset.streaming_encoding is True # sentry forces streaming diff --git a/tests/utils/test_stdin_input.py b/tests/utils/test_stdin_input.py new file mode 100644 index 000000000..3771dc103 --- /dev/null +++ b/tests/utils/test_stdin_input.py @@ -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()