mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
feat(rollout): interactive v1
This commit is contained in:
@@ -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 |
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
@@ -63,12 +63,25 @@ class RolloutStrategy(abc.ABC):
|
||||
self._interpolator = ActionInterpolator(multiplier=ctx.runtime.cfg.interpolation_multiplier)
|
||||
self._engine = ctx.policy.inference
|
||||
logger.info("Starting inference engine...")
|
||||
self._engine.reset()
|
||||
self.reset_control_state()
|
||||
self._engine.start()
|
||||
self._warmup_flushed = False
|
||||
self._cached_obs_processed = None
|
||||
logger.info("Inference engine started")
|
||||
|
||||
def reset_control_state(self) -> None:
|
||||
"""Clear episode-scoped control state so a paused session can restart cleanly.
|
||||
|
||||
Resets the inference engine (policy hidden state, action queues), the
|
||||
action interpolator, and the cached processed observation. Used by the
|
||||
interactive session between run segments; only call while the control
|
||||
loop is not running.
|
||||
"""
|
||||
if self._engine is not None:
|
||||
self._engine.reset()
|
||||
if self._interpolator is not None:
|
||||
self._interpolator.reset()
|
||||
self._cached_obs_processed = None
|
||||
|
||||
def _process_observation_and_notify(self, processors: ProcessorContext, obs_raw: dict) -> dict:
|
||||
"""Run the observation processor and notify the engine — throttled to policy ticks.
|
||||
|
||||
@@ -125,7 +138,7 @@ class RolloutStrategy(abc.ABC):
|
||||
if robot.is_connected:
|
||||
if return_to_initial_position and hw.initial_position:
|
||||
logger.info("Returning robot to initial position before shutdown...")
|
||||
self._return_to_initial_position(hw)
|
||||
self.return_to_initial_position(hw)
|
||||
elif not return_to_initial_position:
|
||||
logger.info(
|
||||
"Skipping return-to-initial-position (disabled by config); leaving robot in final pose."
|
||||
@@ -138,7 +151,7 @@ class RolloutStrategy(abc.ABC):
|
||||
teleop.disconnect()
|
||||
|
||||
@staticmethod
|
||||
def _return_to_initial_position(hw: HardwareContext, duration_s: float = 3.0, fps: int = 50) -> None:
|
||||
def return_to_initial_position(hw: HardwareContext, duration_s: float = 3.0, fps: int = 50) -> None:
|
||||
"""Smoothly interpolate the robot back to its initial position."""
|
||||
robot = hw.robot_wrapper
|
||||
target = hw.initial_position
|
||||
|
||||
@@ -165,7 +165,7 @@ class EpisodicStrategy(RolloutStrategy):
|
||||
|
||||
elif self.config.reset_to_initial_position:
|
||||
# No teleop: return the robot to its startup position.
|
||||
self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
|
||||
self.return_to_initial_position(hw=ctx.hardware, duration_s=1)
|
||||
|
||||
self._reset_loop(
|
||||
ctx=ctx,
|
||||
@@ -187,7 +187,7 @@ class EpisodicStrategy(RolloutStrategy):
|
||||
|
||||
# returns to its initial joint positions captured at startup
|
||||
if not teleop and self.config.reset_to_initial_position:
|
||||
self._return_to_initial_position(hw=ctx.hardware, duration_s=1)
|
||||
self.return_to_initial_position(hw=ctx.hardware, duration_s=1)
|
||||
|
||||
continue
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
Reference in New Issue
Block a user