admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-16 22:41:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

clean up sonic-side

Browse Source
This commit is contained in:
Martino Russi
2026-07-15 16:40:56 +02:00
parent 69a3edfa33
commit 1c88e26c6d
4 changed files with 59 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py
+3 -110
View File
@@ -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
+27 -16
View File
@@ -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.
src/lerobot/robots/unitree_g1/g1_utils.py
+12
View File
@@ -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
+17 -18
View File
@@ -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: