fix lint

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
 
 

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)

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.

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

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.