diff --git a/docs/source/unitree_g1.mdx b/docs/source/unitree_g1.mdx index 69965a561..d6a95ab95 100644 --- a/docs/source/unitree_g1.mdx +++ b/docs/source/unitree_g1.mdx @@ -8,6 +8,15 @@ The Unitree G1 humanoid is now supported in LeRobot! You can teleoperate, train locomanipulation policies, test in sim, and more. Both 29 and 23 DoF variants are supported. + + **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 @@ -59,7 +68,7 @@ lerobot-teleoperate \ --robot.controller=GrootLocomotionController ``` -This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support both [HolosomaLocomotionController](https://github.com/amazon-far/holosoma) and [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl) via `--robot.controller`. +This will launch a [MuJoCo sim instance](https://huggingface.co/lerobot/unitree-g1-mujoco/tree/main) for the G1. You can connect a gamepad to your machine before launching in order to control the robot's locomotion in sim. We support [HolosomaLocomotionController](https://github.com/amazon-far/holosoma), [GrootLocomotionController](https://github.com/NVlabs/GR00T-WholeBodyControl), and [SonicWholeBodyController](https://huggingface.co/nvidia/GEAR-SONIC) via `--robot.controller`. - Press `9` to release the robot - Press `7` / `8` to increase / decrease waist height @@ -290,6 +299,53 @@ 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; `--robot.sonic_token_action=true` 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.sonic_token_action=true \ + --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/README.md b/src/lerobot/robots/unitree_g1/README.md deleted file mode 100644 index 16a8f5381..000000000 --- a/src/lerobot/robots/unitree_g1/README.md +++ /dev/null @@ -1,103 +0,0 @@ -# Unitree G1 — SONIC decoder whole-body control - -This package runs NVIDIA's **SONIC** decoder on the Unitree G1, in MuJoCo simulation or -on real hardware, driven by a **64-D latent motion token**. It is a pure-Python/ONNX -reimplementation of the decode half of the SONIC deploy stack (no `gear_sonic`/torch -dependency): the decoder maps a 64-D latent token + proprioception history into 50 Hz -joint-position targets for the robot's PD controller. The encoder is bypassed — a policy -(e.g. `nepyope/sonic_walk`) emits the token directly. - -## Controllers - -Selected with `--robot.controller=`: - -| Controller | Purpose | -| ------------------------------ | --------------------------------------------------- | -| `SonicWholeBodyController` | SONIC decoder driven by a 64-D latent motion token | -| `GrootLocomotionController` | GR00T locomotion policy | -| `HolosomaLocomotionController` | Holosoma locomotion policy | - -The rest of this document covers the SONIC token path. - -Each tick the `SonicWholeBodyController` takes a 64-D latent token -(`motion_token.0.pos … motion_token.63.pos`) and decodes it directly (encoder bypassed). -Before the first token arrives it holds a captured **neutral token** (a stable standing -pose), then holds the last token received between ticks (the ~30 Hz token stream vs. the -~50 Hz control loop). On startup the controller **interpolates** from the robot's measured -pose into the policy's commanded target over ~3 s (no snap). - -## Requirements - -- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended). Verify with: - ```bash - python -c "import onnxruntime as ort; print(ort.get_available_providers())" - ``` -- `mujoco` for simulation (`is_simulation=True`). -- The SONIC encoder/decoder ONNX models download automatically from the - `nvidia/GEAR-SONIC` Hub repo. - -## Architecture: controller always runs onboard - -The controller runs **on the robot**, never on the laptop. The laptop is a thin client: -it negotiates the controller with `run_g1_server` (handshake), then PUSHes the 64-D token -and reads back the `observation.state` echo + camera frames over ZMQ. - -## Running a rollout (real robot) - -On the robot — host the SONIC decoder + camera onboard: - -```bash -python -m lerobot.robots.unitree_g1.run_g1_server --handshake \ - --cameras "ego_view:/dev/v4l/by-path/platform-3610000.usb-usb-0:2.1:1.3-video-index0:640x480" -``` - -On the laptop — `lerobot-rollout` drives the thin client: - -```bash -lerobot-rollout \ - --policy.path=nepyope/sonic_walk \ - --robot.type=unitree_g1 \ - --robot.is_simulation=false --robot.onboard=false \ - --robot.robot_ip= \ - --robot.controller=SonicWholeBodyController --robot.sonic_token_action=true \ - --robot.cameras='{ego_view: {type: zmq, server_address: , port: 5555, camera_name: ego_view, width: 640, height: 480, fps: 30}}' \ - --task="walk back and forth" --device=cuda -``` - -## Training a token policy (no pi05 code patch) - -The SONIC token interface needs **no modeling changes** to pi05. A 64-D token action is -handled entirely by config: pi05 builds its action projections straight from config -(`action_in_proj = nn.Linear(max_action_dim, …)`, `action_out_proj = nn.Linear(…, -max_action_dim)`), pads the action to `max_action_dim`, then slices back to the dataset's -action dim. Set both dims to 64 and the pad/slice is a no-op, so the full 64-D token is -supervised. - -Requirements: - -1. The dataset carries a 64-D `action` and 64-D `observation.state` (the motion tokens). -2. Pass the dims to `lerobot-train`: - -```bash -lerobot-train \ - --dataset.repo_id=nepyope/walk_back_and_forth \ - --policy.type=pi05 \ - --policy.max_action_dim=64 \ - --policy.max_state_dim=64 \ - --policy.chunk_size=50 --policy.n_action_steps=50 -``` - -`nepyope/sonic_walk` was trained exactly this way (`config.json`: `max_action_dim=64`, -`max_state_dim=64`, `output_features.action.shape=[64]`). Same stock code path for train -and inference — the checkpoint's 64-wide `Linear`s load with unmodified pi05. - -## Observation / action interface (token mode) - -With `--robot.sonic_token_action=true` the robot advertises: - -- action: 64-D `motion_token.{i}.pos` (the decoder consumes it directly), -- `observation.state`: 64-D `motion_token_state.{i}.pos` (the last commanded token, - echoed so a token-output VLA closes the loop on its own previous token), - -plus the ego camera image. The controller always runs onboard (or in sim); it is never -built on the laptop client. diff --git a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py index c171c9179..73e73758f 100644 --- a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py @@ -16,11 +16,9 @@ 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 @@ -28,6 +26,7 @@ from ..g1_utils import ( REMOTE_AXES, G1_29_JointArmIndex, G1_29_JointIndex, + compute_pd_gains, get_gravity_orientation, ) @@ -59,12 +58,23 @@ POLICY_FILES = { "ppo": "ppo_g1_29dof.onnx", } +# Per-joint motor model in Holosoma's joint order, plus the joints that get a x2 +# stiffness/damping factor. These reproduce the kp/kd that used to be read from the +# policy's ONNX metadata exactly (both fastsac and ppo), so gains are now derived from +# the shared motor model (see g1_utils.compute_pd_gains) instead. +HOLOSOMA_MOTOR_MODELS = ( + ["7520_14", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2 + + ["7520_14", "5020", "5020"] + + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2 +) +HOLOSOMA_DOUBLE = {4, 5, 10, 11, 13, 14} + def load_policy( repo_id: str = DEFAULT_HOLOSOMA_REPO_ID, policy_type: str = "fastsac", ) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray]: - """Load Holosoma locomotion policy and extract KP/KD from metadata. + """Load the Holosoma locomotion policy and its motor-model-derived PD gains. Args: repo_id: Hugging Face Hub repo ID @@ -83,16 +93,7 @@ def load_policy( policy = ort.InferenceSession(policy_path) logger.info(f"Policy loaded: {policy.get_inputs()[0].shape} → {policy.get_outputs()[0].shape}") - # Extract KP/KD from ONNX metadata - model = onnx.load(policy_path, load_external_data=False) - metadata = {prop.key: prop.value for prop in model.metadata_props} - - if "kp" not in metadata or "kd" not in metadata: - raise ValueError("ONNX model must contain 'kp' and 'kd' in metadata") - - kp = np.array(json.loads(metadata["kp"]), dtype=np.float32) - kd = np.array(json.loads(metadata["kd"]), dtype=np.float32) - logger.info(f"Loaded KP/KD from ONNX ({len(kp)} joints)") + kp, kd = compute_pd_gains(HOLOSOMA_MOTOR_MODELS, HOLOSOMA_DOUBLE) return policy, kp, kd diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index ccc4bda84..c0f27f703 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -38,12 +38,14 @@ from huggingface_hub import hf_hub_download from ..g1_utils import ( ISAACLAB_TO_MUJOCO, + MOTOR_ARMATURE, MUJOCO_TO_ISAACLAB, + NATURAL_FREQ, G1_29_JointIndex, + compute_pd_gains, get_gravity_orientation, lowstate_to_obs, make_ort_session_options, - ort_providers, ) logger = logging.getLogger(__name__) @@ -65,16 +67,14 @@ DEFAULT_ANGLES = np.array( dtype=np.float32, ) -# Per-motor-type parameters used to derive action scaling and PD gains. Keys are Unitree -# motor model names; ARMATURE = rotor inertia, EFFORT = torque limit (N·m). -NATURAL_FREQ = 10.0 * 2.0 * np.pi # target closed-loop stiffness bandwidth (rad/s) -ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425} +# Per-motor torque limits (N·m), used only for SONIC's residual-action scaling. The +# armature / bandwidth constants and the PD-gain formula are shared (see g1_utils). EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0} def _action_scale(k): """Per-motor residual-action scale (maps policy output to joint-angle delta).""" - return 0.25 * EFFORT[k] / (ARMATURE[k] * NATURAL_FREQ**2) + return 0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2) # Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of truth @@ -101,23 +101,13 @@ def _to_mujoco(a): DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES) +# Ankle + waist joint indices (IsaacLab order) that get a x2 stiffness/damping factor. +_SONIC_DOUBLE = {4, 5, 10, 11, 13, 14} + + def compute_kp_kd(): - """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth. - - Ankle and waist joints get a x2 factor for extra stiffness. Returns two (29,) float32 - arrays in IsaacLab joint order. - """ - - def s(k): - return ARMATURE[k] * NATURAL_FREQ**2 - - def d(k): - return 2.0 * 2.0 * ARMATURE[k] * NATURAL_FREQ - - _double = {4, 5, 10, 11, 13, 14} # ankle + waist indices with factor 2 - kp = np.array([2 * s(k) if i in _double else s(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32) - kd = np.array([2 * d(k) if i in _double else d(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32) - return kp, kd + """SONIC per-joint PD gains (kp, kd), (29,) float32 in IsaacLab joint order.""" + return compute_pd_gains(MOTOR_MODELS, _SONIC_DOUBLE) # Action-feature prefix for the latent-token interface (see _extract_token_from_action). @@ -303,24 +293,11 @@ class SonicRuntime: latent token supplied directly by the policy. """ - def __init__(self, force_cpu: bool = False): + def __init__(self): decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx") - providers = ort_providers(force_cpu=force_cpu) so = make_ort_session_options() - decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers) - - # Report the provider actually bound, not the one requested: ORT silently falls back - # to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a CPU decoder - # drifts the closed-loop heading. Warn loudly so it can't hide. - self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider" - if not force_cpu and not self.use_gpu: - print( - "[SONIC] WARNING: decoder bound to CPUExecutionProvider (CUDA unavailable). " - "Closed-loop replay/control will drift. Ensure libcudnn is on LD_LIBRARY_PATH " - "(site-packages/nvidia/*/lib).", - flush=True, - ) + decoder_sess = ort.InferenceSession(decoder_path, sess_options=so) self.kp, self.kd = compute_kp_kd() self.controller = SonicDecoder(decoder_sess) @@ -342,9 +319,9 @@ class SonicWholeBodyController: control_dt = CONTROL_DT full_body = True - def __init__(self, force_cpu: bool = False): + def __init__(self): logger.info("Loading SONIC whole-body controller...") - self._runtime = SonicRuntime(force_cpu=force_cpu) + self._runtime = SonicRuntime() self.kp = self._runtime.kp self.kd = self._runtime.kd self.controller = self._runtime.controller diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 5901391ba..68c5697fe 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -119,15 +119,33 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray: return gravity_orientation -def ort_providers(force_cpu: bool = False) -> list[str]: - """ONNX Runtime providers, preferring CUDA when available (shared by the ONNX - controllers: SONIC decoder, GR00T). Falls back to CPU.""" - import onnxruntime as ort +# Unitree motor-model parameters shared by the controllers that derive their PD gains +# from motor physics rather than hand-tuning (SONIC decoder, Holosoma). NATURAL_FREQ is +# the target closed-loop stiffness bandwidth (rad/s); MOTOR_ARMATURE is per-model rotor +# inertia (keys are Unitree motor model names). From these: kp = armature * w**2 and +# kd = 4 * armature * w, with an optional x2 factor on stiff joints (ankles/waist). +NATURAL_FREQ = 10.0 * 2.0 * np.pi +MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425} - avail = ort.get_available_providers() - if not force_cpu and "CUDAExecutionProvider" in avail: - return ["CUDAExecutionProvider", "CPUExecutionProvider"] - return ["CPUExecutionProvider"] + +def compute_pd_gains(motor_models, double_indices=()) -> tuple[np.ndarray, np.ndarray]: + """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth. + + ``motor_models`` is a per-joint sequence of Unitree motor model names (in the + controller's own joint order); joints whose index is in ``double_indices`` get a + x2 stiffness/damping factor. Returns two (N,) float32 arrays in that same order. + """ + double = set(double_indices) + + def s(k): + return MOTOR_ARMATURE[k] * NATURAL_FREQ**2 + + def d(k): + return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ + + kp = np.array([2 * s(k) if i in double else s(k) for i, k in enumerate(motor_models)], dtype=np.float32) + kd = np.array([2 * d(k) if i in double else d(k) for i, k in enumerate(motor_models)], dtype=np.float32) + return kp, kd def make_ort_session_options(intra_op_num_threads: int | None = None, inter_op_num_threads: int | None = None):