admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-11 12:01:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(runtime): real-robot interactive mode + rerun live camera view

Browse Source 
Add physical-robot support to the language runtime, plus a live rerun viewer.

- runtime/rerun_viz.py: headless rerun gRPC + web viewer; logs camera frames
  (every control tick) and joint state. Prints an auto-connect ?url= view URL.
- runtime/cli.py:
  - _run_robot_interactive: real-time control loop (background thread) with a
    clean chat prompt — a typed command switches task/subtask immediately and
    regenerates. Starts running as soon as a task is set (via --task or the
    picker); otherwise paused until the first command. No flag needed.
  - --rerun (+ --rerun.web_port / --rerun.grpc_port): live camera view; the
    robot obs provider and action executor log frames to rerun.
  - --direct_subtask (general, sim or robot): the typed text is the subtask fed
    to the action expert; the LM subtask generator is disabled.
  - Inference overrides: force compile_model=False and gradient_checkpointing
    =False (torch.compile recompiled on every prompt-length change -> >1min per
    chunk; grad checkpointing only slows the forward pass).

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn223
2026-07-09 11:52:09 +02:00
parent cd15a66286
commit 1e94a4f62d
2 changed files with 276 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/runtime/cli.py
+167 -5
View File
@@ -228,6 +228,22 @@ def _parse_args(argv: list[str] | None = None, *, prog: str | None = None) -> ar
"``lerobot-record`` uses. Default ``None`` = no clipping."
),
)
p.add_argument(
"--rerun",
action="store_true",
help="Live rerun viewer for the robot cameras (real-robot mode). Serves a "
"headless web viewer; forward --rerun.web_port and --rerun.grpc_port over SSH.",
)
p.add_argument("--rerun.web_port", dest="rerun_web_port", type=int, default=9090,
help="rerun web-viewer port (default 9090).")
p.add_argument("--rerun.grpc_port", dest="rerun_grpc_port", type=int, default=9876,
help="rerun gRPC data port (default 9876).")
p.add_argument(
"--direct_subtask",
action="store_true",
help="Direct-subtask mode (sim OR robot): your typed text IS the subtask "
"fed to the action expert; the LM subtask generator is disabled.",
)
p.add_argument(
"--auto_start",
action="store_true",
@@ -465,6 +481,15 @@ def _load_policy_and_preprocessor(
cfg = PreTrainedConfig.from_pretrained(policy_path)
cfg.pretrained_path = policy_path
# Inference-only overrides (mirror lerobot-eval). torch.compile recompiles
# whenever the prompt length changes (every subtask switch) — catastrophic
# in the interactive runtime — and gradient checkpointing only slows the
# forward pass. Neither is wanted for serving.
if getattr(cfg, "compile_model", False):
cfg.compile_model = False
if getattr(cfg, "gradient_checkpointing", False):
cfg.gradient_checkpointing = False
ds_meta = None
preprocessor = None
postprocessor = None
@@ -821,6 +846,7 @@ def _build_robot_observation_provider(
device: str,
task: str | None,
ds_features: dict[str, Any] | None,
rerun_log: bool = False,
) -> Callable[[], dict | None]:
"""Closure reading from the robot each call: ``robot.get_observation()`` →
``build_inference_frame`` (state vector + image tensors, batched, on device)
@@ -870,6 +896,21 @@ def _build_robot_observation_provider(
logger.warning("robot.get_observation failed: %s", exc)
return None
# Live camera view: log the raw frames + joint state to rerun before any
# resize (natural camera resolution). Best-effort — never blocks control.
if rerun_log:
from lerobot.runtime import rerun_viz # noqa: PLC0415
cam_keys = list(target_image_shapes.keys()) or [
k for k, v in raw.items() if hasattr(v, "ndim") and getattr(v, "ndim", 0) == 3
]
state = {
k: v
for k, v in raw.items()
if isinstance(v, (int, float)) and k not in cam_keys
}
rerun_viz.log_robot_frame(raw, cam_keys, state=state, task=task)
# The runtime supplies messages itself; strip any language
# columns the robot stream may carry through.
_strip_runtime_owned_language_cols(raw)
@@ -977,6 +1018,7 @@ def _build_robot_action_executor(
robot,
postprocessor: Any,
ds_meta: Any,
rerun_log: bool = False,
) -> Callable[[Any], None]:
"""Closure that postprocesses an action and dispatches to the robot.
@@ -1005,6 +1047,12 @@ def _build_robot_action_executor(
logger.warning("unsupported action type %r — skipping", type(action))
return
robot.send_action(action_dict)
# Smooth live view: log the cameras every control tick (buffered
# async_read is cheap). Best-effort — never blocks control.
if rerun_log:
from lerobot.runtime import rerun_viz # noqa: PLC0415
rerun_viz.log_cameras(robot)
except Exception as exc: # noqa: BLE001
logger.error("robot.send_action failed: %s", exc, exc_info=True)
@@ -1554,6 +1602,14 @@ def run(
if sim_stream_server is not None:
sim_backend.attach_stream_server(sim_stream_server)
elif autonomous_mode:
if args.rerun:
from lerobot.runtime.rerun_viz import start_rerun # noqa: PLC0415
start_rerun(
app_name=f"lerobot_{policy_type or 'runtime'}",
grpc_port=args.rerun_grpc_port,
web_port=args.rerun_web_port,
)
print(
f"[runtime] connecting to robot.type={args.robot_type} port={args.robot_port}",
flush=True,
@@ -1571,11 +1627,13 @@ def run(
device=str(getattr(policy.config, "device", "cpu")),
task=args.task,
ds_features=ds_meta.features if ds_meta is not None else None,
rerun_log=bool(args.rerun),
)
robot_executor = _build_robot_action_executor(
robot=robot,
postprocessor=postprocessor,
ds_meta=ds_meta,
rerun_log=bool(args.rerun),
)
elif args.dataset_repo_id is not None:
print(
@@ -1604,7 +1662,7 @@ def run(
top_p=float(args.text_top_p or 1.0),
chunks_per_regen=max(1, int(args.subtask_chunks_per_gen or 1)),
enable_memory=not bool(getattr(args, "disable_memory", False)),
enable_subtask=not bool(getattr(args, "sim_direct_subtask", False)),
enable_subtask=not _direct_subtask_enabled(args),
)
runtime = LanguageConditionedRuntime(
policy_adapter=adapter_factory(policy, gen_config),
@@ -1644,16 +1702,16 @@ def run(
initial_task=args.task,
max_ticks=args.max_ticks,
panel_label=panel_label,
direct_subtask=bool(args.sim_direct_subtask),
direct_subtask=_direct_subtask_enabled(args),
)
if autonomous_mode:
return _run_autonomous(
return _run_robot_interactive(
runtime,
robot=robot,
auto_start=args.auto_start,
robot,
initial_task=args.task,
max_ticks=args.max_ticks,
direct_subtask=_direct_subtask_enabled(args),
panel_label=panel_label,
)
# Fire one full pipeline tick at startup so the obs diagnostic
@@ -1674,6 +1732,11 @@ def run(
return _run_repl(runtime, initial_task=args.task, max_ticks=args.max_ticks, panel_label=panel_label)
def _direct_subtask_enabled(args: Any) -> bool:
"""Direct-subtask mode via either the general or sim-scoped flag."""
return bool(getattr(args, "direct_subtask", False) or getattr(args, "sim_direct_subtask", False))
def _run_sim_interactive(
runtime: Any,
sim_backend: Any,
@@ -1800,6 +1863,105 @@ def _run_sim_interactive(
return 0
def _run_robot_interactive(
runtime: Any,
robot: Any,
*,
initial_task: str | None,
max_ticks: int | None,
direct_subtask: bool = False,
panel_label: str = "Runtime",
) -> int:
"""Real-robot interactive loop.
The control loop runs at real-time rates in a background thread
(``runtime.run()`` — a robot must be driven at a steady ``ctrl_hz``), while
the foreground is a clean chat prompt: type a command to run it (generate- or
direct-subtask mode), ``/pause`` / ``/resume`` / ``stop``. Starts PAUSED so
the arm doesn't move until you issue a command.
"""
import threading # noqa: PLC0415
import time # noqa: PLC0415
if initial_task:
runtime.set_task(initial_task)
runtime.state["current_subtask"] = initial_task if direct_subtask else None
# A task was given (via --task or the startup picker) => start running it
# immediately. Without an initial task we stay paused until the first
# typed command (which switches to action). No flag needed.
runtime.state["mode"] = "action"
mode_line = (
"DIRECT subtask (your text drives the action expert)"
if direct_subtask
else "task (the model generates a subtask from your text)"
)
starting_action = runtime.state.get("mode", "paused") == "action"
start_line = (
f" Starting in ACTION — the ARM WILL MOVE NOW on: {initial_task!r}\n"
if starting_action
else " Starts PAUSED. Type a command + Enter to run it — the ARM WILL MOVE.\n"
)
print(
f"\n{'=' * 64}\n"
f" {panel_label} — OMX robot runtime · Mode: {mode_line}\n"
f"{start_line}"
f" Commands: /pause · /resume · stop\n"
f"{'=' * 64}",
flush=True,
)
thread = threading.Thread(
target=runtime.run, kwargs={"max_ticks": max_ticks}, name="runtime-loop", daemon=True
)
thread.start()
try:
while thread.is_alive():
try:
line = input("\n> ").strip()
except EOFError:
break
if not line:
continue
low = line.lower()
if low in {"stop", "quit", "exit"}:
break
elif low in {"/pause", "pause", "/p"}:
runtime.state["mode"] = "paused"
_clear_action_queue(runtime)
print("[paused] robot holding", flush=True)
elif low in {"/resume", "resume", "/run"}:
runtime.state["mode"] = "action"
print("[running]", flush=True)
else:
# New command: switch task/subtask immediately and regenerate.
runtime.set_task(line)
runtime.state["current_subtask"] = line if direct_subtask else None
_clear_action_queue(runtime)
adapter = getattr(runtime, "policy_adapter", None)
if adapter is not None and hasattr(adapter, "_chunks_until_regen"):
adapter._chunks_until_regen = 0
gate = getattr(runtime, "_language_gate", None)
if gate is not None and hasattr(gate, "rearm"):
gate.rearm()
runtime.state["mode"] = "action"
print(f"[running] {line}", flush=True)
except KeyboardInterrupt:
print("\n[stopping]", flush=True)
finally:
runtime.stop()
for _ in range(10):
if not thread.is_alive():
break
time.sleep(0.1)
try:
robot.disconnect()
print("[runtime] robot disconnected", flush=True)
except Exception as exc: # noqa: BLE001
print(f"[runtime] WARNING: robot.disconnect raised {exc}", flush=True)
return 0
def _run_repl(
runtime: Any, *, initial_task: str | None, max_ticks: int | None, panel_label: str = "Runtime"
) -> int:
src/lerobot/runtime/rerun_viz.py
+109
View File
@@ -0,0 +1,109 @@
# 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.
"""Rerun live visualisation for the interactive runtime (real-robot camera view).
Starts a headless rerun gRPC server + web viewer so a remote operator can watch
the robot's cameras (and state / subtask) over SSH by forwarding two ports and
opening the web viewer in a browser. Logging is best-effort — a rerun failure
never interrupts robot control.
"""
from __future__ import annotations
import logging
from typing import Any
logger = logging.getLogger(__name__)
_ENABLED = False
def start_rerun(app_name: str = "lerobot_runtime", grpc_port: int = 9876, web_port: int = 9090) -> bool:
"""Init rerun and serve a headless gRPC + web viewer. Returns True on success."""
global _ENABLED
try:
import rerun as rr # noqa: PLC0415
rr.init(app_name)
url = rr.serve_grpc(grpc_port=grpc_port)
rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url)
_ENABLED = True
# Open the viewer with the data URL as a query param so it auto-connects
# to the gRPC stream (plain http://host:web_port shows only the welcome
# screen — the web app needs the ?url= to know where the data is).
view_url = f"http://localhost:{web_port}/?url={url}"
print(
f"[runtime] rerun live view: {view_url}\n"
f" (over SSH forward both ports: "
f"ssh -L {web_port}:localhost:{web_port} -L {grpc_port}:localhost:{grpc_port} <host>)",
flush=True,
)
return True
except Exception as exc: # noqa: BLE001
logger.warning("[runtime] could not start rerun: %s", exc)
print(f"[runtime] WARNING: rerun unavailable ({exc})", flush=True)
return False
def log_cameras(robot: Any) -> None:
"""Log every robot camera's latest buffered frame (cheap async_read).
Called every control tick for a smooth live view (best-effort)."""
if not _ENABLED:
return
try:
import rerun as rr # noqa: PLC0415
cams = getattr(robot, "cameras", None) or {}
for name, cam in cams.items():
try:
frame = cam.async_read(timeout_ms=1)
rr.log(f"cameras/{name}", rr.Image(frame))
except Exception: # noqa: BLE001
pass
except Exception as exc: # noqa: BLE001
logger.debug("[runtime] rerun camera log failed: %s", exc)
def log_robot_frame(
raw_obs: dict[str, Any],
camera_keys: list[str],
state: dict[str, float] | None = None,
task: str | None = None,
subtask: str | None = None,
) -> None:
"""Log camera images + optional state/task/subtask for one step (best-effort)."""
if not _ENABLED:
return
try:
import numpy as np # noqa: PLC0415
import rerun as rr # noqa: PLC0415
for cam in camera_keys:
img = raw_obs.get(cam)
if isinstance(img, np.ndarray) and img.ndim == 3 and img.shape[-1] == 3:
rr.log(f"cameras/{cam}", rr.Image(img))
if state:
for name, val in state.items():
try:
rr.log(f"state/{name}", rr.Scalars(float(val)))
except Exception: # noqa: BLE001
pass
if task:
rr.log("prompt/task", rr.TextLog(task))
if subtask:
rr.log("prompt/subtask", rr.TextLog(subtask))
except Exception as exc: # noqa: BLE001
logger.debug("[runtime] rerun log failed: %s", exc)