From 1c88e26c6d3f7eb56a95f920400cd9c2045903d5 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Wed, 15 Jul 2026 16:40:56 +0200 Subject: [PATCH] clean up sonic-side --- .../unitree_g1/controllers/sonic_pipeline.py | 113 +----------------- .../controllers/sonic_whole_body.py | 43 ++++--- src/lerobot/robots/unitree_g1/g1_utils.py | 12 ++ src/lerobot/robots/unitree_g1/unitree_g1.py | 35 +++--- 4 files changed, 59 insertions(+), 144 deletions(-) diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index c63cb783a..3bfb378f2 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -57,7 +57,6 @@ from __future__ import annotations import logging import math -import os import queue import threading import time @@ -70,6 +69,8 @@ import numpy as np from lerobot.utils.import_utils import _onnxruntime_available from ..g1_utils import ( + ISAACLAB_TO_MUJOCO, + MUJOCO_TO_ISAACLAB, G1_29_JointIndex, get_gravity_orientation, ) @@ -158,77 +159,6 @@ BLEND_FRAMES = 8 # cross-fade length when swapping in a freshly planned motion # Seconds between automatic replans, per motion class (faster for dynamic motions). REPLAN_INTERVAL = {"running": 0.1, "crawling": 0.2, "boxing": 1.0, "default": 1.0} -# Joint-order permutations between IsaacLab (policy) and MuJoCo (deploy) layouts. -ISAACLAB_TO_MUJOCO = np.array( - [ - 0, - 3, - 6, - 9, - 13, - 17, - 1, - 4, - 7, - 10, - 14, - 18, - 2, - 5, - 8, - 11, - 15, - 19, - 21, - 23, - 25, - 27, - 12, - 16, - 20, - 22, - 24, - 26, - 28, - ], - dtype=np.int32, -) - -MUJOCO_TO_ISAACLAB = np.array( - [ - 0, - 6, - 12, - 1, - 7, - 13, - 2, - 8, - 14, - 3, - 9, - 15, - 22, - 4, - 10, - 16, - 23, - 5, - 11, - 17, - 24, - 18, - 25, - 19, - 26, - 20, - 27, - 21, - 28, - ], - dtype=np.int32, -) - def _to_mujoco(a): """Reorder a 29-vector from IsaacLab order into MuJoCo/deploy order.""" @@ -505,23 +435,10 @@ 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.""" + """Build ONNX Runtime SessionOptions (quiet logging, default 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 @@ -806,30 +723,6 @@ class StandingEncoderDecoder: ) return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex} - def print_input_diagnostics(self): - """Log sanity checks on the reference/anchor/gravity terms (debugging aid).""" - names = {0: "g1", 1: "teleop", 2: "smpl"} - anchor = self._anchor_6d(np.array([1, 0, 0, 0], np.float32), np.array([1, 0, 0, 0], np.float32)) - gravity = get_gravity_orientation(np.array([1, 0, 0, 0], np.float32)) - dec0 = self.build_decoder_obs() - logger.debug( - "[Diag] Standing reference checks\n" - " encoder mode: %d (%s)\n" - " DEFAULT_ANGLES range: [%+.4f, %+.4f]\n" - " anchor_6d(identity): %s\n" - " gravity(identity): %s (expect [0,0,-1])\n" - " decoder q-delta max: %.6f\n" - " decoder dq max: %.6f", - self.encode_mode, - names.get(self.encode_mode, "unknown"), - DEFAULT_ANGLES.min(), - DEFAULT_ANGLES.max(), - anchor, - gravity, - np.max(np.abs(dec0[94:384])), - np.max(np.abs(dec0[384:674])), - ) - # ── Planner motion buffer ───────────────────────────────────────────────────── 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 e1ca458ea..6272c706f 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -20,7 +20,6 @@ from __future__ import annotations import logging import math -import os from typing import TYPE_CHECKING import numpy as np @@ -179,7 +178,15 @@ class SonicWholeBodyController: control_dt = CONTROL_DT full_body = True - def __init__(self, force_cpu: bool = False): + def __init__( + self, + force_cpu: bool = False, + *, + enable_smpl_root: bool = False, + enable_smpl_stream: bool = False, + smpl_host: str | None = None, + smpl_port: int | None = None, + ): logger.info("Loading SONIC whole-body controller...") self._runtime = SonicRuntime(force_cpu=force_cpu) self.kp = self._runtime.kp @@ -187,6 +194,12 @@ class SonicWholeBodyController: self.controller = self._runtime.controller self.ms = self._runtime.ms + # When True, the per-frame SMPL root quaternion steers the mode-2 anchor. + # Off by default: feeding it currently produces root-acceleration spikes + # (NaN QACC at DOF 0) until the reference root trajectory is smoothed and + # rate-matched (30 Hz dataset -> 50 Hz control). + self.enable_smpl_root = enable_smpl_root + # 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. @@ -194,10 +207,11 @@ class SonicWholeBodyController: # 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 - # with SONIC_SMPL_STREAM=1; override endpoint via SONIC_SMPL_HOST/PORT. + # estop/joystick) before the dedicated pico_headset teleop exists. + self._smpl_host = smpl_host + self._smpl_port = smpl_port self._smpl_stream = None - if os.environ.get("SONIC_SMPL_STREAM", "0") not in ("0", "", "false", "False"): + if enable_smpl_stream: self._init_smpl_stream() logger.info( @@ -215,8 +229,8 @@ class SonicWholeBodyController: SmplStream, ) - host = os.environ.get("SONIC_SMPL_HOST", DEFAULT_SMPL_HOST) - port = int(os.environ.get("SONIC_SMPL_PORT", DEFAULT_SMPL_PORT)) + host = self._smpl_host or DEFAULT_SMPL_HOST + port = self._smpl_port or DEFAULT_SMPL_PORT self._smpl_stream = SmplStream(host=host, port=port) logger.info("SONIC subscribed to rt/smpl @ tcp://%s:%d", host, port) @@ -329,7 +343,7 @@ class SonicWholeBodyController: 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 + # direct rt/smpl subscription when enabled (enable_smpl_stream). A stale # stream (headset silent past its timeout) is treated as "no SMPL" so the # robot doesn't stay frozen tracking the last pose. smpl = _extract_smpl_from_action(action) @@ -345,14 +359,11 @@ class SonicWholeBodyController: if self.controller.encode_mode != 2: self._enter_wholebody() self.controller.smpl_joints_10frame_step1 = smpl - # Root orientation (if provided) steers the mode-2 anchor/heading. - # Temporarily disabled: feeding the per-frame SMPL root quaternion produced - # root-acceleration spikes (NaN QACC at DOF 0, sim unstable) mid-episode. - # Keep the anchor self-driven until the reference root trajectory is - # smoothed/rate-matched (30 Hz dataset -> 50 Hz control). See - # docs/SONIC_REPLAY_DEBUGGING.md. - self.controller.smpl_root_quat = None - _ = root_quat + # Root orientation steers the mode-2 anchor/heading, but only when + # explicitly enabled (see enable_smpl_root); otherwise the anchor stays + # self-driven to avoid root-acceleration spikes from an unsmoothed + # reference trajectory. + self.controller.smpl_root_quat = root_quat if self.enable_smpl_root else None return self._runtime.tick(obs, debug=False, use_joystick=False) # No (or stale) SMPL: fall back to locomotion so the robot stays balanced. diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index f3dace89c..8861676aa 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -23,6 +23,18 @@ import numpy as np NUM_MOTORS = 29 +# Joint-order permutations between the two 29-DoF layouts used across the G1 stack: +# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]`` +# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa. +ISAACLAB_TO_MUJOCO = np.array( + [0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18, 2, 5, 8, 11, 15, 19, 21, 23, 25, 27, 12, 16, 20, 22, 24, 26, 28], + dtype=np.int32, +) +MUJOCO_TO_ISAACLAB = np.array( + [0, 6, 12, 1, 7, 13, 2, 8, 14, 3, 9, 15, 22, 4, 10, 16, 23, 5, 11, 17, 24, 18, 25, 19, 26, 20, 27, 21, 28], + dtype=np.int32, +) + 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 diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 3c11b70b8..3d862d4f0 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -35,7 +35,6 @@ from .g1_kinematics import G1_29_ArmIK from .g1_utils import ( KEYBOARD_KEYS_FIELD, REMOTE_AXES, - REMOTE_KEYS, G1_29_JointArmIndex, G1_29_JointIndex, default_remote_input, @@ -477,29 +476,29 @@ class UnitreeG1(Robot): return action def _update_controller_action(self, action: RobotAction) -> None: - """Update controller input state from incoming teleop action.""" + """Update controller input state from an incoming teleop action. + + Controller-agnostic: every value-carrying key is forwarded verbatim into + ``controller_input`` (joystick ``remote.*``, whole-body ``smpl.*``/``root.*`` + from pico_headset, or whatever a future controller expects), and each + controller extracts only the keys it understands. The robot deliberately does + not enumerate any controller's key schema here. + + KeyboardTeleop is the one special case: it emits the currently-pressed keys as + bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so + those are collected into a single held-key set under ``KEYBOARD_KEYS_FIELD``, + rebuilt each tick so releases clear. Special keys arrive as pynput objects and + are normalised to their name ("space", ...). + """ with self._controller_action_lock: - for key in REMOTE_KEYS: - if key in action: - self.controller_input[key] = action[key] - # Forward the whole-body SMPL reference (pico_headset teleop) to the - # controller. Carried as flat smpl.{i} floats so it passes untouched - # through the standard action pipeline; SonicWholeBodyController - # reassembles it into the 720-vec encoder window. - for key in action: - 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 } + for key, value in action.items(): + if isinstance(key, str) and value is not None: + self.controller_input[key] = value @property def is_calibrated(self) -> bool: