feat(unitree_g1): make SONIC token interface implicit

Drop the ``sonic_token_action`` config flag; the 64-D latent-token
action/observation interface now switches on automatically whenever the
SONIC whole-body controller is selected (``controller == "SonicWholeBodyController"``).
Keyed via a ``_sonic_token`` property so client, onboard and sim roles agree.

docs/source/unitree_g1.mdx

@@ -8,6 +8,15 @@
 
 The Unitree G1 humanoid is now supported in LeRobot! You can teleoperate, train locomanipulation policies, test in sim, and more. Both 29 and 23 DoF variants are supported.
 
+<Tip>
+  **New: SONIC whole-body control.** The `SonicWholeBodyController` runs NVIDIA's
+  [GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC) decoder on the G1, turning a
+  64-D latent motion token into full-body joint targets at 50 Hz. This lets you drive
+  the robot from a VLA policy trained on SONIC motion tokens (token in → whole-body
+  motion out) with `lerobot-rollout`, in sim or on the physical robot. See
+  [Whole-body control with SONIC](#whole-body-control-with-sonic) below.
+</Tip>
+
 ---
 
 ## Part 1: Getting Started
@@ -59,7 +68,7 @@ lerobot-teleoperate \
   --robot.controller=GrootLocomotionController
 ```
 
-This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support both [HolosomaLocomotionController](https://github.com/amazon-far/holosoma) and [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl) via `--robot.controller`.
+This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support [HolosomaLocomotionController](https://github.com/amazon-far/holosoma), [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl), and [SonicWholeBodyController](https://huggingface.co/nvidia/GEAR-SONIC) via `--robot.controller`.
 
 - Press `9` to release the robot
 - Press `7` / `8` to increase / decrease waist height
@@ -290,6 +299,52 @@ lerobot-rollout \
 
 ---
 
+## Whole-body control with SONIC
+
+The `SonicWholeBodyController` runs NVIDIA's [GEAR-SONIC](https://huggingface.co/nvidia/GEAR-SONIC)
+decoder on the G1. Each 50 Hz tick it consumes a **64-D latent motion token** and emits
+full-body joint targets — the encoder is bypassed, so a policy feeds tokens in and the
+decoder turns them into motion. Before the first token arrives the controller holds a
+neutral (idle) pose.
+
+This makes the G1 drivable by a VLA policy trained to output SONIC motion tokens (token
+as both `observation.state` and `action`, e.g. [`nepyope/sonic_walk`](https://huggingface.co/nepyope/sonic_walk))
+using the standard `lerobot-rollout`. The controller always runs **onboard** the robot;
+the laptop is a thin client that streams tokens and receives camera frames over ZMQ.
+
+**On the robot** — start the server in handshake mode so it instantiates and runs the
+controller onboard against local DDS at full rate:
+
+```bash
+cd ~/lerobot
+python src/lerobot/robots/unitree_g1/run_g1_server.py --handshake --camera
+```
+
+**From your laptop** — run the token policy; selecting `--robot.controller=SonicWholeBodyController`
+implicitly switches the robot to the 64-D latent-token action/observation interface:
+
+```bash
+lerobot-rollout \
+  --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}}' \
+  --task="walk back and forth" \
+  --duration=1000 \
+  --fps=30
+```
+
+<Tip>
+  SONIC is a token-only decoder in LeRobot: the only input path is the 64-D latent
+  vector. To train your own token policy, expose the 64-D token as the action (a config
+  choice, e.g. `pi05` with a 64-D action dim) — no policy code changes are needed.
+</Tip>
+
+---
+
 ## Additional Resources
 
 - [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python)

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -62,12 +62,34 @@ class UnitreeG1Config(RobotConfig):
     # Socket config for ZMQ bridge
     robot_ip: str = "192.168.123.164"  # default G1 IP
 
+    # Run the locomotion / whole-body controller ONBOARD the robot (policy on the G1
+    # itself, against local DDS at full rate) instead of on the laptop over the ZMQ
+    # socket bridge. In this mode the robot object uses the real Unitree SDK channels
+    # and expects high-level actions (arm targets + joystick axes, or 64-D SONIC
+    # tokens) fed via send_action -- e.g. by run_g1_server's serve_onboard_controller,
+    # which receives them from the laptop over ZMQ. Mutually exclusive with is_simulation.
+    onboard: bool = False
+    # DDS network interface for onboard mode (None = SDK default, matching
+    # run_g1_server.py's ChannelFactoryInitialize(0)).
+    dds_interface: str | None = None
+    # Onboard sub-flags. On a real G1 both are True: the built-in motion services
+    # must be released before we can write lowcmd, and locomotion axes are read from
+    # the physical wireless remote. Against a DDS sim neither applies (no
+    # MotionSwitcher, no physical remote), so set both False so the controller takes
+    # its locomotion axes purely from send_action (ZMQ) input.
+    release_motion_control: bool = True
+    physical_remote: bool = True
+
     # Cameras (ZMQ-based remote cameras)
     cameras: dict[str, CameraConfig] = field(default_factory=dict)
 
     # 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.
+    # 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).
     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

@@ -14,6 +14,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import logging
 from collections import deque
 
@@ -21,7 +23,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,
@@ -68,9 +70,15 @@ def load_groot_policies(
         filename="GR00T-WholeBodyControl-Walk.onnx",
     )
 
-    # Load ONNX policies
-    policy_balance = ort.InferenceSession(balance_path)
-    policy_walk = ort.InferenceSession(walk_path)
+    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
+    # background thread alongside the (torch) upper-body policy, IK and sim; letting
+    # ORT grab every core starves those and makes the whole rollout stutter. These
+    # are small MLPs, so 1 thread is both enough and lowest-latency.
+    from ..g1_utils import make_ort_session_options
+
+    so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
+    policy_balance = ort.InferenceSession(balance_path, sess_options=so)
+    policy_walk = ort.InferenceSession(walk_path, sess_options=so)
 
     logger.info("GR00T policies loaded successfully")
 
@@ -196,6 +204,16 @@ class GrootLocomotionController:
         # Transform action back to target joint positions
         target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
 
+        # Waist override: an external upper-body IK can command the 3 waist joints
+        # (indices 12/13/14) via ``kWaist{Yaw,Roll,Pitch}.q`` in the action dict. When
+        # present, we substitute the balance policy's waist target so the torso tracks
+        # the IK while the policy keeps only the legs balanced. Single-publisher stays
+        # intact (this thread still owns joints 0-14).
+        for idx in (G1_29_JointIndex.kWaistYaw, G1_29_JointIndex.kWaistRoll, G1_29_JointIndex.kWaistPitch):
+            key = f"{idx.name}.q"
+            if key in action and action[key] is not None:
+                target_dof_pos_15[idx.value] = float(action[key])
+
         # Build action dict
         action_dict = {}
         for i in range(15):

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

@@ -14,18 +14,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
+from __future__ import annotations
+
 import logging
 
 import numpy as np
-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,
+    compute_pd_gains,
     get_gravity_orientation,
 )
 
@@ -57,12 +58,23 @@ POLICY_FILES = {
     "ppo": "ppo_g1_29dof.onnx",
 }
 
+# Per-joint motor model in Holosoma's joint order, plus the joints that get a x2
+# stiffness/damping factor. These reproduce the kp/kd that used to be read from the
+# policy's ONNX metadata exactly (both fastsac and ppo), so gains are now derived from
+# the shared motor model (see g1_utils.compute_pd_gains) instead.
+HOLOSOMA_MOTOR_MODELS = (
+    ["7520_14", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+    + ["7520_14", "5020", "5020"]
+    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
+)
+HOLOSOMA_DOUBLE = {4, 5, 10, 11, 13, 14}
+
 
 def load_policy(
     repo_id: str = DEFAULT_HOLOSOMA_REPO_ID,
     policy_type: str = "fastsac",
 ) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray]:
-    """Load Holosoma locomotion policy and extract KP/KD from metadata.
+    """Load the Holosoma locomotion policy and its motor-model-derived PD gains.
 
     Args:
         repo_id: Hugging Face Hub repo ID
@@ -81,16 +93,7 @@ def load_policy(
     policy = ort.InferenceSession(policy_path)
     logger.info(f"Policy loaded: {policy.get_inputs()[0].shape} → {policy.get_outputs()[0].shape}")
 
-    # Extract KP/KD from ONNX metadata
-    model = onnx.load(policy_path, load_external_data=False)
-    metadata = {prop.key: prop.value for prop in model.metadata_props}
-
-    if "kp" not in metadata or "kd" not in metadata:
-        raise ValueError("ONNX model must contain 'kp' and 'kd' in metadata")
-
-    kp = np.array(json.loads(metadata["kp"]), dtype=np.float32)
-    kd = np.array(json.loads(metadata["kd"]), dtype=np.float32)
-    logger.info(f"Loaded KP/KD from ONNX ({len(kp)} joints)")
+    kp, kd = compute_pd_gains(HOLOSOMA_MOTOR_MODELS, HOLOSOMA_DOUBLE)
 
     return policy, kp, kd
 

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

@@ -0,0 +1,401 @@
+#!/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 ``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 logging
+
+import numpy as np
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+
+from ..g1_utils import (
+    ISAACLAB_TO_MUJOCO,
+    MOTOR_ARMATURE,
+    MUJOCO_TO_ISAACLAB,
+    NATURAL_FREQ,
+    G1_29_JointIndex,
+    compute_pd_gains,
+    get_gravity_orientation,
+    lowstate_to_obs,
+    make_ort_session_options,
+)
+
+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
+
+# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals
+# added on top of this.
+DEFAULT_ANGLES = np.array(
+    [
+        -0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
+        -0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
+        0.0, 0.0, 0.0,
+        0.2, 0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
+        0.2, -0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
+    ],
+    dtype=np.float32,
+)
+
+# Per-motor torque limits (N·m), used only for SONIC's residual-action scaling. The
+# armature / bandwidth constants and the PD-gain formula are shared (see g1_utils).
+EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
+
+
+def _action_scale(k):
+    """Per-motor residual-action scale (maps policy output to joint-angle delta)."""
+    return 0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2)
+
+
+# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of truth
+# for both ACTION_SCALE and compute_kp_kd().
+MOTOR_MODELS = (
+    ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+    + ["7520_14", "5020", "5020"]
+    + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
+)
+ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32)  # (29,) IsaacLab
+
+
+def _to_mujoco(a):
+    """Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
+
+    NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
+    permutation direction are a fixed convention validated against the deployed SONIC ONNX
+    policy (the decoder consumes vectors in this order). Do not "correct" the table or
+    rename toward the opposite direction without re-validating on hardware.
+    """
+    return a[MUJOCO_TO_ISAACLAB]
+
+
+DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
+
+
+# Ankle + waist joint indices (IsaacLab order) that get a x2 stiffness/damping factor.
+_SONIC_DOUBLE = {4, 5, 10, 11, 13, 14}
+
+
+def compute_kp_kd():
+    """SONIC per-joint PD gains (kp, kd), (29,) float32 in IsaacLab joint order."""
+    return compute_pd_gains(MOTOR_MODELS, _SONIC_DOUBLE)
+
+
+# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
+TOKEN_ACTION_PREFIX = "motion_token"
+# 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"
+
+
+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
+
+# Neutral ("zero pose") SONIC token, held by token_mode until the first real token arrives.
+# Captured from the encoder's own output while the robot stood idle in sim: the encoder is
+# an FSQ bottleneck (~5 bit/dim, Div(16)), so its tokens live on the 1/16 grid. We store the
+# integer FSQ codes and rescale by 1/16, giving an exact on-grid token -- unlike the literal
+# all-zero token, which is off the learned manifold and decodes to a slightly goofy stance.
+# This one decodes to a stable, natural standing pose.
+_NEUTRAL_TOKEN_CODES = np.array(
+    [-1, 3, 1, -1, 1, -3, 6, 1, 1, 1, -2, -4, -2, 0, -3, -1,
+     2, -1, -3, -5, 3, 1, 1, -4, -1, -1, 1, -7, 0, 1, 2, -2,
+     5, -2, -2, -4, 0, -1, 3, -1, 0, -5, -1, 0, -4, 0, 0, -1,
+     -1, 2, -2, 1, 3, 3, 1, 0, 0, 6, 0, -7, 3, 0, 2, -2],
+    dtype=np.float32,
+)
+NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0  # FSQ Div(16): integer codes -> on-grid token
+
+
+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).
+    """
+
+    def __init__(self, decoder):
+        self.decoder = decoder
+        self.decoder_input = decoder.get_inputs()[0].name
+        dec_dim = int(decoder.get_inputs()[0].shape[1])
+        if dec_dim != 994:
+            raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected 994)")
+        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)
+        q_mj = _to_mujoco(q)
+        dq_mj = _to_mujoco(dq)
+        self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + 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", 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 = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * 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_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
+
+        so = make_ort_session_options()
+        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so)
+
+        self.kp, self.kd = compute_kp_kd()
+        self.controller = SonicDecoder(decoder_sess)
+
+    @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
+
+        # 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. ``token_mode`` is set True by the robot whenever a SONIC
+        # whole-body controller is selected (token-driven deploy): the controller then 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). This lives here (not in the entry-point script) so it applies
+        # uniformly to run_g1_server, lerobot-rollout and the sim replays.
+        self.token_mode = False
+        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", 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, lowstate) -> dict:
+        if lowstate is None:
+            return {}
+        obs = lowstate_to_obs(lowstate)
+
+        # 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 and self.token_mode:
+            # Token-driven deploy, but no token has arrived yet: hold the captured neutral
+            # token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural standing pose.
+            self._last_token = NEUTRAL_TOKEN.copy()
+        if self._last_token is None:
+            # No token yet and not in token_mode: hold (keep last target).
+            return {}
+        # 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 token_mode re-seeds the neutral token after a reset.
+        self._last_token = None
+
+    def shutdown(self):
+        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -23,11 +23,82 @@ 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,
+)
+MUJOCO_TO_ISAACLAB = np.array(
+    [
+        0,
+        6,
+        12,
+        1,
+        7,
+        13,
+        2,
+        8,
+        14,
+        3,
+        9,
+        15,
+        22,
+        4,
+        10,
+        16,
+        23,
+        5,
+        11,
+        17,
+        24,
+        18,
+        25,
+        19,
+        26,
+        20,
+        27,
+        21,
+        28,
+    ],
+    dtype=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
 
-
 def default_remote_input() -> dict[str, float]:
     """Return a zeroed-out remote input dict (axes + buttons)."""
     return dict.fromkeys(REMOTE_KEYS, 0.0)
@@ -43,6 +114,53 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
     return gravity_orientation
 
 
+# Unitree motor-model parameters shared by the controllers that derive their PD gains
+# from motor physics rather than hand-tuning (SONIC decoder, Holosoma). NATURAL_FREQ is
+# the target closed-loop stiffness bandwidth (rad/s); MOTOR_ARMATURE is per-model rotor
+# inertia (keys are Unitree motor model names). From these: kp = armature * w**2 and
+# kd = 4 * armature * w, with an optional x2 factor on stiff joints (ankles/waist).
+NATURAL_FREQ = 10.0 * 2.0 * np.pi
+MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
+
+
+def compute_pd_gains(motor_models, double_indices=()) -> tuple[np.ndarray, np.ndarray]:
+    """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth.
+
+    ``motor_models`` is a per-joint sequence of Unitree motor model names (in the
+    controller's own joint order); joints whose index is in ``double_indices`` get a
+    x2 stiffness/damping factor. Returns two (N,) float32 arrays in that same order.
+    """
+    double = set(double_indices)
+
+    def s(k):
+        return MOTOR_ARMATURE[k] * NATURAL_FREQ**2
+
+    def d(k):
+        return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ
+
+    kp = np.array([2 * s(k) if i in double else s(k) for i, k in enumerate(motor_models)], dtype=np.float32)
+    kd = np.array([2 * d(k) if i in double else d(k) for i, k in enumerate(motor_models)], dtype=np.float32)
+    return kp, kd
+
+
+def make_ort_session_options(intra_op_num_threads: int | None = None, inter_op_num_threads: int | None = None):
+    """Build quiet ONNX Runtime SessionOptions, optionally capping the CPU thread pool.
+
+    These tiny MLP policies are latency-bound, not throughput-bound, so letting ORT grab
+    every core starves the real-time control loop / torch policy and causes stutter. Pass
+    1 intra + 1 inter thread for lowest-latency per-step inference.
+    """
+    import onnxruntime as ort
+
+    so = ort.SessionOptions()
+    so.log_severity_level = 3
+    if intra_op_num_threads is not None:
+        so.intra_op_num_threads = intra_op_num_threads
+    if inter_op_num_threads is not None:
+        so.inter_op_num_threads = inter_op_num_threads
+    return so
+
+
 class G1_29_JointArmIndex(IntEnum):
     # Left arm
     kLeftShoulderPitch = 15
@@ -63,13 +181,55 @@ class G1_29_JointArmIndex(IntEnum):
     kRightWristYaw = 28
 
 
+def lowstate_to_obs(lowstate) -> dict:
+    """Build a robot observation dict from a Unitree lowstate.
+
+    Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
+    lowstate -> obs mapping lives in exactly one place. Keys match the
+    ``<joint>.q``/``imu.*`` schema consumed across the controllers.
+    """
+    obs: dict = {}
+
+    for motor in G1_29_JointIndex:
+        idx = motor.value
+        obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
+        obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
+        obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
+
+    imu = lowstate.imu_state
+    if imu.gyroscope:
+        obs["imu.gyro.x"] = imu.gyroscope[0]
+        obs["imu.gyro.y"] = imu.gyroscope[1]
+        obs["imu.gyro.z"] = imu.gyroscope[2]
+    if imu.accelerometer:
+        obs["imu.accel.x"] = imu.accelerometer[0]
+        obs["imu.accel.y"] = imu.accelerometer[1]
+        obs["imu.accel.z"] = imu.accelerometer[2]
+    if imu.quaternion:
+        obs["imu.quat.w"] = imu.quaternion[0]
+        obs["imu.quat.x"] = imu.quaternion[1]
+        obs["imu.quat.y"] = imu.quaternion[2]
+        obs["imu.quat.z"] = imu.quaternion[3]
+    if imu.rpy:
+        obs["imu.rpy.roll"] = imu.rpy[0]
+        obs["imu.rpy.pitch"] = imu.rpy[1]
+        obs["imu.rpy.yaw"] = imu.rpy[2]
+
+    wr = getattr(lowstate, "wireless_remote", None)
+    if wr:
+        obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
+
+    return obs
+
+
 def make_locomotion_controller(name: str | None):
     """Instantiate a locomotion controller by class name. Returns None if name is 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/run_g1_server.py

@@ -22,16 +22,33 @@ This server runs on the robot and forwards:
 - Robot commands (LowCmd) from ZMQ to DDS (from remote clients)
 
 Uses JSON for secure serialization instead of pickle.
+
+Controller-negotiation handshake
+--------------------------------
+The first message from a client agrees on which controller the server will run onboard
+(``serve_onboard_controller``); the controller NEVER runs on the laptop client.
+Test the handshake in isolation (no DDS, runs on a laptop) in two terminals::
+
+    # terminal A: handshake-only server
+    python -m lerobot.robots.unitree_g1.run_g1_server --handshake-only
+
+    # terminal B: client proposes a controller
+    python -m lerobot.robots.unitree_g1.run_g1_server \\
+        --handshake-client SonicWholeBodyController --sonic-token-action --server-ip 127.0.0.1
+
+On the real robot, add ``--handshake`` to the normal bridge to require agreement first.
 """
 
 import argparse
 import base64
 import contextlib
 import json
+import signal
 import threading
 import time
 from typing import Any
 
+import numpy as np
 import zmq
 from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
 from unitree_sdk2py.core.channel import ChannelFactoryInitialize, ChannelPublisher, ChannelSubscriber
@@ -50,6 +67,253 @@ LOWCMD_PORT = 6000
 LOWSTATE_PORT = 6001
 NUM_MOTORS = 35
 
+# Onboard high-level channels (serve_onboard_controller): compact actions in, state out.
+ACTION_PORT = 6004
+STATE_PORT = 6005
+
+# Controller-negotiation handshake (REQ/REP). The client's first message agrees on
+# which controller the server will run before any control data flows.
+HANDSHAKE_PORT = 6002
+PROTOCOL_VERSION = 1
+
+# Controllers that can run ONBOARD (must match g1_utils.make_locomotion_controller).
+# ``None`` (a.k.a. "bridge") means no onboard controller: the laptop owns control and
+# streams raw lowcmd over the ZMQ DDS bridge (the legacy run_g1_server behavior).
+VALID_CONTROLLERS = (
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicWholeBodyController",
+)
+# SONIC latent-token dimensionality (mirrors sonic_whole_body.TOKEN_DIM; kept local so
+# the handshake can run without importing the heavy controller / onnxruntime).
+TOKEN_DIM = 64
+_BRIDGE_ALIASES = {"", "none", "null", "bridge", "raw"}
+
+
+def _normalize_controller(name: str | None) -> str | None:
+    """Map a requested controller name to a canonical value (or None for raw bridge)."""
+    if name is None:
+        return None
+    low = str(name).strip().lower()
+    if low in _BRIDGE_ALIASES:
+        return None
+    for c in VALID_CONTROLLERS:
+        if c.lower() == low:
+            return c
+    raise ValueError(f"Unknown controller {name!r}. Available: {list(VALID_CONTROLLERS)} or 'bridge'")
+
+
+def _capabilities(controller: str | None, sonic_token_action: bool) -> dict[str, Any]:
+    """The interface the server advertises for an agreed controller."""
+    caps: dict[str, Any] = {
+        "controller": controller,
+        "sonic_token_action": bool(sonic_token_action),
+        "protocol": PROTOCOL_VERSION,
+    }
+    if controller is None:
+        # Raw DDS bridge: the laptop runs the controller and streams lowcmd.
+        caps["mode"] = "bridge"
+        caps["lowcmd_port"] = LOWCMD_PORT
+        caps["lowstate_port"] = LOWSTATE_PORT
+    else:
+        # Onboard: the controller runs here; the laptop ships compact high-level actions.
+        caps["mode"] = "onboard"
+        caps["action_port"] = ACTION_PORT
+        caps["state_port"] = STATE_PORT
+        if sonic_token_action:
+            caps["action_space"] = "motion_token"
+            caps["action_dim"] = TOKEN_DIM
+    return caps
+
+
+def negotiate_controller(sock: zmq.Socket, shutdown_event: threading.Event) -> dict[str, Any]:
+    """Server side of the handshake: block on one REP socket until a client sends a
+    valid ``hello``, then reply with the negotiated capabilities and return them.
+
+    Rejects malformed / unknown-controller requests with an error reply and keeps
+    waiting (a rejected client can retry). Honors ``shutdown_event`` so Ctrl-C works.
+    """
+    poller = zmq.Poller()
+    poller.register(sock, zmq.POLLIN)
+    while not shutdown_event.is_set():
+        if not dict(poller.poll(timeout=200)):
+            continue
+        raw = sock.recv()
+        try:
+            hello = json.loads(raw.decode("utf-8"))
+        except (ValueError, UnicodeDecodeError) as e:
+            sock.send_json({"type": "error", "ok": False, "error": f"bad hello: {e}"})
+            continue
+        try:
+            controller = _normalize_controller(hello.get("controller"))
+        except ValueError as e:
+            sock.send_json(
+                {"type": "error", "ok": False, "error": str(e), "available": list(VALID_CONTROLLERS)}
+            )
+            continue
+        reply = {"type": "welcome", "ok": True, **_capabilities(controller, hello.get("sonic_token_action", False))}
+        sock.send_json(reply)
+        return reply
+    raise KeyboardInterrupt
+
+
+def request_controller(
+    server_ip: str,
+    controller: str | None,
+    *,
+    sonic_token_action: bool = False,
+    port: int = HANDSHAKE_PORT,
+    timeout_s: float = 5.0,
+) -> dict[str, Any]:
+    """Client side of the handshake: propose a controller, return the server's agreed
+    capabilities (or raise on rejection / timeout)."""
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.REQ)
+    sock.setsockopt(zmq.LINGER, 0)
+    sock.setsockopt(zmq.RCVTIMEO, int(timeout_s * 1000))
+    sock.setsockopt(zmq.SNDTIMEO, int(timeout_s * 1000))
+    sock.connect(f"tcp://{server_ip}:{port}")
+    hello = {
+        "type": "hello",
+        "controller": controller,
+        "sonic_token_action": bool(sonic_token_action),
+        "protocol": PROTOCOL_VERSION,
+    }
+    try:
+        sock.send_json(hello)
+        reply = sock.recv_json()
+    except zmq.Again as e:
+        raise TimeoutError(f"no handshake reply from {server_ip}:{port} within {timeout_s}s") from e
+    finally:
+        sock.close(linger=0)
+    if not reply.get("ok"):
+        raise RuntimeError(f"handshake rejected: {reply.get('error')} (available: {reply.get('available')})")
+    return reply
+
+
+def serve_onboard_controller(
+    *,
+    controller: str,
+    sonic_token_action: bool,
+    dds_interface: str | None = None,
+    sim: bool = False,
+    cameras: dict | None = None,
+    camera_fps: int = 30,
+    camera_port: int = 5555,
+    action_port: int = ACTION_PORT,
+    state_port: int = STATE_PORT,
+    state_fps: float = 30.0,
+    stop: threading.Event | None = None,
+) -> None:
+    """Run the negotiated controller ONBOARD -- the single control path on the robot.
+
+    Builds ``UnitreeG1(onboard=True, controller=...)`` so the controller/balance loop runs
+    locally against DDS at full rate (the 50 Hz ``_controller_loop`` thread lives in
+    UnitreeG1), then receives compact high-level actions from the laptop over ZMQ
+    (:action_port), decodes them via the controller, publishes ``observation.state``
+    (:state_port), and optionally serves the ego camera. The controller NEVER runs on the
+    laptop; the laptop (lerobot-rollout thin-client) only ships tokens/axes and reads back
+    state + camera frames.
+    """
+    # Imported lazily: UnitreeG1 imports request_controller from this module, so a
+    # top-level import here would be circular.
+    from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
+    from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
+
+    if stop is None:
+        stop = threading.Event()
+        signal.signal(signal.SIGINT, lambda *_: stop.set())
+        signal.signal(signal.SIGTERM, lambda *_: stop.set())
+
+    cfg = UnitreeG1Config(
+        is_simulation=False,
+        onboard=True,
+        controller=controller,
+        dds_interface=dds_interface,
+        release_motion_control=not sim,
+        physical_remote=not sim,
+        cameras={},
+    )
+
+    # Optional camera server (background daemon thread; independent of DDS).
+    if cameras:
+        camera_server = ImageServer({"fps": camera_fps, "cameras": cameras}, port=camera_port)
+        threading.Thread(target=camera_server.run, daemon=True).start()
+        cam_summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
+        print(f"Camera server started on :{camera_port}: {cam_summary}")
+
+    robot = UnitreeG1(cfg)
+    print(f"Connecting onboard robot (controller={controller}, token={sonic_token_action})...")
+    robot.connect()
+
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.PULL)
+    sock.setsockopt(zmq.CONFLATE, 1)  # only ever act on the freshest command
+    sock.setsockopt(zmq.RCVTIMEO, 200)  # keeps the loop responsive to the stop event
+    sock.bind(f"tcp://0.0.0.0:{action_port}")
+    print(f"Onboard controller live. Waiting for laptop actions on :{action_port} ...")
+    print("Ctrl-C for graceful shutdown.")
+
+    state_sock = None
+    if state_fps > 0:
+        state_sock = ctx.socket(zmq.PUB)
+        state_sock.setsockopt(zmq.SNDHWM, 2)
+        state_sock.setsockopt(zmq.LINGER, 0)
+        state_sock.bind(f"tcp://0.0.0.0:{state_port}")
+        print(f"Publishing observation.state on :{state_port} at {state_fps:.0f} Hz")
+
+        def publish_state() -> None:
+            period = 1.0 / state_fps
+            while not stop.is_set():
+                t0 = time.time()
+                obs = robot.get_observation()
+                if obs:
+                    # Forward every scalar proprio key the robot exposes (29 joint .q, IMU,
+                    # and the SONIC token echo: 64-D motion_token_state.*). Camera arrays are
+                    # streamed separately by the ImageServer, so drop ndarrays here. This
+                    # makes the laptop thin-client a pure relay.
+                    state = {
+                        k: float(v)
+                        for k, v in obs.items()
+                        if isinstance(v, (bool, int, float, np.floating, np.integer))
+                    }
+                    with contextlib.suppress(zmq.Again):
+                        state_sock.send_json(state, zmq.NOBLOCK)
+                time.sleep(max(0.0, period - (time.time() - t0)))
+
+        threading.Thread(target=publish_state, daemon=True).start()
+    else:
+        print("observation.state PUB disabled (state_fps<=0)")
+
+    n = 0
+    try:
+        while not stop.is_set():
+            try:
+                payload = sock.recv()
+            except zmq.Again:
+                continue
+            except zmq.ContextTerminated:
+                break
+
+            try:
+                action = json.loads(payload.decode("utf-8"))
+            except (ValueError, UnicodeDecodeError) as e:
+                print(f"Dropping malformed action: {e}")
+                continue
+
+            robot.send_action(action)
+
+            n += 1
+            if n % 60 == 0:
+                print(f"Applied {n} actions")
+    finally:
+        print("Shutting down onboard controller...")
+        stop.set()
+        if state_sock is not None:
+            with contextlib.suppress(Exception):
+                state_sock.close(linger=0)
+        robot.disconnect()
+
 
 def lowstate_to_dict(msg: hg_LowState) -> dict[str, Any]:
     """Convert LowState SDK message to a JSON-serializable dictionary."""
@@ -160,8 +424,86 @@ def main() -> None:
     parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
     parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
     parser.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port (default: 5555)")
+    # Controller-negotiation handshake (first message agrees on the controller).
+    parser.add_argument("--handshake", action="store_true",
+                        help="Wait for a client to negotiate the controller before bridging")
+    parser.add_argument("--handshake-port", type=int, default=HANDSHAKE_PORT,
+                        help=f"Handshake REQ/REP port (default: {HANDSHAKE_PORT})")
+    parser.add_argument("--handshake-only", action="store_true",
+                        help="Run ONLY the handshake server (no DDS/cameras) to test negotiation")
+    parser.add_argument("--handshake-client", default=None, metavar="CONTROLLER",
+                        help="Act as a client: propose CONTROLLER (or 'bridge') to --server-ip and print the reply")
+    parser.add_argument("--server-ip", default="127.0.0.1", help="[--handshake-client] server IP")
+    parser.add_argument("--sonic-token-action", action="store_true",
+                        help="[handshake] negotiate the 64-D SONIC token action interface")
     args = parser.parse_args()
 
+    # --- Isolated handshake test paths (no DDS, safe to run on a laptop) ---
+    if args.handshake_client is not None:
+        controller = None if args.handshake_client.strip().lower() in _BRIDGE_ALIASES else args.handshake_client
+        reply = request_controller(
+            args.server_ip, controller,
+            sonic_token_action=args.sonic_token_action, port=args.handshake_port,
+        )
+        print(json.dumps(reply, indent=2))
+        return
+
+    if args.handshake_only:
+        ctx = zmq.Context.instance()
+        rep = ctx.socket(zmq.REP)
+        rep.bind(f"tcp://0.0.0.0:{args.handshake_port}")
+        print(f"[handshake] server listening on :{args.handshake_port} (no DDS). Ctrl-C to stop.")
+        shutdown = threading.Event()
+        try:
+            while True:
+                reply = negotiate_controller(rep, shutdown)
+                print(f"[handshake] agreed: controller={reply['controller']} mode={reply['mode']} "
+                      f"sonic_token_action={reply['sonic_token_action']}")
+        except KeyboardInterrupt:
+            print("\n[handshake] stopping")
+        finally:
+            rep.close(linger=0)
+            ctx.term()
+        return
+
+    # Controller-negotiation handshake: the client's first message agrees on the
+    # controller, which we then run ONBOARD (the controller NEVER runs on the laptop).
+    # Bridge/None falls through to the legacy raw DDS forward (deprecated laptop control).
+    if args.handshake:
+        ctx = zmq.Context.instance()
+        hs = ctx.socket(zmq.REP)
+        hs.bind(f"tcp://0.0.0.0:{args.handshake_port}")
+        print(f"[handshake] waiting for client controller agreement on :{args.handshake_port} ...")
+        shutdown = threading.Event()
+        try:
+            agreed = negotiate_controller(hs, shutdown)
+        except KeyboardInterrupt:
+            print("[handshake] interrupted before agreement; exiting")
+            hs.close(linger=0)
+            ctx.term()
+            return
+        hs.close(linger=0)
+        if agreed["controller"] is not None:
+            print(f"[handshake] running controller ONBOARD: {agreed['controller']} "
+                  f"(sonic_token_action={agreed['sonic_token_action']})")
+            cameras = None
+            if args.camera:
+                cameras = {
+                    "head_camera": {
+                        "device_id": args.camera_device,
+                        "shape": [args.camera_height, args.camera_width],
+                    }
+                }
+            serve_onboard_controller(
+                controller=agreed["controller"],
+                sonic_token_action=bool(agreed["sonic_token_action"]),
+                cameras=cameras,
+                camera_fps=args.camera_fps,
+                camera_port=args.camera_port,
+            )
+            return
+        print("[handshake] client selected raw DDS bridge (laptop owns control) -> legacy forward.")
+
     # Optionally start camera server in background thread
     camera_thread = None
     if args.camera:
@@ -205,6 +547,7 @@ def main() -> None:
 
     # initialize ZMQ
     ctx = zmq.Context.instance()
+    shutdown_event = threading.Event()
 
     # receive commands from remote client
     lowcmd_sock = ctx.socket(zmq.PULL)
@@ -215,7 +558,6 @@ def main() -> None:
     lowstate_sock.bind(f"tcp://0.0.0.0:{LOWSTATE_PORT}")
 
     state_period = 0.002  # ~500 hz
-    shutdown_event = threading.Event()
 
     # start observation forwarding in background thread
     t_state = threading.Thread(

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -16,6 +16,8 @@
 
 from __future__ import annotations
 
+import contextlib
+import json
 import logging
 import threading
 import time
@@ -26,6 +28,7 @@ from typing import TYPE_CHECKING, Protocol, runtime_checkable
 import numpy as np
 
 from lerobot.cameras import make_cameras_from_configs
+from lerobot.utils.errors import DeviceNotConnectedError
 from lerobot.types import RobotAction, RobotObservation
 from lerobot.utils.import_utils import _unitree_sdk_available, require_package
 
@@ -34,10 +37,10 @@ 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,
+    lowstate_to_obs,
     make_locomotion_controller,
 )
 
@@ -47,7 +50,9 @@ if TYPE_CHECKING or _unitree_sdk_available:
         ChannelPublisher as _SDKChannelPublisher,
         ChannelSubscriber as _SDKChannelSubscriber,
     )
-    from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
+    from unitree_sdk2py.idl.default import (
+        unitree_hg_msg_dds__LowCmd_,
+    )
     from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
         LowCmd_ as hg_LowCmd,
         LowState_ as hg_LowState,
@@ -79,6 +84,14 @@ class LocomotionController(Protocol):
 kTopicLowCommand_Debug = "rt/lowcmd"
 kTopicLowState = "rt/lowstate"
 
+# Wireless-remote button byte layout, mapped to the positional button indices the
+# locomotion controllers expect. Used in onboard mode to read the physical Unitree
+# remote from lowstate (mirrors the exo teleoperator's RemoteController).
+_REMOTE_BUTTON_MAP: list[str] = [
+    "RB", "LB", "start", "back", "RT", "LT", "", "",
+    "A", "B", "X", "Y", "up", "right", "down", "left",
+]
+
 
 @dataclass
 class MotorState:
@@ -119,24 +132,37 @@ class UnitreeG1(Robot):
         self.config = config
         self.control_dt = config.control_dt
 
+        # Three mutually-exclusive roles:
+        #   * simulation : local DDS + controller run in-process against a MuJoCo world.
+        #   * onboard    : local DDS + controller run in-process on the robot NX.
+        #   * client     : thin laptop client. No DDS, no controller. It negotiates a
+        #                  controller with ``run_g1_server`` (which runs it onboard),
+        #                  PUSHes high-level actions and reads back state + cameras over
+        #                  ZMQ. The controller *always* runs on the robot, never here.
+        self._client = not config.is_simulation and not config.onboard
+
         # Initialize cameras config (ZMQ-based) - actual connection in connect()
         self._cameras = make_cameras_from_configs(config.cameras)
 
-        # Import channel classes based on mode
-        if config.is_simulation:
+        # DDS channel classes are only needed by the in-process control roles. The thin
+        # client never touches DDS, so we don't import the socket shim at all.
+        if config.is_simulation or config.onboard:
             self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
             self._ChannelPublisher = _SDKChannelPublisher
             self._ChannelSubscriber = _SDKChannelSubscriber
         else:
-            from .unitree_sdk2_socket import (
-                ChannelFactoryInitialize,
-                ChannelPublisher,
-                ChannelSubscriber,
-            )
+            self._ChannelFactoryInitialize = None
+            self._ChannelPublisher = None
+            self._ChannelSubscriber = None
 
-            self._ChannelFactoryInitialize = ChannelFactoryInitialize
-            self._ChannelPublisher = ChannelPublisher
-            self._ChannelSubscriber = ChannelSubscriber
+        # Client-side ZMQ handles / negotiated capabilities (populated in connect()).
+        self._client_action_sock = None
+        self._client_state_sock = None
+        self._client_state_latest: dict[str, float] = {}
+        self._client_caps: dict | None = None
+
+        # Optional arm gravity compensation (feed-forward torque via the arm IK solver).
+        self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
 
         # Initialize state variables
         self.sim_env = None
@@ -146,24 +172,69 @@ class UnitreeG1(Robot):
         self._shutdown_event = threading.Event()
         self.subscribe_thread = None
 
-        self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
+        # Lower-body controller loaded dynamically. GUARDRAIL: the controller must never
+        # be built or run on the laptop client -- it always runs onboard (or in sim).
+        if self._client:
+            self.controller: LocomotionController | None = None
+        else:
+            self.controller = make_locomotion_controller(config.controller)
 
-        # Lower-body controller loaded dynamically
-        self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
+            # Token-driven deploy: a SONIC whole-body controller always runs in token
+            # mode -- it holds a neutral token until the first real one arrives, then
+            # holds the last token between control ticks.
+            if hasattr(self.controller, "token_mode"):
+                self.controller.token_mode = True
 
         # 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 = {}
 
+        # Onboard-only: parser for the physical Unitree wireless remote (read straight
+        # from local lowstate so joystick locomotion works without a laptop round-trip).
+        self._joystick = None
+
+        # 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:
-                self.sim_env.step()
+                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
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -231,15 +302,46 @@ 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._cameras_ft}
+        return {
+            **self._motors_ft,
+            **self._token_state_ft,
+            **self._cameras_ft,
+        }
 
     @cached_property
     def action_features(self) -> dict[str, type]:
-        if self.controller is None:
+        # Role-agnostic: the schema is a pure function of the controller name. The thin
+        # client advertises the same schema as the onboard robot so the exact same
+        # policy output routes straight through.
+
+        # No controller configured at all: raw 29-DoF joint teleop.
+        if self.config.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}
@@ -255,6 +357,11 @@ 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
 
@@ -271,6 +378,13 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
+                # Onboard: the physical Unitree remote (in local lowstate) takes
+                # priority for locomotion when active; otherwise laptop/ZMQ axes stand.
+                if self.config.onboard:
+                    wl = self._wireless_remote_input(lowstate)
+                    if wl is not None:
+                        controller_input.update(wl)
+
                 # Run controller step
                 controller_action = self.controller.run_step(controller_input, lowstate)
 
@@ -293,7 +407,163 @@ class UnitreeG1(Robot):
     def configure(self) -> None:
         pass
 
+    def _wireless_remote_input(self, lowstate) -> dict | None:
+        """Parse the physical Unitree remote from lowstate into controller inputs.
+
+        Onboard only. Returns None when the remote is idle so the laptop-provided
+        (ZMQ) axes keep control; otherwise the physical remote takes priority.
+        """
+        js = self._joystick
+        if js is None:
+            return None
+        wr = getattr(lowstate, "wireless_remote", None)
+        if not wr or len(wr) < 24:
+            return None
+        try:
+            js.extract(wr)
+        except Exception:  # noqa: BLE001
+            return None
+
+        axes = {
+            "remote.lx": float(js.lx.data),
+            "remote.ly": float(js.ly.data),
+            "remote.rx": float(js.rx.data),
+            "remote.ry": float(js.ry.data),
+        }
+        active = any(abs(v) > 1e-2 for v in axes.values())
+        out = dict(axes)
+        for i, name in enumerate(_REMOTE_BUTTON_MAP):
+            if name:
+                val = float(getattr(js, name).data)
+                out[f"remote.button.{i}"] = val
+                if val:
+                    active = True
+        return out if active else None
+
+    def _release_motion_control(self) -> None:
+        """Release the robot's built-in motion services so we can send raw lowcmd.
+
+        Onboard-only. Mirrors run_g1_server.py: on the real robot the factory
+        locomotion/hand services must relinquish control before our controller can
+        write to ``rt/lowcmd``, otherwise commands are ignored or fought.
+        """
+        from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
+
+        msc = MotionSwitcherClient()
+        msc.SetTimeout(5.0)
+        msc.Init()
+        _, result = msc.CheckMode()
+        while result is not None and "name" in result and result["name"]:
+            logger.info("[UnitreeG1] Releasing built-in mode '%s'...", result["name"])
+            msc.ReleaseMode()
+            _, result = msc.CheckMode()
+            time.sleep(1.0)
+
+    # ------------------------------------------------------------------ #
+    # Thin-client role (laptop): no DDS, no controller. Talks to run_g1_server
+    # over ZMQ. The controller ALWAYS runs onboard; we only relay high-level
+    # actions and read back the state echo + camera frames.
+    # ------------------------------------------------------------------ #
+    def _connect_client(self) -> None:
+        import zmq
+
+        from .run_g1_server import ACTION_PORT, HANDSHAKE_PORT, STATE_PORT, request_controller
+
+        server_ip = self.config.robot_ip
+        if not server_ip:
+            raise ValueError("client mode requires config.robot_ip (the G1 running run_g1_server)")
+
+        # 1) Handshake: agree with the server on which controller it will run onboard.
+        logger.info(
+            "[client] handshaking with %s:%d (controller=%s, token=%s)...",
+            server_ip, HANDSHAKE_PORT, self.config.controller, self._sonic_token,
+        )
+        self._client_caps = request_controller(
+            server_ip,
+            self.config.controller,
+            sonic_token_action=self._sonic_token,
+            port=HANDSHAKE_PORT,
+        )
+        logger.info("[client] server agreed: %s", self._client_caps)
+
+        ctx = zmq.Context.instance()
+
+        # 2) Action PUSH: ship compact high-level actions to the onboard controller.
+        self._client_action_sock = ctx.socket(zmq.PUSH)
+        self._client_action_sock.setsockopt(zmq.SNDHWM, 2)
+        self._client_action_sock.setsockopt(zmq.LINGER, 0)
+        self._client_action_sock.connect(f"tcp://{server_ip}:{ACTION_PORT}")
+
+        # 3) State SUB: read the onboard observation.state echo (last token / joints).
+        self._client_state_sock = ctx.socket(zmq.SUB)
+        self._client_state_sock.setsockopt(zmq.CONFLATE, 1)
+        self._client_state_sock.setsockopt_string(zmq.SUBSCRIBE, "")
+        self._client_state_sock.connect(f"tcp://{server_ip}:{STATE_PORT}")
+
+        # 4) Cameras (ZMQ ImageServer served by run_g1_server) - same as any client.
+        for cam in self._cameras.values():
+            if not cam.is_connected:
+                cam.connect()
+        logger.info("[client] connected: actions ->:%d, state <-:%d, %d camera(s).",
+                    ACTION_PORT, STATE_PORT, len(self._cameras))
+
+    def _recv_client_state(self) -> None:
+        """Drain the state SUB (CONFLATE keeps only the freshest) into the latest cache."""
+        import zmq
+
+        if self._client_state_sock is None:
+            return
+        while True:
+            try:
+                state = self._client_state_sock.recv_json(flags=zmq.NOBLOCK)
+            except zmq.Again:
+                break
+            except (ValueError, zmq.ZMQError):
+                break
+            if isinstance(state, dict):
+                self._client_state_latest = {k: float(v) for k, v in state.items()}
+
+    def _get_observation_client(self) -> RobotObservation:
+        self._recv_client_state()
+        obs: dict = dict(self._client_state_latest)
+        for cam_name, cam in self._cameras.items():
+            if getattr(cam, "use_rgb", True):
+                obs[cam_name] = cam.read_latest()
+            if getattr(cam, "use_depth", False):
+                obs[f"{cam_name}_depth"] = cam.read_latest_depth()
+        return obs
+
+    def _send_action_client(self, action: RobotAction) -> RobotAction:
+        """Relay the raw action straight to the onboard controller. NO processing here:
+        the controller negotiated in the handshake interprets it (token / wb / arm)."""
+        import zmq
+
+        if self._client_action_sock is None:
+            raise DeviceNotConnectedError("UnitreeG1 client is not connected")
+        payload = json.dumps({k: float(v) for k, v in action.items()}).encode("utf-8")
+        with contextlib.suppress(zmq.Again):
+            self._client_action_sock.send(payload, zmq.NOBLOCK)
+        return action
+
+    def _disconnect_client(self) -> None:
+        for sock in (self._client_action_sock, self._client_state_sock):
+            if sock is not None:
+                with contextlib.suppress(Exception):
+                    sock.close(linger=0)
+        self._client_action_sock = None
+        self._client_state_sock = None
+        for cam in self._cameras.values():
+            with contextlib.suppress(Exception):
+                cam.disconnect()
+
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
+        # Thin-client role: no DDS, no controller. Negotiate the controller with
+        # run_g1_server (which runs it onboard), then open the high-level ZMQ links:
+        # PUSH actions on :ACTION_PORT, SUB state echo on :STATE_PORT, cameras via ZMQ.
+        if self._client:
+            self._connect_client()
+            return
+
         # Initialize DDS channel and simulation environment
         if self.config.is_simulation:
             from lerobot.envs import make_env
@@ -302,6 +572,28 @@ class UnitreeG1(Robot):
             self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
             # Extract the actual gym env from the dict structure
             self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
+        elif self.config.onboard:
+            # Real robot, controller running onboard against local DDS. Initialize the
+            # real SDK channel factory on the robot's DDS interface and take low-level
+            # control from the built-in services before we start writing lowcmd.
+            if self.config.dds_interface:
+                self._ChannelFactoryInitialize(0, self.config.dds_interface)
+            else:
+                self._ChannelFactoryInitialize(0)
+            # Real robot: hand low-level control over from the built-in services.
+            # A DDS sim has no MotionSwitcher, so this is skipped there.
+            if self.config.release_motion_control:
+                self._release_motion_control()
+            # Real robot: read the physical wireless remote from lowstate for
+            # locomotion. A sim has no physical remote, so leave _joystick=None and
+            # let send_action (ZMQ) drive the locomotion axes instead.
+            if self.config.physical_remote:
+                from unitree_sdk2py.utils.joystick import Joystick
+
+                self._joystick = Joystick()
+                for axis in (self._joystick.lx, self._joystick.ly, self._joystick.rx, self._joystick.ry):
+                    axis.smooth = 1.0
+                    axis.deadzone = 0.0
         else:
             self._ChannelFactoryInitialize(0, config=self.config)
 
@@ -343,6 +635,9 @@ 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
@@ -350,7 +645,8 @@ 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
 
-        # Start controller thread if enabled
+        # Start the 50 Hz controller thread (runs the locomotion/whole-body policy and
+        # publishes low commands to DDS).
         if self.controller is not None:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
             self._controller_thread.start()
@@ -372,12 +668,34 @@ class UnitreeG1(Robot):
             logger.warning(f"Failed to send zero-torque on disconnect: {e}")
 
     def disconnect(self):
-        # Put robot in passive mode before stopping threads
-        if not self.config.is_simulation:
+        if self._client:
+            self._disconnect_client()
+            return
+
+        # Stop the controller loop first so it isn't fighting the shutdown ramp.
+        self._shutdown_event.set()
+        controller_stopped = True
+        if self._controller_thread is not None:
+            # Wait long enough for any in-flight inference tick to finish and the loop
+            # to observe the shutdown flag, so no stray low command is published while
+            # the ramp runs (the shutdown routine must be the single publisher).
+            self._controller_thread.join(timeout=5.0)
+            if self._controller_thread.is_alive():
+                controller_stopped = False
+                logger.error(
+                    "Controller thread did not stop; skipping graceful ramp to avoid "
+                    "concurrent low commands (fail-safe: joints keep last command until exit)"
+                )
+
+        # Put the robot in passive mode (zero-torque) before stopping the rest (real
+        # robot only; the subscribe thread is still alive here to supply the current
+        # pose). Only publish once the controller thread has definitely exited so the
+        # two aren't publishing at once.
+        if not self.config.is_simulation and controller_stopped:
             self._send_zero_torque()
 
-        # Signal thread to stop and unblock any waits
-        self._shutdown_event.set()
+        if self.controller is not None and hasattr(self.controller, "shutdown"):
+            self.controller.shutdown()
 
         # Wait for subscribe thread to finish
         if self.subscribe_thread is not None:
@@ -385,12 +703,6 @@ class UnitreeG1(Robot):
             if self.subscribe_thread.is_alive():
                 logger.warning("Subscribe thread did not stop cleanly")
 
-        # Wait for controller thread to finish
-        if self._controller_thread is not None:
-            self._controller_thread.join(timeout=2.0)
-            if self._controller_thread.is_alive():
-                logger.warning("Controller thread did not stop cleanly")
-
         # Close simulation environment
         if self.config.is_simulation and self.sim_env is not None:
             try:
@@ -417,49 +729,25 @@ class UnitreeG1(Robot):
             cam.disconnect()
 
     def get_observation(self) -> RobotObservation:
+        if self._client:
+            return self._get_observation_client()
+
         with self._lowstate_lock:
             lowstate = self._lowstate
         if lowstate is None:
             return {}
 
-        obs = {}
+        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
+        obs = lowstate_to_obs(lowstate)
 
-        # Motors - q, dq, tau for all joints
-        for motor in G1_29_JointIndex:
-            name = motor.name
-            idx = motor.value
-            obs[f"{name}.q"] = lowstate.motor_state[idx].q
-            obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
-            obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
+        # 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
 
-        # IMU - gyroscope
-        if lowstate.imu_state.gyroscope:
-            obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
-            obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
-            obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
-
-        # IMU - accelerometer
-        if lowstate.imu_state.accelerometer:
-            obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
-            obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
-            obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
-
-        # IMU - quaternion
-        if lowstate.imu_state.quaternion:
-            obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
-            obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
-            obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
-            obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
-
-        # IMU - rpy
-        if lowstate.imu_state.rpy:
-            obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
-            obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
-            obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
-
-        # Wireless remote (raw bytes for teleoperator)
-        if lowstate.wireless_remote:
-            obs["wireless_remote"] = lowstate.wireless_remote
+            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():
@@ -471,11 +759,22 @@ class UnitreeG1(Robot):
         return obs
 
     def send_action(self, action: RobotAction) -> RobotAction:
+        if self._client:
+            return self._send_action_client(action)
+
         action_to_publish = action
         if self.controller is not None:
+            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)
+            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
@@ -503,11 +802,17 @@ class UnitreeG1(Robot):
         return action
 
     def _update_controller_action(self, action: RobotAction) -> None:
-        """Update controller input state from incoming teleop action."""
+        """Update controller input state from an incoming teleop action.
+
+        Controller-agnostic: every value-carrying key (e.g. locomotion ``remote.*``
+        axes/buttons) 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 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:
@@ -515,6 +820,8 @@ class UnitreeG1(Robot):
 
     @property
     def is_connected(self) -> bool:
+        if self._client:
+            return self._client_action_sock is not None
         with self._lowstate_lock:
             return self._lowstate is not None
 
@@ -537,43 +844,64 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
-        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)
+        # Full-body controllers (SONIC / OpenHLM) 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
 
-            # get current state
-            obs = self.get_observation()
+        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)
 
-            # 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"]
+                # get current state
+                obs = self.get_observation()
 
-            # Interpolate to default position
-            for step in range(num_steps):
-                start_time = time.time()
-
-                alpha = step / num_steps
-                action_dict = {}
+                # record current positions
+                init_dof_pos = np.zeros(29, dtype=np.float32)
                 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)
+                    init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
-                self.send_action(action_dict)
+                # Interpolate to default position
+                for step in range(num_steps):
+                    start_time = time.time()
 
-                # Maintain constant control rate
-                elapsed = time.time() - start_time
-                sleep_time = max(0, control_dt - elapsed)
-                time.sleep(sleep_time)
+                    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)
 
-        # Reset controller internal state (gait phase, obs history, etc.)
-        if self.controller is not None and hasattr(self.controller, "reset"):
-            self.controller.reset()
+                    # 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()
 
         logger.info("Reached default position")