Apply ruff-format

Co-authored-by: Cursor <cursoragent@cursor.com>

pyproject.toml

@@ -374,7 +374,11 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
 
 [tool.setuptools.package-data]
-lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
+lerobot = [
+    "envs/*.json",
+    "annotations/steerable_pipeline/prompts/*.txt",
+    "teleoperators/pico_headset/assets/*.npz",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -68,6 +68,6 @@ class UnitreeG1Config(RobotConfig):
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 
-    # Lower-body controller class name, e.g. "GrootLocomotionController" or
-    # "HolosomaLocomotionController". None disables it.
+    # Locomotion controller class name, e.g. "GrootLocomotionController",
+    # "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
     controller: str | None = None

src/lerobot/robots/unitree_g1/controllers/__init__.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
+
+__all__ = [
+    "GrootLocomotionController",
+    "HolosomaLocomotionController",
+    "SonicWholeBodyController",
+    "SonicRuntime",
+]

src/lerobot/robots/unitree_g1/controllers/gr00t_locomotion.py

@@ -21,7 +21,7 @@ import numpy as np
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from lerobot.robots.unitree_g1.g1_utils import (
     REMOTE_AXES,
     REMOTE_BUTTONS,
     G1_29_JointIndex,

src/lerobot/robots/unitree_g1/controllers/holosoma_locomotion.py

@@ -22,7 +22,7 @@ import onnx
 import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 
-from .g1_utils import (
+from lerobot.robots.unitree_g1.g1_utils import (
     REMOTE_AXES,
     G1_29_JointArmIndex,
     G1_29_JointIndex,

src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py

@@ -0,0 +1,262 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SONIC full-body controller for Unitree G1."""
+
+import logging
+import os
+
+import numpy as np
+import onnxruntime as ort
+from huggingface_hub import hf_hub_download
+
+from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
+    CONTROL_DT,
+    DEBUG_PRINT_EVERY,
+    DEFAULT_ANGLES,
+    ENCODER_UPDATE_EVERY,
+    LM,
+    MOTION_SETS,
+    MovementState,
+    PlannerController,
+    SonicPlanner,
+    _ort_providers,
+    _snapshot_ms,
+    clamp_mode_params,
+    compute_kp_kd,
+    process_joystick,
+    should_replan_request,
+)
+from lerobot.robots.unitree_g1.g1_utils import lowstate_to_obs
+
+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.
+
+    The pico_headset teleoperator emits the whole-body reference as flat floats so
+    it flows unchanged through the standard lerobot action pipeline.
+    """
+    if not action or f"{SMPL_ACTION_PREFIX}0" not in action:
+        return None
+    arr = np.fromiter(
+        (float(action.get(f"{SMPL_ACTION_PREFIX}{i}", 0.0)) for i in range(SMPL_ACTION_DIM)),
+        dtype=np.float32,
+        count=SMPL_ACTION_DIM,
+    )
+    return arr
+
+
+def _extract_root_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a (4,) SMPL root quaternion (wxyz) from ``root.{i}`` keys, or None."""
+    if not action or f"{ROOT_ACTION_PREFIX}0" not in action:
+        return None
+    q = np.fromiter(
+        (float(action.get(f"{ROOT_ACTION_PREFIX}{i}", 0.0)) for i in range(ROOT_ACTION_DIM)),
+        dtype=np.float32,
+        count=ROOT_ACTION_DIM,
+    )
+    n = float(np.linalg.norm(q))
+    if n < 1e-6:
+        return None
+    return q / n
+
+
+class SonicRuntime:
+    """Shared SONIC control loop state (standalone demo + locomotion controller)."""
+
+    def __init__(self, force_cpu: bool = False):
+        planner_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="planner_sonic.onnx")
+        encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx")
+        decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
+
+        providers = _ort_providers(force_cpu=force_cpu)
+        self.use_gpu = providers[0] == "CUDAExecutionProvider"
+        so = ort.SessionOptions()
+        so.log_severity_level = 3
+
+        planner_sess = ort.InferenceSession(planner_path, sess_options=so, providers=providers)
+        encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
+        decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
+
+        self.kp, self.kd = compute_kp_kd()
+        self.ms = MovementState()
+        self.planner = SonicPlanner(planner_sess, planner_path)
+        self.controller = PlannerController(self.planner, encoder_sess, decoder_sess)
+
+        motion = self.planner.initialize(DEFAULT_ANGLES, self.ms)
+        self.controller.load_initial_motion(motion)
+        self.planner.start_subprocess(self.controller, use_gpu=self.use_gpu)
+
+        self.step = 0
+        self.replan_timer = 0.0
+        self.last_ms = _snapshot_ms(self.ms)
+
+    @property
+    def pipeline(self):
+        return self.controller
+
+    def tick(self, obs: dict, *, debug: bool | None = None, use_joystick: bool = True) -> dict:
+        if not obs:
+            self.step += 1
+            return {}
+
+        if use_joystick:
+            process_joystick(obs, self.ms, self.controller)
+        clamp_mode_params(self.ms)
+
+        if self.step > 0:
+            self.replan_timer += CONTROL_DT
+        if should_replan_request(self.ms, self.last_ms, self.replan_timer, self.step):
+            self.planner.request_replan(self.controller.ref_cursor, self.ms)
+            self.replan_timer = 0.0
+            self.ms.needs_replan = False
+            self.last_ms = _snapshot_ms(self.ms)
+
+        do_enc = self.step % ENCODER_UPDATE_EVERY == 0
+        if debug is None:
+            debug = self.step % DEBUG_PRINT_EVERY == 0
+        action = self.controller.step(obs, update_encoder=do_enc, debug=debug)
+
+        result = self.planner.try_get_new_motion()
+        if result:
+            self.controller.blend_new_motion(*result)
+
+        self.controller.advance_cursor()
+        self.step += 1
+        return action
+
+    def reset(self):
+        self.ms = MovementState()
+        self.controller.reinit_heading = True
+        self.controller.playing = True
+        self.step = 0
+        self.replan_timer = 0.0
+        self.last_ms = _snapshot_ms(self.ms)
+
+    def shutdown(self):
+        self.planner.stop_subprocess()
+
+
+class SonicWholeBodyController:
+    """Full-body SONIC controller for UnitreeG1's background controller thread."""
+
+    control_dt = CONTROL_DT
+    full_body = True
+
+    def __init__(self, force_cpu: bool = False):
+        logger.info("Loading SONIC whole-body controller...")
+        self._runtime = SonicRuntime(force_cpu=force_cpu)
+        self.kp = self._runtime.kp
+        self.kd = self._runtime.kd
+        self.controller = self._runtime.controller
+        self.ms = self._runtime.ms
+
+        # 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.
+        self._smpl_stream = None
+        if os.environ.get("SONIC_SMPL_STREAM", "0") not in ("0", "", "false", "False"):
+            self._init_smpl_stream()
+
+        logger.info(
+            "SONIC ready: %s (default mode: %s, smpl_stream=%s)",
+            MOTION_SETS[0][0],
+            LM(self.ms.mode).name,
+            self._smpl_stream is not None,
+        )
+
+    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 (
+            DEFAULT_SMPL_HOST,
+            DEFAULT_SMPL_PORT,
+            SmplStream,
+        )
+
+        host = os.environ.get("SONIC_SMPL_HOST", DEFAULT_SMPL_HOST)
+        port = int(os.environ.get("SONIC_SMPL_PORT", DEFAULT_SMPL_PORT))
+        self._smpl_stream = SmplStream(host=host, port=port)
+        logger.info("SONIC subscribed to rt/smpl @ tcp://%s:%d", host, port)
+
+    def _enter_wholebody(self) -> None:
+        """Switch into SMPL whole-body tracking (encode_mode 2)."""
+        self.controller.encode_mode = 2
+        self.controller.reinit_heading = True
+        logger.info("SONIC: SMPL stream active -> whole-body tracking (mode 2)")
+
+    def _exit_wholebody(self) -> None:
+        """Revert to locomotion/standing (encode_mode 0) after SMPL is lost.
+
+        Mirrors the 'M' toggle in sonic.py so the handoff is clean: the robot holds
+        a standing reference and (if a joystick teleop is attached) can be driven.
+        """
+        self.controller.encode_mode = 0
+        self.controller.playing = True
+        self.controller.reinit_heading = True
+        self.ms.needs_replan = True
+        logger.warning("SONIC: SMPL stream lost/stale -> reverting to locomotion (standing)")
+
+    def run_step(self, action: dict, lowstate) -> dict:
+        if lowstate is None:
+            return {}
+        obs = lowstate_to_obs(lowstate)
+
+        # 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
+        # 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)
+        root_quat = _extract_root_from_action(action)
+        if smpl is None and self._smpl_stream is not None:
+            window = self._smpl_stream.step()
+            if self._smpl_stream.has_data and not self._smpl_stream.is_stale:
+                smpl = window
+                root_quat = np.asarray(self._smpl_stream.root_quat, np.float32)
+
+        if smpl is not None:
+            # Full-body whole-body tracking: SMPL drives the reference, not joystick.
+            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.
+            self.controller.smpl_root_quat = 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.
+        if self.controller.encode_mode == 2:
+            self.controller.smpl_root_quat = None
+            self._exit_wholebody()
+        return self._runtime.tick(obs, debug=False)
+
+    def reset(self):
+        self._runtime.reset()
+
+    def shutdown(self):
+        if self._smpl_stream is not None:
+            self._smpl_stream.close()
+        self._runtime.shutdown()

src/lerobot/robots/unitree_g1/g1_utils.py

@@ -63,13 +63,55 @@ class G1_29_JointArmIndex(IntEnum):
     kRightWristYaw = 28
 
 
+def lowstate_to_obs(lowstate) -> dict:
+    """Build a robot observation dict from a Unitree lowstate.
+
+    Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
+    lowstate -> obs mapping lives in exactly one place. Keys match the
+    ``<joint>.q``/``imu.*`` schema consumed across the controllers.
+    """
+    obs: dict = {}
+
+    for motor in G1_29_JointIndex:
+        idx = motor.value
+        obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
+        obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
+        obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
+
+    imu = lowstate.imu_state
+    if imu.gyroscope:
+        obs["imu.gyro.x"] = imu.gyroscope[0]
+        obs["imu.gyro.y"] = imu.gyroscope[1]
+        obs["imu.gyro.z"] = imu.gyroscope[2]
+    if imu.accelerometer:
+        obs["imu.accel.x"] = imu.accelerometer[0]
+        obs["imu.accel.y"] = imu.accelerometer[1]
+        obs["imu.accel.z"] = imu.accelerometer[2]
+    if imu.quaternion:
+        obs["imu.quat.w"] = imu.quaternion[0]
+        obs["imu.quat.x"] = imu.quaternion[1]
+        obs["imu.quat.y"] = imu.quaternion[2]
+        obs["imu.quat.z"] = imu.quaternion[3]
+    if imu.rpy:
+        obs["imu.rpy.roll"] = imu.rpy[0]
+        obs["imu.rpy.pitch"] = imu.rpy[1]
+        obs["imu.rpy.yaw"] = imu.rpy[2]
+
+    wr = getattr(lowstate, "wireless_remote", None)
+    if wr:
+        obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
+
+    return obs
+
+
 def make_locomotion_controller(name: str | None):
     """Instantiate a locomotion controller by class name. Returns None if name is None."""
     if name is None:
         return None
     controllers = {
-        "GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
-        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
+        "GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
+        "HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
+        "SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
     }
     module_path = controllers.get(name)
     if module_path is None:

src/lerobot/robots/unitree_g1/smpl_stream.py

@@ -0,0 +1,203 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Live SMPL stream as a SONIC reference motion (drop-in for ``SmplMotion``).
+
+Instead of reading an ``.npz`` clip, this pulls per-frame SMPL joints live off the
+``rt/smpl`` ZMQ channel published by the GEAR PICO teleop
+(``gear_sonic/scripts/pico_manager_thread_server.py``). Each message carries one
+frame of **canonical** (root-orientation-removed) SMPL local joints ``(24, 3)`` --
+the exact per-frame format ``SmplMotion`` emits -- so this class exposes the same
+``step() -> (720,)`` window interface and can be handed to ``sonic.py`` (or the
+``pico_headset`` teleoperator) wherever a reference motion is expected
+(``encode_mode == 2``).
+
+Transport mirrors the Unitree SDK socket bridge (``unitree_sdk2_socket.py``): a ZMQ
+``SUB`` socket with ``CONFLATE`` (keep only the latest frame) subscribed to the
+``rt/smpl`` topic, JSON payloads.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from collections import deque
+
+import numpy as np
+import zmq
+
+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.
+
+    Args:
+        host: publisher host (the laptop running pico_manager_thread_server.py).
+        port: publisher port for the ``rt/smpl`` channel.
+        fps: nominal source rate, only used for status/reporting.
+        stale_after_s: log a warning if no fresh frame arrives within this window.
+        loop: accepted for API parity with ``SmplMotion`` (ignored; a stream never ends).
+    """
+
+    def __init__(
+        self,
+        host: str = DEFAULT_SMPL_HOST,
+        port: int = DEFAULT_SMPL_PORT,
+        fps: float = 50.0,
+        stale_after_s: float = 0.5,
+        loop: bool = True,
+    ):
+        self.host = host
+        self.port = port
+        self.fps = float(fps)
+        self.loop = loop
+        self.stale_after_s = stale_after_s
+
+        self._ctx = zmq.Context.instance()
+        self._sock = self._ctx.socket(zmq.SUB)
+        self._sock.setsockopt(zmq.CONFLATE, 1)  # keep only the most recent frame
+        self._sock.connect(f"tcp://{host}:{port}")
+        # Single-frame JSON messages (topic embedded in payload); CONFLATE does not
+        # support multipart, so subscribe to everything on this dedicated port.
+        self._sock.setsockopt_string(zmq.SUBSCRIBE, "")
+        self._poller = zmq.Poller()
+        self._poller.register(self._sock, zmq.POLLIN)
+
+        # Rolling window, oldest -> newest (matches SmplMotion.window layout).
+        self._buf: deque[np.ndarray] = deque(maxlen=WINDOW)
+        self._last_frame = np.zeros((N_JOINTS, JOINT_DIM), np.float32)
+        # Latest root/torso pose (updated every received frame).
+        self.root_quat = np.array([1.0, 0.0, 0.0, 0.0], np.float32)  # (w, x, y, z)
+        self.root_transl = np.zeros(3, np.float32)
+        self._last_index = -1
+        self._last_recv_t = 0.0
+        self._warned_stale = False
+        self._got_first = False
+
+    # -- SmplMotion-compatible attributes ------------------------------------
+    @property
+    def num_frames(self) -> int:
+        """Streams are unbounded; report 0 (kept for API parity)."""
+        return 0
+
+    @property
+    def done(self) -> bool:
+        """A live stream never finishes."""
+        return False
+
+    @property
+    def has_data(self) -> bool:
+        """True once at least one real frame has been received."""
+        return self._got_first
+
+    @property
+    def seconds_since_last(self) -> float:
+        """Wall-clock seconds since the last real frame (inf before the first)."""
+        if not self._got_first:
+            return float("inf")
+        return time.time() - self._last_recv_t
+
+    @property
+    def is_stale(self) -> bool:
+        """True when the stream has gone silent past ``stale_after_s``.
+
+        Consumers use this to stop feeding a frozen pose and let the controller
+        fall back to a safe standing/locomotion mode.
+        """
+        if not self._got_first or not self.stale_after_s:
+            return False
+        return self.seconds_since_last > self.stale_after_s
+
+    def reset(self):
+        self._buf.clear()
+        self._got_first = False
+
+    # -- core ----------------------------------------------------------------
+    def _drain_latest(self) -> np.ndarray | None:
+        """Return the newest available (24, 3) frame, or None if nothing new.
+
+        CONFLATE already keeps only the last message, but we poll non-blocking so
+        the 50 Hz control loop never stalls waiting on the headset.
+        """
+        frame = None
+        while dict(self._poller.poll(0)).get(self._sock) == zmq.POLLIN:
+            payload = self._sock.recv()
+            data = json.loads(payload.decode("utf-8")).get("data", {})
+            joints = np.asarray(data.get("smpl_joints_local", []), np.float32)
+            if joints.size != N_JOINTS * JOINT_DIM:
+                continue
+            frame = joints.reshape(N_JOINTS, JOINT_DIM)
+            self._last_index = int(data.get("frame_index", self._last_index + 1))
+            rq = data.get("root_quat")
+            if rq is not None and len(rq) == 4:
+                self.root_quat = np.asarray(rq, np.float32)
+            rt = data.get("root_transl")
+            if rt is not None and len(rt) == 3:
+                self.root_transl = np.asarray(rt, np.float32)
+        return frame
+
+    def step(self) -> np.ndarray:
+        """Advance one control tick, returning the current 720-vec window.
+
+        If no new headset frame arrived this tick we hold the last one, so the
+        policy keeps tracking the latest pose rather than snapping to zero.
+        """
+        frame = self._drain_latest()
+        now = time.time()
+
+        if frame is not None:
+            self._last_frame = frame
+            self._last_recv_t = now
+            self._warned_stale = False
+            if not self._got_first:
+                # Pre-fill the window so the first send is a full, coherent clip.
+                self._buf.extend([frame.copy() for _ in range(WINDOW)])
+                self._got_first = True
+            else:
+                self._buf.append(frame)
+        elif self._got_first:
+            # No fresh frame: repeat the most recent to keep the window moving.
+            self._buf.append(self._last_frame.copy())
+            if (
+                self.stale_after_s
+                and not self._warned_stale
+                and (now - self._last_recv_t) > self.stale_after_s
+            ):
+                print(
+                    f"[SmplStream] no {SMPL_TOPIC} frame for "
+                    f"{now - self._last_recv_t:.2f}s (holding last pose)"
+                )
+                self._warned_stale = True
+
+        if not self._got_first:
+            return np.zeros(SMPL_OBS_DIM, np.float32)
+        # Flatten to (720,): frames oldest->newest, joint-major within a frame
+        # [f0_j0_xyz, f0_j1_xyz, ..., f9_j23_xyz] — matches SmplMotion.window.
+        return np.concatenate(list(self._buf), dtype=np.float32).reshape(-1)
+
+    def close(self):
+        with contextlib.suppress(Exception):
+            self._poller.unregister(self._sock)
+        self._sock.close(0)

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -38,6 +38,7 @@ from .g1_utils import (
     G1_29_JointArmIndex,
     G1_29_JointIndex,
     default_remote_input,
+    lowstate_to_obs,
     make_locomotion_controller,
 )
 
@@ -343,6 +344,9 @@ class UnitreeG1(Robot):
 
         self.kp = np.array(self.config.kp, dtype=np.float32)
         self.kd = np.array(self.config.kd, dtype=np.float32)
+        if self.controller is not None and hasattr(self.controller, "kp"):
+            self.kp = np.array(self.controller.kp, dtype=np.float32)
+            self.kd = np.array(self.controller.kd, dtype=np.float32)
 
         for joint in G1_29_JointIndex:
             self.msg.motor_cmd[joint].mode = 1
@@ -379,6 +383,9 @@ class UnitreeG1(Robot):
         # Signal thread to stop and unblock any waits
         self._shutdown_event.set()
 
+        if self.controller is not None and hasattr(self.controller, "shutdown"):
+            self.controller.shutdown()
+
         # Wait for subscribe thread to finish
         if self.subscribe_thread is not None:
             self.subscribe_thread.join(timeout=2.0)
@@ -422,44 +429,8 @@ class UnitreeG1(Robot):
         if lowstate is None:
             return {}
 
-        obs = {}
-
-        # Motors - q, dq, tau for all joints
-        for motor in G1_29_JointIndex:
-            name = motor.name
-            idx = motor.value
-            obs[f"{name}.q"] = lowstate.motor_state[idx].q
-            obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
-            obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
-
-        # IMU - gyroscope
-        if lowstate.imu_state.gyroscope:
-            obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
-            obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
-            obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
-
-        # IMU - accelerometer
-        if lowstate.imu_state.accelerometer:
-            obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
-            obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
-            obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
-
-        # IMU - quaternion
-        if lowstate.imu_state.quaternion:
-            obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
-            obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
-            obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
-            obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
-
-        # IMU - rpy
-        if lowstate.imu_state.rpy:
-            obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
-            obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
-            obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
-
-        # Wireless remote (raw bytes for teleoperator)
-        if lowstate.wireless_remote:
-            obs["wireless_remote"] = lowstate.wireless_remote
+        # Motors + IMU + wireless remote (shared lowstate -> obs mapping)
+        obs = lowstate_to_obs(lowstate)
 
         # Cameras - read images from ZMQ cameras
         for cam_name, cam in self._cameras.items():
@@ -473,9 +444,11 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
+            self._update_controller_action(action)
+            if getattr(self.controller, "full_body", False):
+                return action
             # Controller thread owns legs/waist. Here we only update joystick inputs
             # and publish arm targets from the teleoperator.
-            self._update_controller_action(action)
             arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
             action_to_publish = {
                 key: value
@@ -508,6 +481,13 @@ class UnitreeG1(Robot):
             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 key.startswith("smpl.") or key.startswith("root."):
+                    self.controller_input[key] = action[key]
 
     @property
     def is_calibrated(self) -> bool:

src/lerobot/scripts/lerobot_teleoperate.py

@@ -116,6 +116,7 @@ from lerobot.teleoperators import (  # noqa: F401
     omx_leader,
     openarm_leader,
     openarm_mini,
+    pico_headset,
     reachy2_teleoperator,
     rebot_102_leader,
     so_leader,

src/lerobot/teleoperators/pico_headset/__init__.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .config_pico_headset import PicoHeadsetConfig
+from .pico_headset import PicoHeadset
+
+__all__ = ["PicoHeadset", "PicoHeadsetConfig"]

src/lerobot/teleoperators/pico_headset/config_pico_headset.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+
+from ..config import TeleoperatorConfig
+
+
+@TeleoperatorConfig.register_subclass("pico_headset")
+@dataclass
+class PicoHeadsetConfig(TeleoperatorConfig):
+    """PICO full-body headset teleop: live SMPL over the rt/smpl ZMQ stream.
+
+    Consumes the ``rt/smpl`` channel published by the GEAR PICO manager
+    (``gear_sonic/scripts/pico_manager_thread_server.py``) and emits the whole-body
+    SONIC reference window (``encode_mode == 2``) for SonicWholeBodyController.
+    """
+
+    smpl_host: str = "127.0.0.1"
+    """Host of the pico_manager rt/smpl publisher (the laptop bridging the PICO)."""
+    smpl_port: int = 5560
+    """Port of the rt/smpl publisher."""
+    stale_after_s: float = 0.5
+    """Warn if no fresh headset frame arrives within this many seconds."""

src/lerobot/teleoperators/pico_headset/pico_headset.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""PICO full-body headset teleoperator (live SMPL -> SONIC whole-body)."""
+
+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
+
+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.
+
+    Subscribes to the ``rt/smpl`` ZMQ channel and, once real frames are flowing,
+    emits the 720-element encoder window as ``smpl.{i}`` floats. Before the first
+    frame arrives it emits no SMPL keys, so the robot stays in safe locomotion mode
+    rather than tracking a zero pose.
+    """
+
+    config_class = PicoHeadsetConfig
+    name = "pico_headset"
+
+    def __init__(self, config: PicoHeadsetConfig):
+        super().__init__(config)
+        self.config = config
+        self._stream: SmplStream | None = None
+
+    @property
+    def action_features(self) -> dict:
+        feats = {f"{SMPL_ACTION_PREFIX}{i}": float for i in range(SMPL_OBS_DIM)}
+        feats.update({f"{ROOT_ACTION_PREFIX}{i}": float for i in range(ROOT_ACTION_DIM)})
+        return feats
+
+    @property
+    def feedback_features(self) -> dict:
+        return {}
+
+    @property
+    def is_connected(self) -> bool:
+        return self._stream is not None
+
+    def connect(self, calibrate: bool = True) -> None:
+        if self._stream is not None:
+            raise RuntimeError(f"{self} already connected")
+        self._stream = SmplStream(
+            host=self.config.smpl_host,
+            port=self.config.smpl_port,
+            stale_after_s=self.config.stale_after_s,
+        )
+        logger.info(
+            "PicoHeadset subscribed to rt/smpl @ tcp://%s:%d",
+            self.config.smpl_host,
+            self.config.smpl_port,
+        )
+
+    @property
+    def is_calibrated(self) -> bool:
+        return True
+
+    def calibrate(self) -> None:
+        # Calibration happens on the headset / pico_manager side, not here.
+        pass
+
+    def configure(self) -> None:
+        pass
+
+    def get_action(self) -> RobotAction:
+        if self._stream is None:
+            raise RuntimeError(f"{self} is not connected")
+        window = self._stream.step()
+        # Emit SMPL only while the headset is actively streaming: hold back before
+        # the first frame (so the controller doesn't track an all-zero collapsed
+        # pose) and once the stream goes stale (so the controller falls back to a
+        # safe standing/locomotion mode instead of freezing on the last pose).
+        if not self._stream.has_data or self._stream.is_stale:
+            return {}
+        action = {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)}
+        action.update({f"{ROOT_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.root_quat)})
+        return action
+
+    def send_feedback(self, feedback: dict[str, Any]) -> None:
+        pass
+
+    def disconnect(self) -> None:
+        if self._stream is not None:
+            self._stream.close()
+            self._stream = None

src/lerobot/teleoperators/pico_headset/pico_publisher.py

@@ -0,0 +1,205 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Standalone ``rt/smpl`` publisher for the PICO headset (no gear_sonic / torch).
+
+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.
+
+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
+(plus numpy/scipy/pyzmq), not the full ``gear_sonic`` stack.
+
+Usage:
+    # Real headset (XRoboToolkit PC Service must be running and connected):
+    python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --port 5560
+
+    # No hardware — emit a synthetic waving motion to test the consumer end-to-end:
+    python -m lerobot.teleoperators.pico_headset.pico_publisher --fake
+
+    # Replay a canned SMPL clip to the robot through the same rt/smpl -> SONIC path:
+    python -m lerobot.teleoperators.pico_headset.pico_publisher \
+        --motion-file examples/unitree_g1/motions/walk_forward.npz
+"""
+
+from __future__ import annotations
+
+import argparse
+import contextlib
+import json
+import time
+
+import numpy as np
+import zmq
+
+from lerobot.teleoperators.pico_headset.smpl_fk import (
+    SmplForwardKinematics,
+    canonicalize_smpl_joints,
+    root_quats_from_aa,
+)
+
+SMPL_TOPIC = "rt/smpl"
+DEFAULT_SMPL_PORT = 5560
+
+
+def pack_message(
+    smpl_joints_local: np.ndarray,
+    frame_index: int,
+    stamp_ns: int,
+    root_quat: np.ndarray | None = None,
+    root_transl: np.ndarray | None = None,
+) -> bytes:
+    """Build the rt/smpl JSON message (single frame, topic embedded in payload)."""
+    data = {
+        "smpl_joints_local": np.asarray(smpl_joints_local, np.float32).reshape(-1).tolist(),
+        "frame_index": int(frame_index),
+        "stamp_ns": int(stamp_ns),
+    }
+    if root_quat is not None:
+        data["root_quat"] = np.asarray(root_quat, np.float32).reshape(-1).tolist()
+    if root_transl is not None:
+        data["root_transl"] = np.asarray(root_transl, np.float32).reshape(-1).tolist()
+    return json.dumps({"topic": SMPL_TOPIC, "data": data}).encode("utf-8")
+
+
+def _fake_body_poses(t: float) -> np.ndarray:
+    """Synthetic (24, 7) body poses: identity rotations + a gently waving right arm."""
+    poses = np.zeros((24, 7), np.float32)
+    poses[:, 6] = 1.0  # unit quaternion (qw = 1), scalar-last
+    poses[:, 1] = 1.0  # ~1 m pelvis height (positions only matter for root_transl)
+    # Wave the right shoulder (SMPL body joint 17) about Z.
+    ang = 0.5 * np.sin(2.0 * np.pi * 0.5 * t)
+    poses[17, 3:7] = [0.0, 0.0, np.sin(ang / 2), np.cos(ang / 2)]
+    return poses
+
+
+def _load_motion_clip(path: str) -> dict:
+    """Load an SMPL ``.npz`` clip and canonicalize it for rt/smpl streaming.
+
+    Expects the same keys as ``motion_loader.SmplMotion``:
+        smpl_joints (T, 24, 3), pose_aa (T, 72) optional, transl (T, 3) optional.
+    Returns per-frame joints already in the encoder's root-removed convention,
+    plus optional per-frame root quat/translation.
+    """
+    data = np.load(path)
+    joints = data["smpl_joints"].astype(np.float32)
+    if joints.ndim != 3 or joints.shape[1:] != (24, 3):
+        raise ValueError(f"Expected smpl_joints (T, 24, 3), got {joints.shape}")
+
+    pose_aa = data["pose_aa"].astype(np.float32) if "pose_aa" in data.files else None
+    root_quat = None
+    if pose_aa is not None:
+        joints = canonicalize_smpl_joints(joints, pose_aa[:, :3])
+        root_quat = root_quats_from_aa(pose_aa[:, :3])
+    transl = data["transl"].astype(np.float32) if "transl" in data.files else None
+    return {"joints": joints, "root_quat": root_quat, "transl": transl}
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--port", type=int, default=DEFAULT_SMPL_PORT, help="ZMQ PUB port for rt/smpl")
+    p.add_argument("--fps", type=float, default=50.0, help="Target publish rate (Hz)")
+    p.add_argument("--skeleton", type=str, default=None, help="Path to smpl_skeleton.npz")
+    src = p.add_mutually_exclusive_group()
+    src.add_argument("--fake", action="store_true", help="Publish synthetic motion (no headset)")
+    src.add_argument("--motion-file", type=str, default=None, help="Replay an SMPL .npz clip over rt/smpl")
+    p.add_argument("--no-loop", action="store_true", help="Play a --motion-file once, then stop")
+    args = p.parse_args()
+
+    clip = _load_motion_clip(args.motion_file) if args.motion_file else None
+
+    # FK is only needed for live/synthetic (24,7) body poses; clips are pre-canonical.
+    fk = None
+    if clip is None:
+        fk = SmplForwardKinematics(args.skeleton) if args.skeleton else SmplForwardKinematics()
+
+    xrt = None
+    if clip is None and not args.fake:
+        try:
+            import xrobotoolkit_sdk as xrt  # noqa: PLC0415
+        except ImportError as e:
+            raise SystemExit(
+                "xrobotoolkit_sdk not available. Install it, or run with --fake / --motion-file "
+                "to test the pipeline without a headset."
+            ) from e
+        xrt.init()
+        print("[pico_publisher] XRoboToolkit initialized")
+
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.PUB)
+    sock.bind(f"tcp://*:{args.port}")
+    src_desc = f"motion-file {args.motion_file}" if clip else ("fake" if args.fake else "headset")
+    print(
+        f"[pico_publisher] '{SMPL_TOPIC}' bound to tcp://*:{args.port} @ {args.fps:.0f} Hz "
+        f"[source: {src_desc}]"
+    )
+    if clip is not None:
+        print(f"[pico_publisher] clip frames={clip['joints'].shape[0]} loop={not args.no_loop}")
+
+    period = 1.0 / max(1.0, args.fps)
+    frame_index = 0
+    t0 = time.time()
+    try:
+        while True:
+            loop_start = time.time()
+            if clip is not None:
+                n = clip["joints"].shape[0]
+                if args.no_loop and frame_index >= n:
+                    print("\n[pico_publisher] clip finished")
+                    break
+                i = frame_index % n
+                joints = clip["joints"][i]
+                root_quat = None if clip["root_quat"] is None else clip["root_quat"][i]
+                root_transl = None if clip["transl"] is None else clip["transl"][i]
+                stamp_ns = time.time_ns()
+            else:
+                if args.fake:
+                    body_poses = _fake_body_poses(loop_start - t0)
+                    stamp_ns = time.time_ns()
+                else:
+                    body_poses = np.asarray(xrt.get_body_joints_pose(), np.float32)
+                    stamp_ns = int(xrt.get_time_stamp_ns())
+                    if body_poses.shape != (24, 7):
+                        time.sleep(0.005)
+                        continue
+                out = fk.compute(body_poses)
+                joints = out["smpl_joints_local"]
+                root_quat = out["root_quat"]
+                root_transl = out["root_transl"]
+
+            sock.send(
+                pack_message(joints, frame_index, stamp_ns, root_quat=root_quat, root_transl=root_transl)
+            )
+            frame_index += 1
+            if frame_index % int(max(1, args.fps)) == 0:
+                print(f"[pico_publisher] sent {frame_index} frames", end="\r")
+
+            dt = time.time() - loop_start
+            if dt < period:
+                time.sleep(period - dt)
+    except KeyboardInterrupt:
+        print("\n[pico_publisher] stopping")
+    finally:
+        sock.close(0)
+        if xrt is not None:
+            with contextlib.suppress(Exception):
+                xrt.close()
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/teleoperators/pico_headset/smpl_fk.py

@@ -0,0 +1,227 @@
+#!/usr/bin/env python
+
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Standalone SMPL forward kinematics + canonicalization in pure numpy/scipy.
+
+This mirrors the ``rt/smpl`` producer path in ``gear_sonic`` (``compute_from_body_poses``
+-> ``process_smpl_joints`` -> ``compute_human_joints``) without depending on torch,
+the SMPL mesh model, or ``gear_sonic``. It only needs a small fixed skeleton table
+(rest-pose joints + kinematic tree), vendored in ``assets/smpl_skeleton.npz``.
+
+Given the 24 body-joint poses reported by the XRoboToolkit headset SDK
+(``xrt.get_body_joints_pose()`` -> (24, 7) of ``[x, y, z, qx, qy, qz, qw]``), it
+produces the root-orientation-removed 24x3 SMPL joints the SONIC encoder expects,
+plus the root orientation quaternion and pelvis translation.
+
+Quaternions are scalar-first (w, x, y, z) unless noted.
+"""
+
+from pathlib import Path
+
+import numpy as np
+from scipy.spatial.transform import Rotation as R  # noqa: N817
+
+# 24-joint parent tree used by the headset body-pose stream (SMPL-X body subset).
+# Matches PoseStreamer.parent_indices in gear_sonic's pico_manager_thread_server.py.
+BODY24_PARENTS = np.array(
+    [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14, 16, 17, 18, 19, 20, 22],
+    dtype=np.int64,
+)
+
+# FK output joints: first 22 SMPL body joints + two thumb tips (SMPL-X indices 39, 54).
+OUTPUT_JOINT_INDEX = np.concatenate([np.arange(22), np.array([39, 54])])
+
+_SKELETON_PATH = Path(__file__).parent / "assets" / "smpl_skeleton.npz"
+
+
+# ── quaternion helpers (scalar-first w, x, y, z) ─────────────────────────────
+
+
+def aa_to_quat(aa: np.ndarray) -> np.ndarray:
+    aa = np.asarray(aa, np.float64)
+    angle = np.linalg.norm(aa, axis=-1, keepdims=True)
+    small = angle < 1e-8
+    safe = np.where(small, 1.0, angle)
+    axis = np.where(small, 0.0, aa / safe)
+    half = angle * 0.5
+    return np.concatenate([np.cos(half), axis * np.sin(half)], axis=-1)
+
+
+def quat_to_aa(q: np.ndarray) -> np.ndarray:
+    q = np.asarray(q, np.float64)
+    w = q[..., 0:1]
+    xyz = q[..., 1:]
+    n = np.linalg.norm(xyz, axis=-1, keepdims=True)
+    angle = 2.0 * np.arctan2(n, w)
+    small = n < 1e-8
+    safe = np.where(small, 1.0, n)
+    axis = np.where(small, 0.0, xyz / safe)
+    return axis * angle
+
+
+def quat_mul(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    aw, ax, ay, az = a[..., 0], a[..., 1], a[..., 2], a[..., 3]
+    bw, bx, by, bz = b[..., 0], b[..., 1], b[..., 2], b[..., 3]
+    return np.stack(
+        [
+            aw * bw - ax * bx - ay * by - az * bz,
+            aw * bx + ax * bw + ay * bz - az * by,
+            aw * by - ax * bz + ay * bw + az * bx,
+            aw * bz + ax * by - ay * bx + az * bw,
+        ],
+        axis=-1,
+    )
+
+
+def quat_conj(q: np.ndarray) -> np.ndarray:
+    return np.concatenate([q[..., 0:1], -q[..., 1:]], axis=-1)
+
+
+def quat_inv(q: np.ndarray) -> np.ndarray:
+    c = quat_conj(q)
+    return c / (np.linalg.norm(c, axis=-1, keepdims=True) + 1e-12)
+
+
+def quat_apply(q: np.ndarray, v: np.ndarray) -> np.ndarray:
+    xyz = q[..., 1:]
+    w = q[..., 0:1]
+    t = 2.0 * np.cross(xyz, v)
+    return v + w * t + np.cross(xyz, t)
+
+
+def smpl_root_ytoz_up(root_quat: np.ndarray) -> np.ndarray:
+    """Rotate the root quaternion 90 deg about X to map SMPL Y-up to robot Z-up."""
+    base = aa_to_quat(np.array([np.pi / 2.0, 0.0, 0.0]))
+    return quat_mul(base, root_quat)
+
+
+def remove_smpl_base_rot(root_quat: np.ndarray) -> np.ndarray:
+    """Conjugate out SMPL's default rest orientation ([0.5, 0.5, 0.5, 0.5])."""
+    base_conj = quat_conj(np.array([0.5, 0.5, 0.5, 0.5]))
+    return quat_mul(root_quat, base_conj)
+
+
+# ── forward kinematics ───────────────────────────────────────────────────────
+
+
+def canonicalize_smpl_joints(smpl_joints: np.ndarray, root_aa: np.ndarray) -> np.ndarray:
+    """Remove per-frame root orientation -> SONIC ``smpl_joints_local`` format.
+
+    Mirrors the deploy transform (and ``motion_loader.canonicalize_smpl_joints``):
+    reference clips store world-frame joints, but the encoder wants each frame's
+    joints with the body root orientation removed.
+
+    Args:
+        smpl_joints: (T, 24, 3) world-frame (z-up) SMPL joint positions.
+        root_aa: (T, 3) SMPL global-orient axis-angle (y-up convention).
+
+    Returns:
+        (T, 24, 3) per-frame root-orientation-removed joints.
+    """
+    rx90 = R.from_euler("x", 90, degrees=True)  # smpl_root_ytoz_up
+    base120 = R.from_quat([0.5, 0.5, 0.5, 0.5])  # remove_smpl_base_rot
+    a = rx90 * R.from_rotvec(root_aa)
+    b_inv = base120 * a.inv()
+    return np.einsum("tij,tkj->tki", b_inv.as_matrix(), smpl_joints).astype(np.float32)
+
+
+def root_quats_from_aa(root_aa: np.ndarray) -> np.ndarray:
+    """Per-frame root orientation as (T, 4) wxyz, matching the live ``root_quat``.
+
+    Same convention as the headset stream: ytoz-up then base-rotation removed.
+    """
+    rx90 = R.from_euler("x", 90, degrees=True)
+    base120 = R.from_quat([0.5, 0.5, 0.5, 0.5])
+    r = (rx90 * R.from_rotvec(root_aa)) * base120.inv()
+    q = r.as_quat()  # xyzw
+    return np.concatenate([q[:, 3:4], q[:, :3]], axis=1).astype(np.float32)  # -> wxyz
+
+
+class SmplForwardKinematics:
+    """Rest-skeleton SMPL forward kinematics (no mesh, no torch)."""
+
+    def __init__(self, skeleton_path: str | Path = _SKELETON_PATH):
+        data = np.load(skeleton_path)
+        self.J = data["J"].astype(np.float64)  # (55, 3) rest joint positions
+        self.parents = data["parents"].astype(np.int64)  # (55,) kinematic tree
+        self.n_joints = self.J.shape[0]
+
+    def _fk(self, full_pose_aa: np.ndarray) -> np.ndarray:
+        """full_pose_aa: (n_joints, 3) axis-angle (joint 0 = global). Returns (24, 3)."""
+        rot = R.from_rotvec(full_pose_aa).as_matrix()  # (n, 3, 3)
+        rel = self.J.copy()
+        rel[1:] -= self.J[self.parents[1:]]
+
+        transforms = np.zeros((self.n_joints, 4, 4), np.float64)
+        transforms[:, :3, :3] = rot
+        transforms[:, :3, 3] = rel
+        transforms[:, 3, 3] = 1.0
+
+        chain = [transforms[0]]
+        for i in range(1, self.n_joints):
+            chain.append(chain[self.parents[i]] @ transforms[i])
+        joints = np.stack(chain)[:, :3, 3]
+        return joints[OUTPUT_JOINT_INDEX]
+
+    def compute(self, body_poses_np: np.ndarray) -> dict:
+        """Convert (24, 7) headset body poses to canonical SMPL joints.
+
+        Args:
+            body_poses_np: (24, 7) rows of [x, y, z, qx, qy, qz, qw] (scalar-last).
+
+        Returns:
+            dict with:
+              - smpl_joints_local: (24, 3) root-orientation-removed joints
+              - root_quat: (4,) root/torso orientation (w, x, y, z)
+              - root_transl: (3,) pelvis translation
+        """
+        body_poses_np = np.asarray(body_poses_np, np.float64)
+        positions = body_poses_np[:, :3]
+
+        # Global joint rotations from the headset (scalar-last), with the SMPL
+        # +180 deg-about-Y frame fix, converted to per-joint local axis-angle.
+        global_rots = R.from_quat(body_poses_np[:, 3:7]) * R.from_euler("y", 180, degrees=True)
+        gm = global_rots.as_matrix()  # (24, 3, 3)
+
+        local_aa = np.zeros((24, 3), np.float64)
+        for i in range(24):
+            p = BODY24_PARENTS[i]
+            m = gm[i] if p == -1 else gm[p].T @ gm[i]
+            local_aa[i] = R.from_matrix(m).as_rotvec()
+
+        global_orient = local_aa[0]
+        body_pose = local_aa[1:].reshape(-1)[:63]  # 21 body joints
+
+        # Root: Y-up -> Z-up, then run FK with the transformed root.
+        root_quat = aa_to_quat(global_orient)
+        root_quat = smpl_root_ytoz_up(root_quat)
+        global_orient_new = quat_to_aa(root_quat)
+
+        full_pose = np.concatenate([global_orient_new, body_pose, np.zeros(3 * self.n_joints - 66)]).reshape(
+            self.n_joints, 3
+        )
+        joints = self._fk(full_pose)  # (24, 3)
+
+        # Canonicalize: strip SMPL base rot and the root orientation.
+        root_quat = remove_smpl_base_rot(root_quat)
+        inv = np.broadcast_to(quat_inv(root_quat), (joints.shape[0], 4))
+        smpl_joints_local = quat_apply(inv, joints)
+
+        return {
+            "smpl_joints_local": smpl_joints_local.astype(np.float32),
+            "root_quat": root_quat.astype(np.float32),
+            "root_transl": positions[0].astype(np.float32),
+        }