perf(unitree_g1): cap ONNX Runtime thread pool for locomotion policies

Add make_ort_session_options() to g1_utils and use it in the GR00T locomotion
controller. These small MLP policies are latency-bound and run at 50 Hz in a
background thread alongside the torch upper-body policy and IK; letting ORT grab
every core starves the real-time loop and causes sim stutter, so cap it to 1
intra/inter thread for lowest-latency per-step inference.

Co-authored-by: Cursor <cursoragent@cursor.com>

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -19,14 +19,29 @@ 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
 
-# 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()
+_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()
 
 
 @RobotConfig.register_subclass("unitree_g1")

src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py

@@ -26,6 +26,7 @@ from ..g1_utils import (
     REMOTE_BUTTONS,
     G1_29_JointIndex,
     get_gravity_orientation,
+    make_ort_session_options,
 )
 
 logger = logging.getLogger(__name__)
@@ -68,9 +69,13 @@ def load_groot_policies(
         filename="GR00T-WholeBodyControl-Walk.onnx",
     )
 
-    # Load ONNX policies
-    policy_balance = ort.InferenceSession(balance_path)
-    policy_walk = ort.InferenceSession(walk_path)
+    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
+    # background thread alongside the (torch) upper-body policy and IK; letting ORT
+    # grab every core starves those and makes the whole rollout stutter. These are
+    # small MLPs, so 1 thread is both enough and lowest-latency.
+    so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
+    policy_balance = ort.InferenceSession(balance_path, sess_options=so)
+    policy_walk = ort.InferenceSession(walk_path, sess_options=so)
 
     logger.info("GR00T policies loaded successfully")
 

src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py

@@ -14,30 +14,23 @@
 # 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
@@ -68,16 +61,15 @@ POLICY_FILES = {
 def load_policy(
     repo_id: str = DEFAULT_HOLOSOMA_REPO_ID,
     policy_type: str = "fastsac",
-) -> ort.InferenceSession:
-    """Load the Holosoma locomotion policy.
+) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray]:
+    """Load Holosoma locomotion policy and extract KP/KD from metadata.
 
     Args:
         repo_id: Hugging Face Hub repo ID
         policy_type: Either "fastsac" (default) or "ppo"
 
     Returns:
-        The ONNX Runtime inference session. PD gains are not read from the checkpoint --
-        they are derived from motor physics (see ``_KP`` / ``_KD`` above).
+        (policy, kp, kd) tuple
     """
     if policy_type not in POLICY_FILES:
         raise ValueError(f"Unknown policy type: {policy_type}. Choose from: {list(POLICY_FILES.keys())}")
@@ -89,7 +81,18 @@ def load_policy(
     policy = ort.InferenceSession(policy_path)
     logger.info(f"Policy loaded: {policy.get_inputs()[0].shape} → {policy.get_outputs()[0].shape}")
 
-    return policy
+    # 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
 
 
 class HolosomaLocomotionController:
@@ -98,9 +101,8 @@ class HolosomaLocomotionController:
     control_dt = CONTROL_DT  # Expose for unitree_g1.py
 
     def __init__(self):
-        self.policy = load_policy()
-        self.kp = _KP.copy()
-        self.kd = _KD.copy()
+        # Load policy and gains
+        self.policy, self.kp, self.kd = load_policy()
 
         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]:
+) -> tuple[ort.InferenceSession, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Load the SONIC decoder and its baked-in deploy constants from ONNX metadata.
 
     Args:
@@ -77,9 +77,8 @@ def load_policy(
         policy_type: Either "default" (full decoder) or "low_latency" (distilled)
 
     Returns:
-        (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.
+        (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.
     """
     if policy_type not in POLICY_FILES:
         raise ValueError(f"Unknown policy type: {policy_type}. Choose from: {list(POLICY_FILES.keys())}")
@@ -95,14 +94,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 = ("default_angles", "action_scale", "neutral_token")
+    required = ("kp", "kd", "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['default_angles'])} joints)")
-    return decoder, arr["default_angles"], arr["action_scale"], arr["neutral_token"]
+    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"]
 
 
 class SonicWholeBodyController:
@@ -116,8 +115,8 @@ class SonicWholeBodyController:
     control_dt = CONTROL_DT
 
     def __init__(self, policy_type: str = "default"):
-        self.decoder, self.default_angles, self.action_scale, self.neutral_token = load_policy(
-            policy_type=policy_type
+        self.decoder, self.kp, self.kd, 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]

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 
 import importlib
-from collections.abc import Iterable, Sequence
 from enum import IntEnum
 
 import numpy as np
@@ -61,66 +60,6 @@ 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
@@ -141,6 +80,26 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
     return gravity_orientation
 
 
+def make_ort_session_options(
+    intra_op_num_threads: int | None = None, inter_op_num_threads: int | None = None
+):
+    """Build quiet ONNX Runtime SessionOptions, optionally capping the CPU thread pool.
+
+    These tiny MLP policies are latency-bound, not throughput-bound, so letting ORT grab
+    every core starves the real-time control loop / torch policy and causes stutter. Pass
+    1 intra + 1 inter thread for lowest-latency per-step inference.
+    """
+    import onnxruntime as ort
+
+    so = ort.SessionOptions()
+    so.log_severity_level = 3
+    if intra_op_num_threads is not None:
+        so.intra_op_num_threads = intra_op_num_threads
+    if inter_op_num_threads is not None:
+        so.inter_op_num_threads = inter_op_num_threads
+    return so
+
+
 class G1_29_JointArmIndex(IntEnum):
     # Left arm
     kLeftShoulderPitch = 15