feat(dataset-viz): render generic custom scalar columns

Salvage of huggingface/lerobot#3918 by @nathon-lee — rebased to main.

lerobot-dataset-viz only logged known scalars; custom float/bool fields
now appear in the Rerun blueprint and frame stream when shape is scalar.

refactor(viz): add new non-default colums viz

Co-authored-by: nathon-lee <leejianwoo@gmail.com>

src/lerobot/scripts/lerobot_dataset_viz.py

@@ -87,6 +87,11 @@ import tqdm
 from lerobot.configs import DEPTH_MILLIMETER_UNIT
 from lerobot.datasets import LeRobotDataset
 from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
+from lerobot.utils.dataset_visualization_utils import (
+    get_extra_scalar_keys,
+    is_scalar_like,
+    scalar_to_float,
+)
 from lerobot.utils.utils import init_logging
 
 logger = logging.getLogger(__name__)
@@ -153,6 +158,8 @@ def build_blueprint_from_dataset(dataset: LeRobotDataset):
     for key in (DONE, REWARD, SUCCESS):
         if key in dataset.features:
             views.append(rrb.TimeSeriesView(origin=key, name=key))
+    for key in get_extra_scalar_keys(dataset):
+        views.append(rrb.TimeSeriesView(origin=key, name=key))
 
     return rrb.Blueprint(rrb.Grid(*views))
 
@@ -244,6 +251,8 @@ def visualize_dataset(
         hi = stats["q99"] if "q99" in stats else stats["max"]
         depth_ranges[key] = (float(np.asarray(lo).item()), float(np.asarray(hi).item()))
 
+    extra_scalar_keys = get_extra_scalar_keys(dataset)
+
     first_index = None
     for batch in tqdm.tqdm(dataloader, total=len(dataloader)):
         if first_index is None:
@@ -287,6 +296,10 @@ def visualize_dataset(
             if SUCCESS in batch:
                 rr.log(SUCCESS, rr.Scalars(batch[SUCCESS][i].item()))
 
+            for key in extra_scalar_keys:
+                if key in batch and is_scalar_like(batch[key][i]):
+                    rr.log(key, rr.Scalars(scalar_to_float(batch[key][i])))
+
     # save .rrd locally
     if mode == "local" and save:
         output_dir.mkdir(parents=True, exist_ok=True)

src/lerobot/utils/dataset_visualization_utils.py

@@ -0,0 +1,99 @@
+# Copyright 2024 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.
+
+"""Shared helpers for visualizing scalar features from a LeRobot dataset."""
+
+from collections.abc import Iterable, Mapping
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+
+from .constants import ACTION, DEFAULT_FEATURES, DONE, OBS_STATE, REWARD, SUCCESS
+
+if TYPE_CHECKING:
+    from lerobot.datasets import LeRobotDataset
+
+
+METADATA_KEYS = {*DEFAULT_FEATURES, "task"}
+KNOWN_SCALAR_KEYS = {DONE, REWARD, SUCCESS}
+SCALAR_DTYPE_KINDS = {"b", "i", "u", "f"}
+
+
+def is_scalar_feature(feature: Mapping) -> bool:
+    """Return whether a feature schema describes a numeric or boolean scalar."""
+
+    dtype = feature.get("dtype")
+    if not isinstance(dtype, str):
+        return False
+    try:
+        dtype_kind = np.dtype(dtype).kind
+    except (TypeError, ValueError):
+        return False
+    if dtype_kind not in SCALAR_DTYPE_KINDS:
+        return False
+
+    shape = feature.get("shape")
+    if shape is None:
+        return True
+    if isinstance(shape, int):
+        return shape == 1
+    if not isinstance(shape, (list, tuple)):
+        return False
+    return len(shape) == 0 or (len(shape) == 1 and shape[0] == 1)
+
+
+def get_extra_scalar_keys(dataset: "LeRobotDataset", additional_known_keys: Iterable[str] = ()) -> list[str]:
+    """Return scalar feature keys not handled by the visualizer's standard paths."""
+
+    known_keys = {
+        ACTION,
+        OBS_STATE,
+        *KNOWN_SCALAR_KEYS,
+        *METADATA_KEYS,
+        *additional_known_keys,
+        *dataset.meta.camera_keys,
+    }
+    return [
+        key
+        for key, feature in dataset.features.items()
+        if key not in known_keys and is_scalar_feature(feature)
+    ]
+
+
+def is_scalar_like(value: object) -> bool:
+    """Return whether a runtime value contains exactly one numeric or boolean scalar."""
+
+    if isinstance(value, torch.Tensor):
+        return value.numel() == 1 and not value.is_complex()
+    if isinstance(value, np.ndarray):
+        return value.size == 1 and value.dtype.kind in SCALAR_DTYPE_KINDS
+    return np.isscalar(value) and np.asarray(value).dtype.kind in SCALAR_DTYPE_KINDS
+
+
+def scalar_to_float(value: object) -> float:
+    """Convert a scalar-like tensor, array, or Python value to ``float``."""
+
+    return float(value.item() if hasattr(value, "item") else value)
+
+
+def get_scalar_values(sample: Mapping, keys: Iterable[str]) -> dict[str, float]:
+    """Select and convert scalar-like values from ``sample`` for the requested keys."""
+
+    values = {}
+    for key in keys:
+        value = sample.get(key)
+        if value is not None and is_scalar_like(value):
+            values[key] = scalar_to_float(value)
+    return values

src/lerobot/utils/foxglove_visualization.py

@@ -23,6 +23,7 @@ importing from here directly. Requires the ``viz`` extra (``pip install 'lerobot
 import logging
 import numbers
 import time
+from collections.abc import Iterable, Mapping
 
 import cv2
 import numpy as np
@@ -41,6 +42,7 @@ from .constants import (
     SUCCESS,
     TRUNCATED,
 )
+from .dataset_visualization_utils import get_extra_scalar_keys, get_scalar_values
 from .import_utils import require_package
 
 # Static schema shared by all scalar topics. Each message carries a flat list of ``{label, value}``
@@ -405,6 +407,24 @@ def _frame_to_scalars(sample: dict, key: str, labels: list[str] | None = None) -
     return _labeled_scalars(name, arr.flatten(), labels)
 
 
+def _dataset_frame_scalar_groups(
+    sample: Mapping, extra_scalar_keys: Iterable[str]
+) -> tuple[dict[str, float], dict[str, float]]:
+    """Return standard episode scalars and custom scalar features for one dataset frame."""
+
+    episode_scalars = {}
+    for feature, label in (
+        (DONE, "done"),
+        (TRUNCATED, "truncated"),
+        (REWARD, "reward"),
+        (SUCCESS, "success"),
+    ):
+        value = sample.get(feature)
+        if value is not None:
+            episode_scalars[label] = float(value)
+    return episode_scalars, get_scalar_values(sample, extra_scalar_keys)
+
+
 def serve_foxglove_dataset_playback(
     dataset,
     episode_index: int,
@@ -452,6 +472,7 @@ def serve_foxglove_dataset_playback(
         raise ValueError("Cannot visualize an empty episode.")
     first_ns, last_ns = times_ns[0], times_ns[-1]
     camera_keys = list(dataset.meta.camera_keys)
+    extra_scalar_keys = get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,))
     # Dataset-wide q01/q99 depth bounds (fallback min/max) used to normalize depth to [0, 1].
     depth_ranges: dict[str, tuple[float, float]] = {}
     for key in dataset.meta.depth_keys:
@@ -500,17 +521,14 @@ def serve_foxglove_dataset_playback(
             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)
+        episode_scalars, extra_scalars = _dataset_frame_scalar_groups(sample, extra_scalar_keys)
         _log_foxglove_scalars("/episode/state", episode_scalars, channels=channels, log_time=log_time)
+        _log_foxglove_scalars(
+            "/episode/extras",
+            extra_scalars,
+            channels=channels,
+            log_time=log_time,
+        )
 
     lock = threading.Lock()
     stop_event = threading.Event()

tests/datasets/test_visualize_dataset.py

@@ -13,11 +13,78 @@
 # 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 sys
+from types import SimpleNamespace
+
+import numpy as np
 import pytest
+import torch
 
 pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
 
 from lerobot.scripts.lerobot_dataset_viz import visualize_dataset
+from lerobot.utils import import_utils
+from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS, TRUNCATED
+from lerobot.utils.dataset_visualization_utils import (
+    get_extra_scalar_keys,
+    is_scalar_feature,
+    is_scalar_like,
+    scalar_to_float,
+)
+
+
+class DummyMeta:
+    camera_keys = ["observation.images.front"]
+
+
+class DummyFeatureDataset:
+    meta = DummyMeta()
+    features = {
+        "index": {"dtype": "int64", "shape": [1]},
+        "timestamp": {"dtype": "float32", "shape": [1]},
+        "episode_index": {"dtype": "int64", "shape": [1]},
+        "frame_index": {"dtype": "int64", "shape": [1]},
+        "task_index": {"dtype": "int64", "shape": [1]},
+        ACTION: {"dtype": "float32", "shape": [6]},
+        OBS_STATE: {"dtype": "float32", "shape": [6]},
+        DONE: {"dtype": "bool", "shape": [1]},
+        REWARD: {"dtype": "float32", "shape": [1]},
+        SUCCESS: {"dtype": "bool", "shape": [1]},
+        TRUNCATED: {"dtype": "bool", "shape": [1]},
+        "observation.images.front": {"dtype": "video", "shape": [3, 480, 640]},
+        "q_target": {"dtype": "float32", "shape": [1]},
+        "quality": {"dtype": "double", "shape": [1]},
+        "intervention": {"dtype": "bool", "shape": []},
+        "embedding": {"dtype": "float32", "shape": [32]},
+        "comment": {"dtype": "string", "shape": [1]},
+    }
+
+
+class DummyVizDataset:
+    repo_id = "dummy/custom-scalars"
+    depth_output_unit = "m"
+    meta = SimpleNamespace(camera_keys=[], depth_keys=[], stats=None)
+    features = {
+        "index": {"dtype": "int64", "shape": [1]},
+        "timestamp": {"dtype": "float32", "shape": [1]},
+        ACTION: {"dtype": "float32", "shape": [2], "names": ["x", "y"]},
+        "q_target": {"dtype": "float32", "shape": [1]},
+        "intervention": {"dtype": "bool", "shape": [1]},
+        "embedding": {"dtype": "float32", "shape": [2]},
+    }
+
+    def __len__(self):
+        return 2
+
+    def __getitem__(self, index):
+        return {
+            "index": torch.tensor(index),
+            "timestamp": torch.tensor(index / 10),
+            ACTION: torch.tensor([index, index + 1], dtype=torch.float32),
+            "q_target": torch.tensor([index + 0.5]),
+            "intervention": torch.tensor([index == 0]),
+            "embedding": torch.tensor([index, index + 1], dtype=torch.float32),
+        }
 
 
 @pytest.mark.skip("TODO: add dummy videos")
@@ -33,3 +100,93 @@ def test_visualize_local_dataset(tmp_path, lerobot_dataset_factory):
         output_dir=output_dir,
     )
     assert rrd_path.exists()
+
+
+def test_get_extra_scalar_keys_skips_known_metadata_and_non_scalars():
+    dataset = DummyFeatureDataset()
+
+    assert get_extra_scalar_keys(dataset) == [TRUNCATED, "q_target", "quality", "intervention"]
+    assert get_extra_scalar_keys(dataset, additional_known_keys=(TRUNCATED,)) == [
+        "q_target",
+        "quality",
+        "intervention",
+    ]
+
+
+@pytest.mark.parametrize(
+    ("feature", "expected"),
+    [
+        ({"dtype": "float32", "shape": (1,)}, True),
+        ({"dtype": "double", "shape": [1]}, True),
+        ({"dtype": "bool", "shape": []}, True),
+        ({"dtype": "int64", "shape": 1}, True),
+        ({"dtype": "float32", "shape": (3,)}, False),
+        ({"dtype": "complex64", "shape": [1]}, False),
+        ({"dtype": "datetime64[ns]", "shape": [1]}, False),
+        ({"dtype": "string", "shape": [1]}, False),
+        ({"shape": [1]}, False),
+    ],
+)
+def test_is_scalar_feature(feature, expected):
+    assert is_scalar_feature(feature) is expected
+
+
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        (torch.tensor(1.5), True),
+        (torch.tensor([1.5]), True),
+        (torch.tensor([1.5, 2.5]), False),
+        (np.array(2.0), True),
+        (np.array([2.0]), True),
+        (np.array([2.0, 3.0]), False),
+        (True, True),
+        ("not numeric", False),
+    ],
+)
+def test_is_scalar_like(value, expected):
+    assert is_scalar_like(value) is expected
+
+
+def test_scalar_to_float():
+    assert scalar_to_float(torch.tensor(3.0)) == 3.0
+    assert scalar_to_float(np.array([4.0])) == 4.0
+
+
+def test_visualize_dataset_logs_extra_scalars(monkeypatch):
+    logged = []
+    initialized = []
+
+    dummy_rrb = SimpleNamespace(
+        Spatial2DView=lambda origin=None, name=None: SimpleNamespace(
+            kind="spatial", origin=origin, name=name
+        ),
+        TimeSeriesView=lambda origin=None, name=None, overrides=None: SimpleNamespace(
+            kind="time_series", origin=origin, name=name, overrides=overrides
+        ),
+        Grid=lambda *views: SimpleNamespace(views=views),
+        Blueprint=lambda root: SimpleNamespace(root=root),
+    )
+    dummy_rr = SimpleNamespace(
+        SeriesLines=lambda names=None: SimpleNamespace(names=names),
+        Scalars=lambda value: SimpleNamespace(value=value),
+        init=lambda *args, **kwargs: initialized.append((args, kwargs)),
+        log=lambda key, entity: logged.append((key, entity)),
+        set_time=lambda *args, **kwargs: None,
+        blueprint=dummy_rrb,
+    )
+    monkeypatch.setitem(sys.modules, "rerun", dummy_rr)
+    monkeypatch.setitem(sys.modules, "rerun.blueprint", dummy_rrb)
+    monkeypatch.setattr(import_utils, "require_package", lambda *args, **kwargs: None)
+
+    visualize_dataset(DummyVizDataset(), episode_index=0, batch_size=2)
+
+    logged_keys = [key for key, _ in logged]
+    assert logged_keys.count("q_target") == 2
+    assert logged_keys.count("intervention") == 2
+    assert "embedding" not in logged_keys
+
+    blueprint = initialized[0][1]["default_blueprint"]
+    time_series_origins = {view.origin for view in blueprint.root.views if view.kind == "time_series"}
+    assert {"q_target", "intervention"} <= time_series_origins
+    assert "embedding" not in time_series_origins

tests/utils/test_foxglove_visualization.py

@@ -24,7 +24,7 @@ the functions that talk to the server, so the helpers below run in the base test
 import numpy as np
 
 from lerobot.utils import foxglove_visualization as fv
-from lerobot.utils.constants import ACTION, OBS_STATE
+from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS
 
 
 def test_foxglove_safe_name_collapses_dots():
@@ -93,6 +93,24 @@ def test_feature_dim_names_formats():
     assert fv._feature_dim_names({"shape": [2]}) is None
 
 
+def test_dataset_frame_scalar_groups_include_custom_scalars():
+    sample = {
+        DONE: np.array(True),
+        REWARD: np.array(0.5),
+        SUCCESS: np.array(False),
+        "q_target": np.array([0.75]),
+        "intervention": np.array([True]),
+        "embedding": np.array([1.0, 2.0]),
+    }
+
+    episode_scalars, extra_scalars = fv._dataset_frame_scalar_groups(
+        sample, ("q_target", "intervention", "embedding")
+    )
+
+    assert episode_scalars == {"done": 1.0, "reward": 0.5, "success": 0.0}
+    assert extra_scalars == {"q_target": 0.75, "intervention": 1.0}
+
+
 def test_is_scalar():
     assert fv._is_scalar(1.0)
     assert fv._is_scalar(np.float32(2.0))