refactor(runtime): reuse shared rerun visualization

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.

src/lerobot/runtime/rerun_viz.py

@@ -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} <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():
-            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)

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)
 

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