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/robots/unitree_g1/config_unitree_g1.py

@@ -68,10 +68,6 @@ class UnitreeG1Config(RobotConfig):
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Locomotion controller class name, e.g. "GrootLocomotionController",
-    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
-    # Selecting "SonicWholeBodyController" implicitly switches the robot to the 64-D
-    # latent-token action/observation interface (``motion_token.{i}.pos`` action and a
-    # ``motion_token_state.{i}.pos`` state echo) so ``lerobot-rollout`` can drive a
-    # policy trained on SONIC motion tokens (e.g. nepyope/sonic_walk).
+    # Lower-body controller class name, e.g. "GrootLocomotionController" or
+    # "HolosomaLocomotionController". None disables it.
     controller: str | None = None

src/lerobot/robots/unitree_g1/controllers/__init__.py

@@ -1,27 +0,0 @@
-#!/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.
-
-"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
-
-from .gr00t_locomotion import GrootLocomotionController
-from .holosoma_locomotion import HolosomaLocomotionController
-from .sonic_whole_body import SonicWholeBodyController
-
-__all__ = [
-    "GrootLocomotionController",
-    "HolosomaLocomotionController",
-    "SonicWholeBodyController",
-]

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -1,360 +0,0 @@
-#!/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.
-
-"""SONIC decoder whole-body controller for the Unitree G1 (token-only).
-
-Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy stack.
-The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``)
-supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder**
-maps ``token + recent proprioception history`` to a residual action that is scaled and
-added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
-for the robot's PD controller.
-
-Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
-**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils)
-convert between them. Quaternions are scalar-first ``(w, x, y, z)``.
-"""
-
-from __future__ import annotations
-
-import json
-import logging
-
-import numpy as np
-import onnx
-import onnxruntime as ort
-from huggingface_hub import hf_hub_download
-
-from ..g1_utils import (
-    ISAACLAB_TO_MUJOCO,
-    MUJOCO_TO_ISAACLAB,
-    G1_29_JointIndex,
-    get_gravity_orientation,
-)
-
-logger = logging.getLogger(__name__)
-
-# ── Constants (hardware-validated; see the NVIDIA SONIC deploy reference) ──────
-CONTROL_DT = 0.02  # 50 Hz control period (s)
-TOKEN_DIM = 64  # decoder latent size
-
-# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
-# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
-# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the
-# model *and* all of these straight from the checkpoint (the Holosoma convention), so no
-# motor-physics math happens at deploy time.
-DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
-DECODER_FILENAME = "model_decoder.onnx"
-DECODER_INPUT_DIM = 994  # token(64) + 10-frame proprio history + gravity
-
-
-def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
-    """Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
-
-    Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
-    gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
-    (64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
-    recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
-    """
-    decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
-    so = ort.SessionOptions()
-    so.log_severity_level = 3  # quiet ORT logs
-    session = ort.InferenceSession(decoder_path, sess_options=so)
-    dec_dim = int(session.get_inputs()[0].shape[1])
-    if dec_dim != DECODER_INPUT_DIM:
-        raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
-
-    meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
-    required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
-    missing = [k for k in required if k not in meta]
-    if missing:
-        raise ValueError(
-            f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
-            "re-run upload_sonic_decoder.py to (re)provision the checkpoint."
-        )
-    arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
-    logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
-    return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
-
-
-# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
-TOKEN_ACTION_PREFIX = "motion_token"  # nosec B105 - feature-key prefix, not a secret
-# Proprio-state prefix for the token interface: the robot echoes the last commanded token
-# here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
-TOKEN_STATE_PREFIX = "motion_token_state"  # nosec B105 - feature-key prefix, not a secret
-
-
-def token_action_key(i: int) -> str:
-    """Action-dict key for the i-th component of the 64-D SONIC latent token.
-
-    The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``, which
-    only routes ``.pos`` scalar features onto the policy action vector.
-    """
-    return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
-
-
-def token_state_key(i: int) -> str:
-    """Observation key for the i-th component of the 64-D SONIC latent token state."""
-    return f"{TOKEN_STATE_PREFIX}.{i}.pos"
-
-
-# Startup blend duration: over the first control ticks, linearly interpolate every joint
-# from the robot's initial measured pose into the policy's commanded target, so control
-# eases in without a snap on the first command.
-INIT_RAMP_S = 3.0
-
-
-def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
-    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
-
-    The token-only interface: the caller supplies the 64-D encoder latent directly (e.g. a
-    token-output VLA's action), which the decoder consumes with the encoder bypassed.
-    Requires the full dense token; a partial one is ignored (returns None).
-    """
-    if not action:
-        return None
-    keys = [token_action_key(i) for i in range(TOKEN_DIM)]
-    if any(key not in action for key in keys):
-        return None
-    return np.fromiter((float(action[key]) for key in keys), dtype=np.float32, count=TOKEN_DIM)
-
-
-class SonicDecoder:
-    """Runs the SONIC decoder ONNX model and owns the proprioception history.
-
-    Each tick it appends the latest robot state to 10-frame history buffers, then maps the
-    supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
-    The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and
-    ``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint.
-    """
-
-    def __init__(self, decoder, default_angles, action_scale):
-        self.decoder = decoder
-        self.decoder_input = decoder.get_inputs()[0].name
-        self.default_angles = np.asarray(default_angles, np.float32)
-        self.action_scale = np.asarray(action_scale, np.float32)
-        self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-
-    def reset(self):
-        """Clear the token and 10-frame proprioception history.
-
-        ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new episode
-        are not contaminated by the previous episode's state.
-        """
-        self.token = np.zeros(TOKEN_DIM, np.float32)
-        self.last_action_mj = np.zeros(29, np.float32)
-        self.h_q_mj = [np.zeros(29, np.float32)] * 10
-        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
-        self.h_ang = [np.zeros(3, np.float32)] * 10
-        self.h_act_mj = [np.zeros(29, np.float32)] * 10
-        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
-
-    def update_history(self, q, dq, ang, quat):
-        """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
-        quat = quat / (np.linalg.norm(quat) + 1e-8)
-        # Reorder IsaacLab-order state into the MuJoCo order the decoder consumes. This
-        # permutation direction is validated against the deployed SONIC ONNX; don't flip it.
-        q_mj = q[MUJOCO_TO_ISAACLAB]
-        dq_mj = dq[MUJOCO_TO_ISAACLAB]
-        self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1]
-        self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
-        self.h_ang = [ang.copy()] + self.h_ang[:-1]
-        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
-        self.h_quat = [quat.copy()] + self.h_quat[:-1]
-
-    def build_decoder_obs(self):
-        """Assemble the 994-D decoder input: token + 10-frame proprioception history + gravity."""
-        obs = np.zeros(994, np.float32)
-        off = 0
-        obs[off : off + 64] = self.token
-        off += 64
-        for h, sz in [
-            (list(reversed(self.h_ang)), 3),
-            (list(reversed(self.h_q_mj)), 29),
-            (list(reversed(self.h_dq_mj)), 29),
-            (list(reversed(self.h_act_mj)), 29),
-        ]:
-            for f in range(10):
-                obs[off : off + sz] = h[f]
-                off += sz
-        for q in reversed(self.h_quat):
-            obs[off : off + 3] = get_gravity_orientation(q)
-            off += 3
-        assert off == 994, f"Decoder obs mismatch: {off}"
-        return obs
-
-    def step(self, robot_obs, token, debug=False):
-        """One control tick: read robot obs, decode the supplied token -> joint targets.
-
-        Args:
-            robot_obs: dict with ``<joint>.q``/``.dq`` and ``imu.*`` fields.
-            token: 64-D latent supplied by the policy (encoder bypassed).
-            debug: log action/delta norms.
-
-        Returns:
-            dict of ``<joint>.q`` target positions (rad) in IsaacLab joint order.
-        """
-        self.token = np.asarray(token, np.float32)
-        jnames = [m.name for m in G1_29_JointIndex]
-        q = np.array(
-            [
-                robot_obs.get(f"{n}.q", self.default_angles[m.value])
-                for m, n in zip(G1_29_JointIndex, jnames, strict=False)
-            ],
-            np.float32,
-        )
-        dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
-        quat = np.array(
-            [
-                robot_obs.get("imu.quat.w", 1),
-                robot_obs.get("imu.quat.x", 0),
-                robot_obs.get("imu.quat.y", 0),
-                robot_obs.get("imu.quat.z", 0),
-            ],
-            np.float32,
-        )
-        ang = np.array([robot_obs.get(f"imu.gyro.{a}", 0) for a in "xyz"], np.float32)
-        self.update_history(q, dq, ang, quat)
-        action_mj = (
-            self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1, -1)})[0]
-            .squeeze()
-            .astype(np.float32)
-        )
-        self.last_action_mj = action_mj.copy()
-        target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale
-        if debug:
-            delta = target - q
-            logger.debug(
-                "token_norm=%.4f action_norm=%.4f delta_max=%.4f delta_rms=%.4f",
-                np.linalg.norm(self.token),
-                np.linalg.norm(action_mj),
-                np.max(np.abs(delta)),
-                np.sqrt(np.mean(delta**2)),
-            )
-        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
-
-
-class SonicRuntime:
-    """Loads the SONIC decoder ONNX model and owns the decode controller.
-
-    Token-only deploy: the encoder is bypassed; each tick the decoder consumes a 64-D
-    latent token supplied directly by the policy.
-    """
-
-    def __init__(self):
-        decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
-        self.default_angles = default_angles
-        self.neutral_token = neutral_token
-        self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
-
-    @property
-    def pipeline(self):
-        return self.controller
-
-    def reset(self):
-        self.controller.reset()
-
-    def shutdown(self):
-        pass
-
-
-class SonicWholeBodyController:
-    """Full-body SONIC controller for UnitreeG1's background controller thread."""
-
-    control_dt = CONTROL_DT
-    full_body = True
-
-    def __init__(self):
-        logger.info("Loading SONIC whole-body controller...")
-        self._runtime = SonicRuntime()
-        self.kp = self._runtime.kp
-        self.kd = self._runtime.kd
-        self.controller = self._runtime.controller
-        self._default_angles = self._runtime.default_angles
-        self._neutral_token = self._runtime.neutral_token
-
-        # Startup blend: ease from the robot's initial pose into the first commanded policy
-        # targets over INIT_RAMP_S (captured on the first control tick).
-        self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
-        self._init_step = 0
-        self._start_pose: dict[str, float] = {}
-
-        # Token-interface state. The controller holds a stable *neutral* token until the first
-        # real token arrives, and afterwards holds the *last* token received between ticks (the
-        # async controller runs ~50 Hz while a token VLA streams ~30 Hz).
-        self._last_token: np.ndarray | None = None
-
-        logger.info("SONIC ready (decoder, 64-D token command path)")
-
-    def _startup_blend(self, obs: dict, out: dict) -> dict:
-        """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
-        interpolate between the robot's pose captured on the first tick and the policy's
-        live commanded target, so the handoff has no snap.
-
-        ``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
-        climbs 0->1 over the ramp, after which the raw policy target passes through.
-        """
-        if self._init_step >= self._init_ramp_steps or not out:
-            return out
-        if self._init_step == 0:
-            # Capture the robot's actual pose as the interpolation start point.
-            self._start_pose = {
-                f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value]))
-                for m in G1_29_JointIndex
-            }
-        self._init_step += 1
-        ratio = min(1.0, self._init_step / self._init_ramp_steps)
-        blended = {
-            k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
-            for k, tgt in out.items()
-        }
-        if self._init_step >= self._init_ramp_steps:
-            logger.info("SONIC startup blend complete -> full policy control")
-        return blended
-
-    def run_step(self, action: dict, obs: dict) -> dict:
-        if not obs:
-            return {}
-
-        # Token-only interface (token-output VLA): a dense 64-D ``motion_token.{i}`` command
-        # is decoded directly, encoder bypassed.
-        token = _extract_token_from_action(action)
-        if token is not None:
-            self._last_token = token
-        elif self._last_token is None:
-            # No token has arrived yet: hold the checkpoint's neutral token, which the decoder
-            # maps to a stable, natural standing pose.
-            self._last_token = self._neutral_token.copy()
-        # Either a fresh token this tick or the last one received (held between the ~30 Hz
-        # token stream and the ~50 Hz control loop).
-        return self._startup_blend(obs, self.controller.step(obs, self._last_token))
-
-    def reset(self):
-        self._runtime.reset()
-        self._init_step = 0  # re-run the startup blend after a reset
-        self._start_pose = {}
-        # Drop the held token so the neutral token is re-seeded after a reset.
-        self._last_token = None
-
-    def shutdown(self):
-        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,47 +23,6 @@ import numpy as np
 
 NUM_MOTORS = 29
 
-# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
-# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
-# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
-ISAACLAB_TO_MUJOCO = np.array(
-    [
-        0,
-        3,
-        6,
-        9,
-        13,
-        17,
-        1,
-        4,
-        7,
-        10,
-        14,
-        18,
-        2,
-        5,
-        8,
-        11,
-        15,
-        19,
-        21,
-        23,
-        25,
-        27,
-        12,
-        16,
-        20,
-        22,
-        24,
-        26,
-        28,
-    ],
-    dtype=np.int32,
-)
-# The two orderings are inverses of each other, so derive one from the other (argsort) to
-# guarantee they can never drift out of sync.
-MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
-
 REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
 REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
 REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
@@ -109,9 +68,8 @@ def make_locomotion_controller(name: str | None):
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
-        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/gr00t_locomotion.py

@@ -21,7 +21,7 @@ import numpy as np
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,

src/lerobot/robots/unitree_g1/holosoma_locomotion.py

@@ -22,7 +22,7 @@ import onnx
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from ..g1_utils import (
+from .g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -34,6 +34,7 @@ from .config_unitree_g1 import UnitreeG1Config
 from .g1_kinematics import G1_29_ArmIK
 from .g1_utils import (
     REMOTE_AXES,
+    REMOTE_KEYS,
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
@@ -147,54 +148,22 @@ class UnitreeG1(Robot):
 
         self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
 
-        # Lower-body / whole-body controller loaded dynamically
+        # Lower-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
+
         # Controller thread state
         self._controller_thread = None
-        # When set, the controller loop stops publishing low commands so reset() can
-        # drive the joints directly without two publishers fighting (single-publisher).
-        self._controller_paused = threading.Event()
         self._controller_action_lock = threading.Lock()
         self.controller_input = default_remote_input()
         self.controller_output = {}
 
-        # Token-mode state: last 64-D SONIC latent token commanded by the policy,
-        # echoed back as ``observation.state`` so a token-output VLA closes the loop
-        # on its own previous token. Implicit whenever the SONIC whole-body controller
-        # is active. Seeded to zeros; the controller's startup blend eases joints in.
-        self._last_token: np.ndarray | None = None
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import TOKEN_DIM
-
-            self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
-
-    @property
-    def _sonic_token(self) -> bool:
-        """Whether the SONIC whole-body decoder is active.
-
-        A SONIC controller consumes a 64-D latent motion token as its action and echoes
-        the last commanded token as ``observation.state``. Keyed purely off the selected
-        controller so the token interface is implicit -- no separate config flag.
-        """
-        return self.config.controller == "SonicWholeBodyController"
-
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
             # Step simulation if in simulation mode
             if self.config.is_simulation and self.sim_env is not None:
-                try:
-                    self.sim_env.step()
-                except ValueError as e:
-                    # Startup race: the sim thread can step once before reset() has
-                    # written a valid base pose, giving a zero-norm pelvis quaternion
-                    # (scipy>=1.11 raises instead of normalizing). Skip and retry so
-                    # the thread survives instead of dying and freezing the sim.
-                    if "zero norm" not in str(e).lower():
-                        raise
-                    time.sleep(self.control_dt)
-                    continue
+                self.sim_env.step()
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -262,38 +231,15 @@ class UnitreeG1(Robot):
                 features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
         return features
 
-    @property
-    def _token_state_ft(self) -> dict[str, type]:
-        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
-
-        Exposed only when a SONIC whole-body controller is active; aggregated by the
-        rollout into a 64-D ``observation.state`` (the last token the policy commanded).
-        """
-        if not self._sonic_token:
-            return {}
-        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
-
-        return {token_state_key(i): float for i in range(TOKEN_DIM)}
-
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft}
+        return {**self._motors_ft, **self._cameras_ft}
 
     @cached_property
     def action_features(self) -> dict[str, type]:
-        # No controller configured at all: raw 29-DoF joint teleop.
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
-        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
-        # (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
-        # straight onto the decoder, bypassing the encoder.
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
-
-            return {token_action_key(i): float for i in range(TOKEN_DIM)}
-
-        # Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
         arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
         remote_features = dict.fromkeys(REMOTE_AXES, float)
         return {**arm_features, **remote_features}
@@ -309,11 +255,6 @@ class UnitreeG1(Robot):
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
-            # Paused during reset() so the reset routine is the sole low-cmd publisher.
-            if self._controller_paused.is_set():
-                time.sleep(control_dt)
-                continue
-
             with self._lowstate_lock:
                 lowstate = self._lowstate
 
@@ -330,12 +271,8 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
-                # Full-body controllers (SONIC) consume the full observation dict; others
-                # take the raw lowstate. get_observation() is the single lowstate -> obs builder.
-                controller_state = (
-                    self.get_observation() if getattr(self.controller, "full_body", False) else lowstate
-                )
-                controller_action = self.controller.run_step(controller_input, controller_state)
+                # Run controller step
+                controller_action = self.controller.run_step(controller_input, lowstate)
 
                 # Write controller output snapshot
                 with self._controller_action_lock:
@@ -406,9 +343,6 @@ class UnitreeG1(Robot):
 
         self.kp = np.array(self.config.kp, dtype=np.float32)
         self.kd = np.array(self.config.kd, dtype=np.float32)
-        if self.controller is not None and hasattr(self.controller, "kp"):
-            self.kp = np.array(self.controller.kp, dtype=np.float32)
-            self.kd = np.array(self.controller.kd, dtype=np.float32)
 
         for joint in G1_29_JointIndex:
             self.msg.motor_cmd[joint].mode = 1
@@ -457,10 +391,6 @@ class UnitreeG1(Robot):
             if self._controller_thread.is_alive():
                 logger.warning("Controller thread did not stop cleanly")
 
-        # Release controller resources (e.g. SONIC decoder sessions).
-        if self.controller is not None and hasattr(self.controller, "shutdown"):
-            self.controller.shutdown()
-
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:
@@ -531,15 +461,6 @@ class UnitreeG1(Robot):
         if lowstate.wireless_remote:
             obs["wireless_remote"] = lowstate.wireless_remote
 
-        # Token mode: echo the last commanded latent token as observation.state so a
-        # token-output VLA closes the loop on its own previous token.
-        if self._sonic_token:
-            from .controllers.sonic_whole_body import token_state_key
-
-            token = self._last_token if self._last_token is not None else []
-            for i, v in enumerate(token):
-                obs[token_state_key(i)] = float(v)
-
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
             if getattr(cam, "use_rgb", True):
@@ -552,22 +473,9 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
-            # SONIC decoder: pull the 64-D latent token out of the action and remember it
-            # for the observation.state echo. The controller thread reads it back from
-            # controller_input (populated below) and decodes it into a 29-DoF command.
-            if self._sonic_token:
-                from .controllers.sonic_whole_body import _extract_token_from_action
-
-                token = _extract_token_from_action(action)
-                if token is not None:
-                    self._last_token = token
-            self._update_controller_action(action)
-            # Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
-            # publish here (the controller thread is the sole publisher).
-            if getattr(self.controller, "full_body", False):
-                return action
             # Controller thread owns legs/waist. Here we only update joystick inputs
             # and publish arm targets from the teleoperator.
+            self._update_controller_action(action)
             arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
             action_to_publish = {
                 key: value
@@ -595,17 +503,11 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from an incoming teleop action.
-
-        Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
-        SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
-        and each controller extracts only the keys it understands. The robot deliberately
-        does not enumerate any controller's key schema here.
-        """
+        """Update controller input state from incoming teleop action."""
         with self._controller_action_lock:
-            for key, value in action.items():
-                if isinstance(key, str) and value is not None:
-                    self.controller_input[key] = value
+            for key in REMOTE_KEYS:
+                if key in action:
+                    self.controller_input[key] = action[key]
 
     @property
     def is_calibrated(self) -> bool:
@@ -635,64 +537,43 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
-        # Full-body controllers (SONIC) own the whole 29-DoF command and ignore
-        # ``<joint>.q`` in send_action(), so reset() must publish the default pose
-        # directly. Pause the background controller first so the two aren't both writing
-        # low commands while the robot moves to the default pose.
-        full_body = getattr(self.controller, "full_body", False)
-        paused = False
-        if full_body and self._controller_thread is not None:
-            self._controller_paused.set()
-            paused = True
-            time.sleep(control_dt)  # let any in-flight controller tick settle
+        if self.config.is_simulation and self.sim_env is not None:
+            self.sim_env.reset()
+            self.publish_lowcmd(
+                {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
+            )
+        else:
+            total_time = 3.0
+            num_steps = int(total_time / control_dt)
 
-        try:
-            if self.config.is_simulation and self.sim_env is not None:
-                self.sim_env.reset()
-                self.publish_lowcmd(
-                    {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
-                )
-            else:
-                total_time = 3.0
-                num_steps = int(total_time / control_dt)
+            # get current state
+            obs = self.get_observation()
 
-                # get current state
-                obs = self.get_observation()
+            # record current positions
+            init_dof_pos = np.zeros(29, dtype=np.float32)
+            for motor in G1_29_JointIndex:
+                init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
-                # record current positions
-                init_dof_pos = np.zeros(29, dtype=np.float32)
+            # Interpolate to default position
+            for step in range(num_steps):
+                start_time = time.time()
+
+                alpha = step / num_steps
+                action_dict = {}
                 for motor in G1_29_JointIndex:
-                    init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
+                    target_pos = default_positions[motor.value]
+                    interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
+                    action_dict[f"{motor.name}.q"] = float(interp_pos)
 
-                # Interpolate to default position
-                for step in range(num_steps):
-                    start_time = time.time()
+                self.send_action(action_dict)
 
-                    alpha = step / num_steps
-                    action_dict = {}
-                    for motor in G1_29_JointIndex:
-                        target_pos = default_positions[motor.value]
-                        interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
-                        action_dict[f"{motor.name}.q"] = float(interp_pos)
+                # Maintain constant control rate
+                elapsed = time.time() - start_time
+                sleep_time = max(0, control_dt - elapsed)
+                time.sleep(sleep_time)
 
-                    # Full-body controllers no-op in send_action(); publish the pose
-                    # directly (arm-only controllers keep the send_action() path).
-                    if full_body:
-                        self.publish_lowcmd(action_dict)
-                    else:
-                        self.send_action(action_dict)
-
-                    # Maintain constant control rate
-                    elapsed = time.time() - start_time
-                    sleep_time = max(0, control_dt - elapsed)
-                    time.sleep(sleep_time)
-
-            # Reset controller internal state (gait phase, obs history, etc.) before
-            # resuming so its buffers reflect the post-reset pose.
-            if self.controller is not None and hasattr(self.controller, "reset"):
-                self.controller.reset()
-        finally:
-            if paused:
-                self._controller_paused.clear()
+        # Reset controller internal state (gait phase, obs history, etc.)
+        if self.controller is not None and hasattr(self.controller, "reset"):
+            self.controller.reset()
 
         logger.info("Reached default position")

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