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