From 95d902903914b2b7df815167a85fa67fc1cf2844 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Thu, 30 Jul 2026 11:40:01 +0200 Subject: [PATCH] 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). --- docs/source/unitree_g1.mdx | 57 +-- .../robots/unitree_g1/config_unitree_g1.py | 18 - .../controllers/gr00t_locomotion.py | 24 +- .../controllers/holosoma_locomotion.py | 29 +- .../robots/unitree_g1/run_g1_server.py | 344 +--------------- src/lerobot/robots/unitree_g1/unitree_g1.py | 381 ++++-------------- 6 files changed, 105 insertions(+), 748 deletions(-) diff --git a/docs/source/unitree_g1.mdx b/docs/source/unitree_g1.mdx index 7e2585afc..69965a561 100644 --- a/docs/source/unitree_g1.mdx +++ b/docs/source/unitree_g1.mdx @@ -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. - - **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. - - --- ## 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.controller=SonicWholeBodyController \ - --robot.cameras='{"ego_view": {"type": "zmq", "server_address": "", "port": 5555, "camera_name": "head_camera", "width": 640, "height": 480, "fps": 30}}' \ - --task="walk back and forth" \ - --duration=1000 \ - --fps=30 -``` - - - 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. - - ---- - ## Additional Resources - [Unitree SDK Documentation](https://github.com/unitreerobotics/unitree_sdk2_python) diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index adcc108d2..702faad2b 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -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) diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index 2cc39051a..c9d08c870 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py @@ -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 - # 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) + # Load ONNX policies + policy_balance = ort.InferenceSession(balance_path) + policy_walk = ort.InferenceSession(walk_path) 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): diff --git a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py index 73e73758f..0e2d1c92c 100644 --- a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py @@ -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 diff --git a/src/lerobot/robots/unitree_g1/run_g1_server.py b/src/lerobot/robots/unitree_g1/run_g1_server.py index c1b5aeb56..b5bd0baf8 100644 --- a/src/lerobot/robots/unitree_g1/run_g1_server.py +++ b/src/lerobot/robots/unitree_g1/run_g1_server.py @@ -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( diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 718a216e2..35548d221 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -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 ( - 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, @@ -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 - # client never touches DDS, so we don't import the socket shim at all. - if config.is_simulation or config.onboard: + # Import channel classes based on mode + if config.is_simulation: self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize self._ChannelPublisher = _SDKChannelPublisher self._ChannelSubscriber = _SDKChannelSubscriber else: - self._ChannelFactoryInitialize = None - self._ChannelPublisher = None - self._ChannelSubscriber = None + from .unitree_sdk2_socket import ( + ChannelFactoryInitialize, + ChannelPublisher, + 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 + self._ChannelFactoryInitialize = ChannelFactoryInitialize + self._ChannelPublisher = ChannelPublisher + self._ChannelSubscriber = ChannelSubscriber # 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 - # 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) + self.arm_ik = G1_29_ArmIK() if config.gravity_compensation else None - # 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 + # Lower-body / whole-body controller loaded dynamically + self.controller: LocomotionController | None = make_locomotion_controller(config.controller) + + # 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 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 { - **self._motors_ft, - **self._token_state_ft, - **self._cameras_ft, - } + return {**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 - # publishes low commands to DDS). + # Start controller thread if enabled 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: - 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: + # Put robot in passive mode before stopping threads + if not self.config.is_simulation: self._send_zero_torque() - if self.controller is not None and hasattr(self.controller, "shutdown"): - self.controller.shutdown() + # Signal thread to stop and unblock any waits + 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 = lowstate_to_obs(lowstate) + obs = {} + + # 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.*`` - 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. + Controller-agnostic: every value-carrying key (locomotion ``remote.*`` axes or + SONIC ``motion_token.*`` values) is forwarded verbatim into ``controller_input`` + and each controller extracts only the keys it understands. The robot deliberately + does not enumerate any controller's key schema here. """ with self._controller_action_lock: for key, value in action.items(): @@ -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 - # ignore ``.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 controllers (SONIC) own the whole 29-DoF command and ignore + # ``.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: