admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-30 13:09:40 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(unitree_g1): load SONIC constants from onnx

Browse Source
This commit is contained in:
Martino Russi
2026-07-30 13:44:24 +02:00
parent f33089b027
commit fffa42cc5e
2 changed files with 60 additions and 229 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py
+60 -169
View File
@@ -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``) 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** 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 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 added onto the standing pose (``default_angles``) to produce 50 Hz joint-position targets
controller. for the robot's PD controller.
Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and
**MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` (in g1_utils) **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 from __future__ import annotations
import json
import logging import logging
import numpy as np import numpy as np
import onnx
import onnxruntime as ort import onnxruntime as ort
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
from ..g1_utils import ( from ..g1_utils import (
ISAACLAB_TO_MUJOCO, ISAACLAB_TO_MUJOCO,
MOTOR_ARMATURE,
MUJOCO_TO_ISAACLAB, MUJOCO_TO_ISAACLAB,
NATURAL_FREQ,
G1_29_JointIndex, G1_29_JointIndex,
compute_pd_gains,
get_gravity_orientation, get_gravity_orientation,
) )
from ..unitree_g1 import lowstate_to_obs from ..unitree_g1 import lowstate_to_obs
@@ -53,61 +52,43 @@ logger = logging.getLogger(__name__)
CONTROL_DT = 0.02 # 50 Hz control period (s) CONTROL_DT = 0.02 # 50 Hz control period (s)
TOKEN_DIM = 64 # decoder latent size TOKEN_DIM = 64 # decoder latent size
# Nominal standing pose (rad), 29 joints in IsaacLab order. Decoder actions are residuals # SONIC decoder checkpoint: NVIDIA's decoder ONNX re-packaged with its deploy constants
# added on top of this. # (kp/kd PD gains, the standing pose default_angles, and the residual action_scale) embedded
DEFAULT_ANGLES = np.array( # 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
-0.312, # motor-physics math happens at deploy time.
0.0, DEFAULT_SONIC_REPO_ID = "lerobot/sonic_decoder"
0.0, DECODER_FILENAME = "model_decoder.onnx"
0.669, DECODER_INPUT_DIM = 994 # token(64) + 10-frame proprio history + gravity
-0.363,
0.0,
-0.312, def load_sonic_decoder(repo_id: str = DEFAULT_SONIC_REPO_ID):
0.0, """Load the SONIC decoder ONNX and its baked-in deploy constants from the checkpoint.
0.0,
0.669, Returns ``(decoder_session, kp, kd, default_angles, action_scale, neutral_token)``. The
-0.363, gains/pose/scale are (29,) float32 in IsaacLab joint order and ``neutral_token`` is the
0.0, (64,) float32 idle latent -- all read from the ONNX ``metadata_props`` rather than
0.0, recomputed/hardcoded at deploy time (mirrors ``holosoma_locomotion.load_policy``).
0.0, """
0.0, decoder_path = hf_hub_download(repo_id=repo_id, filename=DECODER_FILENAME)
0.2, so = ort.SessionOptions()
0.2, so.log_severity_level = 3 # quiet ORT logs
0.0, session = ort.InferenceSession(decoder_path, sess_options=so)
0.6, dec_dim = int(session.get_inputs()[0].shape[1])
0.0, if dec_dim != DECODER_INPUT_DIM:
0.0, raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected {DECODER_INPUT_DIM})")
0.0,
0.2, meta = {p.key: p.value for p in onnx.load(decoder_path, load_external_data=False).metadata_props}
-0.2, required = ("kp", "kd", "default_angles", "action_scale", "neutral_token")
0.0, missing = [k for k in required if k not in meta]
0.6, if missing:
0.0, raise ValueError(
0.0, f"SONIC decoder ONNX at {repo_id} is missing metadata {missing}; "
0.0, "re-run upload_sonic_decoder.py to (re)provision the checkpoint."
],
dtype=np.float32,
) )
arr = {k: np.array(json.loads(meta[k]), dtype=np.float32) for k in required}
# Per-motor torque limits (N·m), used only for SONIC's residual-action scaling. The logger.info("Loaded SONIC deploy constants from %s (%d joints)", repo_id, len(arr["kp"]))
# armature / bandwidth constants and the PD-gain formula are shared (see g1_utils). return session, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
def _action_scale(k):
"""Per-motor residual-action scale (maps policy output to joint-angle delta)."""
return 0.25 * EFFORT[k] / (MOTOR_ARMATURE[k] * NATURAL_FREQ**2)
# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of truth
# for both ACTION_SCALE and compute_kp_kd().
MOTOR_MODELS = (
["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2
+ ["7520_14", "5020", "5020"]
+ ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2
)
ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32) # (29,) IsaacLab
def _to_mujoco(a): def _to_mujoco(a):
@@ -121,18 +102,6 @@ def _to_mujoco(a):
return a[MUJOCO_TO_ISAACLAB] 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). # 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 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 # 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. # eases in without a snap on the first command.
INIT_RAMP_S = 3.0 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: def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
"""Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or 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. """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 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 supplied 64-D ``token`` + that history to a residual action added onto ``default_angles``.
``DEFAULT_ANGLES``. The encoder is bypassed entirely (token supplied by the policy). 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 = decoder
self.decoder_input = decoder.get_inputs()[0].name self.decoder_input = decoder.get_inputs()[0].name
dec_dim = int(decoder.get_inputs()[0].shape[1]) self.default_angles = np.asarray(default_angles, np.float32)
if dec_dim != 994: self.action_scale = np.asarray(action_scale, np.float32)
raise RuntimeError(f"Unexpected decoder input dim {dec_dim} (expected 994)") self.default_angles_mj = _to_mujoco(self.default_angles)
self.token = np.zeros(TOKEN_DIM, np.float32) self.token = np.zeros(TOKEN_DIM, np.float32)
self.last_action_mj = np.zeros(29, np.float32) self.last_action_mj = np.zeros(29, np.float32)
self.h_q_mj = [np.zeros(29, np.float32)] * 10 self.h_q_mj = [np.zeros(29, np.float32)] * 10
@@ -293,7 +186,7 @@ class SonicDecoder:
quat = quat / (np.linalg.norm(quat) + 1e-8) quat = quat / (np.linalg.norm(quat) + 1e-8)
q_mj = _to_mujoco(q) q_mj = _to_mujoco(q)
dq_mj = _to_mujoco(dq) 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_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
self.h_ang = [ang.copy()] + self.h_ang[:-1] self.h_ang = [ang.copy()] + self.h_ang[:-1]
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1] self.h_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] jnames = [m.name for m in G1_29_JointIndex]
q = np.array( 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) for m, n in zip(G1_29_JointIndex, jnames, strict=False)
], ],
np.float32, np.float32,
@@ -358,7 +251,7 @@ class SonicDecoder:
.astype(np.float32) .astype(np.float32)
) )
self.last_action_mj = action_mj.copy() 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: if debug:
delta = target - q delta = target - q
logger.debug( logger.debug(
@@ -379,14 +272,10 @@ class SonicRuntime:
""" """
def __init__(self): def __init__(self):
decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx") decoder_sess, self.kp, self.kd, default_angles, action_scale, neutral_token = load_sonic_decoder()
self.default_angles = default_angles
so = ort.SessionOptions() self.neutral_token = neutral_token
so.log_severity_level = 3 # quiet ORT logs self.controller = SonicDecoder(decoder_sess, default_angles, action_scale)
decoder_sess = ort.InferenceSession(decoder_path, sess_options=so)
self.kp, self.kd = compute_kp_kd()
self.controller = SonicDecoder(decoder_sess)
@property @property
def pipeline(self): def pipeline(self):
@@ -411,6 +300,8 @@ class SonicWholeBodyController:
self.kp = self._runtime.kp self.kp = self._runtime.kp
self.kd = self._runtime.kd self.kd = self._runtime.kd
self.controller = self._runtime.controller 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 # 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). # targets over INIT_RAMP_S (captured on the first control tick).
@@ -442,7 +333,7 @@ class SonicWholeBodyController:
if self._init_step == 0: if self._init_step == 0:
# Capture the robot's actual pose as the interpolation start point. # Capture the robot's actual pose as the interpolation start point.
self._start_pose = { 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 for m in G1_29_JointIndex
} }
self._init_step += 1 self._init_step += 1
@@ -466,9 +357,9 @@ class SonicWholeBodyController:
if token is not None: if token is not None:
self._last_token = token self._last_token = token
elif self._last_token is None and self.token_mode: elif self._last_token is None and self.token_mode:
# Token-driven deploy, but no token has arrived yet: hold the captured neutral # Token-driven deploy, but no token has arrived yet: hold the checkpoint's neutral
# token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural standing pose. # token, which the decoder maps to a stable, natural standing pose.
self._last_token = NEUTRAL_TOKEN.copy() self._last_token = self._neutral_token.copy()
if self._last_token is None: if self._last_token is None:
# No token yet and not in token_mode: hold (keep last target). # No token yet and not in token_mode: hold (keep last target).
return {} return {}
src/lerobot/robots/unitree_g1/g1_utils.py
+3 -63
View File
@@ -60,40 +60,9 @@ ISAACLAB_TO_MUJOCO = np.array(
], ],
dtype=np.int32, dtype=np.int32,
) )
MUJOCO_TO_ISAACLAB = np.array( # The two orderings are inverses of each other, so derive one from the other (argsort) to
[ # guarantee they can never drift out of sync.
0, MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32)
6,
12,
1,
7,
13,
2,
8,
14,
3,
9,
15,
22,
4,
10,
16,
23,
5,
11,
17,
24,
18,
25,
19,
26,
20,
27,
21,
28,
],
dtype=np.int32,
)
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry") REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16)) 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 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): class G1_29_JointArmIndex(IntEnum):
# Left arm # Left arm
kLeftShoulderPitch = 15 kLeftShoulderPitch = 15