refactor(unitree_g1): reduce branch to SONIC-only diff vs main

Split the onboard-controller server/handshake/thin-client work out to feat/g1_onboard_controller and revert it here: - unitree_g1.py: drop client + onboard roles, wireless-remote parsing and motion-service release; restore main's sim/socket-bridge transport. Keep only the SONIC integration (implicit token action/state, full-body reset/pause, controller kp/kd + shutdown). - config: drop onboard/dds_interface/release_motion_control/physical_remote. - run_g1_server.py + unitree_g1.mdx: reverted to main. - gr00t/holosoma: keep the move into controllers/ but revert their content to main (only the package-relative import changes).
2026-07-30 13:09:40 +00:00 · 2026-07-30 11:40:01 +02:00
parent bbfc4ff443
commit 95d9029039
6 changed files with 105 additions and 748 deletions
@@ -8,15 +8,6 @@
 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
@@ -68,7 +59,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 [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`.
+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`.
 - Press `9` to release the robot
 - Press `7` / `8` to increase / decrease waist height
@@ -299,52 +290,6 @@ 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)
@@ -62,24 +62,6 @@ 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)
@@ -14,8 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import annotations
 import logging
 from collections import deque
@@ -70,15 +68,9 @@ def load_groot_policies(
        filename="GR00T-WholeBodyControl-Walk.onnx",
    )
-    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
+    # Load ONNX policies
-    # background thread alongside the (torch) upper-body policy, IK and sim; letting
+    policy_balance = ort.InferenceSession(balance_path)
-    # ORT grab every core starves those and makes the whole rollout stutter. These
+    policy_walk = ort.InferenceSession(walk_path)
    # 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")
@@ -204,16 +196,6 @@ 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):
@@ -14,11 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
-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
@@ -26,7 +26,6 @@ from ..g1_utils import (
    REMOTE_AXES,
    G1_29_JointArmIndex,
    G1_29_JointIndex,
    compute_pd_gains,
    get_gravity_orientation,
 )
@@ -58,23 +57,12 @@ 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 the Holosoma locomotion policy and its motor-model-derived PD gains.
+    """Load Holosoma locomotion policy and extract KP/KD from metadata.
    Args:
        repo_id: Hugging Face Hub repo ID
@@ -93,7 +81,16 @@ def load_policy(
    policy = ort.InferenceSession(policy_path)
    logger.info(f"Policy loaded: {policy.get_inputs()[0].shape} → {policy.get_outputs()[0].shape}")
-    kp, kd = compute_pd_gains(HOLOSOMA_MOTOR_MODELS, HOLOSOMA_DOUBLE)
+    # 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)")
    return policy, kp, kd
@@ -22,33 +22,16 @@ 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
@@ -67,253 +50,6 @@ 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."""
@@ -424,86 +160,8 @@ 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:
@@ -547,7 +205,6 @@ def main() -> None:
    # initialize ZMQ
    ctx = zmq.Context.instance()
    shutdown_event = threading.Event()
    # receive commands from remote client
    lowcmd_sock = ctx.socket(zmq.PULL)
@@ -558,6 +215,7 @@ 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(
@@ -16,8 +16,6 @@
 from __future__ import annotations
 import contextlib
 import json
 import logging
 import threading
 import time
@@ -28,7 +26,6 @@ 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
@@ -40,7 +37,6 @@ from .g1_utils import (
    G1_29_JointArmIndex,
    G1_29_JointIndex,
    default_remote_input,
    lowstate_to_obs,
    make_locomotion_controller,
 )
@@ -50,9 +46,7 @@ if TYPE_CHECKING or _unitree_sdk_available:
        ChannelPublisher as _SDKChannelPublisher,
        ChannelSubscriber as _SDKChannelSubscriber,
    )
-    from unitree_sdk2py.idl.default import (
+    from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
        unitree_hg_msg_dds__LowCmd_,
    )
    from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
        LowCmd_ as hg_LowCmd,
        LowState_ as hg_LowState,
@@ -84,14 +78,6 @@ 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:
@@ -132,37 +118,24 @@ 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)
-        # DDS channel classes are only needed by the in-process control roles. The thin
+        # Import channel classes based on mode
-        # client never touches DDS, so we don't import the socket shim at all.
+        if config.is_simulation:
        if config.is_simulation or config.onboard:
            self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
            self._ChannelPublisher = _SDKChannelPublisher
            self._ChannelSubscriber = _SDKChannelSubscriber
        else:
-            self._ChannelFactoryInitialize = None
+            from .unitree_sdk2_socket import (
-            self._ChannelPublisher = None
+                ChannelFactoryInitialize,
-            self._ChannelSubscriber = None
+                ChannelPublisher,
                ChannelSubscriber,
            )
-        # Client-side ZMQ handles / negotiated capabilities (populated in connect()).
+            self._ChannelFactoryInitialize = ChannelFactoryInitialize
-        self._client_action_sock = None
+            self._ChannelPublisher = ChannelPublisher
-        self._client_state_sock = None
+            self._ChannelSubscriber = ChannelSubscriber
        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
@@ -172,18 +145,15 @@ class UnitreeG1(Robot):
        self._shutdown_event = threading.Event()
        self.subscribe_thread = None
-        # Lower-body controller loaded dynamically. GUARDRAIL: the controller must never
+        self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None
        # 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)
-            # Token-driven deploy: a SONIC whole-body controller always runs in token
+        # Lower-body / whole-body controller loaded dynamically
-            # mode -- it holds a neutral token until the first real one arrives, then
+        self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
-            # holds the last token between control ticks.
+
-            if hasattr(self.controller, "token_mode"):
+        # A SONIC whole-body controller always runs in token mode: it holds a neutral
-                self.controller.token_mode = True
+        # token until the first real one arrives, then holds the last token between ticks.
        if self.controller is not None and hasattr(self.controller, "token_mode"):
            self.controller.token_mode = True
        # Controller thread state
        self._controller_thread = None
@@ -194,10 +164,6 @@ class UnitreeG1(Robot):
        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
@@ -317,20 +283,12 @@ class UnitreeG1(Robot):
    @cached_property
    def observation_features(self) -> dict[str, type | tuple]:
-        return {
+        return {**self._motors_ft, **self._token_state_ft, **self._cameras_ft}
            **self._motors_ft,
            **self._token_state_ft,
            **self._cameras_ft,
        }
    @cached_property
    def action_features(self) -> dict[str, type]:
        # 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:
+        if self.controller is None:
            return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
@@ -378,13 +336,6 @@ 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)
@@ -407,163 +358,7 @@ 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
@@ -572,28 +367,6 @@ 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)
@@ -645,8 +418,7 @@ 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 the 50 Hz controller thread (runs the locomotion/whole-body policy and
+        # Start controller thread if enabled
        # 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()
@@ -668,34 +440,12 @@ class UnitreeG1(Robot):
            logger.warning(f"Failed to send zero-torque on disconnect: {e}")
    def disconnect(self):
-        if self._client:
+        # Put robot in passive mode before stopping threads
-            self._disconnect_client()
+        if not self.config.is_simulation:
            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()
-        if self.controller is not None and hasattr(self.controller, "shutdown"):
+        # Signal thread to stop and unblock any waits
-            self.controller.shutdown()
+        self._shutdown_event.set()
        # Wait for subscribe thread to finish
        if self.subscribe_thread is not None:
@@ -703,6 +453,16 @@ 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")
        # Release controller resources (e.g. SONIC decoder sessions).
        if self.controller is not None and hasattr(self.controller, "shutdown"):
            self.controller.shutdown()
        # Close simulation environment
        if self.config.is_simulation and self.sim_env is not None:
            try:
@@ -729,16 +489,49 @@ 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 {}
-        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
+        obs = {}
-        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
        # 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 mode: echo the last commanded latent token as observation.state so a
        # token-output VLA closes the loop on its own previous token.
@@ -759,11 +552,11 @@ 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:
            # SONIC decoder: pull the 64-D latent token out of the action and remember it
            # for the observation.state echo. The controller thread reads it back from
            # controller_input (populated below) and decodes it into a 29-DoF command.
            if self._sonic_token:
                from .controllers.sonic_whole_body import _extract_token_from_action
@@ -771,6 +564,8 @@ class UnitreeG1(Robot):
                if token is not None:
                    self._last_token = token
            self._update_controller_action(action)
            # Full-body controllers (SONIC) own the whole 29-DoF command; nothing to
            # publish here (the controller thread is the sole publisher).
            if getattr(self.controller, "full_body", False):
                return action
            # Controller thread owns legs/waist. Here we only update joystick inputs
@@ -804,10 +599,10 @@ class UnitreeG1(Robot):
    def _update_controller_action(self, action: RobotAction) -> None:
        """Update controller input state from an incoming teleop action.
-        Controller-agnostic: every value-carrying key (e.g. locomotion ``remote.*``
+        Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or
-        axes/buttons) is forwarded verbatim into ``controller_input`` and each
+        SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input``
-        controller extracts only the keys it understands. The robot deliberately does
+        and each controller extracts only the keys it understands. The robot deliberately
-        not enumerate any controller's key schema here.
+        does not enumerate any controller's key schema here.
        """
        with self._controller_action_lock:
            for key, value in action.items():
@@ -820,8 +615,6 @@ 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
@@ -844,10 +637,10 @@ class UnitreeG1(Robot):
        if default_positions is None:
            default_positions = np.array(self.config.default_positions, dtype=np.float32)
-        # Full-body controllers (SONIC / OpenHLM) own the whole 29-DoF command and
+        # Full-body controllers (SONIC) own the whole 29-DoF command and ignore
-        # ignore ``<joint>.q`` in send_action(), so reset() must publish the default
+        # ``<joint>.q`` in send_action(), so reset() must publish the default pose
-        # pose directly. Pause the background controller first so the two aren't both
+        # directly. Pause the background controller first so the two aren't both writing
-        # writing low commands while the robot moves to the default pose.
+        # 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: