From fffa42cc5e080e0a4c23102e9c48811c0d8a977e Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Thu, 30 Jul 2026 13:44:24 +0200 Subject: [PATCH] refactor(unitree_g1): load SONIC constants from onnx --- .../controllers/sonic_whole_body.py | 223 +++++------------- src/lerobot/robots/unitree_g1/g1_utils.py | 66 +----- 2 files changed, 60 insertions(+), 229 deletions(-) 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 d1429d2b2..91cf8e42f 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -20,8 +20,8 @@ Pure-Python/ONNX re-implementation of the *decode* half of NVIDIA's SONIC deploy The encoder is intentionally absent: a token-output VLA (e.g. ``nepyope/sonic_walk``) supplies the 64-D latent ``motion_token`` directly each tick, and the SONIC **decoder** maps ``token + recent proprioception history`` to a residual action that is scaled and -added onto ``DEFAULT_ANGLES`` to produce 50 Hz joint-position targets for the robot's PD -controller. +added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets +for the robot's PD controller. Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and **MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils) @@ -30,19 +30,18 @@ convert between them. Quaternions are scalar-first ``(w, x, y, z)``. from __future__ import annotations +import json import logging import numpy as np +import onnx import onnxruntime as ort from huggingface_hub import hf_hub_download from ..g1_utils import ( ISAACLAB_TO_MUJOCO, - MOTOR_ARMATURE, MUJOCO_TO_ISAACLAB, - NATURAL_FREQ, G1_29_JointIndex, - compute_pd_gains, get_gravity_orientation, ) from ..unitree_g1 import lowstate_to_obs @@ -53,61 +52,43 @@ logger = logging.getLogger(__name__) CONTROL_DT = 0.02 # 50 Hz control period (s) TOKEN_DIM = 64 # decoder latent size -# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals -# added on top of this. -DEFAULT_ANGLES = np.array( - [ - -0.312, - 0.0, - 0.0, - 0.669, - -0.363, - 0.0, - -0.312, - 0.0, - 0.0, - 0.669, - -0.363, - 0.0, - 0.0, - 0.0, - 0.0, - 0.2, - 0.2, - 0.0, - 0.6, - 0.0, - 0.0, - 0.0, - 0.2, - -0.2, - 0.0, - 0.6, - 0.0, - 0.0, - 0.0, - ], - dtype=np.float32, -) - -# Per-motor torque limits (N·m), used only for SONIC's residual-action scaling. The -# armature / bandwidth constants and the PD-gain formula are shared (see g1_utils). -EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0} +# SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants +# (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded +# in the ONNX metadata; see upload_sonic_decoder.py for provisioning. The runtime loads the +# model *and* all of these straight from the checkpoint (the Holosoma convention), so no +# motor-physics math happens at deploy time. +DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder" +DECODER_FILENAME = "model_decoder.onnx" +DECODER_INPUT_DIM = 994 # token(64) + 10-frame proprio history + gravity -def _action_scale(k): - """Per-motor residual-action scale (maps policy output to joint-angle delta).""" - return 0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2) +def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID): + """Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint. + Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The + gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the + (64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than + recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``). + """ + decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME) + so = ort.SessionOptions() + so.log_severity_level = 3 # quiet ORT logs + session = ort.InferenceSession(decoder_path, sess_options=so) + dec_dim = int(session.get_inputs()[0].shape[1]) + if dec_dim != DECODER_INPUT_DIM: + raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})") -# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of truth -# for both ACTION_SCALE and compute_kp_kd(). -MOTOR_MODELS = ( - ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2 - + ["7520_14", "5020", "5020"] - + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2 -) -ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32) # (29,) IsaacLab + meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props} + required = ("kp", "kd", "default_angles", "action_scale", "neutral_token") + missing = [k for k in required if k not in meta] + if missing: + raise ValueError( + f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; " + "re-run upload_sonic_decoder.py to (re)provision the checkpoint." + ) + arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required} + logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"])) + return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"] def _to_mujoco(a): @@ -121,18 +102,6 @@ def _to_mujoco(a): return a[MUJOCO_TO_ISAACLAB] -DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES) - - -# Ankle + waist joint indices (IsaacLab order) that get a x2 stiffness/damping factor. -_SONIC_DOUBLE = {4, 5, 10, 11, 13, 14} - - -def compute_kp_kd(): - """SONIC per-joint PD gains (kp, kd), (29,) float32 in IsaacLab joint order.""" - return compute_pd_gains(MOTOR_MODELS, _SONIC_DOUBLE) - - # Action-feature prefix for the latent-token interface (see _extract_token_from_action). TOKEN_ACTION_PREFIX = "motion_token" # nosec B105 - feature-key prefix, not a secret # Proprio-state prefix for the token interface: the robot echoes the last commanded token @@ -159,83 +128,6 @@ def token_state_key(i: int) -> str: # eases in without a snap on the first command. INIT_RAMP_S = 3.0 -# Neutral ("zero pose") SONIC token, held by token_mode until the first real token arrives. -# Captured from the encoder's own output while the robot stood idle in sim: the encoder is -# an FSQ bottleneck (~5 bit/dim, Div(16)), so its tokens live on the 1/16 grid. We store the -# integer FSQ codes and rescale by 1/16, giving an exact on-grid token -- unlike the literal -# all-zero token, which is off the learned manifold and decodes to a slightly goofy stance. -# This one decodes to a stable, natural standing pose. -_NEUTRAL_TOKEN_CODES = np.array( - [ - -1, - 3, - 1, - -1, - 1, - -3, - 6, - 1, - 1, - 1, - -2, - -4, - -2, - 0, - -3, - -1, - 2, - -1, - -3, - -5, - 3, - 1, - 1, - -4, - -1, - -1, - 1, - -7, - 0, - 1, - 2, - -2, - 5, - -2, - -2, - -4, - 0, - -1, - 3, - -1, - 0, - -5, - -1, - 0, - -4, - 0, - 0, - -1, - -1, - 2, - -2, - 1, - 3, - 3, - 1, - 0, - 0, - 6, - 0, - -7, - 3, - 0, - 2, - -2, - ], - dtype=np.float32, -) -NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0 # FSQ Div(16): integer codes -> on-grid token - def _extract_token_from_action(action: dict | None) -> np.ndarray | None: """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None. @@ -256,16 +148,17 @@ class SonicDecoder: """Runs the SONIC decoder ONNX model and owns the proprioception history. Each tick it appends the latest robot state to 10-frame history buffers, then maps the - supplied 64-D ``token`` + that history to a residual action added onto - ``DEFAULT_ANGLES``. The encoder is bypassed entirely (token supplied by the policy). + supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``. + The encoder is bypassed entirely (token supplied by the policy). ``default_angles`` and + ``action_scale`` are (29,) float32 in IsaacLab order, loaded from the checkpoint. """ - def __init__(self, decoder): + def __init__(self, decoder, default_angles, action_scale): self.decoder = decoder self.decoder_input = decoder.get_inputs()[0].name - dec_dim = int(decoder.get_inputs()[0].shape[1]) - if dec_dim != 994: - raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected 994)") + self.default_angles = np.asarray(default_angles, np.float32) + self.action_scale = np.asarray(action_scale, np.float32) + self.default_angles_mj = _to_mujoco(self.default_angles) self.token = np.zeros(TOKEN_DIM, np.float32) self.last_action_mj = np.zeros(29, np.float32) self.h_q_mj = [np.zeros(29, np.float32)] * 10 @@ -293,7 +186,7 @@ class SonicDecoder: quat = quat / (np.linalg.norm(quat) + 1e-8) q_mj = _to_mujoco(q) dq_mj = _to_mujoco(dq) - self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + self.h_q_mj[:-1] + self.h_q_mj = [q_mj - self.default_angles_mj] + self.h_q_mj[:-1] self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1] self.h_ang = [ang.copy()] + self.h_ang[:-1] self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1] @@ -335,7 +228,7 @@ class SonicDecoder: jnames = [m.name for m in G1_29_JointIndex] q = np.array( [ - robot_obs.get(f"{n}.q", DEFAULT_ANGLES[m.value]) + robot_obs.get(f"{n}.q", self.default_angles[m.value]) for m, n in zip(G1_29_JointIndex, jnames, strict=False) ], np.float32, @@ -358,7 +251,7 @@ class SonicDecoder: .astype(np.float32) ) self.last_action_mj = action_mj.copy() - target = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * ACTION_SCALE + target = self.default_angles + action_mj[ISAACLAB_TO_MUJOCO] * self.action_scale if debug: delta = target - q logger.debug( @@ -379,14 +272,10 @@ class SonicRuntime: """ def __init__(self): - decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx") - - so = ort.SessionOptions() - so.log_severity_level = 3 # quiet ORT logs - decoder_sess = ort.InferenceSession(decoder_path, sess_options=so) - - self.kp, self.kd = compute_kp_kd() - self.controller = SonicDecoder(decoder_sess) + decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder() + self.default_angles = default_angles + self.neutral_token = neutral_token + self.controller = SonicDecoder(decoder_sess, default_angles, action_scale) @property def pipeline(self): @@ -411,6 +300,8 @@ class SonicWholeBodyController: self.kp = self._runtime.kp self.kd = self._runtime.kd self.controller = self._runtime.controller + self._default_angles = self._runtime.default_angles + self._neutral_token = self._runtime.neutral_token # Startup blend: ease from the robot's initial pose into the first commanded policy # targets over INIT_RAMP_S (captured on the first control tick). @@ -442,7 +333,7 @@ class SonicWholeBodyController: if self._init_step == 0: # Capture the robot's actual pose as the interpolation start point. self._start_pose = { - f"{m.name}.q": float(obs.get(f"{m.name}.q", DEFAULT_ANGLES[m.value])) + f"{m.name}.q": float(obs.get(f"{m.name}.q", self._default_angles[m.value])) for m in G1_29_JointIndex } self._init_step += 1 @@ -466,9 +357,9 @@ class SonicWholeBodyController: if token is not None: self._last_token = token elif self._last_token is None and self.token_mode: - # Token-driven deploy, but no token has arrived yet: hold the captured neutral - # token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural standing pose. - self._last_token = NEUTRAL_TOKEN.copy() + # Token-driven deploy, but no token has arrived yet: hold the checkpoint's neutral + # token, which the decoder maps to a stable, natural standing pose. + self._last_token = self._neutral_token.copy() if self._last_token is None: # No token yet and not in token_mode: hold (keep last target). return {} diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index c92c890f4..83f5eae31 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -60,40 +60,9 @@ ISAACLAB_TO_MUJOCO = np.array( ], dtype=np.int32, ) -MUJOCO_TO_ISAACLAB = np.array( - [ - 0, - 6, - 12, - 1, - 7, - 13, - 2, - 8, - 14, - 3, - 9, - 15, - 22, - 4, - 10, - 16, - 23, - 5, - 11, - 17, - 24, - 18, - 25, - 19, - 26, - 20, - 27, - 21, - 28, - ], - dtype=np.int32, -) +# The two orderings are inverses of each other, so derive one from the other (argsort) to +# guarantee they can never drift out of sync. +MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32) REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry") REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16)) @@ -115,35 +84,6 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray: return gravity_orientation -# Unitree motor-model parameters shared by the controllers that derive their PD gains -# from motor physics rather than hand-tuning (SONIC decoder, Holosoma). NATURAL_FREQ is -# the target closed-loop stiffness bandwidth (rad/s); MOTOR_ARMATURE is per-model rotor -# inertia (keys are Unitree motor model names). From these: kp = armature * w**2 and -# kd = 4 * armature * w, with an optional x2 factor on stiff joints (ankles/waist). -NATURAL_FREQ = 10.0 * 2.0 * np.pi -MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425} - - -def compute_pd_gains(motor_models, double_indices=()) -> tuple[np.ndarray, np.ndarray]: - """Derive per-joint PD gains (kp, kd) from motor armature and target bandwidth. - - ``motor_models`` is a per-joint sequence of Unitree motor model names (in the - controller's own joint order); joints whose index is in ``double_indices`` get a - x2 stiffness/damping factor. Returns two (N,) float32 arrays in that same order. - """ - double = set(double_indices) - - def s(k): - return MOTOR_ARMATURE[k] * NATURAL_FREQ**2 - - def d(k): - return 4.0 * MOTOR_ARMATURE[k] * NATURAL_FREQ - - kp = np.array([2 * s(k) if i in double else s(k) for i, k in enumerate(motor_models)], dtype=np.float32) - kd = np.array([2 * d(k) if i in double else d(k) for i, k in enumerate(motor_models)], dtype=np.float32) - return kp, kd - - class G1_29_JointArmIndex(IntEnum): # Left arm kLeftShoulderPitch = 15