From 6662431923c95ad8f8965d2256e167bfbec444f7 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 31 Jul 2026 13:48:27 +0200 Subject: [PATCH] feat(unitree_g1): derive PD gains from motor physics Add compute_pd_gains / compute_action_scale to g1_utils (kp = armature*w^2, kd = 4*armature*w, with a 2x factor on stiff ankle/waist joints), plus the G1 motor-model layout and shared NATURAL_FREQ/armature/effort constants. The config now derives its default kp/kd from these instead of hand-tuned values. Both controllers stop reading kp/kd from ONNX metadata: SONIC uses the config gains (identical to its baked values), and Holosoma derives them with a hip-pitch override (7520_14 instead of 7520_22) to match its checkpoint. Co-authored-by: Cursor --- .../robots/unitree_g1/config_unitree_g1.py | 29 +++------ .../controllers/holosoma_locomotion.py | 36 ++++++----- .../controllers/sonic_whole_body.py | 17 +++--- src/lerobot/robots/unitree_g1/g1_utils.py | 61 +++++++++++++++++++ 4 files changed, 94 insertions(+), 49 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/config_unitree_g1.py b/src/lerobot/robots/unitree_g1/config_unitree_g1.py index 3fbeb37fe..0bb2efe57 100644 --- a/src/lerobot/robots/unitree_g1/config_unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/config_unitree_g1.py @@ -19,29 +19,14 @@ from dataclasses import dataclass, field from lerobot.cameras import CameraConfig from ..config import RobotConfig +from .g1_utils import G1_MOTOR_MODELS, STIFF_JOINT_INDICES, compute_pd_gains -_GAINS: dict[str, dict[str, list[float]]] = { - "left_leg": { - "kp": [150, 150, 150, 300, 40, 40], - "kd": [2, 2, 2, 4, 2, 2], - }, # pitch, roll, yaw, knee, ankle_pitch, ankle_roll - "right_leg": {"kp": [150, 150, 150, 300, 40, 40], "kd": [2, 2, 2, 4, 2, 2]}, - "waist": {"kp": [250, 250, 250], "kd": [5, 5, 5]}, # yaw, roll, pitch - "left_arm": {"kp": [50, 50, 80, 80], "kd": [3, 3, 3, 3]}, # shoulder_pitch/roll/yaw, elbow - "left_wrist": {"kp": [40, 40, 40], "kd": [1.5, 1.5, 1.5]}, # roll, pitch, yaw - "right_arm": {"kp": [50, 50, 80, 80], "kd": [3, 3, 3, 3]}, - "right_wrist": {"kp": [40, 40, 40], "kd": [1.5, 1.5, 1.5]}, -} - - -def _build_gains() -> tuple[list[float], list[float]]: - """Build kp and kd lists from body-part groupings.""" - kp = [v for g in _GAINS.values() for v in g["kp"]] - kd = [v for g in _GAINS.values() for v in g["kd"]] - return kp, kd - - -_DEFAULT_KP, _DEFAULT_KD = _build_gains() +# Default PD gains are derived from Unitree motor physics (armature + target natural +# frequency; see g1_utils.compute_pd_gains) rather than hand-tuned. Controllers may tweak +# individual joints on top of these. +_kp, _kd = compute_pd_gains(G1_MOTOR_MODELS, STIFF_JOINT_INDICES) +_DEFAULT_KP: list[float] = _kp.tolist() +_DEFAULT_KD: list[float] = _kd.tolist() @RobotConfig.register_subclass("unitree_g1") diff --git a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py index da9ef2ddb..12c67256f 100644 --- a/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py +++ b/src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py @@ -14,23 +14,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -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 ( + G1_MOTOR_MODELS, REMOTE_AXES, + STIFF_JOINT_INDICES, G1_29_JointArmIndex, G1_29_JointIndex, + compute_pd_gains, get_gravity_orientation, ) logger = logging.getLogger(__name__) +# Holosoma runs the hip-pitch joints (left=0, right=6) at the softer 7520_14 gain rather than +# the 7520_22 default; every other joint uses the derived motor-physics gains. +_HOLOSOMA_MOTOR_MODELS = list(G1_MOTOR_MODELS) +_HOLOSOMA_MOTOR_MODELS[0] = _HOLOSOMA_MOTOR_MODELS[6] = "7520_14" +_KP, _KD = compute_pd_gains(_HOLOSOMA_MOTOR_MODELS, STIFF_JOINT_INDICES) + DEFAULT_ANGLES = np.zeros(29, dtype=np.float32) DEFAULT_ANGLES[[0, 6]] = -0.312 # Hip pitch DEFAULT_ANGLES[[3, 9]] = 0.669 # Knee @@ -61,15 +68,16 @@ POLICY_FILES = { 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. +) -> ort.InferenceSession: + """Load the Holosoma locomotion policy. Args: repo_id: Hugging Face Hub repo ID policy_type: Either "fastsac" (default) or "ppo" Returns: - (policy, kp, kd) tuple + The ONNX Runtime inference session. PD gains are not read from the checkpoint -- + they are derived from motor physics (see ``_KP`` / ``_KD`` above). """ if policy_type not in POLICY_FILES: raise ValueError(f"Unknown policy type: {policy_type}. Choose from: {list(POLICY_FILES.keys())}") @@ -81,18 +89,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)") - - return policy, kp, kd + return policy class HolosomaLocomotionController: @@ -101,8 +98,9 @@ class HolosomaLocomotionController: control_dt = CONTROL_DT # Expose for unitree_g1.py def __init__(self): - # Load policy and gains - self.policy, self.kp, self.kd = load_policy() + self.policy = load_policy() + self.kp = _KP.copy() + self.kd = _KD.copy() self.default_angles = DEFAULT_ANGLES self.cmd = np.zeros(3, dtype=np.float32) 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 bf96eaf51..fc12a029a 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -69,7 +69,7 @@ POLICY_FILES = { def load_policy( repo_id: str = DEFAULT_SONIC_REPO_ID, policy_type: str = "default", -) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: +) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray, np.ndarray]: """Load the SONIC decoder and its baked-in deploy constants from ONNX metadata. Args: @@ -77,8 +77,9 @@ def load_policy( policy_type: Either "default" (full decoder) or "low_latency" (distilled) Returns: - (decoder, kp, kd, default_angles, action_scale, neutral_token) tuple. The gains/pose/ - scale are (29,) float32 in IsaacLab joint order; neutral_token is the (64,) idle latent. + (decoder, default_angles, action_scale, neutral_token) tuple. The pose/scale are (29,) + float32 in IsaacLab joint order; neutral_token is the (64,) idle latent. PD gains are + not read here -- SONIC uses the robot's config gains, which equal the derived defaults. """ if policy_type not in POLICY_FILES: raise ValueError(f"Unknown policy type: {policy_type}. Choose from: {list(POLICY_FILES.keys())}") @@ -94,14 +95,14 @@ def load_policy( model = onnx.load(decoder_path, load_external_data=False) metadata = {prop.key: prop.value for prop in model.metadata_props} - required = ("kp", "kd", "default_angles", "action_scale", "neutral_token") + required = ("default_angles", "action_scale", "neutral_token") missing = [k for k in required if k not in metadata] if missing: raise ValueError(f"ONNX model must contain {list(required)} in metadata (missing {missing})") arr = {k: np.array(json.loads(metadata[k]), dtype=np.float32) for k in required} - logger.info(f"Loaded SONIC deploy constants from ONNX ({len(arr['kp'])} joints)") - return decoder, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"] + logger.info(f"Loaded SONIC deploy constants from ONNX ({len(arr['default_angles'])} joints)") + return decoder, arr["default_angles"], arr["action_scale"], arr["neutral_token"] class SonicWholeBodyController: @@ -115,8 +116,8 @@ class SonicWholeBodyController: control_dt = CONTROL_DT def __init__(self, policy_type: str = "default"): - self.decoder, self.kp, self.kd, self.default_angles, self.action_scale, self.neutral_token = ( - load_policy(policy_type=policy_type) + self.decoder, self.default_angles, self.action_scale, self.neutral_token = load_policy( + policy_type=policy_type ) self.decoder_input = self.decoder.get_inputs()[0].name self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB] diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 5d1d4a018..c284b2c1b 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -15,6 +15,7 @@ # limitations under the License. import importlib +from collections.abc import Iterable, Sequence from enum import IntEnum import numpy as np @@ -60,6 +61,66 @@ ISAACLAB_TO_MUJOCO = np.array( ) MUJOCO_TO_ISAACLAB = np.argsort(ISAACLAB_TO_MUJOCO).astype(np.int32) +# ── G1 PD-gain derivation from Unitree motor physics (shared by SONIC + Holosoma) ── +# Rather than hand-tuning, the PD gains (kp/kd) and the residual action_scale are derived +# from motor physics: each joint is one of four Unitree motor models with a known rotor +# armature (inertia) and peak effort. For a target closed-loop natural frequency w, a +# critically-damped second-order response gives kp = armature * w**2 and kd = 4 * armature +# * w; stiff joints (ankles/waist) optionally get a 2x factor. Callers pass their own +# controller's per-joint motor-model sequence (SONIC uses IsaacLab order, Holosoma its +# own), so the returned arrays come back in that same joint order. + +# Target closed-loop natural frequency (rad/s), i.e. 10 Hz. +NATURAL_FREQ = 10.0 * 2.0 * np.pi + +# Rotor armature (kg·m²) and peak effort (N·m) per Unitree motor model. +G1_MOTOR_ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425} +G1_MOTOR_EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0} + +# Per-joint Unitree motor model in standard G1_29_JointIndex order (legs, waist, arms), and +# the stiff joints (ankles + waist roll/pitch) that receive a 2x gain factor. Controllers may +# override individual joints (e.g. Holosoma softens the hip-pitch joints) on top of these. +G1_MOTOR_MODELS = ( + ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] # left leg + + ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] # right leg + + ["7520_14", "5020", "5020"] # waist: yaw, roll, pitch + + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] # left arm + + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] # right arm +) +STIFF_JOINT_INDICES = frozenset({4, 5, 10, 11, 13, 14}) + + +def compute_pd_gains( + motor_models: Sequence[str], double_indices: Iterable[int] = () +) -> tuple[np.ndarray, np.ndarray]: + """Derive PD gains (kp, kd) from Unitree motor physics. + + ``motor_models`` is a per-joint sequence of Unitree motor-model names in the caller's + own joint order; joints whose index is in ``double_indices`` get a 2x stiffness/damping + factor (the stiff ankle/waist joints). ``kp = armature * w**2`` and ``kd = 4 * armature + * w`` (critically damped). Returns ``(kp, kd)``, each an (N,) float32 array in the same + order as ``motor_models``. + """ + double = set(double_indices) + armature = np.array([G1_MOTOR_ARMATURE[m] for m in motor_models], dtype=np.float32) + factor = np.array([2.0 if i in double else 1.0 for i in range(len(motor_models))], dtype=np.float32) + kp = factor * armature * NATURAL_FREQ**2 + kd = factor * 4.0 * armature * NATURAL_FREQ + return kp, kd + + +def compute_action_scale(motor_models: Sequence[str]) -> np.ndarray: + """Derive the residual action scale from Unitree motor physics. + + ``action_scale = 0.25 * effort / (armature * w**2)``, returned as an (N,) float32 array + in the same order as ``motor_models``. Maps a policy's residual output to a joint-angle + delta added on top of the standing pose (``default_angles``). + """ + armature = np.array([G1_MOTOR_ARMATURE[m] for m in motor_models], dtype=np.float32) + effort = np.array([G1_MOTOR_EFFORT[m] for m in motor_models], dtype=np.float32) + return 0.25 * effort / (armature * NATURAL_FREQ**2) + + REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry") REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16)) REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS