feat(rollout): interactive v1

This commit is contained in:
Steven Palma
2026-08-07 15:22:44 +02:00
parent 266be2bd17
commit 072c697c0e
9 changed files with 1121 additions and 29 deletions
+2 -1
View File
@@ -244,7 +244,7 @@ See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.
## 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 | -- |
@@ -256,6 +256,7 @@ See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters.
| `--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 (`/start`, `/reset`, `/stop`); 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 |
+14 -2
View File
@@ -47,6 +47,13 @@ from .inference import (
SyncInferenceEngine,
create_inference_engine,
)
from .interactive import (
InteractiveCommand,
InteractiveSession,
LinkedEvent,
StdinCommandListener,
parse_command,
)
from .strategies import (
BaseStrategy,
DAggerStrategy,
@@ -65,13 +72,16 @@ __all__ = [
"DAggerStrategy",
"DAggerStrategyConfig",
"DatasetContext",
"EpisodicStrategy",
"EpisodicStrategyConfig",
"HardwareContext",
"HighlightStrategy",
"HighlightStrategyConfig",
"EpisodicStrategy",
"EpisodicStrategyConfig",
"InferenceEngine",
"InferenceEngineConfig",
"InteractiveCommand",
"InteractiveSession",
"LinkedEvent",
"PolicyContext",
"ProcessorContext",
"RTCInferenceConfig",
@@ -83,9 +93,11 @@ __all__ = [
"RuntimeContext",
"SentryStrategy",
"SentryStrategyConfig",
"StdinCommandListener",
"SyncInferenceConfig",
"SyncInferenceEngine",
"build_rollout_context",
"create_inference_engine",
"create_strategy",
"parse_command",
]
+16
View File
@@ -239,6 +239,11 @@ class RolloutConfig:
# Runtime
fps: float = 30.0
duration: float = 0.0 # 0 = infinite (24/7 mode)
# Interactive session: control the rollout from stdin with chat-style
# commands (/start, /reset, /stop) while hardware and policy stay warm.
# The robot does not move until /start is received. Currently limited to
# --strategy.type=base.
interactive: bool = False
interpolation_multiplier: int = 1
device: str | None = None
task: str = ""
@@ -294,6 +299,17 @@ class RolloutConfig:
"Base strategy does not record data. Use sentry, highlight, or dagger for recording."
)
# Interactive mode drives strategy.run() in restartable segments and reads
# commands from stdin. 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):
raise ValueError(
f"--interactive=true currently supports only --strategy.type=base "
f"(got '{self.strategy.type}')."
)
# Sentry MUST use streaming encoding to avoid disk I/O blocking the control loop
if (
isinstance(self.strategy, SentryStrategyConfig)
+23 -1
View File
@@ -140,6 +140,9 @@ class RTCInferenceEngine(InferenceEngine):
self._action_queue: ActionQueue | None = None
self._obs_holder: dict[str, Any] = {}
self._obs_lock = Lock()
# Bumped by reset() (under _obs_lock) so chunks whose inference started
# before a reset are discarded instead of merged into the fresh queue.
self._reset_epoch = 0
self._policy_active = Event()
self._compile_warmup_done = Event()
self._shutdown_event = Event()
@@ -235,13 +238,26 @@ class RTCInferenceEngine(InferenceEngine):
self._policy_active.set()
def reset(self) -> None:
"""Reset the policy, processors, and action queue."""
"""Reset the policy, processors, and action queue.
Call while the engine is paused (both DAgger transitions and the
interactive session do): the RTC thread may still be finishing an
inference started before the pause, so ``reset`` also drops the last
published observation — it can be arbitrarily stale by the time the
engine resumes (e.g. the robot was returned to its initial position
in the meantime), and a chunk computed from it would jerk the robot
toward the old pose — and bumps the reset epoch so any in-flight
chunk is discarded instead of merged into the cleared queue.
"""
logger.info("Resetting RTC inference state (policy + processors + queue)")
self._policy.reset()
self._preprocessor.reset()
self._postprocessor.reset()
if self._action_queue is not None:
self._action_queue.clear()
with self._obs_lock:
self._obs_holder["obs"] = None
self._reset_epoch += 1
# ------------------------------------------------------------------
# Action production (called from main thread)
@@ -281,6 +297,7 @@ class RTCInferenceEngine(InferenceEngine):
queue = self._action_queue
with self._obs_lock:
obs = self._obs_holder.get("obs")
epoch_before = self._reset_epoch
if queue is None or obs is None:
time.sleep(_RTC_IDLE_SLEEP_S)
continue
@@ -339,7 +356,12 @@ class RTCInferenceEngine(InferenceEngine):
else:
latency_tracker.add(new_latency)
with self._obs_lock:
epoch_unchanged = epoch_before == self._reset_epoch
if epoch_unchanged:
queue.merge(original, processed, new_delay, idx_before)
else:
logger.info("Discarding action chunk computed before an engine reset")
if (
is_warmup
+474
View File
@@ -0,0 +1,474 @@
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Interactive rollout session: chat-style stdin commands for ``lerobot-rollout``.
Enabled with ``--interactive=true``, this module lets the operator control a
rollout from the terminal while hardware and policy stay connected and warm:
/start start (or restart) the policy control loop
/reset stop movement and return the robot to its initial position
/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 sets
thread-safe flags — it never touches hardware or the inference engine. 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)
propagate through the linked event's parent, so Ctrl-C behaves exactly as in
non-interactive runs.
The command table is intentionally a name → handler mapping so future
commands (``/subtask``, ``/ask`` — see the language-runtime work in
PR #4183/#4234) can be registered without restructuring the parser or the
session loop.
"""
from __future__ import annotations
import logging
import os
import select
import sys
import time
from collections.abc import Callable
from dataclasses import dataclass
from threading import Event, Thread
from typing import IO, TYPE_CHECKING
from lerobot.utils.utils import log_say
if TYPE_CHECKING:
from .context import RolloutContext
from .strategies import RolloutStrategy
logger = logging.getLogger(__name__)
_BANNER_RULE = "" * 60
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
@dataclass(frozen=True)
class InteractiveCommand:
"""A parsed ``/name args`` line from the interactive prompt."""
name: str
args: 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 ``/``).
"""
line = line.strip()
if not line.startswith("/"):
return None
head, *rest = line.split(maxsplit=1)
name = head[1:].lower()
if not name:
return None
return InteractiveCommand(name=name, args=rest[0].strip() if rest else "")
class StdinCommandListener:
"""Daemon thread that reads input lines and forwards them to a callback.
On POSIX the reader polls the stream with ``select`` so ``stop()`` can
end the thread promptly; elsewhere (or for file-like objects without a
file descriptor) it falls back to a blocking ``readline`` daemon thread
that dies with the process. Blank lines are skipped; end-of-file and
unexpected read errors trigger ``on_eof`` (an interactive Ctrl-D or an
exhausted piped script both mean "no more commands" — the session must
not keep the robot running with no way to command it).
Unlike :class:`lerobot.utils.keyboard_input.TerminalKeyListener`, this
reader leaves the terminal in canonical (line-buffered, echoing) mode —
the operator is typing chat-style commands, not pressing hotkeys.
"""
def __init__(
self,
on_line: Callable[[str], None],
on_eof: Callable[[], None] | None = None,
stream: IO[str] | None = None,
poll_interval_s: float = 0.2,
) -> None:
self._on_line = on_line
self._on_eof = on_eof
self._stream = stream if stream is not None else sys.stdin
self._poll_interval_s = poll_interval_s
self._running = False
self._thread: Thread | None = None
self._use_select = False
if os.name == "posix":
try:
self._stream.fileno()
self._use_select = True
except (OSError, ValueError, AttributeError):
pass
def start(self) -> None:
"""Start the reader thread (idempotent)."""
if self._thread is not None:
return
self._running = True
self._thread = Thread(target=self._run, daemon=True, name="InteractiveStdin")
self._thread.start()
if not self._use_select:
logger.info("stdin listener running in blocking mode (select unavailable for this stream)")
def stop(self) -> None:
"""Stop the reader thread.
Blocking-mode threads may be stuck inside ``readline`` and cannot be
joined; they are daemons and die with the process. Late lines are
ignored via the ``_running`` flag either way.
"""
self._running = False
thread = self._thread
self._thread = None
if thread is not None and thread.is_alive() and self._use_select:
thread.join(timeout=1.0)
def _run(self) -> None:
if self._use_select:
self._run_select()
else:
self._run_blocking()
def _run_select(self) -> None:
"""Poll the file descriptor and split lines from raw bytes.
Reading raw bytes (instead of ``stream.readline()``) matters: a
buffered file object can slurp several lines off the descriptor at
once, after which ``select`` reports the drained fd as not-ready and
the buffered lines would never be delivered — breaking pasted or
piped command sequences.
"""
fd = self._stream.fileno()
buffer = b""
while self._running:
try:
ready, _, _ = select.select([fd], [], [], self._poll_interval_s)
except (OSError, ValueError): # stream closed underneath us
self._emit_read_error()
return
if not ready:
continue
try:
chunk = os.read(fd, 4096)
except OSError:
self._emit_read_error()
return
if not self._running:
return
if chunk == b"": # EOF: Ctrl-D or the piped input ended
self._emit_line(buffer) # a final command without trailing newline still counts
self._emit_eof()
return
buffer += chunk
while b"\n" in buffer:
raw, buffer = buffer.split(b"\n", 1)
self._emit_line(raw)
def _run_blocking(self) -> None:
while self._running:
try:
line = self._stream.readline()
except (OSError, ValueError):
self._emit_read_error()
return
if not self._running:
return
if line == "": # EOF
self._emit_eof()
return
self._emit_line(line.encode() if isinstance(line, str) else line)
def _emit_line(self, raw: bytes) -> None:
line = raw.decode(errors="replace").strip()
if not line:
return
try:
self._on_line(line)
except Exception: # never let a handler error kill the reader thread
logger.exception("Error while handling interactive input %r", line)
def _emit_eof(self) -> None:
logger.info("Interactive input stream closed (EOF)")
if self._on_eof is not None:
try:
self._on_eof()
except Exception:
logger.exception("Error while handling interactive input EOF")
def _emit_read_error(self) -> None:
"""Treat an unexpected read failure like EOF so the session shuts down.
A dead command channel must not leave the robot running with no way
to stop it. Deliberate ``stop()`` calls clear ``_running`` first and
do not reach this path.
"""
if self._running:
logger.warning("Interactive input stream failed — treating as EOF")
self._emit_eof()
class InteractiveSession:
"""Drive a rollout strategy from chat-style stdin commands.
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.
Commands are last-write-wins: ``/reset`` and ``/stop`` cancel a pending
``/start`` so the robot never starts moving after the operator's final
command asked it not to. End-of-file on the command stream stops the
session (a closed stdin means there is no way left to command the
robot), so piped scripts must keep stdin open for the intended session
duration, e.g. ``(printf '/start\\n'; sleep 60; printf '/stop\\n') |
lerobot-rollout ... --interactive=true``.
"""
_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
self._listener = StdinCommandListener(self._handle_line, on_eof=self._handle_eof, stream=input_stream)
# Written by the listener thread, consumed by the main loop.
self._start_requested = Event()
self._reset_requested = Event()
self._stop_requested = Event()
self._wake = Event()
self._running = Event()
# name -> (handler, help line); /help and the banner render from this
# table, so future commands (/subtask, /ask) stay documented for free.
self._commands: dict[str, tuple[Callable[[InteractiveCommand], None], str]] = {
"start": (self._cmd_start, "start (or restart) the policy control loop"),
"reset": (self._cmd_reset, "stop movement and return the robot to its initial position"),
"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
self._print(self._render_banner())
self._listener.start()
try:
while not self._global_shutdown.is_set():
if self._ctx.policy.inference.failed:
self._print("Inference engine failed — shutting down. See the log for the error.")
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()
finally:
self._listener.stop()
log_say("Interactive session ended", play_sounds)
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("Rollout running — /reset to pause and 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 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)
# ------------------------------------------------------------------
def _handle_line(self, line: str) -> None:
cmd = parse_command(line)
if cmd is None:
self._print("Input not recognized — commands start with '/'. Type /help for the list.")
return
entry = self._commands.get(cmd.name)
if entry is None:
self._print(f"Unknown command '/{cmd.name}'. Type /help for the list.")
return
handler, _ = entry
handler(cmd)
def _handle_eof(self) -> None:
self._print("Input stream closed — stopping the session.")
self._request_stop()
def _cmd_start(self, cmd: InteractiveCommand) -> None:
if self._running.is_set():
self._print("Already running — /reset to pause first, or /stop to shut down.")
return
self._start_requested.set()
self._wake.set()
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()
self._reset_requested.set()
self._segment_stop.set()
self._wake.set()
def _cmd_stop(self, cmd: InteractiveCommand) -> None:
self._request_stop()
def _request_stop(self) -> None:
self._start_requested.clear() # last command wins, see _cmd_reset
self._stop_requested.set()
self._segment_stop.set()
self._wake.set()
def _cmd_help(self, cmd: InteractiveCommand) -> None:
self._print(self._render_help())
# ------------------------------------------------------------------
# Rendering
# ------------------------------------------------------------------
def _render_help(self) -> str:
width = max(len(name) for name in self._commands)
lines = [f" /{name:<{width}} {help_line}" for name, (_, help_line) in self._commands.items()]
return "Available commands:\n" + "\n".join(lines)
def _render_banner(self) -> str:
return (
f"{_BANNER_RULE}\n"
"Interactive rollout session — the robot will NOT move until you type /start.\n"
f"{self._render_help()}\n"
f"{_BANNER_RULE}"
)
@staticmethod
def _print(message: str) -> None:
"""User-facing chat output; logging stays on stderr, replies on stdout."""
print(message, flush=True)
+17 -4
View File
@@ -63,12 +63,25 @@ class RolloutStrategy(abc.ABC):
self._interpolator = ActionInterpolator(multiplier=ctx.runtime.cfg.interpolation_multiplier)
self._engine = ctx.policy.inference
logger.info("Starting inference engine...")
self._engine.reset()
self.reset_control_state()
self._engine.start()
self._warmup_flushed = False
self._cached_obs_processed = None
logger.info("Inference engine started")
def reset_control_state(self) -> None:
"""Clear episode-scoped control state so a paused session can restart cleanly.
Resets the inference engine (policy hidden state, action queues), the
action interpolator, and the cached processed observation. Used by the
interactive session between run segments; only call while the control
loop is not running.
"""
if self._engine is not None:
self._engine.reset()
if self._interpolator is not None:
self._interpolator.reset()
self._cached_obs_processed = None
def _process_observation_and_notify(self, processors: ProcessorContext, obs_raw: dict) -> dict:
"""Run the observation processor and notify the engine — throttled to policy ticks.
@@ -125,7 +138,7 @@ class RolloutStrategy(abc.ABC):
if robot.is_connected:
if return_to_initial_position and hw.initial_position:
logger.info("Returning robot to initial position before shutdown...")
self._return_to_initial_position(hw)
self.return_to_initial_position(hw)
elif not return_to_initial_position:
logger.info(
"Skipping return-to-initial-position (disabled by config); leaving robot in final pose."
@@ -138,7 +151,7 @@ class RolloutStrategy(abc.ABC):
teleop.disconnect()
@staticmethod
def _return_to_initial_position(hw: HardwareContext, duration_s: float = 3.0, fps: int = 50) -> None:
def return_to_initial_position(hw: HardwareContext, duration_s: float = 3.0, fps: int = 50) -> None:
"""Smoothly interpolate the robot back to its initial position."""
robot = hw.robot_wrapper
target = hw.initial_position
+2 -2
View File
@@ -165,7 +165,7 @@ class EpisodicStrategy(RolloutStrategy):
elif self.config.reset_to_initial_position:
# No teleop: return the robot to its startup position.
self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
self.return_to_initial_position(hw=ctx.hardware, duration_s=1)
self._reset_loop(
ctx=ctx,
@@ -187,7 +187,7 @@ class EpisodicStrategy(RolloutStrategy):
# returns to its initial joint positions captured at startup
if not teleop and self.config.reset_to_initial_position:
self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
self.return_to_initial_position(hw=ctx.hardware, duration_s=1)
continue
+26 -1
View File
@@ -44,6 +44,17 @@ Usage examples
--robot.port=/dev/ttyACM0 \\
--task="pick up cube" --duration=30
# Base mode — interactive session: the robot stays idle until /start is
# typed; /reset returns it to the initial position (hardware and policy
# stay warm); /stop shuts down gracefully
lerobot-rollout \\
--strategy.type=base \\
--policy.path=lerobot/act_koch_real \\
--robot.type=koch_follower \\
--robot.port=/dev/ttyACM0 \\
--task="pick up cube" \\
--interactive=true
# Base mode — RTC inference for slow VLAs (Pi0, Pi0.5, SmolVLA)
lerobot-rollout \\
--strategy.type=base \\
@@ -173,7 +184,13 @@ from lerobot.robots import ( # noqa: F401
so_follower,
unitree_g1 as unitree_g1_robot,
)
from lerobot.rollout import RolloutConfig, build_rollout_context, create_strategy
from lerobot.rollout import (
InteractiveSession,
LinkedEvent,
RolloutConfig,
build_rollout_context,
create_strategy,
)
from lerobot.teleoperators import ( # noqa: F401
Teleoperator,
TeleoperatorConfig,
@@ -215,6 +232,10 @@ def rollout(cfg: RolloutConfig):
signal_handler = ProcessSignalHandler(use_threads=True, display_pid=False)
shutdown_event = signal_handler.shutdown_event
if cfg.interactive:
# Session commands (/reset, /stop) end the running control loop by setting
# the local flag; process signals still propagate through the parent event.
shutdown_event = LinkedEvent(shutdown_event)
logger.info("Building rollout context...")
ctx = build_rollout_context(cfg, shutdown_event)
@@ -230,6 +251,10 @@ def rollout(cfg: RolloutConfig):
try:
strategy.setup(ctx)
if cfg.interactive:
logger.info("Rollout setup complete — starting interactive session (robot idle until /start)")
InteractiveSession(strategy, ctx).run()
else:
logger.info("Rollout setup complete, starting rollout...")
strategy.run(ctx)
except KeyboardInterrupt:
+529
View File
@@ -0,0 +1,529 @@
# 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 interactive rollout session (--interactive=true)."""
from __future__ import annotations
import contextlib
import os
import time
from threading import Event, Thread
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.rollout import ( # noqa: E402
InteractiveCommand,
InteractiveSession,
LinkedEvent,
StdinCommandListener,
parse_command,
)
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 session's stdin listener never sees EOF."""
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()
# ---------------------------------------------------------------------------
# Command parser
# ---------------------------------------------------------------------------
def test_parse_command_basic():
cmd = parse_command("/start")
assert cmd is not None
assert cmd.name == "start"
assert cmd.args == ""
def test_parse_command_case_whitespace_and_args():
cmd = parse_command(" /SubTask Grab the red cube ")
assert cmd is not None
assert cmd.name == "subtask"
assert cmd.args == "Grab the red cube"
def test_parse_command_tab_separated():
assert parse_command("/subtask\tgrab the cube") == InteractiveCommand(
name="subtask", args="grab the cube"
)
def test_parse_command_non_commands():
assert parse_command("hello robot") is None
assert parse_command("") is None
assert parse_command(" ") is None
assert parse_command("/") is None
assert parse_command("/ start") is None
# ---------------------------------------------------------------------------
# LinkedEvent
# ---------------------------------------------------------------------------
def test_linked_event_local_flag():
parent = Event()
event = LinkedEvent(parent)
assert not event.is_set()
event.set()
assert event.is_set()
assert not parent.is_set()
event.clear()
assert not event.is_set()
def test_linked_event_reflects_parent():
parent = Event()
event = LinkedEvent(parent)
parent.set()
assert event.is_set()
# Clearing the local flag never masks the parent.
event.clear()
assert event.is_set()
def test_linked_event_wait():
parent = Event()
event = LinkedEvent(parent)
assert event.wait(timeout=0.05) is False
parent.set()
assert event.wait(timeout=0.05) is True
parent.clear()
event.set()
assert event.wait(timeout=0.05) is True
def test_linked_event_wait_wakes_on_parent_set():
parent = Event()
event = LinkedEvent(parent)
Thread(target=lambda: (time.sleep(0.05), parent.set()), daemon=True).start()
assert event.wait(timeout=2.0) is True
# ---------------------------------------------------------------------------
# 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."""
import io
lines: list[str] = []
eof = Event()
listener = StdinCommandListener(lines.append, on_eof=eof.set, stream=io.StringIO("/start\n\n/help\n"))
assert not listener._use_select
listener.start()
assert _wait_for(eof.is_set)
assert lines == ["/start", "/help"]
listener.stop()
# ---------------------------------------------------------------------------
# InteractiveSession
# ---------------------------------------------------------------------------
def _make_session(input_stream, run_behavior=None):
"""Build a session around a mock strategy and a minimal fake context."""
parent = Event()
stop_event = LinkedEvent(parent)
engine = MagicMock()
engine.failed = False
ctx = SimpleNamespace(
runtime=SimpleNamespace(
cfg=SimpleNamespace(play_sounds=False),
shutdown_event=stop_event,
),
policy=SimpleNamespace(inference=engine),
hardware=SimpleNamespace(initial_position={"joint.pos": 0.0}),
)
strategy = MagicMock()
run_started = Event()
def default_run(c):
run_started.set()
while not c.runtime.shutdown_event.is_set():
time.sleep(0.005)
strategy.run.side_effect = run_behavior or default_run
session = InteractiveSession(strategy, ctx, input_stream=input_stream)
return session, strategy, engine, parent, run_started
def _start_session_thread(session) -> Thread:
thread = Thread(target=session.run, daemon=True)
thread.start()
return thread
def test_session_requires_linked_event():
ctx = SimpleNamespace(runtime=SimpleNamespace(shutdown_event=Event()))
with pytest.raises(TypeError, match="LinkedEvent"):
InteractiveSession(MagicMock(), ctx)
def test_session_start_reset_restart_stop_flow():
with _pipe_stream() as (reader, _writer):
session, strategy, engine, _parent, run_started = _make_session(reader)
thread = _start_session_thread(session)
# Idle until /start: the strategy loop must not run on its own.
time.sleep(0.05)
strategy.run.assert_not_called()
session._handle_line("/start")
assert _wait_for(run_started.is_set)
assert strategy.reset_control_state.call_count == 1
# /reset ends the segment, pauses the engine, and returns to the initial position.
session._handle_line("/reset")
assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
assert engine.pause.call_count >= 1
assert thread.is_alive()
# /start again runs a fresh segment with freshly reset control state.
run_started.clear()
session._handle_line("/start")
assert _wait_for(run_started.is_set)
assert strategy.run.call_count == 2
assert strategy.reset_control_state.call_count == 2
# /stop ends the session; teardown stays with the caller (the CLI script).
session._handle_line("/stop")
thread.join(timeout=2.0)
assert not thread.is_alive()
strategy.teardown.assert_not_called()
def test_session_stop_while_idle():
with _pipe_stream() as (reader, _writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader)
thread = _start_session_thread(session)
session._handle_line("/stop")
thread.join(timeout=2.0)
assert not thread.is_alive()
strategy.run.assert_not_called()
def test_session_reset_while_idle_returns_to_initial_position():
with _pipe_stream() as (reader, _writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader)
thread = _start_session_thread(session)
session._handle_line("/reset")
assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
assert thread.is_alive()
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_start_while_running_is_rejected():
with _pipe_stream() as (reader, _writer):
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)
session._handle_line("/start")
time.sleep(0.05)
assert strategy.run.call_count == 1
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_reset_cancels_pending_start():
"""Last command wins: a queued /start must not fire after a later /reset."""
with _pipe_stream() as (reader, _writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader)
# Queue both commands before the session loop starts servicing them.
session._handle_line("/start")
session._handle_line("/reset")
thread = _start_session_thread(session)
assert _wait_for(lambda: strategy.return_to_initial_position.call_count == 1)
time.sleep(0.05)
strategy.run.assert_not_called()
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_stop_cancels_pending_start():
with _pipe_stream() as (reader, _writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader)
session._handle_line("/start")
session._handle_line("/stop")
thread = _start_session_thread(session)
thread.join(timeout=2.0)
assert not thread.is_alive()
strategy.run.assert_not_called()
def test_session_exits_on_parent_shutdown():
with _pipe_stream() as (reader, _writer):
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)
parent.set() # SIGINT/SIGTERM path
thread.join(timeout=2.0)
assert not thread.is_alive()
def test_session_stops_on_engine_failure():
def failing_run(c):
# Mimic the RTC thread's fatal-error path: flag the failure and set
# the shutdown event the engine was built with (the LinkedEvent).
c.policy.inference.failed = True
c.runtime.shutdown_event.set()
with _pipe_stream() as (reader, _writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader, run_behavior=failing_run)
thread = _start_session_thread(session)
session._handle_line("/start")
thread.join(timeout=2.0)
assert not thread.is_alive()
# A failed engine ends the session instead of returning to idle.
strategy.return_to_initial_position.assert_not_called()
def test_session_stops_on_engine_failure_while_idle():
"""A fatal engine error while idle ends the session instead of being masked by /start."""
with _pipe_stream() as (reader, _writer):
session, strategy, engine, _parent, _run_started = _make_session(reader)
thread = _start_session_thread(session)
time.sleep(0.05)
engine.failed = True
thread.join(timeout=2.0)
assert not thread.is_alive()
strategy.run.assert_not_called()
def test_session_returns_to_idle_when_run_ends_naturally():
def finite_run(c):
return None # e.g. --duration elapsed
with _pipe_stream() as (reader, _writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader, run_behavior=finite_run)
thread = _start_session_thread(session)
session._handle_line("/start")
assert _wait_for(lambda: strategy.run.call_count == 1)
time.sleep(0.05)
assert thread.is_alive() # back to idle, not shut down
# The session accepts another /start after a natural end.
session._handle_line("/start")
assert _wait_for(lambda: strategy.run.call_count == 2)
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_unknown_input_does_not_start(capsys):
with _pipe_stream() as (reader, _writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader)
thread = _start_session_thread(session)
session._handle_line("/frobnicate")
session._handle_line("hello robot")
session._handle_line("/help")
time.sleep(0.05)
strategy.run.assert_not_called()
out = capsys.readouterr().out
assert "/frobnicate" in out
assert "commands start with '/'" in out
session._handle_line("/stop")
thread.join(timeout=2.0)
def test_session_eof_stops_session():
with _pipe_stream() as (reader, writer):
session, strategy, _engine, _parent, _run_started = _make_session(reader)
thread = _start_session_thread(session)
writer.close() # EOF on the command stream
thread.join(timeout=2.0)
assert not thread.is_alive()
strategy.run.assert_not_called()
def test_session_commands_via_stream():
"""End-to-end: commands flow through the pipe and the listener thread."""
with _pipe_stream() as (reader, writer):
session, strategy, _engine, _parent, run_started = _make_session(reader)
thread = _start_session_thread(session)
writer.write("/start\n")
writer.flush()
assert _wait_for(run_started.is_set)
writer.write("/stop\n")
writer.flush()
thread.join(timeout=2.0)
assert not thread.is_alive()
assert strategy.run.call_count == 1
def test_session_drives_real_base_strategy():
"""End-to-end with a real BaseStrategy control loop (only hardware/engine mocked)."""
from lerobot.rollout import BaseStrategy, BaseStrategyConfig
parent = Event()
stop_event = LinkedEvent(parent)
engine = MagicMock()
engine.failed = False
engine.get_action.return_value = None # no action ready; the loop still ticks
robot = MagicMock()
robot.get_observation.return_value = {"joint.pos": 0.0}
def identity(x):
return x
ctx = SimpleNamespace(
runtime=SimpleNamespace(
cfg=SimpleNamespace(
play_sounds=False,
fps=100.0,
duration=0.0,
use_torch_compile=False,
interpolation_multiplier=1,
display_data=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=None, dataset_features={}, hw_features={}, ordered_action_keys=[]),
)
strategy = BaseStrategy(BaseStrategyConfig())
strategy.setup(ctx)
strategy.return_to_initial_position = MagicMock() # skip the 3s hardware sweep
with _pipe_stream() as (reader, _writer):
session = InteractiveSession(strategy, ctx, input_stream=reader)
thread = _start_session_thread(session)
session._handle_line("/start")
assert _wait_for(lambda: engine.resume.called)
assert _wait_for(lambda: robot.get_observation.call_count >= 3)
session._handle_line("/reset")
assert _wait_for(lambda: strategy.return_to_initial_position.called)
assert engine.pause.called
assert thread.is_alive()
session._handle_line("/start")
assert _wait_for(lambda: engine.resume.call_count >= 2)
session._handle_line("/stop")
thread.join(timeout=2.0)
assert not thread.is_alive()
# ---------------------------------------------------------------------------
# Config validation
# ---------------------------------------------------------------------------
def test_interactive_requires_base_strategy():
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,
)