clean up sonic-side

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 ─────────────────────────────────────────────────────
 

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
+        # estop/joystick) before the dedicated pico_headset teleop exists.
-        # with SONIC_SMPL_STREAM=1; override endpoint via SONIC_SMPL_HOST/PORT.
+        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)
+        host = self._smpl_host or DEFAULT_SMPL_HOST
-        port = int(os.environ.get("SONIC_SMPL_PORT", DEFAULT_SMPL_PORT))
+        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.
+            # Root orientation steers the mode-2 anchor/heading, but only when
-            # Temporarily disabled: feeding the per-frame SMPL root quaternion produced
+            # explicitly enabled (see enable_smpl_root); otherwise the anchor stays
-            # root-acceleration spikes (NaN QACC at DOF 0, sim unstable) mid-episode.
+            # self-driven to avoid root-acceleration spikes from an unsmoothed
-            # Keep the anchor self-driven until the reference root trajectory is
+            # reference trajectory.
-            # smoothed/rate-matched (30 Hz dataset -> 50 Hz control). See
+            self.controller.smpl_root_quat = root_quat if self.enable_smpl_root else None
-            # docs/SONIC_REPLAY_DEBUGGING.md.
-            self.controller.smpl_root_quat = None
-            _ = root_quat
             return self._runtime.tick(obs, debug=False, use_joystick=False)
 
         # No (or stale) SMPL: fall back to locomotion so the robot stays balanced.

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

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: