diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index ab92eab49..c63cb783a 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -134,13 +134,14 @@ def _action_scale(k): return 0.25 * EFFORT[k] / (ARMATURE[k] * NATURAL_FREQ**2) -# Per-joint motor model (IsaacLab order): legs, waist, then arms. -_J = ( +# Per-joint motor model (IsaacLab order): legs, waist, then arms. Single source of +# truth for both ACTION_SCALE and compute_kp_kd(). +MOTOR_MODELS = ( ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2 + ["7520_14", "5020", "5020"] + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2 ) -ACTION_SCALE = np.array([_action_scale(k) for k in _J], dtype=np.float32) # (29,) IsaacLab order +ACTION_SCALE = np.array([_action_scale(k) for k in MOTOR_MODELS], dtype=np.float32) # (29,) IsaacLab order CONTROL_DT = 0.02 # 50 Hz control period (s) DEFAULT_HEIGHT = 0.788740 # nominal pelvis height (m) @@ -260,15 +261,9 @@ def compute_kp_kd(): def d(k): return 2.0 * 2.0 * ARMATURE[k] * NATURAL_FREQ - _kp_keys = ( - ["7520_22", "7520_22", "7520_14", "7520_22", "5020", "5020"] * 2 - + ["7520_14", "5020", "5020"] - + ["5020", "5020", "5020", "5020", "5020", "4010", "4010"] * 2 - ) - _kd_keys = _kp_keys _double = {4, 5, 10, 11, 13, 14} # ankle + waist indices with factor 2 - kp = np.array([2 * s(k) if i in _double else s(k) for i, k in enumerate(_kp_keys)], dtype=np.float32) - kd = np.array([2 * d(k) if i in _double else d(k) for i, k in enumerate(_kd_keys)], dtype=np.float32) + kp = np.array([2 * s(k) if i in _double else s(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32) + kd = np.array([2 * d(k) if i in _double else d(k) for i, k in enumerate(MOTOR_MODELS)], dtype=np.float32) return kp, kd 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 1e637d1b6..e1ca458ea 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -26,6 +26,12 @@ from typing import TYPE_CHECKING import numpy as np from huggingface_hub import hf_hub_download +from lerobot.teleoperators.pico_headset.smpl_constants import ( + ROOT_ACTION_DIM, + ROOT_ACTION_PREFIX, + SMPL_ACTION_PREFIX, + SMPL_OBS_DIM as SMPL_ACTION_DIM, +) from lerobot.utils.import_utils import _onnxruntime_available, require_package from ..g1_utils import KEYBOARD_KEYS_FIELD, lowstate_to_obs @@ -56,15 +62,6 @@ else: logger = logging.getLogger(__name__) -# Length of the flattened SONIC whole-body reference window -# (10 frames x 24 SMPL joints x 3 coords). Matches smpl_joints_10frame_step1. -SMPL_ACTION_DIM = 720 -# Prefix for per-element SMPL floats carried on the teleop action dict. -SMPL_ACTION_PREFIX = "smpl." -# Optional per-frame SMPL root orientation (wxyz) for the mode-2 anchor. -ROOT_ACTION_DIM = 4 -ROOT_ACTION_PREFIX = "root." - def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None: """Reassemble a (720,) SMPL window from ``smpl.{i}`` action keys, or None. @@ -212,7 +209,7 @@ class SonicWholeBodyController: def _init_smpl_stream(self) -> None: # Lazy import so the zmq dependency is only required when streaming is on. - from lerobot.robots.unitree_g1.smpl_stream import ( + from lerobot.teleoperators.pico_headset.smpl_stream import ( DEFAULT_SMPL_HOST, DEFAULT_SMPL_PORT, SmplStream, @@ -352,7 +349,8 @@ class SonicWholeBodyController: # 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). + # smoothed/rate-matched (30 Hz dataset -> 50 Hz control). See + # docs/SONIC_REPLAY_DEBUGGING.md. self.controller.smpl_root_quat = None _ = root_quat return self._runtime.tick(obs, debug=False, use_joystick=False) diff --git a/src/lerobot/teleoperators/pico_headset/pico_headset.py b/src/lerobot/teleoperators/pico_headset/pico_headset.py index d2f7006e5..bbf828388 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_headset.py +++ b/src/lerobot/teleoperators/pico_headset/pico_headset.py @@ -19,22 +19,15 @@ import logging from typing import Any -from lerobot.robots.unitree_g1.smpl_stream import SMPL_OBS_DIM, SmplStream from lerobot.types import RobotAction from ..teleoperator import Teleoperator from .config_pico_headset import PicoHeadsetConfig +from .smpl_constants import ROOT_ACTION_DIM, ROOT_ACTION_PREFIX, SMPL_ACTION_PREFIX, SMPL_OBS_DIM +from .smpl_stream import SmplStream logger = logging.getLogger(__name__) -# Flat action keys carrying the 720-vec SONIC whole-body reference window. Kept as -# scalar floats so the reference flows unchanged through the standard lerobot action -# pipeline; SonicWholeBodyController reassembles them into smpl_joints_10frame_step1. -SMPL_ACTION_PREFIX = "smpl." -# Per-frame SMPL root orientation (wxyz) that steers the SONIC mode-2 anchor/heading. -ROOT_ACTION_PREFIX = "root." -ROOT_ACTION_DIM = 4 - class PicoHeadset(Teleoperator): """Streams full-body SMPL from a PICO headset as a SONIC whole-body reference. diff --git a/src/lerobot/teleoperators/pico_headset/pico_publisher.py b/src/lerobot/teleoperators/pico_headset/pico_publisher.py index 772870960..09f6e147c 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_publisher.py +++ b/src/lerobot/teleoperators/pico_headset/pico_publisher.py @@ -19,7 +19,7 @@ Reads 24 body-joint poses from the XRoboToolkit SDK, runs pure-numpy SMPL forward kinematics + canonicalization (``smpl_fk.py``), and publishes one canonical ``(24, 3)`` SMPL frame per tick over ZMQ on the ``rt/smpl`` topic — the exact -message ``lerobot.robots.unitree_g1.smpl_stream.SmplStream`` consumes. +message ``lerobot.teleoperators.pico_headset.smpl_stream.SmplStream`` consumes. This makes the LeRobot side self-contained: the only runtime dependency to drive SONIC whole-body teleop from the headset is the ``xrobotoolkit_sdk`` Python package diff --git a/src/lerobot/robots/unitree_g1/smpl_stream.py b/src/lerobot/teleoperators/pico_headset/smpl_stream.py similarity index 98% rename from src/lerobot/robots/unitree_g1/smpl_stream.py rename to src/lerobot/teleoperators/pico_headset/smpl_stream.py index cf537223e..908ad37c6 100644 --- a/src/lerobot/robots/unitree_g1/smpl_stream.py +++ b/src/lerobot/teleoperators/pico_headset/smpl_stream.py @@ -41,17 +41,14 @@ from collections import deque import numpy as np import zmq +from .smpl_constants import JOINT_DIM, N_JOINTS, SMPL_OBS_DIM, WINDOW + logger = logging.getLogger(__name__) SMPL_TOPIC = "rt/smpl" DEFAULT_SMPL_HOST = "127.0.0.1" DEFAULT_SMPL_PORT = 5560 -WINDOW = 10 # frames per encoder window (smpl_joints_10frame_step1) -N_JOINTS = 24 -JOINT_DIM = 3 -SMPL_OBS_DIM = WINDOW * N_JOINTS * JOINT_DIM # 720 - class SmplStream: """Live ``rt/smpl`` consumer with the ``SmplMotion`` interface.