From c997abe739810975a34dd53ae718a7286595d0c6 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Tue, 14 Jul 2026 18:11:03 +0200 Subject: [PATCH] (fix) keep num of ORTthreads under core count --- .../unitree_g1/controllers/sonic_pipeline.py | 24 ++++- .../controllers/sonic_whole_body.py | 98 ++++++++++++++++++- src/lerobot/robots/unitree_g1/g1_utils.py | 5 + src/lerobot/robots/unitree_g1/unitree_g1.py | 14 ++- 4 files changed, 135 insertions(+), 6 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 0fdf17220..bb9e331ad 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -56,6 +56,7 @@ worker · ``SonicPlanner`` · ``PlannerController`` · joystick input. from __future__ import annotations import math +import os import queue import threading import time @@ -506,6 +507,26 @@ def ort_providers(force_cpu: bool = False) -> list[str]: return ["CPUExecutionProvider"] +# Number of intra-op threads each ONNX session may use. SONIC runs THREE CPU +# sessions concurrently (async planner thread + per-tick encoder + decoder), and +# ONNX Runtime otherwise defaults to one thread PER CORE per session, which +# massively oversubscribes the CPU and starves the 50 Hz control loop (especially +# alongside the MuJoCo sim). Capping keeps total ORT threads well under the core +# count so the control loop stays real-time. Override with SONIC_ORT_THREADS. +ORT_INTRA_OP_THREADS = int(os.environ.get("SONIC_ORT_THREADS", "2")) + + +def make_ort_session_options(): + """Build ONNX Runtime SessionOptions with capped, non-oversubscribing threading.""" + so = ort.SessionOptions() + so.log_severity_level = 3 + # inter_op=1: these graphs are sequential, parallel op scheduling only adds + # contention. intra_op capped so N concurrent sessions don't fight for cores. + so.inter_op_num_threads = 1 + so.intra_op_num_threads = max(1, ORT_INTRA_OP_THREADS) + return so + + # ── Movement state ──────────────────────────────────────────────────────────── @@ -883,8 +904,7 @@ def _planner_worker(path, req_q, res_q, stop_evt, version, seed, use_gpu): ``req_q``, running inference, resampling to 50 Hz, and putting the newest result on ``res_q`` (dropping stale entries). Runs until ``stop_evt`` is set. """ - so = ort.SessionOptions() - so.log_severity_level = 3 + so = make_ort_session_options() providers = ort_providers(force_cpu=not use_gpu) sess = ort.InferenceSession(path, sess_options=so, providers=providers) while not stop_evt.is_set(): 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 6c915cc02..7fca5f1d1 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -19,6 +19,7 @@ from __future__ import annotations import logging +import math import os from typing import TYPE_CHECKING @@ -27,11 +28,12 @@ from huggingface_hub import hf_hub_download from lerobot.utils.import_utils import _onnxruntime_available, require_package -from ..g1_utils import lowstate_to_obs +from ..g1_utils import KEYBOARD_KEYS_FIELD, lowstate_to_obs from .sonic_pipeline import ( CONTROL_DT, DEBUG_PRINT_EVERY, DEFAULT_ANGLES, + DEFAULT_HEIGHT, ENCODER_UPDATE_EVERY, LM, MOTION_SETS, @@ -40,6 +42,7 @@ from .sonic_pipeline import ( SonicPlanner, clamp_mode_params, compute_kp_kd, + make_ort_session_options, ort_providers, process_joystick, should_replan_request, @@ -108,8 +111,7 @@ class SonicRuntime: providers = ort_providers(force_cpu=force_cpu) self.use_gpu = providers[0] == "CUDAExecutionProvider" - so = ort.SessionOptions() - so.log_severity_level = 3 + so = make_ort_session_options() planner_sess = ort.InferenceSession(planner_path, sess_options=so, providers=providers) encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers) @@ -188,6 +190,11 @@ class SonicWholeBodyController: self.controller = self._runtime.controller self.ms = self._runtime.ms + # Tracks the previous keyboard held-key set so discrete controls (mode, + # motion set, replan, e-stop, WASD direction) fire once per physical press + # instead of every 50 Hz tick while the key is held. + self._prev_keys: set[str] = set() + # Optional: subscribe directly to the rt/smpl headset stream so full-body # teleop works with ANY teleoperator (e.g. --teleop.type=unitree_g1 for the # estop/joystick) before the dedicated pico_headset teleop exists. Enable @@ -234,11 +241,96 @@ class SonicWholeBodyController: self.ms.needs_replan = True logger.warning("SONIC: SMPL stream lost/stale -> reverting to locomotion (standing)") + def _process_keyboard(self, action: dict | None) -> None: + """Translate a native KeyboardTeleop's held-key set into MovementState. + + Mirrors the standalone SONIC demo's keyboard mapping so locomotion (mode 0/1) + can be driven with ``--teleop.type=keyboard`` instead of the PICO SMPL stream. + Discrete controls act on newly-pressed keys (edge-detected against the previous + tick); inherently-continuous controls (facing turn, height, speed) integrate a + small per-tick delta while the key is held so they feel smooth at 50 Hz. + + Controls: WASD move, Q/E turn, 1-8 select mode, 9/0 speed down/up, + -/= height down/up, R replan, Space emergency-stop -> IDLE. + """ + if action is None: + return + keys = action.get(KEYBOARD_KEYS_FIELD) + if keys is None: + return # No KeyboardTeleop attached; leave joystick/SMPL paths untouched. + + ms, controller = self.ms, self.controller + held = {k.lower() if isinstance(k, str) and len(k) == 1 else k for k in keys} + prev = self._prev_keys + pressed = held - prev # newly-pressed this tick (edge) + self._prev_keys = held + + # ── Discrete: fire once per press ──────────────────────────────────── + if "space" in pressed: + ms.mode = LM.IDLE + ms.speed = ms.height = -1.0 + ms.has_movement = False + ms.needs_replan = True + controller.playing = False + controller.reinit_heading = True + logger.info("SONIC keyboard: EMERGENCY STOP -> IDLE") + if "r" in pressed: + ms.needs_replan = True + if "n" in pressed or "p" in pressed: + step = 1 if "n" in pressed else -1 + ms.motion_set_idx = (ms.motion_set_idx + step) % len(MOTION_SETS) + logger.info("SONIC keyboard: motion set -> %s", MOTION_SETS[ms.motion_set_idx][0]) + for digit in ("1", "2", "3", "4", "5", "6", "7", "8"): + if digit in pressed: + idx = int(digit) - 1 + modes = MOTION_SETS[ms.motion_set_idx][1] + if 0 <= idx < len(modes): + ms.mode = modes[idx] + ms.has_movement = False + ms.needs_replan = True + controller.playing = True + controller.reinit_heading = True + logger.info("SONIC keyboard: mode -> %s", LM(ms.mode).name) + # WASD sets the movement direction relative to current facing (press to set, + # Space to stop) to match the standalone demo. + if "w" in pressed: + ms.movement_angle = ms.facing_angle + elif "s" in pressed: + ms.movement_angle = ms.facing_angle + math.pi + elif "a" in pressed: + ms.movement_angle = ms.facing_angle + math.pi / 2 + elif "d" in pressed: + ms.movement_angle = ms.facing_angle - math.pi / 2 + if pressed & {"w", "a", "s", "d"}: + ms.has_movement = True + ms.needs_replan = True + + # ── Continuous: integrate a small delta while held ─────────────────── + if "q" in held: + ms.facing_angle += 0.02 + controller.delta_heading += 0.02 + if "e" in held: + ms.facing_angle -= 0.02 + controller.delta_heading -= 0.02 + if "0" in held: + ms.speed = min(5.0, (ms.speed if ms.speed >= 0 else 1.0) + 0.02) + if "9" in held: + ms.speed = max(0.0, (ms.speed if ms.speed >= 0 else 1.0) - 0.02) + if "=" in held: + ms.height = min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + 0.005) + if "-" in held: + ms.height = max(0.1, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) - 0.005) + def run_step(self, action: dict, lowstate) -> dict: if lowstate is None: return {} obs = lowstate_to_obs(lowstate) + # Keyboard teleop (native KeyboardTeleop) drives the same locomotion intent + # the joystick does; applied before the SMPL check so whole-body tracking + # still takes priority when a headset stream is present. + self._process_keyboard(action) + # Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a # direct rt/smpl subscription when SONIC_SMPL_STREAM is enabled. A stale # stream (headset silent past its timeout) is treated as "no SMPL" so the diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 0ea267d76..f3dace89c 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -27,6 +27,11 @@ 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 +# Reserved action-dict field used to forward the set of currently-pressed keyboard +# keys from a KeyboardTeleop through the standard action pipeline to the SONIC +# whole-body controller (see SonicWholeBodyController._process_keyboard). +KEYBOARD_KEYS_FIELD = "keyboard.keys" + def default_remote_input() -> dict[str, float]: """Return a zeroed-out remote input dict (axes + buttons).""" diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 1ad5818ff..3c11b70b8 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -33,6 +33,7 @@ from ..robot import Robot from .config_unitree_g1 import UnitreeG1Config from .g1_kinematics import G1_29_ArmIK from .g1_utils import ( + KEYBOARD_KEYS_FIELD, REMOTE_AXES, REMOTE_KEYS, G1_29_JointArmIndex, @@ -486,8 +487,19 @@ class UnitreeG1(Robot): # through the standard action pipeline; SonicWholeBodyController # reassembles it into the 720-vec encoder window. for key in action: - if key.startswith("smpl.") or key.startswith("root."): + if isinstance(key, str) and (key.startswith("smpl.") or key.startswith("root.")): self.controller_input[key] = action[key] + # Forward the KeyboardTeleop state. That teleop emits the set of + # currently-pressed keys as bare action keys with a None value + # (dict.fromkeys(pressed, None)); collect them into a single set so the + # SONIC controller sees the full held-key state each tick (and releases + # clear, unlike the persistent merge above). Special keys arrive as + # pynput objects, so normalise them to their name ("space", ...). + self.controller_input[KEYBOARD_KEYS_FIELD] = { + (k if isinstance(k, str) else getattr(k, "name", str(k))) + for k, value in action.items() + if value is None + } @property def is_calibrated(self) -> bool: