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>
2026-07-30 13:09:40 +00:00 · 2026-07-30 12:17:54 +02:00
12 changed files with 365 additions and 612 deletions
@@ -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",
+    # Lower-body controller class name, e.g. "GrootLocomotionController" or
-    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
+    # "HolosomaLocomotionController". 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).
    controller: str | None = None
@@ -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",
 ]
@@ -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()
@@ -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",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
    }
    module_path = controllers.get(name)
    if module_path is None:
@@ -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,
@@ -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,
@@ -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()
                    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
            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
+                # Run controller step
-                # take the raw lowstate. get_observation() is the single lowstate -> obs builder.
+                controller_action = self.controller.run_step(controller_input, lowstate)
                controller_state = (
                    self.get_observation() if getattr(self.controller, "full_body", False) else lowstate
                )
                controller_action = self.controller.run_step(controller_input, controller_state)
                # 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.
+        """Update controller input state from 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.
        """
        with self._controller_action_lock:
-            for key, value in action.items():
+            for key in REMOTE_KEYS:
-                if isinstance(key, str) and value is not None:
+                if key in action:
-                    self.controller_input[key] = value
+                    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
+        if self.config.is_simulation and self.sim_env is not None:
-        # ``<joint>.q`` in send_action(), so reset() must publish the default pose
+            self.sim_env.reset()
-        # directly. Pause the background controller first so the two aren't both writing
+            self.publish_lowcmd(
-        # low commands while the robot moves to the default pose.
+                {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
-        full_body = getattr(self.controller, "full_body", False)
+            )
-        paused = False
+        else:
-        if full_body and self._controller_thread is not None:
+            total_time = 3.0
-            self._controller_paused.set()
+            num_steps = int(total_time / control_dt)
            paused = True
            time.sleep(control_dt)  # let any in-flight controller tick settle
-        try:
+            # get current state
-            if self.config.is_simulation and self.sim_env is not None:
+            obs = self.get_observation()
                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
+            # record current positions
-                obs = self.get_observation()
+            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
+            # Interpolate to default position
-                init_dof_pos = np.zeros(29, dtype=np.float32)
+            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
+                self.send_action(action_dict)
                for step in range(num_steps):
                    start_time = time.time()
-                    alpha = step / num_steps
+                # Maintain constant control rate
-                    action_dict = {}
+                elapsed = time.time() - start_time
-                    for motor in G1_29_JointIndex:
+                sleep_time = max(0, control_dt - elapsed)
-                        target_pos = default_positions[motor.value]
+                time.sleep(sleep_time)
                        interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
                        action_dict[f"{motor.name}.q"] = float(interp_pos)
-                    # Full-body controllers no-op in send_action(); publish the pose
+        # Reset controller internal state (gait phase, obs history, etc.)
-                    # directly (arm-only controllers keep the send_action() path).
+        if self.controller is not None and hasattr(self.controller, "reset"):
-                    if full_body:
+            self.controller.reset()
                        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()
        logger.info("Reached default position")
@@ -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)
@@ -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
@@ -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()
@@ -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
@@ -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))