From a09715121ed45ac6e7c8166786ad39a8488f9178 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Wed, 15 Jul 2026 15:59:13 +0200 Subject: [PATCH] refactor(runtime): reuse shared rerun visualization --- src/lerobot/runtime/cli.py | 37 +++++---- src/lerobot/runtime/rerun_viz.py | 98 ------------------------ src/lerobot/utils/rerun_visualization.py | 10 ++- tests/utils/test_rerun_visualization.py | 30 ++++++++ 4 files changed, 60 insertions(+), 115 deletions(-) delete mode 100644 src/lerobot/runtime/rerun_viz.py diff --git a/src/lerobot/runtime/cli.py b/src/lerobot/runtime/cli.py index dd6c8c307..7ff08079a 100644 --- a/src/lerobot/runtime/cli.py +++ b/src/lerobot/runtime/cli.py @@ -455,6 +455,7 @@ def _build_rollout_runtime_io( prepare_observation_for_inference, ) from lerobot.utils.feature_utils import build_dataset_frame # noqa: PLC0415 + from lerobot.utils.rerun_visualization import log_rerun_data # noqa: PLC0415 robot = ctx.hardware.robot_wrapper device = torch.device(ctx.runtime.cfg.device or "cpu") @@ -466,11 +467,10 @@ def _build_rollout_runtime_io( latest_raw.clear() latest_raw.update(raw) if rerun_log: - from lerobot.runtime import rerun_viz # noqa: PLC0415 - - camera_keys = list(robot.cameras) - state = {k: v for k, v in raw.items() if isinstance(v, (int, float))} - rerun_viz.log_robot_frame(raw, camera_keys, state=state, task=get_task()) + try: + log_rerun_data(observation=raw) + except Exception as exc: # noqa: BLE001 + logger.debug("rerun observation log failed: %s", exc) _strip_runtime_owned_language_cols(raw) processed = ctx.processors.robot_observation_processor(raw) observation = build_dataset_frame(ctx.data.dataset_features, processed, prefix="observation") @@ -501,10 +501,6 @@ def _build_rollout_runtime_io( raw = latest_raw or robot.get_observation() robot_action = ctx.processors.robot_action_processor((action_dict, raw)) robot.send_action(robot_action) - 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 action pipeline failed: %s", exc, exc_info=True) @@ -840,19 +836,28 @@ def run( if sim_stream_server is not None: sim_backend.attach_stream_server(sim_stream_server) elif autonomous_mode: + rerun_log = False if args.rerun: - from lerobot.runtime.rerun_viz import start_rerun # noqa: PLC0415 + from lerobot.utils.rerun_visualization import init_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, - ) + try: + init_rerun( + session_name=f"lerobot_{policy_type or 'runtime'}", + port=args.rerun_grpc_port, + web_port=args.rerun_web_port, + ) + rerun_log = True + print( + f"[runtime] rerun live view: http://localhost:{args.rerun_web_port}", + flush=True, + ) + except Exception as exc: # noqa: BLE001 + logger.warning("could not start rerun: %s", exc) robot = rollout_ctx.hardware.robot_wrapper.inner print(f"[runtime] connected to {robot.name}", flush=True) observation_provider, robot_executor = _build_rollout_runtime_io( rollout_ctx, - rerun_log=bool(args.rerun), + rerun_log=rerun_log, get_task=_live_task, ) # Generation settings belong to the adapter rather than mutable runtime state. diff --git a/src/lerobot/runtime/rerun_viz.py b/src/lerobot/runtime/rerun_viz.py deleted file mode 100644 index c7b81deba..000000000 --- a/src/lerobot/runtime/rerun_viz.py +++ /dev/null @@ -1,98 +0,0 @@ -# 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. - -"""Best-effort Rerun camera visualization for local or SSH-forwarded robot rollouts.""" - -from __future__ import annotations - -import logging -from contextlib import suppress -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 - # Include the stream URL so the web viewer connects automatically. - 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(): - with suppress(Exception): - frame = cam.async_read(timeout_ms=1) - rr.log(f"cameras/{name}", rr.Image(frame)) - 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(): - with suppress(Exception): - rr.log(f"state/{name}", rr.Scalars(float(val))) - 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) diff --git a/src/lerobot/utils/rerun_visualization.py b/src/lerobot/utils/rerun_visualization.py index 46f2c0b4b..5dd3eb174 100644 --- a/src/lerobot/utils/rerun_visualization.py +++ b/src/lerobot/utils/rerun_visualization.py @@ -38,7 +38,10 @@ def _is_scalar(x): def init_rerun( - session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None + session_name: str = "lerobot_control_loop", + ip: str | None = None, + port: int | None = None, + web_port: int | None = None, ) -> None: """ Initializes the Rerun SDK for visualizing the control loop. @@ -47,6 +50,7 @@ def init_rerun( session_name: Name of the Rerun session. ip: Optional IP for connecting to a Rerun server. port: Optional port for connecting to a Rerun server. + web_port: Serve a headless web viewer on this port, using ``port`` for gRPC. """ require_package("rerun-sdk", extra="viz", import_name="rerun") @@ -60,6 +64,10 @@ def init_rerun( memory_limit = os.getenv("LEROBOT_RERUN_MEMORY_LIMIT", "10%") if ip and port: rr.connect_grpc(url=f"rerun+http://{ip}:{port}/proxy") + elif web_port is not None: + grpc_port = port or 9876 + url = rr.serve_grpc(grpc_port=grpc_port) + rr.serve_web_viewer(web_port=web_port, open_browser=False, connect_to=url) else: rr.spawn(memory_limit=memory_limit) diff --git a/tests/utils/test_rerun_visualization.py b/tests/utils/test_rerun_visualization.py index d4c3e6999..a52ef9077 100644 --- a/tests/utils/test_rerun_visualization.py +++ b/tests/utils/test_rerun_visualization.py @@ -129,6 +129,36 @@ def _views_by_kind(blueprint, kind): return [v for v in blueprint.root.views if v.kind == kind] +def test_init_rerun_can_serve_headless_web_viewer(mock_rerun, monkeypatch): + rv, _calls, _blueprints = mock_rerun + rr = sys.modules["rerun"] + served = {} + + def serve_grpc(grpc_port): + served["grpc"] = grpc_port + return "rerun+http://localhost" + + monkeypatch.setattr( + rr, + "serve_grpc", + serve_grpc, + raising=False, + ) + monkeypatch.setattr( + rr, + "serve_web_viewer", + lambda **kwargs: served.setdefault("web", kwargs), + raising=False, + ) + + rv.init_rerun(session_name="runtime", port=8765, web_port=9091) + + assert served["grpc"] == 8765 + assert served["web"]["web_port"] == 9091 + assert served["web"]["open_browser"] is False + assert served["web"]["connect_to"] == "rerun+http://localhost" + + def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun): rv, calls, blueprints = mock_rerun