update docs

Co-authored-by: Cursor <cursoragent@cursor.com>

docs/source/unitree_g1.mdx

@@ -290,6 +290,47 @@ lerobot-rollout \
 
 ---
 
+## Part 5: SONIC Whole-Body Control (Latent-Token Policies)
+
+LeRobot now ports the **decoder** half of NVIDIA's [SONIC](https://github.com/NVIDIA/sonic) whole-body deploy stack as the `SonicWholeBodyController`. Instead of commanding joints, the policy outputs a 64-D **latent motion token** (`motion_token.{i}.pos`) each tick; the SONIC decoder maps that token plus recent proprioception history onto a residual action that is scaled and added to the standing pose, producing 50 Hz joint-position targets for all 29 DoF.
+
+The controller loads its ONNX decoder and all deploy constants (PD gains, `default_angles`, `action_scale`, and the neutral token) from [`lerobot/sonic_decoder`](https://huggingface.co/lerobot/sonic_decoder). The repo also ships a distilled low-latency decoder, selectable with the controller's `policy_type="low_latency"` argument (defaults to `"default"`).
+
+### Test in Simulation
+
+Any token-output policy trained on SONIC motion tokens (e.g. [`nepyope/sonic_walk`](https://huggingface.co/nepyope/sonic_walk)) can drive it:
+
+```bash
+lerobot-rollout \
+  --strategy.type=base \
+  --policy.path=nepyope/sonic_walk \
+  --policy.device=cuda \
+  --robot.type=unitree_g1 \
+  --robot.is_simulation=true \
+  --robot.controller=SonicWholeBodyController \
+  --robot.cameras='{"ego_view": {"type": "zmq", "server_address": "localhost", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30, "warmup_s": 15}}'
+```
+
+### Run on the Physical Robot
+
+Start the robot server (see Part 2), then point the rollout at the robot instead of the sim:
+
+```bash
+lerobot-rollout \
+  --strategy.type=base \
+  --policy.path=nepyope/sonic_walk \
+  --policy.device=cuda \
+  --robot.type=unitree_g1 \
+  --robot.is_simulation=false \
+  --robot.robot_ip=<ROBOT_IP> \
+  --robot.controller=SonicWholeBodyController \
+  --robot.cameras='{"ego_view": {"type": "zmq", "server_address": "<ROBOT_IP>", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30, "warmup_s": 5}}'
+```
+
+On connect, the robot eases every joint to the decoder's `default_angles` before the controller takes over, so the first policy commands don't snap from the connect-time pose.
+
+---
+
 ## Additional Resources
 
 - [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python)
@@ -300,4 +341,4 @@ lerobot-rollout \
 
 ---
 
-_Last updated: March 2026_
+_Last updated: July 2026_

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -68,6 +68,6 @@ class UnitreeG1Config(RobotConfig):
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Lower-body controller class name, e.g. "GrootLocomotionController" or
-    # "HolosomaLocomotionController". None disables it.
+    # Controller class name, e.g. GrootLocomotionController / HolosomaLocomotionController /
+    # SonicWholeBodyController. None disables it.
     controller: str | None = None

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

@@ -0,0 +1,27 @@
+#!/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/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,
@@ -86,6 +86,7 @@ class GrootLocomotionController:
         # Load policies
         self.policy_balance, self.policy_walk = load_groot_policies()
 
+        self.default_angles = GROOT_DEFAULT_ANGLES
         self.cmd = np.array([0.0, 0.0, 0.0], dtype=np.float32)  # vx, vy, theta_dot
 
         # Robot state

src/lerobot/robots/unitree_g1/controllers/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,
@@ -104,6 +104,7 @@ class HolosomaLocomotionController:
         # Load policy and gains
         self.policy, self.kp, self.kd = load_policy()
 
+        self.default_angles = DEFAULT_ANGLES
         self.cmd = np.zeros(3, dtype=np.float32)
 
         # Robot state

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

@@ -0,0 +1,195 @@
+#!/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__)
+
+CONTROL_DT = 0.02  # 50 Hz control period (s)
+TOKEN_DIM = 64  # decoder latent size
+
+# Latent-token feature-key prefixes: action carries the token, obs echoes it back.
+TOKEN_ACTION_PREFIX = "motion_token"  # nosec B105 - feature-key prefix, not a secret
+TOKEN_STATE_PREFIX = "motion_token_state"  # nosec B105 - feature-key prefix, not a secret
+
+# SONIC decoder checkpoint. Deploy constants (kp/kd, default_angles, action_scale,
+# neutral_token) are baked into the ONNX metadata; see upload_sonic_decoder.py.
+DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
+DECODER_INPUT_DIM = 994  # token(64) + 10-frame proprio history + gravity
+
+# Decoder filename mapping: the full decoder (default) or NVIDIA's distilled low-latency one.
+POLICY_FILES = {
+    "default": "model_decoder.onnx",
+    "low_latency": "low_latency/model_decoder.onnx",
+}
+
+
+def load_policy(
+    repo_id: str = DEFAULT_SONIC_REPO_ID,
+    policy_type: str = "default",
+) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Load the SONIC decoder and its baked-in deploy constants from ONNX metadata.
+
+    Args:
+        repo_id: Hugging Face Hub repo ID
+        policy_type: Either "default" (full decoder) or "low_latency" (distilled)
+
+    Returns:
+        (decoder, kp, kd, default_angles, action_scale, neutral_token) tuple. The gains/pose/
+        scale are (29,) float32 in IsaacLab joint order; neutral_token is the (64,) idle latent.
+    """
+    if policy_type not in POLICY_FILES:
+        raise ValueError(f"Unknown policy type: {policy_type}. Choose from: {list(POLICY_FILES.keys())}")
+
+    filename = POLICY_FILES[policy_type]
+    logger.info(f"Loading {policy_type.upper()} SONIC decoder from: {repo_id}/{filename}")
+    decoder_path = hf_hub_download(repo_id=repo_id, filename=filename)
+
+    decoder = ort.InferenceSession(decoder_path)
+    logger.info(f"Decoder loaded: {decoder.get_inputs()[0].shape} → {decoder.get_outputs()[0].shape}")
+
+    # Extract deploy constants from ONNX metadata
+    model = onnx.load(decoder_path, load_external_data=False)
+    metadata = {prop.key: prop.value for prop in model.metadata_props}
+
+    required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
+    missing = [k for k in required if k not in metadata]
+    if missing:
+        raise ValueError(f"ONNX model must contain {list(required)} in metadata (missing {missing})")
+
+    arr = {k: np.array(json.loads(metadata[k]), dtype=np.float32) for k in required}
+    logger.info(f"Loaded SONIC deploy constants from ONNX ({len(arr['kp'])} joints)")
+    return decoder, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
+
+
+class SonicWholeBodyController:
+    """Full-body SONIC decoder controller for UnitreeG1's background controller thread.
+
+    Token-only deploy (encoder bypassed): each tick it appends the latest robot state to
+    10-frame history buffers, then maps the policy-supplied 64-D token + that history to a
+    residual action added onto ``default_angles`` -> 50 Hz joint-position targets.
+    """
+
+    control_dt = CONTROL_DT
+
+    def __init__(self, policy_type: str = "default"):
+        self.decoder, self.kp, self.kd, self.default_angles, self.action_scale, self.neutral_token = (
+            load_policy(policy_type=policy_type)
+        )
+        self.decoder_input = self.decoder.get_inputs()[0].name
+        self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]
+
+        # 64-D latent-token action space; rollout maps the policy's 64-D output onto these keys.
+        self.action_ft = {f"{TOKEN_ACTION_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)}
+        # 64-D token proprio state, aggregated by rollout into observation.state (last token).
+        self.observation_ft = {f"{TOKEN_STATE_PREFIX}.{i}.pos": float for i in range(TOKEN_DIM)}
+
+        self.reset()
+        logger.info("SonicWholeBodyController initialized")
+
+    def reset(self) -> None:
+        """Reset internal state for a new episode: held token and 10-frame history buffers."""
+        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
+        self._last_token = None  # neutral token is re-seeded on the first tick
+
+    def observation_state(self) -> dict[str, float]:
+        """Echo the last decoded token as ``observation.state`` so a token-output VLA closes
+        the loop on its own previous token."""
+        token = self._last_token if self._last_token is not None else np.zeros(TOKEN_DIM, dtype=np.float32)
+        return {f"{TOKEN_STATE_PREFIX}.{i}.pos": float(v) for i, v in enumerate(token)}
+
+    def run_step(self, action: dict, lowstate) -> dict:
+        if lowstate is None:
+            return {}
+
+        # Token: reassemble the dense 64-D latent from motion_token.{i}.pos (all keys required);
+        # else hold the last one (neutral until the first real token, which decodes to a stand).
+        keys = [f"{TOKEN_ACTION_PREFIX}.{i}.pos" for i in range(TOKEN_DIM)]
+        if action and all(k in action for k in keys):
+            self._last_token = np.fromiter(
+                (float(action[k]) for k in keys), dtype=np.float32, count=TOKEN_DIM
+            )
+        elif self._last_token is None:
+            self._last_token = self.neutral_token.copy()
+
+        # Read proprioception from lowstate (IsaacLab joint order).
+        q = np.array([lowstate.motor_state[m.value].q for m in G1_29_JointIndex], np.float32)
+        dq = np.array([lowstate.motor_state[m.value].dq for m in G1_29_JointIndex], np.float32)
+        quat = np.array(lowstate.imu_state.quaternion, np.float32)  # (w, x, y, z)
+        quat = quat / (np.linalg.norm(quat) + 1e-8)
+        ang = np.array(lowstate.imu_state.gyroscope, np.float32)
+
+        # Push into the 10-frame history (newest first). The decoder consumes MuJoCo joint
+        # order, so reorder q/dq via MUJOCO_TO_ISAACLAB (validated against the ONNX; don't flip).
+        self.h_q_mj = [q[MUJOCO_TO_ISAACLAB] - self.default_angles_mj] + self.h_q_mj[:-1]
+        self.h_dq_mj = [dq[MUJOCO_TO_ISAACLAB]] + self.h_dq_mj[:-1]
+        self.h_ang = [ang] + self.h_ang[:-1]
+        self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
+        self.h_quat = [quat] + self.h_quat[:-1]
+
+        # Assemble the 994-D decoder input: token + oldest->newest history + gravity.
+        obs = np.zeros(DECODER_INPUT_DIM, np.float32)
+        obs[:TOKEN_DIM] = self._last_token
+        off = TOKEN_DIM
+        for hist, sz in ((self.h_ang, 3), (self.h_q_mj, 29), (self.h_dq_mj, 29), (self.h_act_mj, 29)):
+            for frame in reversed(hist):
+                obs[off : off + sz] = frame
+                off += sz
+        for hquat in reversed(self.h_quat):
+            obs[off : off + 3] = get_gravity_orientation(hquat)
+            off += 3
+
+        # Decode -> residual action (MuJoCo order) added onto the standing pose.
+        action_mj = (
+            self.decoder.run(None, {self.decoder_input: 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
+        return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,6 +23,43 @@ import numpy as np
 
 NUM_MOTORS = 29
 
+# Joint-order permutation between IsaacLab and Mujoco convention
+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,
+)
+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
@@ -68,8 +105,9 @@ def make_locomotion_controller(name: str | None):
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
+        "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",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -34,7 +34,6 @@ 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,
@@ -148,9 +147,8 @@ class UnitreeG1(Robot):
 
         self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
 
-        # Lower-body controller loaded dynamically
+        # Controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
-
         # Controller thread state
         self._controller_thread = None
         self._controller_action_lock = threading.Lock()
@@ -233,13 +231,24 @@ class UnitreeG1(Robot):
 
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
-        return {**self._motors_ft, **self._cameras_ft}
+        # Controllers may contribute their own proprio features (e.g. SONIC's token state).
+        controller_ft = getattr(self.controller, "observation_ft", {})
+        return {**self._motors_ft, **controller_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}
 
+        # Whole-body controllers (SONIC): 64-D latent token.
+        controller_ft = getattr(self.controller, "action_ft", None)
+        if controller_ft is not None:
+            return dict(controller_ft)
+
+        # Locomotion controllers (GR00T / Holosoma): arm joint targets + joystick axes.
+        # TODO: have GR00T/Holosoma advertise their own action_features too, so every
+        # controller declares its action space and this fallthrough can be dropped.
         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}
@@ -341,8 +350,14 @@ class UnitreeG1(Robot):
         logger.info("[UnitreeG1] Connected to robot.")
         self.msg.mode_machine = lowstate.mode_machine
 
-        self.kp = np.array(self.config.kp, dtype=np.float32)
-        self.kd = np.array(self.config.kd, dtype=np.float32)
+        # Prefer the active controller's gains (e.g. SONIC loads kp/kd from its ONNX);
+        # otherwise fall back to the config defaults.
+        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)
+        else:
+            self.kp = np.array(self.config.kp, dtype=np.float32)
+            self.kd = np.array(self.config.kd, dtype=np.float32)
 
         for joint in G1_29_JointIndex:
             self.msg.motor_cmd[joint].mode = 1
@@ -350,6 +365,11 @@ class UnitreeG1(Robot):
             self.msg.motor_cmd[joint].kd = self.kd[joint.value]
             self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
 
+        # Ease into the controller's home pose before it takes over, so the first commands
+        # don't snap from the connect-time pose.
+        if self.controller is not None and hasattr(self.controller, "default_angles"):
+            self.reset(default_positions=self.controller.default_angles)
+
         # Start controller thread if enabled
         if self.controller is not None:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
@@ -461,6 +481,11 @@ class UnitreeG1(Robot):
         if lowstate.wireless_remote:
             obs["wireless_remote"] = lowstate.wireless_remote
 
+        # Controller-contributed observation (e.g. SONIC echoes its last decoded token as
+        # observation.state so a token-output VLA closes the loop on its own previous token).
+        if self.controller is not None and hasattr(self.controller, "observation_state"):
+            obs.update(self.controller.observation_state())
+
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
             if getattr(cam, "use_rgb", True):
@@ -503,11 +528,12 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from incoming teleop action."""
+        """Forward incoming teleop action values into ``controller_input``; each controller
+        reads only the keys it understands."""
         with self._controller_action_lock:
-            for key in REMOTE_KEYS:
-                if key in action:
-                    self.controller_input[key] = action[key]
+            for key, value in action.items():
+                if isinstance(key, str) and value is not None:
+                    self.controller_input[key] = value
 
     @property
     def is_calibrated(self) -> bool:
@@ -565,7 +591,7 @@ class UnitreeG1(Robot):
                     interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
                     action_dict[f"{motor.name}.q"] = float(interp_pos)
 
-                self.send_action(action_dict)
+                self.publish_lowcmd(action_dict)
 
                 # Maintain constant control rate
                 elapsed = time.time() - start_time