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>

docs/source/cameras.mdx

@@ -164,8 +164,8 @@ includes the range reported by the sensor. Requesting an unsupported control als
 Omitted controls leave the sensor's existing automatic or manual setting unchanged. These options
 require `use_rgb=True`.
 
-Manual color controls require a dedicated RGB module. Cameras without one, such as the RealSense
+On the RealSense D405, the color stream is provided by the Stereo Module, so changing manual
-D405, do not support them and raise an error at connection time.
+exposure or gain also affects the depth stream.
 
 </hfoption>
 </hfoptions>

docs/source/env_processor.mdx

@@ -88,20 +88,6 @@ policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
 
 The same policy can work with different environment processors, and the same environment processor can work with different policies:
 
-````python
-# Use SmolVLA policy with LIBERO environment
-# Use SmolVLA policy with LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=smolvla_cfg,
-)
-smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
-# Or use ACT policy with the same LIBERO environment
-libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
-    env_cfg=libero_cfg,
-    policy_cfg=act_cfg,
-)
-act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
 ```python
 # Use SmolVLA policy with LIBERO environment
 libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
@@ -116,6 +102,7 @@ libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
     policy_cfg=act_cfg,
 )
 act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
+```
 
 ### 3. **Easier Experimentation**
 
@@ -145,7 +132,7 @@ class LiberoVelocityProcessorStep(ObservationProcessorStep):
         state = torch.cat([eef_pos, eef_axisangle, eef_vel,
                           gripper_pos, gripper_vel], dim=-1)  # 14D
         return state
-````
+```
 
 ### 4. **Cleaner Environment Code**
 

docs/source/hope_jr.mdx

@@ -211,7 +211,7 @@ Record, Replay and Train with Hope-JR is still experimental.
 
 ### Record
 
-This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data/settings).
+This step records the dataset, which can be seen as an example [here](https://huggingface.co/datasets/nepyope/hand_record_test_with_video_data).
 
 ```bash
 lerobot-record \

docs/source/integrate_hardware.mdx

@@ -18,7 +18,7 @@ If you're using Feetech or Dynamixel motors, LeRobot provides built-in bus inter
 - [`DynamixelMotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/dynamixel/dynamixel.py) – for controlling Dynamixel servos
 
 Please refer to the [`MotorsBus`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/motors/motors_bus.py) abstract class to learn about its API.
-For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so101_follower/so101_follower.py)
+For a good example of how it can be used, you can have a look at our own [SO101 follower implementation](https://github.com/huggingface/lerobot/blob/main/src/lerobot/robots/so_follower/so_follower.py)
 
 Use these if compatible. Otherwise, you'll need to find or write a Python interface (not covered in this tutorial):
 

docs/source/lekiwi.mdx

@@ -51,7 +51,7 @@ In addition to these instructions, you need to install the Feetech SDK & ZeroMQ
 pip install -e ".[lekiwi]"
 ```
 
-Great :hugs:! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base :robot:.
+Great 🤗! You are now done installing LeRobot, and we can begin assembling the SO100/SO101 arms and the mobile base 🤖.
 Every time you now want to use LeRobot, you can go to the `~/lerobot` folder where we installed LeRobot and run one of the commands.
 
 # Step-by-Step Assembly Instructions

docs/source/pi0fast.mdx

@@ -174,7 +174,7 @@ The model takes images, text instructions, and robot state as input, and outputs
 
 ## Reproducing π₀Fast results
 
-We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40kk steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
+We reproduce the results of π₀Fast on the LIBERO benchmark using the LeRobot implementation. We take the LeRobot PiFast base model [lerobot/pi0fast-base](https://huggingface.co/lerobot/pi0fast-base) and finetune for an additional 40k steps in bfloat16, with batch size of 256 on 8 H100 GPUs using the [HuggingFace LIBERO dataset](https://huggingface.co/datasets/HuggingFaceVLA/libero).
 
 The finetuned model can be found here:
 

docs/source/smolvla.mdx

@@ -93,7 +93,7 @@ lerobot-train --help
 
 ## Evaluate the finetuned model and run it in real-time
 
-Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots).
+Similarly for when recording an episode, it is recommended that you are logged in to the HuggingFace Hub. You can follow the corresponding steps: [Record a dataset](./il_robots#record-a-dataset).
 Once you are logged in, you can run inference in your setup by doing:
 
 ```bash

src/lerobot/cameras/realsense/camera_realsense.py

@@ -365,12 +365,11 @@ class RealSenseCamera(Camera):
         return self._async_read(timeout_ms=10000, read_depth=read_depth)
 
     def _get_color_sensor(self) -> "rs.sensor":
-        """Returns the dedicated "RGB Camera" sensor that controls the color stream.
+        """Returns the sensor that controls the color stream.
 
-        Manual color controls are only applied to a dedicated RGB module. Cameras
+        Most RealSense cameras expose "RGB Camera" for color. The D405 has no
-        without one (e.g. the D405, whose color stream comes from the shared
+        separate RGB module — its color stream comes from "Stereo Module".
-        "Stereo Module") are unsupported, so we never fall back to another sensor
+        We try RGB Camera first, then fall back to Stereo Module.
-        to avoid altering the depth stream.
         """
         if self.rs_profile is None:
             raise RuntimeError(f"{self}: rs_profile must be initialized before use.")
@@ -378,14 +377,12 @@ class RealSenseCamera(Camera):
         device = self.rs_profile.get_device()
         sensors = {s.get_info(rs.camera_info.name): s for s in device.query_sensors()}
 
-        if "RGB Camera" in sensors:
+        for name in ("RGB Camera", "Stereo Module"):
-            return sensors["RGB Camera"]
+            if name in sensors:
+                return sensors[name]
 
         available = list(sensors.keys())
-        raise RuntimeError(
+        raise RuntimeError(f"{self}: no color sensor found. Available sensors: {available}")
-            f"{self}: manual color controls require a dedicated 'RGB Camera' module, which this camera does not have. ",
-            f"Available sensors: {available}.",
-        )
 
     def _set_sensor_option(self, sensor: "rs.sensor", option: "rs.option", value: float, label: str) -> None:
         """Sets a sensor option, re-raising range errors with actionable diagnostics."""

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 = {}
+        episode_scalars, extra_scalars = _dataset_frame_scalar_groups(sample, extra_scalar_keys)
-        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)
+        _log_foxglove_scalars(
+            "/episode/extras",
+            extra_scalars,
+            channels=channels,
+            log_time=log_time,
+        )
 
     lock = threading.Lock()
     stop_event = threading.Event()

tests/cameras/test_realsense.py

@@ -322,16 +322,15 @@ def test_get_color_sensor_prefers_rgb_camera():
     assert camera._get_color_sensor() is rgb
 
 
-def test_get_color_sensor_raises_without_dedicated_rgb_module():
+def test_get_color_sensor_falls_back_to_stereo_module():
-    """D405 has no separate RGB module; we refuse to touch the shared Stereo Module."""
+    """D405 has no separate RGB module; color comes from Stereo Module."""
     config = RealSenseCameraConfig(serial_number_or_name="042")
     camera = RealSenseCamera(config)
 
     stereo = _make_mock_sensor("Stereo Module")
     _attach_mock_color_sensor(camera, stereo)
 
-    with pytest.raises(RuntimeError, match="dedicated 'RGB Camera' module"):
+    assert camera._get_color_sensor() is stereo
-        camera._get_color_sensor()
 
 
 def test_get_color_sensor_raises_with_available_sensors():

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))