From 9e884da74aa1f0965a8a9bce197713230e6ecef2 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Fri, 31 Jul 2026 14:19:51 +0200 Subject: [PATCH] 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 --- .../controllers/gr00t_locomotion.py | 11 +++++++--- src/lerobot/robots/unitree_g1/g1_utils.py | 20 +++++++++++++++++++ 2 files changed, 28 insertions(+), 3 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py b/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py index f293a0b1c..cdc0a10b9 100644 --- a/src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py +++ b/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") diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 5d1d4a018..f4bf92d0b 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -80,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