| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| """Foxglove visualization backend. |
| |
| Live control-loop streaming (:func:`log_foxglove_data`) and seekable dataset playback |
| (:func:`serve_foxglove_dataset_playback`) over a Foxglove WebSocket server. Callers usually select a |
| backend at runtime through the dispatch in :mod:`lerobot.utils.visualization_utils` rather than |
| importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot[viz]'``). |
| """ |
|
|
| import logging |
| import numbers |
| import time |
|
|
| import cv2 |
| import numpy as np |
|
|
| from lerobot.types import RobotAction, RobotObservation |
|
|
| from .constants import ( |
| ACTION, |
| ACTION_PREFIX, |
| DONE, |
| OBS_IMAGES, |
| OBS_PREFIX, |
| OBS_STATE, |
| OBS_STR, |
| REWARD, |
| SUCCESS, |
| TRUNCATED, |
| ) |
| from .import_utils import require_package |
|
|
| |
| |
| |
| |
| |
| _SCALARS_SCHEMA = { |
| "type": "object", |
| "title": "lerobot.Scalars", |
| "properties": { |
| "scalars": { |
| "type": "array", |
| "items": { |
| "type": "object", |
| "properties": { |
| "label": {"type": "string"}, |
| "value": {"type": "number"}, |
| }, |
| }, |
| } |
| }, |
| } |
|
|
|
|
| def _is_scalar(x): |
| return isinstance(x, (float | numbers.Real | np.integer | np.floating)) or ( |
| isinstance(x, np.ndarray) and x.ndim == 0 |
| ) |
|
|
|
|
| def init_foxglove(host: str = "127.0.0.1", port: int | None = 8765) -> None: |
| """ |
| Starts a Foxglove WebSocket server for visualizing the control loop. |
| |
| Connect to it from the Foxglove app at ``ws://<host>:<port>``. Calling this |
| more than once is a no-op while a server is already running. |
| |
| Args: |
| host: Host interface to bind the WebSocket server to. |
| port: Port to bind the WebSocket server to (defaults to 8765). |
| """ |
|
|
| require_package("foxglove-sdk", extra="viz", import_name="foxglove") |
| import foxglove |
|
|
| |
| |
| |
| if getattr(log_foxglove_data, "server", None) is not None: |
| return |
| log_foxglove_data.server = foxglove.start_server(host=host, port=port or 8765) |
| log_foxglove_data.channels = {} |
|
|
|
|
| def shutdown_foxglove() -> None: |
| """Stops the Foxglove WebSocket server and clears cached channels.""" |
|
|
| server = getattr(log_foxglove_data, "server", None) |
| if server is not None: |
| server.stop() |
| log_foxglove_data.server = None |
| log_foxglove_data.channels = {} |
|
|
|
|
| def _foxglove_safe_name(name: str) -> str: |
| """Replace ``.`` with ``_`` so a feature name is a single Foxglove topic-path segment. |
| |
| Foxglove treats ``.`` as a path separator, so an unsanitized name like ``observation.images.front`` |
| would split into nested segments instead of naming one topic. |
| """ |
|
|
| return name.replace(".", "_") |
|
|
|
|
| def _foxglove_topic(key: str, *, is_image: bool = False) -> str: |
| """Build the Foxglove topic for a feature ``key``. |
| |
| Camera features map to a per-source image topic (``/observation/images/<name>``); scalar features |
| share one aggregate topic per source: ``/observation/state`` for observations, ``/action/state`` |
| for actions. |
| """ |
|
|
| if is_image: |
| name = str(key) |
| for prefix in (f"{OBS_IMAGES}.", OBS_PREFIX): |
| if name.startswith(prefix): |
| name = name[len(prefix) :] |
| break |
| return f"/{OBS_STR}/images/{_foxglove_safe_name(name)}" |
| source = ACTION if (str(key).startswith(ACTION_PREFIX) or str(key) == ACTION) else OBS_STR |
| return f"/{source}/state" |
|
|
|
|
| def _log_foxglove_scalars( |
| topic: str, values: dict[str, float], *, channels: dict | None = None, log_time: int | None = None |
| ) -> None: |
| """Log scalars on a typed JSON channel using the static :data:`_SCALARS_SCHEMA`. |
| |
| ``values`` is an ordered mapping of feature name to value; it is emitted as a ``scalars`` array of |
| ``{label, value}`` objects. Insertion order is preserved so series stay stable across messages. |
| |
| ``channels`` is the per-topic channel cache to reuse (defaults to the live-stream cache on |
| :func:`log_foxglove_data`; dataset playback passes its own local cache to stay self-contained). |
| ``log_time`` is the message time in nanoseconds; when ``None`` the server's receive time is used. |
| """ |
|
|
| if not values: |
| return |
|
|
| import foxglove |
|
|
| if channels is None: |
| channels = log_foxglove_data.channels |
| channel = channels.get(topic) |
| if channel is None: |
| channel = channels[topic] = foxglove.Channel(topic, schema=_SCALARS_SCHEMA, message_encoding="json") |
| msg = {"scalars": [{"label": label, "value": value} for label, value in values.items()]} |
| if log_time is None: |
| channel.log(msg) |
| else: |
| channel.log(msg, log_time=log_time) |
|
|
|
|
| def _labeled_scalars(name: str, values, labels: list[str] | None = None) -> dict[str, float]: |
| """Expand a 1D sequence into ``{label: value}`` entries with a consistent fallback.""" |
|
|
| flat = [float(v) for v in values] |
| if labels is None or len(labels) != len(flat): |
| labels = [f"{name}_{i}" for i in range(len(flat))] |
| return dict(zip(labels, flat, strict=True)) |
|
|
|
|
| def _log_foxglove_image( |
| topic: str, |
| frame_id: str, |
| arr: np.ndarray, |
| *, |
| compress_images: bool, |
| channels: dict | None = None, |
| log_time: int | None = None, |
| depth_range: tuple[float, float] | None = None, |
| raw_depth_values: bool = False, |
| ) -> None: |
| """Log an image on a cached per-topic channel. |
| |
| The encoding is chosen from the channel count and dtype: a single-channel ``float`` or ``uint16`` |
| frame is a depth map (``32FC1``/``16UC1``), single-channel ``uint8`` is ``mono8``, 3 => ``rgb8`` |
| (float input assumed in [0, 1], cast to uint8), 4 => ``rgba8``; other counts are skipped with a |
| warning. When ``compress_images`` is set, ``rgb8`` is JPEG-encoded instead. |
| |
| Args: |
| topic: Foxglove topic to log on. |
| frame_id: Frame id stamped on the message. |
| arr: Image as HWC or CHW (CHW is transposed to HWC), any dtype. |
| compress_images: JPEG-encode ``rgb8`` frames; ignored for other encodings. |
| channels: Per-topic channel cache to reuse (see :func:`_log_foxglove_scalars`). |
| log_time: Message time in nanoseconds, also written to the header timestamp; when ``None`` |
| the server's receive time is used. |
| depth_range: ``(lo, hi)`` clip bounds in a depth frame's own input units. Depth frames |
| (``32FC1``/``16UC1``) are rescaled onto Foxglove's default display max for their encoding |
| (``1.0`` / ``10000``) so they show with sensible contrast; ``depth_range`` sets the source |
| range, else the frame's own min/max is used. Ignored for ``mono8``/``rgb8``/``rgba8``. |
| raw_depth_values: If True, depth values are not rescaled and are logged as is. |
| """ |
|
|
| from foxglove.channels import CompressedImageChannel, RawImageChannel |
| from foxglove.messages import CompressedImage, RawImage, Timestamp |
|
|
| if channels is None: |
| channels = log_foxglove_data.channels |
| time_ns = time.time_ns() if log_time is None else log_time |
| timestamp = Timestamp(sec=time_ns // 1_000_000_000, nsec=time_ns % 1_000_000_000) |
| log_kwargs = {} if log_time is None else {"log_time": log_time} |
|
|
| |
| if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): |
| arr = np.transpose(arr, (1, 2, 0)) |
| height, width = arr.shape[0], arr.shape[1] |
| n_channels = 1 if arr.ndim == 2 else arr.shape[2] |
|
|
| if n_channels == 1 and arr.dtype != np.uint8: |
| |
| encoding, target_dtype, value_max = ( |
| ("32FC1", np.float32, 1.0) |
| if np.issubdtype(arr.dtype, np.floating) |
| else ("16UC1", np.uint16, 10000.0) |
| ) |
| if not raw_depth_values: |
| |
| lo, hi = depth_range if depth_range is not None else (float(arr.min()), float(arr.max())) |
| arr = arr.clip(lo, hi).astype(np.float32) |
| arr = (arr - lo) / ((hi - lo) if hi > lo else 1.0) * value_max |
| arr = np.ascontiguousarray(arr, dtype=target_dtype) |
| else: |
| if n_channels == 3 and np.issubdtype(arr.dtype, np.floating): |
| arr = (arr * 255.0).clip(0, 255) |
| arr = np.ascontiguousarray(arr, dtype=np.uint8) |
|
|
| if compress_images and n_channels == 3: |
| buf_src = cv2.cvtColor(arr, cv2.COLOR_RGB2BGR) |
| _, buf = cv2.imencode(".jpg", buf_src) |
| channel = channels.get(topic) |
| if channel is None: |
| channel = channels[topic] = CompressedImageChannel(topic=topic) |
| channel.log( |
| CompressedImage(timestamp=timestamp, frame_id=frame_id, data=buf.tobytes(), format="jpeg"), |
| **log_kwargs, |
| ) |
| return |
|
|
| encoding = {1: "mono8", 3: "rgb8", 4: "rgba8"}.get(n_channels) |
| if encoding is None: |
| logging.warning( |
| "Foxglove: skipping image on topic '%s' with unsupported shape %s (%d channels); " |
| "expected 1 (mono8/16UC1/32FC1), 3 (rgb8), or 4 (rgba8) channels.", |
| topic, |
| tuple(arr.shape), |
| n_channels, |
| ) |
| return |
|
|
| channel = channels.get(topic) |
| if channel is None: |
| channel = channels[topic] = RawImageChannel(topic=topic) |
| channel.log( |
| RawImage( |
| timestamp=timestamp, |
| frame_id=frame_id, |
| width=width, |
| height=height, |
| encoding=encoding, |
| step=width * n_channels * arr.itemsize, |
| data=arr.tobytes(), |
| ), |
| **log_kwargs, |
| ) |
|
|
|
|
| def log_foxglove_data( |
| observation: RobotObservation | None = None, |
| action: RobotAction | None = None, |
| compress_images: bool = False, |
| ) -> None: |
| """ |
| Logs observation and action data to a Foxglove WebSocket server for real-time visualization. |
| |
| Mirrors ``log_rerun_data`` but emits Foxglove messages over the server started by |
| :func:`init_foxglove`. Data is mapped as follows: |
| - Scalars (and elements of 1D arrays) are accumulated per source and logged on the |
| ``/observation/state`` and ``/action/state`` topics as typed JSON messages using the static |
| ``lerobot.Scalars`` schema: a ``scalars`` array of ``{label, value}`` objects (see |
| :data:`_SCALARS_SCHEMA`). The ``label`` field lets Foxglove name each series automatically, so |
| ``/observation/state.scalars[:].value`` plots every feature at once. |
| - 3D NumPy arrays that resemble images are transposed from CHW to HWC when needed and logged on a |
| per-source topic (e.g. ``/observation/images/front``) as a ``RawImage`` (or a JPEG |
| ``CompressedImage`` when ``compress_images`` is True). |
| |
| Args: |
| observation: An optional dictionary containing observation data to log. |
| action: An optional dictionary containing action data to log. |
| compress_images: Whether to JPEG-compress images before logging to save bandwidth in exchange |
| for CPU and quality. |
| """ |
|
|
| require_package("foxglove-sdk", extra="viz", import_name="foxglove") |
|
|
| if getattr(log_foxglove_data, "server", None) is None: |
| raise RuntimeError("init_foxglove() must be called before log_foxglove_data().") |
|
|
| now = time.time_ns() |
|
|
| if observation: |
| obs_scalars: dict[str, float] = {} |
| for k, v in observation.items(): |
| if v is None: |
| continue |
| key = k[len(OBS_PREFIX) :] if str(k).startswith(OBS_PREFIX) else str(k) |
| if _is_scalar(v): |
| obs_scalars[key] = float(v) |
| elif isinstance(v, np.ndarray): |
| if v.ndim == 1: |
| obs_scalars.update(_labeled_scalars(key, v)) |
| else: |
| _log_foxglove_image( |
| _foxglove_topic(k, is_image=True), |
| key, |
| v, |
| compress_images=compress_images, |
| log_time=now, |
| ) |
| _log_foxglove_scalars(_foxglove_topic(OBS_STATE), obs_scalars, log_time=now) |
|
|
| if action: |
| action_scalars: dict[str, float] = {} |
| for k, v in action.items(): |
| if v is None: |
| continue |
| key = k[len(ACTION_PREFIX) :] if str(k).startswith(ACTION_PREFIX) else str(k) |
| if _is_scalar(v): |
| action_scalars[key] = float(v) |
| elif isinstance(v, np.ndarray): |
| action_scalars.update(_labeled_scalars(key, v.flatten())) |
| _log_foxglove_scalars(_foxglove_topic(ACTION), action_scalars, log_time=now) |
|
|
|
|
| |
| |
| |
| |
|
|
|
|
| def _feature_dim_names(feature: dict | None) -> list[str] | None: |
| """Best-effort per-dimension series labels for a 1D feature, or ``None`` to fall back to indices. |
| |
| LeRobot records a feature's ``names`` inconsistently: a flat list (``["x", "y"]``), a category |
| mapping (``{"motors": ["motor_0", "motor_1"]}``), or a name->index mapping |
| (``{"delta_x": 0, "delta_y": 1}``). Each is handled, but labels are only returned when their count |
| matches the feature's 1D shape, so a malformed/mismatched ``names`` can't silently mislabel series. |
| """ |
|
|
| if not feature: |
| return None |
| shape = feature.get("shape") |
| dim = shape[0] if shape and len(shape) == 1 else None |
| names = feature.get("names") |
| labels: list[str] | None = None |
| if isinstance(names, dict): |
| values = list(names.values()) |
| if values and all(isinstance(v, (list, tuple)) for v in values): |
| labels = [str(n) for group in values for n in group] |
| elif values and all(isinstance(v, int) and not isinstance(v, bool) for v in values): |
| labels = [name for name, _ in sorted(names.items(), key=lambda kv: kv[1])] |
| elif isinstance(names, (list, tuple)): |
| labels = [str(n) for n in names] |
| if labels is not None and dim is not None and len(labels) == dim: |
| return labels |
| return None |
|
|
|
|
| def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -> dict[str, float]: |
| """Flatten a frame's vector/scalar feature ``key`` into ``{label: value}`` entries. |
| |
| ``labels`` provides one name per dimension (from the dataset's feature metadata); when absent or |
| the wrong length, dimensions fall back to ``{name}_{i}`` (the short feature name), matching the |
| live stream so series names agree. A scalar feature becomes a single entry. Missing or ``None`` |
| features yield an empty mapping. |
| """ |
|
|
| v = sample.get(key) |
| if v is None: |
| return {} |
| arr = v.numpy() if hasattr(v, "numpy") else np.asarray(v) |
| if key.startswith(OBS_PREFIX): |
| name = key[len(OBS_PREFIX) :] |
| elif key.startswith(ACTION_PREFIX): |
| name = key[len(ACTION_PREFIX) :] |
| else: |
| name = key |
| if arr.ndim == 0: |
| return {name: float(arr)} |
| return _labeled_scalars(name, arr.flatten(), labels) |
|
|
|
|
| def serve_foxglove_dataset_playback( |
| dataset, |
| episode_index: int, |
| *, |
| host: str = "127.0.0.1", |
| port: int = 8765, |
| compress_images: bool = False, |
| autoplay: bool = True, |
| ) -> None: |
| """Serve a single dataset episode to Foxglove as a seekable, scrubbable timeline. |
| |
| Starts a Foxglove WebSocket server advertising the ``PlaybackControl`` capability over the |
| episode's time range. The Foxglove app drives play/pause/seek/speed; a background thread and a |
| ``ServerListener`` read frames from the on-disk ``dataset`` on demand and log them stamped at |
| their dataset timestamps, so the user can scrub anywhere in the episode. Blocks until interrupted. |
| |
| Args: |
| dataset: A ``LeRobotDataset`` loaded for the single episode to visualize. |
| episode_index: Index of the episode being visualized (used only for the session name). |
| host: Host interface to bind the WebSocket server to. |
| port: Port to bind the WebSocket server to. |
| compress_images: Whether to JPEG-compress camera frames before logging. |
| autoplay: If True, start playing automatically as soon as a client connects, instead of |
| waiting for the user to press play in the Foxglove app. |
| """ |
|
|
| require_package("foxglove-sdk", extra="viz", import_name="foxglove") |
| import bisect |
| import threading |
|
|
| import foxglove |
| from foxglove.websocket import ( |
| Capability, |
| PlaybackCommand, |
| PlaybackControlRequest, |
| PlaybackState, |
| PlaybackStatus, |
| ServerListener, |
| ) |
|
|
| |
| times_ns = [int(round(float(t) * 1e9)) for t in dataset.hf_dataset["timestamp"]] |
| n_frames = len(times_ns) |
| if n_frames == 0: |
| raise ValueError("Cannot visualize an empty episode.") |
| first_ns, last_ns = times_ns[0], times_ns[-1] |
| camera_keys = list(dataset.meta.camera_keys) |
| |
| depth_ranges: dict[str, tuple[float, float]] = {} |
| for key in dataset.meta.depth_keys: |
| stats = (dataset.meta.stats or {}).get(key) |
| if not stats: |
| continue |
| lo = stats["q01"] if "q01" in stats else stats["min"] |
| hi = stats["q99"] if "q99" in stats else stats["max"] |
| depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item())) |
| |
| scalar_labels = { |
| OBS_STATE: _feature_dim_names(dataset.meta.features.get(OBS_STATE)), |
| ACTION: _feature_dim_names(dataset.meta.features.get(ACTION)), |
| } |
| |
| channels: dict = {} |
|
|
| def emit_frame(i: int) -> None: |
| """Log every channel for frame ``i`` stamped at its dataset timestamp.""" |
| sample = dataset[i] |
| log_time = times_ns[i] |
| for key in camera_keys: |
| arr = sample.get(key) |
| if arr is None: |
| continue |
| arr = arr.numpy() if hasattr(arr, "numpy") else np.asarray(arr) |
| _log_foxglove_image( |
| _foxglove_topic(key, is_image=True), |
| key, |
| arr, |
| compress_images=compress_images, |
| channels=channels, |
| log_time=log_time, |
| depth_range=depth_ranges.get(key), |
| raw_depth_values=True, |
| ) |
| _log_foxglove_scalars( |
| _foxglove_topic(OBS_STATE), |
| _frame_to_scalars(sample, OBS_STATE, scalar_labels[OBS_STATE]), |
| channels=channels, |
| log_time=log_time, |
| ) |
| _log_foxglove_scalars( |
| _foxglove_topic(ACTION), |
| _frame_to_scalars(sample, ACTION, scalar_labels[ACTION]), |
| channels=channels, |
| log_time=log_time, |
| ) |
| episode_scalars = {} |
| for feat, label in ( |
| (DONE, "done"), |
| (TRUNCATED, "truncated"), |
| (REWARD, "reward"), |
| (SUCCESS, "success"), |
| ): |
| v = sample.get(feat) |
| if v is not None: |
| episode_scalars[label] = float(v) |
| _log_foxglove_scalars("/episode/state", episode_scalars, channels=channels, log_time=log_time) |
|
|
| lock = threading.Lock() |
| stop_event = threading.Event() |
| |
| |
| |
| state = { |
| "status": PlaybackStatus.Paused, |
| "cursor": first_ns, |
| "speed": 1.0, |
| "last_idx": -1, |
| "seek_idx": None, |
| } |
|
|
| def index_at(t_ns: int) -> int: |
| return max(0, min(n_frames - 1, bisect.bisect_right(times_ns, t_ns) - 1)) |
|
|
| |
| autoplay_started = threading.Event() |
|
|
| class _PlaybackListener(ServerListener): |
| def on_subscribe(self, client, channel): |
| |
| |
| |
| if not autoplay: |
| return |
| with lock: |
| if autoplay_started.is_set() or state["status"] != PlaybackStatus.Paused: |
| return |
| autoplay_started.set() |
| state["status"] = PlaybackStatus.Playing |
| cursor, speed = state["cursor"], state["speed"] |
| server.broadcast_playback_state(PlaybackState(PlaybackStatus.Playing, cursor, speed, False, "")) |
|
|
| def on_playback_control_request(self, req: PlaybackControlRequest): |
| |
| with lock: |
| did_seek = False |
| if req.seek_time is not None: |
| cursor = max(first_ns, min(last_ns, req.seek_time)) |
| state["cursor"] = cursor |
| state["last_idx"] = state["seek_idx"] = index_at(cursor) |
| did_seek = True |
| if req.playback_speed and req.playback_speed > 0: |
| state["speed"] = req.playback_speed |
| if req.playback_command == PlaybackCommand.Play: |
| |
| if state["cursor"] >= last_ns: |
| state["cursor"] = first_ns |
| state["last_idx"] = state["seek_idx"] = 0 |
| did_seek = True |
| state["status"] = PlaybackStatus.Playing |
| elif req.playback_command == PlaybackCommand.Pause: |
| state["status"] = PlaybackStatus.Paused |
| status, cursor, speed = state["status"], state["cursor"], state["speed"] |
| request_id = req.request_id or "" |
| return PlaybackState(status, cursor, speed, did_seek, request_id) |
|
|
| server = foxglove.start_server( |
| name=f"{dataset.repo_id}/episode_{episode_index}", |
| host=host, |
| port=port, |
| capabilities=[Capability.PlaybackControl, Capability.Time], |
| server_listener=_PlaybackListener(), |
| playback_time_range=(first_ns, last_ns), |
| ) |
|
|
| def playback_loop() -> None: |
| |
| |
| |
| max_tick_dt_s = 0.25 |
| prev = time.monotonic() |
| while not stop_event.is_set(): |
| time.sleep(1.0 / 60.0) |
| ended = False |
| speed = 1.0 |
| with lock: |
| now = time.monotonic() |
| dt = min(now - prev, max_tick_dt_s) |
| prev = now |
| |
| work = [] |
| seek_idx = state["seek_idx"] |
| if seek_idx is not None: |
| state["seek_idx"] = None |
| work.append(seek_idx) |
| if state["status"] == PlaybackStatus.Playing: |
| cursor = state["cursor"] + int(dt * 1e9 * state["speed"]) |
| start_idx = state["last_idx"] + 1 |
| if cursor >= last_ns: |
| cursor, target, ended = last_ns, n_frames - 1, True |
| else: |
| target = index_at(cursor) |
| state["cursor"] = cursor |
| work.extend(range(start_idx, target + 1)) |
| |
| |
| state["last_idx"] = target |
| if ended: |
| state["status"] = PlaybackStatus.Ended |
| if not work: |
| continue |
| cursor, speed = state["cursor"], state["speed"] |
| |
| |
| for i in work: |
| if stop_event.is_set(): |
| break |
| emit_frame(i) |
| server.broadcast_time(cursor) |
| if ended: |
| server.broadcast_playback_state(PlaybackState(PlaybackStatus.Ended, cursor, speed, False, "")) |
|
|
| |
| |
| emit_frame(0) |
| with lock: |
| state["last_idx"] = 0 |
| server.broadcast_time(first_ns) |
| server.broadcast_playback_state(PlaybackState(PlaybackStatus.Paused, first_ns, 1.0, True, "")) |
|
|
| thread = threading.Thread(target=playback_loop, name="foxglove-playback", daemon=True) |
| thread.start() |
|
|
| print(f"Foxglove server running. Connect the Foxglove app to ws://{host}:{port}") |
| print("Use the playback controls in Foxglove to play/pause and scrub the episode. Ctrl-C to exit.") |
| try: |
| while not stop_event.is_set(): |
| time.sleep(0.5) |
| except KeyboardInterrupt: |
| print("Ctrl-C received. Exiting.") |
| finally: |
| stop_event.set() |
| thread.join(timeout=2.0) |
| server.stop() |
| channels.clear() |
|
|