diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index 74b149e58..746bf2bb4 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -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: diff --git a/src/lerobot/runtime/rerun_viz.py b/src/lerobot/runtime/rerun_viz.py new file mode 100644 index 000000000..ae3098d55 --- /dev/null +++ b/src/lerobot/runtime/rerun_viz.py @@ -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} )", + 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)