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..7e84183f6 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -43,6 +43,7 @@ from ..g1_utils import ( MUJOCO_TO_ISAACLAB, G1_29_JointIndex, get_gravity_orientation, + make_ort_session_options, ) logger = logging.getLogger(__name__) @@ -87,7 +88,12 @@ def load_policy( logger.info(f"Loading {policy_type.upper()} SONIC decoder from: {repo_id}/{filename}") decoder_path = hf_hub_download(repo_id=repo_id, filename=filename) - decoder = ort.InferenceSession(decoder_path) + # Cap the thread pool: onboard, this decoder steps at 50 Hz in the same process as the + # ZMQ camera server (capture + JPEG encode). Default session options let ORT grab every + # core on the NX, which starves the camera thread (stale frames) and jitters the control + # loop (limping gait). It is a small MLP, so 1 thread is enough and lowest-latency. + session_options = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1) + decoder = ort.InferenceSession(decoder_path, sess_options=session_options) logger.info(f"Decoder loaded: {decoder.get_inputs()[0].shape} → {decoder.get_outputs()[0].shape}") # Extract deploy constants from ONNX metadata 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