mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-31 21:49:45 +00:00
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>
This commit is contained in:
@@ -26,6 +26,7 @@ from ..g1_utils import (
|
|||||||
REMOTE_BUTTONS,
|
REMOTE_BUTTONS,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
get_gravity_orientation,
|
get_gravity_orientation,
|
||||||
|
make_ort_session_options,
|
||||||
)
|
)
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -68,9 +69,13 @@ def load_groot_policies(
|
|||||||
filename="GR00T-WholeBodyControl-Walk.onnx",
|
filename="GR00T-WholeBodyControl-Walk.onnx",
|
||||||
)
|
)
|
||||||
|
|
||||||
# Load ONNX policies
|
# Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
|
||||||
policy_balance = ort.InferenceSession(balance_path)
|
# background thread alongside the (torch) upper-body policy and IK; letting ORT
|
||||||
policy_walk = ort.InferenceSession(walk_path)
|
# 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")
|
logger.info("GR00T policies loaded successfully")
|
||||||
|
|
||||||
|
|||||||
@@ -80,6 +80,26 @@ def get_gravity_orientation(quaternion: list[float] | np.ndarray) -> np.ndarray:
|
|||||||
return gravity_orientation
|
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):
|
class G1_29_JointArmIndex(IntEnum):
|
||||||
# Left arm
|
# Left arm
|
||||||
kLeftShoulderPitch = 15
|
kLeftShoulderPitch = 15
|
||||||
|
|||||||
Reference in New Issue
Block a user