Compare commits

..

1 Commits

Author SHA1 Message Date
Steven Palma a56fc0b174 refactor(rollout): integrate feedback -> api, recording, log mut and keyboard 2026-08-07 21:33:17 +02:00
20 changed files with 1912 additions and 1572 deletions
+457
View File
@@ -0,0 +1,457 @@
# Interactive Rollout — Design Notes
Branch: `feat/add_interactive_rollout` · Status: Phases 12 committed; Round 2
(programmatic API, sentry support, muting v2, stdin move) implemented and tested,
uncommitted.
---
## 1. Vision
`lerobot-rollout` runs inference on a real robot: it connects hardware, loads the policy,
builds the processor pipelines, optionally records a dataset, and spins the control loop.
Today that is a **one-shot, fire-and-forget** program. You pass `--task="pick up the cube"`
on the command line, the robot starts moving immediately, and the only interaction left is
Ctrl-C. If you want a different instruction, you kill the process and pay the full startup
cost again — reconnecting motors, re-homing, re-loading a multi-GB VLA onto the GPU.
Since LeRobot gained subtask annotation and language conditioning, that model is the
bottleneck. The **north star** is a chat-style CLI over stdin, where the operator issues
commands *concurrently with the robot moving*:
```
/start begin (or resume) the policy control loop
/subtask Grab the red cube re-instruct the policy on the fly
/ask what's the capital of France? query an LLM while the robot keeps moving
/reset stop movement, return home, clear the subtask —
but keep hardware and policy warm
/stop graceful shutdown
```
The unifying idea: **the expensive things (hardware, policy weights, processors) stay warm
across commands.** Only the cheap things — the instruction, the control loop — start and
stop. That turns a rollout from a batch job into a session you can steer.
## 2. Objective (scoped)
Phased, so each phase lands as a reviewable unit:
| Phase | Scope | Status |
|---|---|---|
| **1** | `--interactive` flag, non-blocking stdin listener, command parser, `/start` `/reset` `/stop` `/help` | ✅ done |
| **1.5** | Mute system logs so they stop fighting the prompt for the terminal | ✅ done |
| **2** | `/subtask <text>` — change the policy's instruction mid-run | ✅ done |
| **2.5** | Round 2: `RolloutController` public API, sentry recording support, muting v2 (errors surface), stdin listener → `lerobot/utils` | ✅ done (see §5) |
| **3** | `/ask` + hierarchical task-vs-subtask semantics (LLM in the loop) | not started |
An explicit constraint through Phases 12: **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 13 (Phases 12):**
- **RTC stale observation (critical).** `RTCInferenceEngine.reset()` never cleared
`_obs_holder["obs"]`. After `/reset` physically moved the robot home, the next `/start`
computed its first chunk from the **pre-reset pose** — a lurch back toward where the arm
used to be. Fixed by clearing the observation and adding a `_reset_epoch` counter so an
in-flight chunk computed across a reset is discarded rather than merged. This also fixes a
pre-existing DAgger staleness path.
- **Muting hid fatal errors** → `failure_traceback` capture + session print (§4.7).
- **Muting scope too narrow** → root-only missed `transformers` / `datasets`; `warnings`
output bypassed logging entirely.
- **Command ordering** → `/reset` and `/stop` didn't cancel a pending `/start` (§4.5); the
`/reset`-then-`/subtask` clobber (§4.8).
- **Flush too heavy** → `policy.reset()``drop_queued_actions()` (§4.8).
- **Empty-task rendering** → `''` replaced with `(none — set one with /subtask <text>)`.
- **Silent switch** → the confirmation now says "(applies from the next policy inference)",
since the explanatory logs are muted.
## 7. Verification
After Round 2:
```
uv run --extra dataset pytest tests/test_interactive_rollout.py \
tests/utils/test_stdin_input.py tests/test_rollout.py -q
→ 81 passed
pre-commit (all changed files)
→ 0 failures
```
Phase 12 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.
+66 -26
View File
@@ -260,7 +260,6 @@ lerobot-rollout \
| ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `/start` | Start (or restart) the policy control loop |
| `/subtask <text>` | Change the instruction the policy follows, without stopping. No argument prints the current task. Only affects policies that condition on language (SmolVLA, π0/π0.5, and similar) |
| `/ask <question>` | Ask a supported policy text head about its latest view. The answer is generated in the background without pausing the session |
| `/reset` | Stop movement, return the robot to its startup position, and restore the `--task` instruction |
| `/stop` | End the session and run the normal shutdown routines |
| `/help` | List the commands |
@@ -270,9 +269,6 @@ lerobot-rollout \
Rollout running — task 'pick up the cube'. /subtask <text> to change it, ...
> /subtask put the cube in the box
Task: 'pick up the cube' → 'put the cube in the box' (applies from the next policy inference)
> /ask where is the red cube?
Question queued: 'where is the red cube?' (the rollout keeps running)
[policy] The red cube is beside the bowl.
> /reset
Task restored to 'pick up the cube'
Resetting — returning the robot to its initial position...
@@ -293,37 +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.
**How `/ask` runs without taking over the rollout.** During an active rollout, the inference engine caches the latest policy-ready observation, so the command reader never touches cameras, processors, or robot hardware. A single background worker sends that snapshot to the optional `PreTrainedPolicy.generate_text(..., kind=TextKind.VQA, user_text=question)` hook and prints the result when ready. Questions are independent turns; there is no conversation history, and a second question is rejected while one is running so stale image tensors cannot accumulate. WALL-OSS (`wall_x`) is the first policy implementing this hook; policies without a compatible text head report that `/ask` is unsupported.
**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.
Text and action calls share one policy safely: the engine gives a pending question priority after the current action inference finishes, while action inference uses a non-blocking gate. The hardware loop therefore keeps ticking and `/ask` never clears an action queue. RTC continues dispatching its buffered actions while text is decoded. Sync keeps the robot on its last commanded target until the policy is available again. Text generation still consumes model/GPU capacity, so response generation can reduce action freshness; RTC is preferred when uninterrupted action buffering matters.
Sessions work over SSH and on headless machines — the command reader uses the terminal (or pipe) directly and needs no display server.
`/stop` suppresses any late answer and gives an active decoder five seconds to finish cleanly. If it is stuck, hardware teardown continues rather than leaving the robot session open indefinitely; the daemon may retain its model/GPU resources until it returns.
**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.
**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.
```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
```
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.
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 |
---
-2
View File
@@ -31,7 +31,6 @@ from .types import (
PipelineFeatureType,
PolicyFeature,
RTCAttentionSchedule,
TextKind,
)
from .video import (
DEFAULT_DEPTH_UNIT,
@@ -55,7 +54,6 @@ __all__ = [
"PipelineFeatureType",
"PolicyFeature",
"RTCAttentionSchedule",
"TextKind",
# Config classes
"DatasetRecordConfig",
"DatasetConfig",
-6
View File
@@ -67,12 +67,6 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno
# Whether the policy employed PEFT for training.
use_peft: bool = False
# Decoding defaults for policies that implement `generate_text`. They live
# in config.json so a text head uses the settings it was trained/evaluated
# with; policy-specific decoding knobs belong on the concrete config.
text_temperature: float = 0.0 # 0.0 = greedy; > 0 enables sampling
text_top_p: float = 1.0
push_to_hub: bool = True # type: ignore[assignment] # TODO: use a different name to avoid override
repo_id: str | None = None
-7
View File
@@ -31,13 +31,6 @@ class PipelineFeatureType(str, Enum):
OBSERVATION = "OBSERVATION"
class TextKind(str, Enum):
"""Text-generation requests understood by interactive policy hooks."""
SUBTASK = "subtask"
VQA = "vqa"
class NormalizationMode(str, Enum):
MIN_MAX = "MIN_MAX"
MEAN_STD = "MEAN_STD"
+1 -23
View File
@@ -28,7 +28,7 @@ from huggingface_hub.errors import HfHubHTTPError
from safetensors.torch import load_model as load_model_as_safetensor
from torch import Tensor, nn
from lerobot.configs import PreTrainedConfig, TextKind
from lerobot.configs import PreTrainedConfig
from lerobot.utils.constants import ACTION
from lerobot.utils.device_utils import resolve_safetensors_device
from lerobot.utils.hub import HubMixin
@@ -234,28 +234,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
if action_queue is not None:
action_queue.clear()
def supports_text_generation(self) -> bool:
"""Whether this policy implements the optional :meth:`generate_text` hook."""
return type(self).generate_text is not PreTrainedPolicy.generate_text
def generate_text(
self,
batch: dict[str, Tensor],
*,
kind: TextKind = TextKind.SUBTASK,
user_text: str | None = None,
) -> str:
"""Generate one string from a policy's optional language head.
Interactive rollout calls this with a policy-ready observation batch.
Implementations must treat the batch as read-only and avoid mutating
action queues or episode state: text generation runs on a background
worker while the control loop remains active.
"""
raise NotImplementedError(
f"{type(self).__name__} has no text head. Implement `generate_text` to support /ask."
)
def supports_rtc(self) -> bool:
"""Whether this policy implements Real-Time Chunking inference semantics."""
return False
@@ -52,7 +52,6 @@ from torch.nn import CrossEntropyLoss
from torchvision.transforms import InterpolationMode
from torchvision.transforms.v2 import functional as tv_functional
from lerobot.configs import TextKind
from lerobot.utils.constants import ACTION, OBS_STATE
from lerobot.utils.import_utils import (
_wallx_deps_available,
@@ -108,7 +107,6 @@ else:
from .utils import (
get_wallx_normal_text,
img_key_mapping,
preprocesser_call,
process_grounding_points,
replace_action_token,
@@ -1587,25 +1585,6 @@ class Qwen2_5_VLMoEForAction(_Qwen2_5_VLForAction_Base): # noqa: N801
- Handles special cases for input_embeds, generation methods, and GPU synchronization
- Manages vision inputs to avoid unnecessary forward passes
"""
if cache_position is None:
past_length = 0
if past_key_values is not None and hasattr(past_key_values, "get_seq_length"):
past_length = int(past_key_values.get_seq_length())
input_length = input_ids.shape[1]
end = input_length if input_length > past_length else past_length + input_length
cache_position = torch.arange(
past_length,
end,
dtype=torch.long,
device=input_ids.device,
)
if cache_position.numel() == 0:
cache_position = torch.arange(
input_length,
dtype=torch.long,
device=input_ids.device,
)
# Initialize MoE token types if not provided
if moe_token_types is None:
moe_token_types = torch.zeros_like(
@@ -1872,23 +1851,6 @@ class WallXPolicy(PreTrainedPolicy):
"""Get parameters for optimization."""
return self.parameters()
@staticmethod
def _observation_prompt(img_keys: list[str]) -> str:
prompt = "Observation:"
for label in img_key_mapping(img_keys):
prompt += f" {label}: <|vision_start|><|image_pad|><|vision_end|>"
return prompt
def _format_text_prompt(self, instruction: str, kind: str, img_keys: list[str]) -> str:
if kind == TextKind.SUBTASK:
instruction = f"{instruction}\nPredict the next action in language."
return (
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
f"<|im_start|>user\n{self._observation_prompt(img_keys)}\n"
f"Instruction: {instruction}<|im_end|>\n"
"<|im_start|>assistant\n"
)
def preprocess_inputs(
self,
batch: dict[str, Any],
@@ -2118,118 +2080,6 @@ class WallXPolicy(PreTrainedPolicy):
return loss, loss_dict
def _build_text_inputs(
self,
batch: dict[str, Any],
*,
kind: str,
user_text: str | list[str] | None,
) -> BatchFeature:
batch_size = batch[OBS_STATE].shape[0]
img_keys = [key for key in self.config.image_features if key in batch]
if not img_keys:
raise ValueError("Wall-X text generation requires at least one image feature.")
image_inputs, dimensions_by_key = _prepare_wall_x_image_inputs(batch, img_keys)
orig_height, orig_width, resized_height, resized_width = dimensions_by_key[img_keys[-1]]
tasks = batch["task"] if isinstance(batch["task"], list) else [batch["task"]] * batch_size
if user_text is None:
instructions = tasks
elif isinstance(user_text, str):
instructions = [user_text] * batch_size
elif len(user_text) == batch_size:
instructions = user_text
else:
raise ValueError(f"Expected one text prompt for each of the {batch_size} samples.")
texts = [
process_grounding_points(
self._format_text_prompt(str(instruction), kind, img_keys),
orig_height,
orig_width,
resized_height,
resized_width,
MODEL_TYPE,
)
for instruction in instructions
]
inputs = preprocesser_call(
processor=self.model.processor,
text=texts,
images=image_inputs,
videos=None,
device=batch[OBS_STATE].device,
padding=True,
truncation=True,
return_tensors="pt",
max_length=TOKENIZER_MAX_LENGTH,
)
inputs.pop("labels", None)
inputs["moe_token_types"] = torch.zeros_like(inputs.input_ids, dtype=torch.bool)
for key, value in inputs.items():
if isinstance(value, torch.Tensor):
inputs[key] = value.to(batch[OBS_STATE].device)
return inputs
@torch.no_grad()
def generate_text(
self,
batch: dict[str, Tensor],
*,
kind: TextKind = TextKind.SUBTASK,
user_text: str | None = None,
) -> str:
"""Generate one grounded language response from the WALL-OSS VLM."""
outputs = self.generate_texts(
batch,
kind=kind,
user_text=user_text,
temperature=self.config.text_temperature,
top_p=self.config.text_top_p,
)
if len(outputs) != 1:
raise ValueError(f"Interactive rollout expected one Wall-X output, got {len(outputs)}.")
return outputs[0]
@torch.no_grad()
def generate_texts(
self,
batch: dict[str, Any],
*,
kind: TextKind = TextKind.VQA,
user_text: str | list[str] | None = None,
max_new_tokens: int = 100,
min_new_tokens: int = 0,
temperature: float = 0.0,
top_p: float = 1.0,
) -> list[str]:
"""Generate grounded Wall-X text for one or more observations."""
self.eval()
if kind not in {TextKind.VQA, TextKind.SUBTASK}:
raise ValueError("Unsupported Wall-X text kind.")
inputs = self._build_text_inputs(batch, kind=kind, user_text=user_text)
prompt_length = inputs.input_ids.shape[1]
sampling = temperature > 0
generation_kwargs: dict[str, Any] = {
"max_new_tokens": max_new_tokens,
"min_new_tokens": min_new_tokens,
"do_sample": sampling,
"eos_token_id": self.model.processor.tokenizer.eos_token_id,
"pad_token_id": self.model.processor.tokenizer.pad_token_id,
"use_cache": True,
}
if sampling:
generation_kwargs.update(temperature=temperature, top_p=top_p)
output_ids = self.model.generate(**inputs, **generation_kwargs)
return [
value.strip()
for value in self.model.processor.tokenizer.batch_decode(
output_ids[:, prompt_length:],
skip_special_tokens=True,
clean_up_tokenization_spaces=True,
)
]
@torch.no_grad()
def predict_action_chunk(self, batch: dict[str, Tensor]) -> Tensor:
"""Predict action chunk for evaluation."""
+7 -7
View File
@@ -38,6 +38,11 @@ from .context import (
RuntimeContext,
build_rollout_context,
)
from .controller import (
LinkedEvent,
RolloutController,
RolloutEvent,
)
from .inference import (
InferenceEngine,
InferenceEngineConfig,
@@ -50,10 +55,6 @@ from .inference import (
from .interactive import (
InteractiveCommand,
InteractiveSession,
LinkedEvent,
StdinCommandListener,
TextQueryRequest,
TextQueryWorker,
parse_command,
)
from .strategies import (
@@ -84,20 +85,19 @@ __all__ = [
"InteractiveCommand",
"InteractiveSession",
"LinkedEvent",
"TextQueryRequest",
"TextQueryWorker",
"PolicyContext",
"ProcessorContext",
"RTCInferenceConfig",
"RTCInferenceEngine",
"RolloutConfig",
"RolloutContext",
"RolloutController",
"RolloutEvent",
"RolloutStrategy",
"RolloutStrategyConfig",
"RuntimeContext",
"SentryStrategy",
"SentryStrategyConfig",
"StdinCommandListener",
"SyncInferenceConfig",
"SyncInferenceEngine",
"build_rollout_context",
+12 -11
View File
@@ -240,11 +240,12 @@ class RolloutConfig:
fps: float = 30.0
duration: float = 0.0 # 0 = infinite (24/7 mode)
# Interactive session: control the rollout from stdin with chat-style
# commands (/start, /subtask <text>, /ask <question>, /reset, /stop) while hardware and
# commands (/start, /subtask <text>, /reset, /stop) while hardware and
# policy stay warm. The robot does not move until /start is received,
# `/subtask` re-instructs the policy mid-run, and 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
+382
View File
@@ -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)
+4 -107
View File
@@ -23,17 +23,10 @@ from __future__ import annotations
import abc
import logging
from copy import copy
from threading import Event, Lock
from typing import TYPE_CHECKING
from threading import Lock
import torch
from lerobot.configs import TextKind
if TYPE_CHECKING:
from lerobot.policies.pretrained import PreTrainedPolicy
logger = logging.getLogger(__name__)
@@ -67,40 +60,19 @@ class InferenceEngine(abc.ABC):
thread via :meth:`_take_task`, so no policy state is ever mutated
across threads.
Text queries
------------
Backends publish their latest policy-ready observation through
:meth:`_publish_text_observation`. The interactive ``/ask`` worker takes
a shallow snapshot and calls :meth:`generate_text`; an engine-owned gate
prevents language decoding from racing an action-model call. Action
inference uses a non-blocking acquire so the hardware loop keeps ticking
while a text response is being decoded.
Optional hooks
--------------
``notify_observation`` / ``pause`` / ``resume`` have a no-op default
so rollout strategies can invoke them unconditionally.
Subclasses must call ``super().__init__(task=..., policy=...)``; the task
holder and optional text-generation plumbing are set up there.
Subclasses must call ``super().__init__(task=...)``; the task holder
is set up there.
"""
def __init__(self, task: str = "", policy: PreTrainedPolicy | None = None) -> None:
def __init__(self, task: str = "") -> None:
self._task = task
self._task_changed = False
self._task_lock = Lock()
self._policy = policy
self._text_observation: dict | None = None
self._text_observation_lock = Lock()
self._text_observation_publication_enabled = True
# A text query gets priority after the current action inference ends.
# Action backends never block on this lock: they return "no action yet"
# and let the control loop keep servicing hardware at its normal rate.
self._policy_call_lock = Lock()
self._text_query_pending = Event()
self._text_query_serial_lock = Lock()
# ------------------------------------------------------------------
# Task (language instruction)
@@ -142,81 +114,6 @@ class InferenceEngine(abc.ABC):
with self._task_lock:
self._task_changed = False
# ------------------------------------------------------------------
# Optional text generation (interactive /ask)
# ------------------------------------------------------------------
def supports_text_generation(self) -> bool:
"""Whether the attached policy implements the optional text hook."""
return self._policy is not None and self._policy.supports_text_generation()
def _publish_text_observation(self, observation: dict) -> None:
"""Cache a policy-ready observation for a future background query."""
with self._text_observation_lock:
if self._text_observation_publication_enabled:
self._text_observation = copy(observation)
def invalidate_text_observation(self) -> None:
"""Discard the cached view and reject publications until the next reset.
``/reset`` calls this before its deferred main-thread homing begins.
Keeping publication disabled matters because action inference may
already be in flight and otherwise republish the pre-reset scene.
"""
with self._text_observation_lock:
self._text_observation = None
self._text_observation_publication_enabled = False
def _reset_text_observation(self) -> None:
"""Clear the cached view and allow the next control segment to publish."""
with self._text_observation_lock:
self._text_observation = None
self._text_observation_publication_enabled = True
def snapshot_text_observation(self) -> dict | None:
"""Return the latest policy-ready view without touching robot hardware."""
with self._text_observation_lock:
if self._text_observation is None:
return None
observation = copy(self._text_observation)
# A /subtask may have arrived since this visual observation was cached.
# Keep the image/state snapshot but pair it with the latest instruction.
observation["task"] = self.task
return observation
def generate_text(
self,
observation: dict,
*,
kind: TextKind,
user_text: str | None = None,
) -> str:
"""Run one policy text query outside the hardware loop.
Only one query is admitted at a time. Setting the pending edge before
acquiring the shared policy gate prevents a busy action backend from
repeatedly winning the lock and starving the question.
"""
if self._policy is None or not self._policy.supports_text_generation():
raise NotImplementedError("This policy does not support text generation.")
with self._text_query_serial_lock:
self._text_query_pending.set()
try:
with self._policy_call_lock, torch.inference_mode():
return self._policy.generate_text(observation, kind=kind, user_text=user_text)
finally:
self._text_query_pending.clear()
def _try_begin_action_inference(self) -> bool:
"""Acquire policy ownership without blocking the hardware loop."""
if self._text_query_pending.is_set():
return False
return self._policy_call_lock.acquire(blocking=False)
def _end_action_inference(self) -> None:
"""Release policy ownership acquired by :meth:`_try_begin_action_inference`."""
self._policy_call_lock.release()
@abc.abstractmethod
def start(self) -> None:
"""Initialise the backend."""
+5 -15
View File
@@ -124,7 +124,8 @@ class RTCInferenceEngine(InferenceEngine):
rtc_queue_threshold: int = 30,
shutdown_event: Event | None = None,
) -> None:
super().__init__(task=task, policy=policy)
super().__init__(task=task)
self._policy = policy
self._preprocessor = preprocessor
self._postprocessor = postprocessor
self._robot = robot_wrapper
@@ -259,16 +260,14 @@ class RTCInferenceEngine(InferenceEngine):
chunk is discarded instead of merged into the cleared queue.
"""
logger.info("Resetting RTC inference state (policy + processors + queue)")
with self._policy_call_lock:
self._policy.reset()
self._preprocessor.reset()
self._postprocessor.reset()
self._policy.reset()
self._preprocessor.reset()
self._postprocessor.reset()
if self._action_queue is not None:
self._action_queue.clear()
with self._obs_lock:
self._obs_holder["obs"] = None
self._reset_epoch += 1
self._reset_text_observation()
# The queue was just cleared, so a pending task change has nothing
# stale left to blend against.
self._discard_task_change()
@@ -317,12 +316,6 @@ class RTCInferenceEngine(InferenceEngine):
continue
if queue.qsize() <= self._rtc_queue_threshold:
if not self._try_begin_action_inference():
# A background /ask is using the policy. The control
# thread can keep draining the already-produced RTC
# actions; do not start another model call meanwhile.
time.sleep(_RTC_IDLE_SLEEP_S)
continue
try:
current_time = time.perf_counter()
idx_before = queue.get_action_index()
@@ -351,7 +344,6 @@ class RTCInferenceEngine(InferenceEngine):
obs_batch["task"] = [task]
preprocessed = self._preprocessor(obs_batch)
self._publish_text_observation(preprocessed)
if prev_actions is not None and self._relative_step is not None:
# Rebase against the raw cached state so the leftover tail stays in
@@ -420,8 +412,6 @@ class RTCInferenceEngine(InferenceEngine):
# Persistent failure: stop retrying and propagate shutdown.
raise
time.sleep(_RTC_ERROR_RETRY_DELAY_S)
finally:
self._end_action_inference()
else:
time.sleep(_RTC_IDLE_SLEEP_S)
+36 -48
View File
@@ -65,7 +65,8 @@ class SyncInferenceEngine(InferenceEngine):
device: str | None,
robot_type: str,
) -> None:
super().__init__(task=task, policy=policy)
super().__init__(task=task)
self._policy = policy
self._preprocessor = preprocessor
self._postprocessor = postprocessor
self._dataset_features = dataset_features
@@ -89,11 +90,9 @@ class SyncInferenceEngine(InferenceEngine):
def reset(self) -> None:
"""Reset the policy and pre/post-processors."""
logger.info("Resetting sync inference state (policy + processors)")
with self._policy_call_lock:
self._policy.reset()
self._preprocessor.reset()
self._postprocessor.reset()
self._reset_text_observation()
self._policy.reset()
self._preprocessor.reset()
self._postprocessor.reset()
# The policy was just reset, so a pending task change has nothing
# stale left to flush.
self._discard_task_change()
@@ -102,46 +101,35 @@ class SyncInferenceEngine(InferenceEngine):
"""Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
if obs_frame is None:
return None
if not self._try_begin_action_inference():
# A background /ask owns (or is waiting for) the policy. Do not
# block the control thread; the robot keeps executing its last
# dispatched target until action inference becomes available.
return None
try:
# Shallow copy is intentional: the caller (`send_next_action`) builds
# ``obs_frame`` fresh per tick via ``build_dataset_frame``, so the
# tensor/array values are not shared with any other reader.
observation = copy(obs_frame)
autocast_ctx = (
torch.autocast(device_type=self._device.type)
if self._device.type == "cuda" and self._policy.config.use_amp
else nullcontext()
)
task, task_changed = self._take_task()
with torch.inference_mode(), autocast_ctx:
if task_changed:
# Chunking policies serve actions from an internal queue filled
# under the previous instruction (up to chunk_size ticks of stale
# behavior), so drop them and let the new instruction take effect
# on this very tick. Deliberately narrower than ``policy.reset``:
# observation history and other episode state are kept, so a
# policy that conditions on them (and one that ignores the task
# entirely) sees no discontinuity. Safe to mutate here — this is
# the thread that calls ``select_action``.
logger.info("Task changed to '%s' — dropping precomputed actions", task)
self._policy.drop_queued_actions()
observation = prepare_observation_for_inference(
observation, self._device, task, self._robot_type
)
observation = self._preprocessor(observation)
self._publish_text_observation(observation)
action = self._policy.select_action(observation)
action = self._postprocessor(action)
action_tensor = action.squeeze(0).cpu()
# Shallow copy is intentional: the caller (`send_next_action`) builds
# ``obs_frame`` fresh per tick via ``build_dataset_frame``, so the
# tensor/array values are not shared with any other reader.
observation = copy(obs_frame)
autocast_ctx = (
torch.autocast(device_type=self._device.type)
if self._device.type == "cuda" and self._policy.config.use_amp
else nullcontext()
)
task, task_changed = self._take_task()
with torch.inference_mode(), autocast_ctx:
if task_changed:
# Chunking policies serve actions from an internal queue filled
# under the previous instruction (up to chunk_size ticks of stale
# behavior), so drop them and let the new instruction take effect
# on this very tick. Deliberately narrower than ``policy.reset``:
# observation history and other episode state are kept, so a
# policy that conditions on them (and one that ignores the task
# entirely) sees no discontinuity. Safe to mutate here — this is
# the thread that calls ``select_action``.
logger.info("Task changed to '%s' — dropping precomputed actions", task)
self._policy.drop_queued_actions()
observation = prepare_observation_for_inference(observation, self._device, task, self._robot_type)
observation = self._preprocessor(observation)
action = self._policy.select_action(observation)
action = self._postprocessor(action)
action_tensor = action.squeeze(0).cpu()
# Reorder to match dataset action ordering so the caller can treat
# the returned tensor uniformly across backends.
action_dict = make_robot_action(action_tensor, self._dataset_features)
return torch.tensor([action_dict[k] for k in self._ordered_action_keys])
finally:
self._end_action_inference()
# Reorder to match dataset action ordering so the caller can treat
# the returned tensor uniformly across backends.
action_dict = make_robot_action(action_tensor, self._dataset_features)
return torch.tensor([action_dict[k] for k in self._ordered_action_keys])
+112 -590
View File
@@ -19,58 +19,55 @@ rollout from the terminal while hardware and policy stay connected and warm:
/start start (or restart) the policy control loop
/subtask <text> change the instruction the policy follows, mid-run
/ask <question> ask the policy about the latest view without stopping
/reset stop movement, return the robot to its initial position,
and restore the instruction passed on the command line
/stop end the session and run the normal shutdown routines
/help show the available commands
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). ``/ask`` snapshots the latest policy-ready observation and hands it
to one background text-query worker; the worker never reads robot hardware.
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 can be registered without restructuring the
parser, the help output, or the session loop.
mapping so further commands (``/ask`` and the rest of the language-runtime
work in PR #4183/#4234) can be registered without restructuring the parser,
the help output, or the session loop.
"""
from __future__ import annotations
import contextlib
import logging
import 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 queue import Empty, Queue
from threading import Event, Lock, Thread
from typing import IO, TYPE_CHECKING
from lerobot.configs import TextKind
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
@@ -80,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)
@@ -170,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("/"):
@@ -185,301 +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()
@dataclass(frozen=True)
class TextQueryRequest:
"""One immutable policy question paired with the view seen at submission."""
question: str
observation: dict
class TextQueryWorker:
"""Single background worker for non-blocking interactive policy questions.
At most one request may be queued or running. This bounds the number of
retained image tensors (the observation can live on the GPU) and gives the
operator an explicit "busy" response instead of accumulating questions
against increasingly stale views.
"""
_STOP = object()
_JOIN_TIMEOUT_S = 5.0
def __init__(
self,
answer: Callable[[TextQueryRequest], str],
on_answer: Callable[[TextQueryRequest, str], None],
on_error: Callable[[TextQueryRequest, Exception], None],
) -> None:
self._answer = answer
self._on_answer = on_answer
self._on_error = on_error
self._queue: Queue[TextQueryRequest | object] = Queue(maxsize=1)
self._state_lock = Lock()
self._busy = False
self._stopping = Event()
self._stop_enqueued = False
self._thread: Thread | None = None
@property
def busy(self) -> bool:
with self._state_lock:
return self._busy
def start(self) -> None:
"""Start the worker (idempotent)."""
with self._state_lock:
if self._thread is not None or self._stopping.is_set():
return
self._thread = Thread(target=self._run, daemon=True, name="InteractiveTextQuery")
self._thread.start()
def submit(self, request: TextQueryRequest) -> bool:
"""Queue ``request`` without blocking; return ``False`` when busy or stopping."""
with self._state_lock:
if self._busy or self._stopping.is_set():
return False
self._busy = True
# stop() takes the same lock before publishing its sentinel, so a
# successful admission cannot race with shutdown and hit Queue.Full.
self._queue.put_nowait(request)
return True
def cancel(self) -> None:
"""Reject new work, discard a queued request, and suppress late callbacks."""
with self._state_lock:
self._stopping.set()
discard_request = not self._stop_enqueued
if discard_request:
self._discard_queued_request()
def stop(self, timeout_s: float = _JOIN_TIMEOUT_S) -> bool:
"""Cancel queued work and give an active model call bounded time to finish.
Returns ``False`` when decoding is still stuck after ``timeout_s``. The
worker is a daemon and callbacks stay suppressed, allowing hardware
teardown to proceed instead of hanging indefinitely.
"""
self.cancel()
with self._state_lock:
thread = self._thread
enqueue_stop = thread is not None and not self._stop_enqueued
self._stop_enqueued = self._stop_enqueued or enqueue_stop
if thread is None:
return True
if enqueue_stop:
self._queue.put(self._STOP)
thread.join(timeout=timeout_s)
stopped = not thread.is_alive()
if stopped:
with self._state_lock:
if self._thread is thread:
self._thread = None
return stopped
def _discard_queued_request(self) -> None:
try:
item = self._queue.get_nowait()
except Empty:
return
self._queue.task_done()
if item is self._STOP:
# A concurrent/repeated cancel must not consume the sentinel that
# an earlier stop() already published for the worker.
self._queue.put_nowait(self._STOP)
return
with self._state_lock:
self._busy = False
def _run(self) -> None:
while True:
item = self._queue.get()
try:
if item is self._STOP:
return
request = item
assert isinstance(request, TextQueryRequest)
if self._stopping.is_set():
continue
try:
answer = self._answer(request)
except Exception as exc: # a language failure must not end robot control
self._deliver(self._on_error, request, exc)
else:
self._deliver(self._on_answer, request, answer)
finally:
if item is not self._STOP:
with self._state_lock:
self._busy = False
self._queue.task_done()
def _deliver(self, callback: Callable[..., None], *args) -> None:
"""Linearize a result callback with cancellation."""
with self._state_lock:
if not self._stopping.is_set():
callback(*args)
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
@@ -487,156 +153,85 @@ 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)
self._text_query = TextQueryWorker(
self._answer_text_query,
self._report_text_answer,
self._report_text_error,
)
# 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 stay documented for free.
# render from this table, so future commands (e.g. /ask) stay
# documented for free.
self._commands: dict[str, tuple[Callable[[InteractiveCommand], None], str, str]] = {
"start": (self._cmd_start, "", "start (or restart) the policy control loop"),
"subtask": (self._cmd_subtask, " <text>", "set the instruction the policy follows"),
"ask": (self._cmd_ask, " <question>", "ask the policy about the latest view"),
"reset": (self._cmd_reset, "", "stop movement, return to initial position, restore the task"),
"stop": (self._cmd_stop, "", "end the session and shut down"),
"help": (self._cmd_help, "", "show this help"),
}
# ------------------------------------------------------------------
# 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._text_query.start()
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()
# A model call cannot be force-cancelled safely. Give it a bounded
# grace period, then prioritize hardware teardown if it is wedged.
if not self._text_query.stop():
self._print(
"Policy question did not finish within 5 seconds — "
"continuing hardware shutdown; its daemon thread will be abandoned."
)
# 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 <text> to change it, /reset to return to initial position, /stop to shut down."
)
elif event is RolloutEvent.SEGMENT_ENDED:
self._print(
"Rollout run ended on its own (duration reached). Robot is holding position — "
"/start to run again, /reset to return to initial position, /stop to shut down."
)
elif event is RolloutEvent.RESET_STARTED:
log_say("Resetting robot to initial position", self._play_sounds)
self._print("Resetting — returning the robot to its initial position...")
elif event is RolloutEvent.RESET_DONE:
self._print("Robot reset — holding at initial position. /start to run.")
elif event is RolloutEvent.RESET_SKIPPED:
self._print("Robot paused — no initial position captured, holding current pose. /start to run.")
elif event is RolloutEvent.ENGINE_FAILED:
self._report_engine_failure()
def _report_engine_failure(self) -> None:
"""Surface a fatal engine error despite the muted console logging."""
self._print("Inference engine failed — shutting down.")
failure_traceback = self._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 <text> to change it, /ask <question> to query the policy, "
"/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:
@@ -653,28 +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
if self._text_query.busy:
self._print("A policy question is still finishing — wait for it before /start.")
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)"
@@ -682,76 +268,12 @@ class InteractiveSession:
else:
self._print(f"Task unchanged: {_format_task(task)}")
def _cmd_ask(self, cmd: InteractiveCommand) -> None:
question = _strip_quotes(cmd.args)
if not question:
self._print("Usage: /ask <question>")
return
engine = self._ctx.policy.inference
if not engine.supports_text_generation():
self._print("This policy does not support /ask (it has no text-generation head).")
return
if not self._running.is_set():
self._print("The rollout is not running — /start it before using /ask.")
return
observation = engine.snapshot_text_observation()
if observation is None:
self._print("No policy observation is available yet — /start the rollout and try again.")
return
request = TextQueryRequest(question=question, observation=observation)
if not self._text_query.submit(request):
self._print("A policy question is already being answered — try again when it finishes.")
return
self._print(f"Question queued: {question!r} (the rollout keeps running)")
def _answer_text_query(self, request: TextQueryRequest) -> str:
return self._ctx.policy.inference.generate_text(
request.observation,
kind=TextKind.VQA,
user_text=request.question,
)
def _report_text_answer(self, request: TextQueryRequest, answer: str) -> None:
if answer:
self._print(f"[policy] {answer}")
else:
self._print(f"The policy returned no answer for {request.question!r}.")
def _report_text_error(self, request: TextQueryRequest, exc: Exception) -> None:
self._print(f"Policy question failed ({type(exc).__name__}): {exc}")
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.
engine = self._ctx.policy.inference
if engine.set_task(self._initial_task):
self._print(f"Task restored to {_format_task(self._initial_task)}")
# Homing changes the scene outside normal inference. Invalidate the
# cached VLM input synchronously so a following /ask cannot capture
# the pre-reset view while the main thread is still unwinding.
engine.invalidate_text_observation()
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:
# Suppress a queued/finishing answer as soon as /stop or EOF is
# observed; stop() in the session's finally block gives the model call
# bounded time to finish before hardware teardown continues.
self._text_query.cancel()
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())
@@ -770,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}"
)
+89 -63
View File
@@ -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."""
+5 -4
View File
@@ -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 <text> re-instructs the policy mid-run; /ask <question>
# queries a supported policy text head in the background; /reset returns
# to the initial position (hardware and policy stay warm); /stop shuts down
# Interactive session (base or sentry strategy): the robot stays idle
# until /start is typed; /subtask <text> re-instructs the policy mid-run;
# /reset returns it to the initial position (hardware and policy stay
# warm); /stop shuts down gracefully. With --strategy.type=sentry the
# session also records continuously, labeling frames with the live task.
lerobot-rollout \\
--strategy.type=base \\
--policy.path=lerobot/act_koch_real \\
+187
View File
@@ -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()
+1 -81
View File
@@ -16,8 +16,6 @@
"""Test script to verify Wall-X policy integration with LeRobot"""
from types import SimpleNamespace
import pytest
import torch
@@ -26,15 +24,11 @@ pytest.importorskip("peft")
pytest.importorskip("transformers")
pytest.importorskip("torchdiffeq")
from lerobot.configs import TextKind # noqa: E402
from lerobot.policies.factory import make_policy_config # noqa: E402
from lerobot.policies.wall_x import (
WallXConfig, # noqa: E402
)
from lerobot.policies.wall_x.modeling_wall_x import ( # noqa: E402
Qwen2_5_VLMoEForAction,
WallXPolicy,
)
from lerobot.policies.wall_x.modeling_wall_x import WallXPolicy # noqa: E402
from lerobot.policies.wall_x.processor_wall_x import make_wall_x_pre_post_processors # noqa: E402
from lerobot.policies.wall_x.qwen_model import Qwen2_5_VLMoEModel, Qwen2_5_VLTextConfig # noqa: E402
from lerobot.utils.random_utils import set_seed # noqa: E402
@@ -82,80 +76,6 @@ def test_moe_model_captures_requested_hidden_states_and_attentions():
assert len(output.attentions) == config.num_hidden_layers
def _make_unloaded_policy():
policy = WallXPolicy.__new__(WallXPolicy)
torch.nn.Module.__init__(policy)
policy.config = SimpleNamespace(
text_temperature=0.0,
text_top_p=1.0,
)
return policy
def test_policy_exposes_grounded_text_generation(monkeypatch):
class Inputs(dict):
__getattr__ = dict.__getitem__
class Tokenizer:
eos_token_id = 2
pad_token_id = 0
@staticmethod
def batch_decode(token_ids, **kwargs):
del kwargs
assert torch.equal(token_ids, torch.tensor([[7, 8]]))
return ["The mug is beside the bowl."]
class Model:
processor = SimpleNamespace(tokenizer=Tokenizer())
@staticmethod
def generate(input_ids, **kwargs):
del kwargs
return torch.cat([input_ids, torch.tensor([[7, 8]])], dim=1)
policy = _make_unloaded_policy()
policy.model = Model()
inputs = Inputs(input_ids=torch.tensor([[1, 2, 3]]), attention_mask=torch.ones(1, 3))
monkeypatch.setattr(policy, "_build_text_inputs", lambda *args, **kwargs: inputs)
batch = {"observation.state": torch.zeros(1, 7), "task": "pick up the cup"}
assert (
policy.generate_text(batch, kind=TextKind.VQA, user_text="Where is the mug?")
== "The mug is beside the bowl."
)
assert policy.supports_text_generation()
prompt = policy._format_text_prompt(
"Where is the mug?",
TextKind.VQA,
["observation.images.face_view"],
)
assert "Observation: front view:" in prompt
assert "Instruction: Where is the mug?" in prompt
assert prompt.endswith("<|im_start|>assistant\n")
def test_text_generation_synthesizes_missing_cache_position():
class Cache:
@staticmethod
def get_seq_length():
return 3
pixel_values = torch.ones(1, 3, 4, 4)
inputs = Qwen2_5_VLMoEForAction.prepare_inputs_for_generation(
object(),
torch.tensor([[9]]),
past_key_values=Cache(),
pixel_values=pixel_values,
)
assert torch.equal(inputs["cache_position"], torch.tensor([3]))
assert torch.equal(inputs["input_ids"], torch.tensor([[9]]))
# A continuation token reuses the KV cache, so the image is not encoded again.
assert inputs["pixel_values"] is None
@require_cuda
@require_hf_token
def test_policy_instantiation():
+412 -432
View File
@@ -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
@@ -21,7 +22,7 @@ import io
import logging
import os
import time
from threading import Event, Thread, current_thread
from threading import Event, Thread
from types import SimpleNamespace
from unittest.mock import MagicMock
@@ -30,15 +31,13 @@ import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.configs import TextKind # noqa: E402
from lerobot.rollout import ( # noqa: E402
InferenceEngine,
InteractiveCommand,
InteractiveSession,
LinkedEvent,
StdinCommandListener,
TextQueryRequest,
TextQueryWorker,
RolloutController,
RolloutEvent,
parse_command,
)
@@ -149,97 +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()
def test_text_query_worker_stop_is_bounded_and_suppresses_callback():
query_started = Event()
release_query = Event()
answers = []
def answer(request):
query_started.set()
assert release_query.wait(timeout=2.0)
return request.question
worker = TextQueryWorker(
answer=answer,
on_answer=lambda request, response: answers.append((request, response)),
on_error=lambda request, exc: pytest.fail(f"unexpected error for {request}: {exc}"),
)
worker.start()
assert worker.submit(TextQueryRequest("question", {}))
assert _wait_for(query_started.is_set)
assert worker.stop(timeout_s=0.01) is False
active_thread = worker._thread
worker.start()
assert worker._thread is active_thread
release_query.set()
assert worker.stop(timeout_s=2.0) is True
assert worker._thread is None
assert answers == []
# ---------------------------------------------------------------------------
# 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.
"""
@@ -253,8 +170,8 @@ class _FakeEngine(InferenceEngine):
def reset(self) -> None: ...
def get_action(self, obs_frame=None): ...
def __init__(self, task: str = "pick up the cube", policy=None) -> None:
super().__init__(task=task, policy=policy)
def __init__(self, task: str = "pick up the cube") -> None:
super().__init__(task=task)
self.start = MagicMock()
self.stop = MagicMock()
self.reset = MagicMock()
@@ -264,11 +181,11 @@ class _FakeEngine(InferenceEngine):
self.get_action = MagicMock(return_value=None)
def _make_session(input_stream, run_behavior=None, policy=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(policy=policy)
engine = _FakeEngine()
ctx = SimpleNamespace(
runtime=SimpleNamespace(
cfg=SimpleNamespace(play_sounds=False),
@@ -287,7 +204,224 @@ def _make_session(input_stream, run_behavior=None, policy=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
@@ -512,58 +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():
@@ -662,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)
@@ -707,220 +842,6 @@ def test_session_reset_restores_initial_task():
thread.join(timeout=2.0)
def test_session_reset_invalidates_text_observation(capsys):
policy = MagicMock()
policy.supports_text_generation.return_value = True
with _pipe_stream() as (reader, _writer):
session, _strategy, engine, _parent, _run_started = _make_session(reader, policy=policy)
engine._publish_text_observation({"observation.state": "before-reset"})
thread = _start_session_thread(session)
session._handle_line("/reset")
# Simulate an action inference that started before /reset but reaches
# publication after the command handler invalidated the old scene.
engine._publish_text_observation({"observation.state": "late-before-reset"})
assert engine.snapshot_text_observation() is None
session._handle_line("/ask what can you see?")
assert "rollout is not running" in capsys.readouterr().out
policy.generate_text.assert_not_called()
engine._reset_text_observation()
engine._publish_text_observation({"observation.state": "after-restart"})
assert engine.snapshot_text_observation()["observation.state"] == "after-restart"
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_ask_reports_usage_unsupported_policy_and_missing_observation(capsys):
with _pipe_stream() as (reader, _writer):
session, _strategy, _engine, _parent, _run_started = _make_session(reader)
thread = _start_session_thread(session)
session._handle_line("/ask")
session._handle_line("/ask what can you see?")
out = capsys.readouterr().out
assert "Usage: /ask <question>" in out
assert "does not support /ask" in out
session._handle_line("/stop")
thread.join(timeout=2.0)
policy = MagicMock()
policy.supports_text_generation.return_value = True
with _pipe_stream() as (reader, _writer):
session, _strategy, _engine, _parent, run_started = _make_session(reader, policy=policy)
thread = _start_session_thread(session)
session._handle_line("/ask what can you see?")
assert "rollout is not running" in capsys.readouterr().out
session._handle_line("/start")
assert _wait_for(run_started.is_set)
session._handle_line("/ask what can you see?")
assert "No policy observation is available yet" in capsys.readouterr().out
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_ask_runs_in_background_without_stopping_rollout(capsys):
query_started = Event()
release_query = Event()
calls = []
class TextPolicy:
@staticmethod
def supports_text_generation() -> bool:
return True
@staticmethod
def generate_text(observation, *, kind, user_text):
calls.append((observation, kind, user_text, current_thread().name))
query_started.set()
assert release_query.wait(timeout=2.0)
return "The red cube is beside the bowl."
with _pipe_stream() as (reader, _writer):
session, strategy, engine, _parent, run_started = _make_session(reader, policy=TextPolicy())
engine._publish_text_observation({"observation.state": "snapshot", "task": "stale task"})
engine.set_task("current task")
thread = _start_session_thread(session)
session._handle_line("/start")
assert _wait_for(run_started.is_set)
session._handle_line('/ask "Where is the red cube?"')
assert _wait_for(query_started.is_set)
# The query worker is blocked in generation, but the rollout segment
# and command listener remain live and the engine has not been paused.
assert session._running.is_set()
assert strategy.run.call_count == 1
engine.pause.assert_not_called()
session._handle_line("/ask another question")
assert "already being answered" in capsys.readouterr().out
release_query.set()
assert _wait_for(lambda: not session._text_query.busy)
out = capsys.readouterr().out
assert "[policy] The red cube is beside the bowl." in out
assert len(calls) == 1
observation, kind, user_text, thread_name = calls[0]
assert observation["observation.state"] == "snapshot"
assert observation["task"] == "current task"
assert kind is TextKind.VQA
assert user_text == "Where is the red cube?"
assert thread_name == "InteractiveTextQuery"
session._handle_line("/stop")
thread.join(timeout=2.0)
assert not thread.is_alive()
def test_session_ask_failure_is_nonfatal(capsys):
class FailingTextPolicy:
@staticmethod
def supports_text_generation() -> bool:
return True
@staticmethod
def generate_text(observation, *, kind, user_text):
del observation, kind, user_text
raise RuntimeError("decoder failed")
with _pipe_stream() as (reader, _writer):
session, _strategy, engine, _parent, run_started = _make_session(reader, policy=FailingTextPolicy())
engine._publish_text_observation({"observation.state": "snapshot"})
thread = _start_session_thread(session)
session._handle_line("/start")
assert _wait_for(run_started.is_set)
session._handle_line("/ask what can you see?")
assert _wait_for(lambda: not session._text_query.busy)
assert "Policy question failed (RuntimeError): decoder failed" in capsys.readouterr().out
assert session._running.is_set()
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_does_not_restart_while_text_query_owns_policy(capsys):
query_started = Event()
release_query = Event()
class TextPolicy:
@staticmethod
def supports_text_generation() -> bool:
return True
@staticmethod
def generate_text(observation, *, kind, user_text):
del observation, kind, user_text
query_started.set()
assert release_query.wait(timeout=2.0)
return "answer"
with _pipe_stream() as (reader, _writer):
session, strategy, engine, _parent, run_started = _make_session(reader, policy=TextPolicy())
engine._publish_text_observation({"observation.state": "snapshot"})
thread = _start_session_thread(session)
session._handle_line("/start")
assert _wait_for(run_started.is_set)
session._handle_line("/ask question")
assert _wait_for(query_started.is_set)
session._handle_line("/reset")
assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
session._handle_line("/start")
assert "question is still finishing" in capsys.readouterr().out
assert strategy.run.call_count == 1
release_query.set()
assert _wait_for(lambda: not session._text_query.busy)
run_started.clear()
session._handle_line("/start")
assert _wait_for(run_started.is_set)
assert strategy.run.call_count == 2
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_stop_suppresses_late_text_answer(capsys):
query_started = Event()
release_query = Event()
class TextPolicy:
@staticmethod
def supports_text_generation() -> bool:
return True
@staticmethod
def generate_text(observation, *, kind, user_text):
del observation, kind, user_text
query_started.set()
assert release_query.wait(timeout=2.0)
return "late answer"
with _pipe_stream() as (reader, _writer):
session, _strategy, engine, _parent, run_started = _make_session(reader, policy=TextPolicy())
engine._publish_text_observation({"observation.state": "snapshot"})
thread = _start_session_thread(session)
session._handle_line("/start")
assert _wait_for(run_started.is_set)
session._handle_line("/ask question")
assert _wait_for(query_started.is_set)
capsys.readouterr()
session._handle_line("/stop")
# Shutdown waits for the model call so teardown cannot race its GPU use.
time.sleep(0.05)
assert thread.is_alive()
release_query.set()
thread.join(timeout=2.0)
assert not thread.is_alive()
assert "[policy] late answer" not in capsys.readouterr().out
# ---------------------------------------------------------------------------
# InferenceEngine task holder (the /subtask plumbing)
# ---------------------------------------------------------------------------
@@ -950,90 +871,6 @@ def test_engine_discard_task_change():
assert engine._take_task() == ("b", False)
def test_engine_text_query_has_priority_without_blocking_action_loop():
query_started = Event()
release_query = Event()
class TextPolicy:
@staticmethod
def supports_text_generation() -> bool:
return True
@staticmethod
def generate_text(observation, *, kind, user_text):
del observation, kind, user_text
query_started.set()
assert release_query.wait(timeout=2.0)
return "answer"
engine = _FakeEngine(policy=TextPolicy())
query_thread = Thread(
target=lambda: engine.generate_text({}, kind=TextKind.VQA, user_text="question"),
daemon=True,
)
query_thread.start()
assert _wait_for(query_started.is_set)
# Action inference probes the gate without waiting for text decoding.
assert engine._try_begin_action_inference() is False
release_query.set()
query_thread.join(timeout=2.0)
assert not query_thread.is_alive()
assert engine._try_begin_action_inference() is True
engine._end_action_inference()
def test_rtc_engine_does_not_race_action_inference_with_text_query():
from lerobot.policies.rtc.configuration_rtc import RTCConfig
from lerobot.rollout.inference import RTCInferenceEngine
query_started = Event()
release_query = Event()
def generate_text(observation, *, kind, user_text):
del observation, kind, user_text
query_started.set()
assert release_query.wait(timeout=2.0)
return "answer"
policy = MagicMock()
policy.supports_text_generation.return_value = True
policy.generate_text.side_effect = generate_text
engine = RTCInferenceEngine(
policy=policy,
preprocessor=SimpleNamespace(steps=[]),
postprocessor=SimpleNamespace(steps=[]),
robot_wrapper=SimpleNamespace(action_features={}, robot_type="test"),
rtc_config=RTCConfig(),
hw_features={},
task="test",
fps=30,
device="cpu",
)
query_thread = Thread(
target=lambda: engine.generate_text({}, kind=TextKind.VQA, user_text="question"),
daemon=True,
)
query_thread.start()
assert _wait_for(query_started.is_set)
engine.start()
try:
engine.notify_observation({"joint.pos": 0.0})
engine.resume()
time.sleep(0.05)
# The RTC worker remains responsive but does not enter the same policy
# while language decoding owns it.
assert engine._rtc_thread is not None and engine._rtc_thread.is_alive()
policy.predict_action_chunk.assert_not_called()
finally:
engine.stop()
release_query.set()
query_thread.join(timeout=2.0)
assert not query_thread.is_alive()
def test_sync_engine_uses_new_task_and_flushes_precomputed_actions():
"""A /subtask switch must reach the policy and drop stale queued actions."""
import torch
@@ -1060,7 +897,6 @@ def test_sync_engine_uses_new_task_and_flushes_precomputed_actions():
engine.get_action({"observation.state": np.zeros(1, dtype=np.float32)})
assert policy.drop_queued_actions.call_count == 0
assert policy.select_action.call_args[0][0]["task"] == "pick up the cube"
assert engine.snapshot_text_observation()["task"] == "pick up the cube"
engine.set_task("fold the towel")
engine.get_action({"observation.state": np.zeros(1, dtype=np.float32)})
@@ -1102,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
+136
View File
@@ -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()