diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index 8013ba8b5..9d10c9e61 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/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) diff --git a/src/lerobot/utils/dataset_visualization_utils.py b/src/lerobot/utils/dataset_visualization_utils.py new file mode 100644 index 000000000..3672dec68 --- /dev/null +++ b/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 diff --git a/src/lerobot/utils/foxglove_visualization.py b/src/lerobot/utils/foxglove_visualization.py index fc4136e12..d41b2e7d9 100644 --- a/src/lerobot/utils/foxglove_visualization.py +++ b/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() diff --git a/tests/datasets/test_visualize_dataset.py b/tests/datasets/test_visualize_dataset.py index 3bf94e6cb..093ba2f01 100644 --- a/tests/datasets/test_visualize_dataset.py +++ b/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 diff --git a/tests/utils/test_foxglove_visualization.py b/tests/utils/test_foxglove_visualization.py index f3e80315c..92cba97ce 100644 --- a/tests/utils/test_foxglove_visualization.py +++ b/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))