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 <cursoragent@cursor.com>

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]]] = {
+# Default PD gains are derived from Unitree motor physics (armature + target natural
-    "left_leg": {
+# frequency; see g1_utils.compute_pd_gains) rather than hand-tuned. Controllers may tweak
-        "kp": [150, 150, 150, 300, 40, 40],
+# individual joints on top of these.
-        "kd": [2, 2, 2, 4, 2, 2],
+_kp, _kd = compute_pd_gains(G1_MOTOR_MODELS, STIFF_JOINT_INDICES)
-    },  # pitch, roll, yaw, knee, ankle_pitch, ankle_roll
+_DEFAULT_KP: list[float] = _kp.tolist()
-    "right_leg": {"kp": [150, 150, 150, 300, 40, 40], "kd": [2, 2, 2, 4, 2, 2]},
+_DEFAULT_KD: list[float] = _kd.tolist()
-    "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()
 
 
 @RobotConfig.register_subclass("unitree_g1")

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]:
+) -> ort.InferenceSession:
-    """Load Holosoma locomotion policy and extract KP/KD from metadata.
+    """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
+    return policy
-    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
 
 
 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 = load_policy()
-        self.policy, self.kp, self.kd = load_policy()
+        self.kp = _KP.copy()
+        self.kd = _KD.copy()
 
         self.default_angles = DEFAULT_ANGLES
         self.cmd = np.zeros(3, dtype=np.float32)

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/
+        (decoder, default_angles, action_scale, neutral_token) tuple. The pose/scale are (29,)
-        scale are (29,) float32 in IsaacLab joint order; neutral_token is the (64,) idle latent.
+        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)")
+    logger.info(f"Loaded SONIC deploy constants from ONNX ({len(arr['default_angles'])} joints)")
-    return decoder, arr["kp"], arr["kd"], arr["default_angles"], arr["action_scale"], arr["neutral_token"]
+    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 = (
+        self.decoder, self.default_angles, self.action_scale, self.neutral_token = load_policy(
-            load_policy(policy_type=policy_type)
+            policy_type=policy_type
         )
         self.decoder_input = self.decoder.get_inputs()[0].name
         self.default_angles_mj = self.default_angles[MUJOCO_TO_ISAACLAB]

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