feat(visualization): add foxglove support (#3902)

* Add Foxglove display mode for teleoperate

Add a --display_mode flag (rerun|foxglove) to lerobot-teleoperate. When set
to foxglove, stream observations/actions over a Foxglove WebSocket server:
images as RawImage/CompressedImage, scalars as typed JSON channels with
schemas generated from the feature names (sanitized so paths don't need
quoting). Adds a `foxglove` extra.

* Add Foxglove display mode to lerobot-record

Wire the --display_mode flag (rerun|foxglove) into lerobot-record, matching
lerobot-teleoperate: route init/log through the backend-agnostic dispatchers
and stop the visualization backend on exit.

* update foxglove-sdk to 0.25.1

* Use static lerobot.Scalars schema for Foxglove state topics

Replace the per-topic JSON schema derived from feature names with a single
static lerobot.Scalars schema: a scalars array of {label, value} objects. The
same schema fits any robot regardless of which observation/action features it
reports, and the label field lets Foxglove name each series automatically so
one filtered path plots every feature.

* add foxglove option to dataset viz

* Make Foxglove dataset playback loop the sole frame emitter

Address review: the listener no longer emits frames, it only mutates
playback state and queues a one-shot seek index that the playback loop
services. The loop is now the only caller of emit_frame, so concurrent
random access into the on-disk dataset / video decoder never overlaps.

Also remove the dead server_holder and tighten the _foxglove_safe_name
docstring to state what it does and why.

* Label Foxglove dataset scalars with feature dimension names

Use the dataset's per-dimension feature names (e.g. joint names) as the
Foxglove series labels for /observation/state and /action/state instead
of bare indices. LeRobot stores `names` inconsistently (flat list,
{category: [...]}, or {name: index}), so _feature_dim_names handles each
and falls back to indices on any unknown format or length mismatch.

* Make Foxglove server host bindable and refactor topic/channel handling

Pass display_ip through as the Foxglove WebSocket bind host (127.0.0.1
for local only, 0.0.0.0 for all interfaces) instead of always binding
locally. In lerobot-dataset-viz, fold the separate --port into --web-port
so one flag covers both the Rerun web viewer and the Foxglove server port.

Add a _foxglove_topic() helper and thread a per-topic channel cache
through the log helpers so dataset playback stays self-contained instead
of mutating the module-global cache. Promote SUCCESS to constants.py.

* feat(viz): add support for foxglove in rollout + add to viz tag

* fix(docs): remove misleading installation note

* fix(visualization): no duplicated prefix, consolidated norm + warnings log

* chore(viz): minor improvements

* refactor(viz): split files + autoplay + updated docs + added minimal tests

* fix(viz): right tags + warning

* feat(deprecated ws-port): removing rerun's depreacted ws-port parameter in dataset visualization

* chore(web ports): adding global variables for default foxglove/rerun web ports

* feat(depth): adding depth support to foxglove visualizer. Because of foxglove limitations (min and max values on RawImage cannot be set from the SDK), depth is normalized between [0,1] when a depth range is provided.

* fix(rerun depth range): making rerun depth range computation safe against missing stats

* chore(foxglove depth): make it simple, and make it work.

* fix(scaling): fixing depth frames scaling

---------

Co-authored-by: Roman Shtylman <roman@foxglove.dev>
Co-authored-by: Caroline Pascal <caroline8.pascal@gmail.com>

tests/utils/test_foxglove_visualization.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+
+"""Tests for the Foxglove backend's pure helpers.
+
+These cover topic naming, series labelling and feature-name parsing. They import
+``foxglove_visualization`` directly and need NO ``foxglove`` extra: the SDK is imported lazily inside
+the functions that talk to the server, so the helpers below run in the base test tier.
+"""
+
+import numpy as np
+
+from lerobot.utils import foxglove_visualization as fv
+from lerobot.utils.constants import ACTION, OBS_STATE
+
+
+def test_foxglove_safe_name_collapses_dots():
+    assert fv._foxglove_safe_name("observation.images.front") == "observation_images_front"
+    assert fv._foxglove_safe_name("plain") == "plain"
+
+
+def test_foxglove_topic_image_strips_prefix_without_doubling_images():
+    # Fully-qualified camera key -> single clean segment (no doubled "images").
+    assert fv._foxglove_topic("observation.images.front", is_image=True) == "/observation/images/front"
+    # A nested camera name keeps its structure via safe-name collapsing.
+    assert (
+        fv._foxglove_topic("observation.images.wrist.left", is_image=True) == "/observation/images/wrist_left"
+    )
+    # Bare camera name (as real robots emit).
+    assert fv._foxglove_topic("front", is_image=True) == "/observation/images/front"
+
+
+def test_foxglove_topic_scalar_sources():
+    assert fv._foxglove_topic(OBS_STATE) == "/observation/state"
+    assert fv._foxglove_topic("observation.environment_state") == "/observation/state"
+    assert fv._foxglove_topic(ACTION) == "/action/state"
+    assert fv._foxglove_topic("action.delta") == "/action/state"
+
+
+def test_labeled_scalars_uses_labels_then_index_fallback():
+    assert fv._labeled_scalars("state", np.array([1.0, 2.0, 3.0])) == {
+        "state_0": 1.0,
+        "state_1": 2.0,
+        "state_2": 3.0,
+    }
+    assert fv._labeled_scalars("state", [1.0, 2.0], ["pan", "lift"]) == {"pan": 1.0, "lift": 2.0}
+    # Wrong-length labels fall back to index naming (never silently mislabels).
+    assert fv._labeled_scalars("q", [1.0, 2.0], ["only_one"]) == {"q_0": 1.0, "q_1": 2.0}
+
+
+def test_frame_to_scalars_matches_live_labeling_and_handles_scalar():
+    frame = {OBS_STATE: np.array([1.0, 2.0])}
+    # No metadata -> {short_name}_{i}, identical to the live-stream fallback.
+    assert fv._frame_to_scalars(frame, OBS_STATE) == fv._labeled_scalars("state", np.array([1.0, 2.0]))
+    assert fv._frame_to_scalars(frame, OBS_STATE) == {"state_0": 1.0, "state_1": 2.0}
+    # Metadata labels are honored.
+    assert fv._frame_to_scalars(frame, OBS_STATE, ["pan", "lift"]) == {"pan": 1.0, "lift": 2.0}
+    # A 0-d scalar becomes a single entry named by the short feature name.
+    assert fv._frame_to_scalars({ACTION: np.array(5.0)}, ACTION) == {"action": 5.0}
+    # A missing feature yields an empty mapping.
+    assert fv._frame_to_scalars({}, OBS_STATE) == {}
+
+
+def test_feature_dim_names_formats():
+    # Flat list of names.
+    assert fv._feature_dim_names({"shape": [2], "names": ["x", "y"]}) == ["x", "y"]
+    # Category mapping (dict of lists).
+    assert fv._feature_dim_names({"shape": [2], "names": {"motors": ["m0", "m1"]}}) == ["m0", "m1"]
+    # name -> index mapping (returned sorted by index).
+    assert fv._feature_dim_names({"shape": [2], "names": {"delta_x": 0, "delta_y": 1}}) == [
+        "delta_x",
+        "delta_y",
+    ]
+    # Bool values must NOT be treated as an index map (bool is a subclass of int).
+    assert fv._feature_dim_names({"shape": [2], "names": {"a": True, "b": False}}) is None
+    # Mismatched length -> None (won't silently mislabel).
+    assert fv._feature_dim_names({"shape": [3], "names": ["x", "y"]}) is None
+    # Missing / absent names -> None.
+    assert fv._feature_dim_names(None) is None
+    assert fv._feature_dim_names({"shape": [2]}) is None
+
+
+def test_is_scalar():
+    assert fv._is_scalar(1.0)
+    assert fv._is_scalar(np.float32(2.0))
+    assert fv._is_scalar(np.array(3.0))  # 0-d array
+    assert not fv._is_scalar(np.array([1.0, 2.0]))
+    assert not fv._is_scalar("x")

tests/utils/test_rerun_visualization.py

@@ -0,0 +1,310 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+
+import importlib
+import sys
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+
+pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])")
+
+from lerobot.types import TransitionKey
+from lerobot.utils.constants import OBS_STATE
+
+
+@pytest.fixture
+def mock_rerun(monkeypatch):
+    """
+    Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
+    depend on the real library. Also reload the module-under-test so it binds to
+    this mock `rr`.
+    """
+    calls = []
+    blueprints = []
+
+    class DummyScalar:
+        def __init__(self, value):
+            # Scalars may be built from a single float or from a 1D array batch.
+            self.value = value
+
+    class DummyImage:
+        def __init__(self, arr):
+            self.arr = arr
+
+        def compress(self, *a, **k):
+            return self
+
+    class DummyDepthImage:
+        def __init__(self, arr, colormap=None):
+            self.arr = arr
+            self.colormap = colormap
+
+    def dummy_log(key, obj=None, **kwargs):
+        # Accept either positional `obj` or keyword `entity` and record remaining kwargs.
+        if obj is None and "entity" in kwargs:
+            obj = kwargs.pop("entity")
+        calls.append((key, obj, kwargs))
+
+    def dummy_send_blueprint(blueprint, *a, **k):
+        blueprints.append(blueprint)
+
+    # Mock the `rerun.blueprint` submodule used to build the layout.
+    dummy_rrb = SimpleNamespace(
+        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
+            kind="Spatial2DView", origin=origin, name=name
+        ),
+        TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
+            kind="TimeSeriesView", name=name, contents=contents
+        ),
+        Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
+        Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
+    )
+
+    dummy_rr = SimpleNamespace(
+        __name__="rerun",
+        __package__="rerun",
+        __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
+        Scalars=DummyScalar,
+        Image=DummyImage,
+        DepthImage=DummyDepthImage,
+        components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
+        log=dummy_log,
+        send_blueprint=dummy_send_blueprint,
+        init=lambda *a, **k: None,
+        spawn=lambda *a, **k: None,
+        blueprint=dummy_rrb,
+    )
+
+    # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
+    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
+    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
+
+    # Now import and reload the module under test, to bind to our rerun mock
+    import lerobot.utils.rerun_visualization as rv
+
+    importlib.reload(rv)
+
+    # Expose the reloaded module, the call recorder and the captured blueprints
+    yield rv, calls, blueprints
+
+
+def _keys(calls):
+    """Helper to extract just the keys logged to rr.log"""
+    return [k for (k, _obj, _kw) in calls]
+
+
+def _obj_for(calls, key):
+    """Find the first object logged under a given key."""
+    for k, obj, _kw in calls:
+        if k == key:
+            return obj
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def _kwargs_for(calls, key):
+    for k, _obj, kw in calls:
+        if k == key:
+            return kw
+    raise KeyError(f"Key {key} not found in calls: {calls}")
+
+
+def _views_by_kind(blueprint, kind):
+    """Return the views of a given kind from the (single) blueprint's grid."""
+    return [v for v in blueprint.root.views if v.kind == kind]
+
+
+def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
+    rv, calls, blueprints = mock_rerun
+
+    # Build EnvTransition dict
+    obs = {
+        f"{OBS_STATE}.temperature": np.float32(25.0),
+        # CHW image should be converted to HWC for rr.Image
+        "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8),
+    }
+    act = {
+        "action.throttle": 0.7,
+        # 1D array should be logged as a single Scalars batch under one entity path
+        "action.vector": np.array([1.0, 2.0], dtype=np.float32),
+    }
+    transition = {
+        TransitionKey.OBSERVATION: obs,
+        TransitionKey.ACTION: act,
+    }
+
+    # Extract observation and action data from transition like in the real call sites
+    obs_data = transition.get(TransitionKey.OBSERVATION, {})
+    action_data = transition.get(TransitionKey.ACTION, {})
+    rv.log_rerun_data(observation=obs_data, action=action_data)
+
+    # We expect:
+    # - observation.state.temperature -> Scalars
+    # - observation.camera -> Image (HWC) with static=True
+    # - action.throttle -> Scalars
+    # - action.vector -> single Scalars batch (no per-element suffix)
+    expected_keys = {
+        f"{OBS_STATE}.temperature",
+        "observation.camera",
+        "action.throttle",
+        "action.vector",
+    }
+    assert set(_keys(calls)) == expected_keys
+
+    # Check scalar types and values
+    temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
+    assert type(temp_obj).__name__ == "DummyScalar"
+    assert float(temp_obj.value) == pytest.approx(25.0)
+
+    throttle_obj = _obj_for(calls, "action.throttle")
+    assert type(throttle_obj).__name__ == "DummyScalar"
+    assert float(throttle_obj.value) == pytest.approx(0.7)
+
+    # 1D vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vector")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
+
+    # Check image handling: CHW -> HWC
+    img_obj = _obj_for(calls, "observation.camera")
+    assert type(img_obj).__name__ == "DummyImage"
+    assert img_obj.arr.shape == (10, 20, 3)  # transposed
+    assert _kwargs_for(calls, "observation.camera").get("static", False) is True  # static=True for images
+
+    # A blueprint should have been built and sent exactly once, and cached on the function.
+    assert len(blueprints) == 1
+    assert rv.log_rerun_data.blueprint is blueprints[0]
+
+    bp = blueprints[0]
+    # One spatial view per image path
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.camera"}
+
+    # One time-series view each for observation and action scalars
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert set(ts_views) == {"observation", "action"}
+    assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vector"]
+
+
+def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
+    rv, calls, blueprints = mock_rerun
+
+    # First dict without prefixes treated as observation
+    # Second dict without prefixes treated as action
+    obs_plain = {
+        "temp": 1.5,
+        # Already HWC image => should stay as-is
+        "img": np.zeros((5, 6, 3), dtype=np.uint8),
+        "none": None,  # should be skipped
+    }
+    act_plain = {
+        "throttle": 0.3,
+        "vec": np.array([9, 8, 7], dtype=np.float32),
+    }
+
+    # Extract observation and action data from list like the old function logic did
+    # First dict was treated as observation, second as action
+    rv.log_rerun_data(observation=obs_plain, action=act_plain)
+
+    # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
+    expected = {
+        "observation.temp",
+        "observation.img",
+        "action.throttle",
+        "action.vec",
+    }
+    logged = set(_keys(calls))
+    assert logged == expected
+
+    # Scalars
+    t = _obj_for(calls, "observation.temp")
+    assert type(t).__name__ == "DummyScalar"
+    assert float(t.value) == pytest.approx(1.5)
+
+    throttle = _obj_for(calls, "action.throttle")
+    assert type(throttle).__name__ == "DummyScalar"
+    assert float(throttle.value) == pytest.approx(0.3)
+
+    # Image stays HWC
+    img = _obj_for(calls, "observation.img")
+    assert type(img).__name__ == "DummyImage"
+    assert img.arr.shape == (5, 6, 3)
+    assert _kwargs_for(calls, "observation.img").get("static", False) is True
+
+    # Vector logged as a single batched Scalars under one entity path
+    vec = _obj_for(calls, "action.vec")
+    assert type(vec).__name__ == "DummyScalar"
+    np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
+
+    # Blueprint sent once with the expected view layout
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    spatial_views = _views_by_kind(bp, "Spatial2DView")
+    assert {v.origin for v in spatial_views} == {"observation.img"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.throttle", "action.vec"]
+
+
+def test_log_rerun_data_kwargs_only(mock_rerun):
+    rv, calls, blueprints = mock_rerun
+
+    rv.log_rerun_data(
+        observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
+        action={"action.a": 1.0},
+    )
+
+    keys = set(_keys(calls))
+    assert "observation.temp" in keys
+    assert "observation.gray" in keys
+    assert "action.a" in keys
+
+    temp = _obj_for(calls, "observation.temp")
+    assert type(temp).__name__ == "DummyScalar"
+    assert float(temp.value) == pytest.approx(10.0)
+
+    img = _obj_for(calls, "observation.gray")
+    assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
+    assert img.arr.shape == (8, 8, 1)  # remains HWC
+    assert _kwargs_for(calls, "observation.gray").get("static", False) is True
+
+    a = _obj_for(calls, "action.a")
+    assert type(a).__name__ == "DummyScalar"
+    assert float(a.value) == pytest.approx(1.0)
+
+    # Blueprint sent once, with a spatial view for the image and time-series views for scalars
+    assert len(blueprints) == 1
+    bp = blueprints[0]
+    assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
+    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
+    assert ts_views["observation"].contents == ["observation.temp"]
+    assert ts_views["action"].contents == ["action.a"]
+
+
+def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
+    """The blueprint is built from the first call and not resent on subsequent calls."""
+    rv, calls, blueprints = mock_rerun
+
+    rv.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
+    assert len(blueprints) == 1
+    first_blueprint = rv.log_rerun_data.blueprint
+
+    rv.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
+    # Still only one blueprint, and the cached one is unchanged.
+    assert len(blueprints) == 1
+    assert rv.log_rerun_data.blueprint is first_blueprint

tests/utils/test_visualization_utils.py

@@ -14,297 +14,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import importlib
-import sys
-from types import SimpleNamespace
+"""Tests for the backend-agnostic visualization dispatch.
+
+These exercise the display-mode routing/validation only; they need neither ``rerun`` nor
+``foxglove`` installed since the unknown-mode branch raises before touching any backend. Backend
+behavior is covered in ``test_rerun_visualization.py`` and ``test_foxglove_visualization.py``.
+"""
 
-import numpy as np
 import pytest
 
-pytest.importorskip("rerun", reason="rerun-sdk is required (install lerobot[viz])")
-
-from lerobot.types import TransitionKey
-from lerobot.utils.constants import OBS_STATE
+from lerobot.utils import visualization_utils as vu
 
 
-@pytest.fixture
-def mock_rerun(monkeypatch):
-    """
-    Provide a mock `rerun` module (and `rerun.blueprint` submodule) so tests don't
-    depend on the real library. Also reload the module-under-test so it binds to
-    this mock `rr`.
-    """
-    calls = []
-    blueprints = []
-
-    class DummyScalar:
-        def __init__(self, value):
-            # Scalars may be built from a single float or from a 1D array batch.
-            self.value = value
-
-    class DummyImage:
-        def __init__(self, arr):
-            self.arr = arr
-
-        def compress(self, *a, **k):
-            return self
-
-    class DummyDepthImage:
-        def __init__(self, arr, colormap=None):
-            self.arr = arr
-            self.colormap = colormap
-
-    def dummy_log(key, obj=None, **kwargs):
-        # Accept either positional `obj` or keyword `entity` and record remaining kwargs.
-        if obj is None and "entity" in kwargs:
-            obj = kwargs.pop("entity")
-        calls.append((key, obj, kwargs))
-
-    def dummy_send_blueprint(blueprint, *a, **k):
-        blueprints.append(blueprint)
-
-    # Mock the `rerun.blueprint` submodule used to build the layout.
-    dummy_rrb = SimpleNamespace(
-        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
-            kind="Spatial2DView", origin=origin, name=name
-        ),
-        TimeSeriesView=lambda name=None, contents=None: SimpleNamespace(
-            kind="TimeSeriesView", name=name, contents=contents
-        ),
-        Grid=lambda *views: SimpleNamespace(kind="Grid", views=list(views)),
-        Blueprint=lambda root: SimpleNamespace(kind="Blueprint", root=root),
-    )
-
-    dummy_rr = SimpleNamespace(
-        __name__="rerun",
-        __package__="rerun",
-        __spec__=SimpleNamespace(name="rerun", submodule_search_locations=None),
-        Scalars=DummyScalar,
-        Image=DummyImage,
-        DepthImage=DummyDepthImage,
-        components=SimpleNamespace(Colormap=SimpleNamespace(Viridis="viridis")),
-        log=dummy_log,
-        send_blueprint=dummy_send_blueprint,
-        init=lambda *a, **k: None,
-        spawn=lambda *a, **k: None,
-        blueprint=dummy_rrb,
-    )
-
-    # Inject fake modules into sys.modules (both `rerun` and `rerun.blueprint`).
-    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
-    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
-
-    # Now import and reload the module under test, to bind to our rerun mock
-    import lerobot.utils.visualization_utils as vu
-
-    importlib.reload(vu)
-
-    # Expose the reloaded module, the call recorder and the captured blueprints
-    yield vu, calls, blueprints
+def test_visualization_modes():
+    assert vu.VISUALIZATION_MODES == ("rerun", "foxglove")
 
 
-def _keys(calls):
-    """Helper to extract just the keys logged to rr.log"""
-    return [k for (k, _obj, _kw) in calls]
-
-
-def _obj_for(calls, key):
-    """Find the first object logged under a given key."""
-    for k, obj, _kw in calls:
-        if k == key:
-            return obj
-    raise KeyError(f"Key {key} not found in calls: {calls}")
-
-
-def _kwargs_for(calls, key):
-    for k, _obj, kw in calls:
-        if k == key:
-            return kw
-    raise KeyError(f"Key {key} not found in calls: {calls}")
-
-
-def _views_by_kind(blueprint, kind):
-    """Return the views of a given kind from the (single) blueprint's grid."""
-    return [v for v in blueprint.root.views if v.kind == kind]
-
-
-def test_log_rerun_data_envtransition_scalars_and_image(mock_rerun):
-    vu, calls, blueprints = mock_rerun
-
-    # Build EnvTransition dict
-    obs = {
-        f"{OBS_STATE}.temperature": np.float32(25.0),
-        # CHW image should be converted to HWC for rr.Image
-        "observation.camera": np.zeros((3, 10, 20), dtype=np.uint8),
-    }
-    act = {
-        "action.throttle": 0.7,
-        # 1D array should be logged as a single Scalars batch under one entity path
-        "action.vector": np.array([1.0, 2.0], dtype=np.float32),
-    }
-    transition = {
-        TransitionKey.OBSERVATION: obs,
-        TransitionKey.ACTION: act,
-    }
-
-    # Extract observation and action data from transition like in the real call sites
-    obs_data = transition.get(TransitionKey.OBSERVATION, {})
-    action_data = transition.get(TransitionKey.ACTION, {})
-    vu.log_rerun_data(observation=obs_data, action=action_data)
-
-    # We expect:
-    # - observation.state.temperature -> Scalars
-    # - observation.camera -> Image (HWC) with static=True
-    # - action.throttle -> Scalars
-    # - action.vector -> single Scalars batch (no per-element suffix)
-    expected_keys = {
-        f"{OBS_STATE}.temperature",
-        "observation.camera",
-        "action.throttle",
-        "action.vector",
-    }
-    assert set(_keys(calls)) == expected_keys
-
-    # Check scalar types and values
-    temp_obj = _obj_for(calls, f"{OBS_STATE}.temperature")
-    assert type(temp_obj).__name__ == "DummyScalar"
-    assert float(temp_obj.value) == pytest.approx(25.0)
-
-    throttle_obj = _obj_for(calls, "action.throttle")
-    assert type(throttle_obj).__name__ == "DummyScalar"
-    assert float(throttle_obj.value) == pytest.approx(0.7)
-
-    # 1D vector logged as a single batched Scalars under one entity path
-    vec = _obj_for(calls, "action.vector")
-    assert type(vec).__name__ == "DummyScalar"
-    np.testing.assert_allclose(np.asarray(vec.value), [1.0, 2.0])
-
-    # Check image handling: CHW -> HWC
-    img_obj = _obj_for(calls, "observation.camera")
-    assert type(img_obj).__name__ == "DummyImage"
-    assert img_obj.arr.shape == (10, 20, 3)  # transposed
-    assert _kwargs_for(calls, "observation.camera").get("static", False) is True  # static=True for images
-
-    # A blueprint should have been built and sent exactly once, and cached on the function.
-    assert len(blueprints) == 1
-    assert vu.log_rerun_data.blueprint is blueprints[0]
-
-    bp = blueprints[0]
-    # One spatial view per image path
-    spatial_views = _views_by_kind(bp, "Spatial2DView")
-    assert {v.origin for v in spatial_views} == {"observation.camera"}
-
-    # One time-series view each for observation and action scalars
-    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
-    assert set(ts_views) == {"observation", "action"}
-    assert ts_views["observation"].contents == [f"{OBS_STATE}.temperature"]
-    assert ts_views["action"].contents == ["action.throttle", "action.vector"]
-
-
-def test_log_rerun_data_plain_list_ordering_and_prefixes(mock_rerun):
-    vu, calls, blueprints = mock_rerun
-
-    # First dict without prefixes treated as observation
-    # Second dict without prefixes treated as action
-    obs_plain = {
-        "temp": 1.5,
-        # Already HWC image => should stay as-is
-        "img": np.zeros((5, 6, 3), dtype=np.uint8),
-        "none": None,  # should be skipped
-    }
-    act_plain = {
-        "throttle": 0.3,
-        "vec": np.array([9, 8, 7], dtype=np.float32),
-    }
-
-    # Extract observation and action data from list like the old function logic did
-    # First dict was treated as observation, second as action
-    vu.log_rerun_data(observation=obs_plain, action=act_plain)
-
-    # Expected keys with auto-prefixes. The 1D vector is a single batched Scalars.
-    expected = {
-        "observation.temp",
-        "observation.img",
-        "action.throttle",
-        "action.vec",
-    }
-    logged = set(_keys(calls))
-    assert logged == expected
-
-    # Scalars
-    t = _obj_for(calls, "observation.temp")
-    assert type(t).__name__ == "DummyScalar"
-    assert float(t.value) == pytest.approx(1.5)
-
-    throttle = _obj_for(calls, "action.throttle")
-    assert type(throttle).__name__ == "DummyScalar"
-    assert float(throttle.value) == pytest.approx(0.3)
-
-    # Image stays HWC
-    img = _obj_for(calls, "observation.img")
-    assert type(img).__name__ == "DummyImage"
-    assert img.arr.shape == (5, 6, 3)
-    assert _kwargs_for(calls, "observation.img").get("static", False) is True
-
-    # Vector logged as a single batched Scalars under one entity path
-    vec = _obj_for(calls, "action.vec")
-    assert type(vec).__name__ == "DummyScalar"
-    np.testing.assert_allclose(np.asarray(vec.value), [9, 8, 7])
-
-    # Blueprint sent once with the expected view layout
-    assert len(blueprints) == 1
-    bp = blueprints[0]
-    spatial_views = _views_by_kind(bp, "Spatial2DView")
-    assert {v.origin for v in spatial_views} == {"observation.img"}
-    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
-    assert ts_views["observation"].contents == ["observation.temp"]
-    assert ts_views["action"].contents == ["action.throttle", "action.vec"]
-
-
-def test_log_rerun_data_kwargs_only(mock_rerun):
-    vu, calls, blueprints = mock_rerun
-
-    vu.log_rerun_data(
-        observation={"observation.temp": 10.0, "observation.gray": np.zeros((8, 8, 1), dtype=np.uint8)},
-        action={"action.a": 1.0},
-    )
-
-    keys = set(_keys(calls))
-    assert "observation.temp" in keys
-    assert "observation.gray" in keys
-    assert "action.a" in keys
-
-    temp = _obj_for(calls, "observation.temp")
-    assert type(temp).__name__ == "DummyScalar"
-    assert float(temp.value) == pytest.approx(10.0)
-
-    img = _obj_for(calls, "observation.gray")
-    assert type(img).__name__ == "DummyDepthImage"  # single-channel -> DepthImage
-    assert img.arr.shape == (8, 8, 1)  # remains HWC
-    assert _kwargs_for(calls, "observation.gray").get("static", False) is True
-
-    a = _obj_for(calls, "action.a")
-    assert type(a).__name__ == "DummyScalar"
-    assert float(a.value) == pytest.approx(1.0)
-
-    # Blueprint sent once, with a spatial view for the image and time-series views for scalars
-    assert len(blueprints) == 1
-    bp = blueprints[0]
-    assert {v.origin for v in _views_by_kind(bp, "Spatial2DView")} == {"observation.gray"}
-    ts_views = {v.name: v for v in _views_by_kind(bp, "TimeSeriesView")}
-    assert ts_views["observation"].contents == ["observation.temp"]
-    assert ts_views["action"].contents == ["action.a"]
-
-
-def test_log_rerun_data_blueprint_sent_only_once(mock_rerun):
-    """The blueprint is built from the first call and not resent on subsequent calls."""
-    vu, calls, blueprints = mock_rerun
-
-    vu.log_rerun_data(observation={"temp": 1.0}, action={"a": 2.0})
-    assert len(blueprints) == 1
-    first_blueprint = vu.log_rerun_data.blueprint
-
-    vu.log_rerun_data(observation={"temp": 3.0}, action={"a": 4.0})
-    # Still only one blueprint, and the cached one is unchanged.
-    assert len(blueprints) == 1
-    assert vu.log_rerun_data.blueprint is first_blueprint
+@pytest.mark.parametrize("func", ["init_visualization", "log_visualization_data", "shutdown_visualization"])
+def test_dispatch_rejects_unknown_mode(func):
+    with pytest.raises(ValueError, match="Unknown display_mode"):
+        getattr(vu, func)("bogus")