feat(unitree_g1): use captured neutral SONIC token instead of zeros

The all-zero token is off the encoder's learned FSQ manifold and decodes to a
slightly goofy stance. Replace it with a NEUTRAL_TOKEN captured from the encoder's
own idle output in sim (stored as integer FSQ codes, rescaled by the encoder's
1/16 quantization step to an exact on-grid token). token_mode now seeds this
neutral, and the onboard sender starts observation.state from it so the first
inference sees the token the decoder is actually holding.

src/lerobot/robots/unitree_g1/README.md

@@ -1,162 +1,89 @@
-# Unitree G1 — SONIC whole-body control
+# Unitree G1 — SONIC encoder/decoder whole-body control
 
-This package runs NVIDIA's **SONIC** whole-body controller (and the GR00T/Holosoma
-locomotion controllers) on the Unitree G1, in MuJoCo simulation or on real hardware.
-SONIC turns a high-level movement intent — or a streamed **SMPL** whole-body pose — into
-50 Hz joint-position targets. It is a pure-Python/ONNX reimplementation of the SONIC
-deploy stack (no `gear_sonic`/torch dependency).
+This package runs NVIDIA's **SONIC** encoder/decoder on the Unitree G1, in MuJoCo
+simulation or on real hardware, driven by a dense **34-D whole-body command** (the
+OpenHLM / pi0.5 action layout). It is a pure-Python/ONNX reimplementation of the
+reference-tracking half of the SONIC deploy stack (no `gear_sonic`/torch dependency, and
+no motion planner): the encoder compresses a reference motion window into a latent token
+and the decoder maps that token + proprioception history into 50 Hz joint-position
+targets for the robot's PD controller.
 
 ## Controllers
 
 Selected with `--robot.controller=<ClassName>`:
 
-| Controller                     | Purpose                                                                                    |
-| ------------------------------ | ------------------------------------------------------------------------------------------ |
-| `SonicWholeBodyController`     | SONIC whole-body: locomotion (mode 0), 3-point VR teleop (mode 1), SMPL imitation (mode 2) |
-| `GrootLocomotionController`    | GR00T locomotion policy                                                                    |
-| `HolosomaLocomotionController` | Holosoma locomotion policy                                                                 |
+| Controller                     | Purpose                                                      |
+| ------------------------------ | ------------------------------------------------------------ |
+| `SonicWholeBodyController`     | SONIC encoder/decoder driven by a 34-D OpenHLM/pi0.5 command |
+| `GrootLocomotionController`    | GR00T locomotion policy                                      |
+| `HolosomaLocomotionController` | Holosoma locomotion policy                                   |
 
-On startup the controller **interpolates** from the robot's measured pose into the
-policy's commanded target over ~3 s (no snap), and on disconnect (Ctrl-C) it performs a
-**graceful damped settle** — holding pose while ramping stiffness to zero over
-`--robot.graceful_stop_s` (default 1.5 s) instead of going instantly limp. Both apply in
-every mode.
+The rest of this document covers the SONIC whole-body path.
+
+Each tick the `SonicWholeBodyController` takes a 34-D command (`wb.0.pos … wb.33.pos`) in the OpenHLM
+layout:
+
+```
+[L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
+ root roll/pitch + yaw-rate(3)]
+```
+
+The 29 joint targets become the SONIC encode-mode-0 reference (accumulated into a rolling
+50-frame trajectory with finite-difference velocities so the encoder's lookahead sees a
+real motion sequence), the root roll/pitch set the anchor orientation, and the two grip
+scalars can drive the Dex3 hands (see below). On startup the controller **interpolates**
+from the robot's measured pose into the policy's commanded target over ~3 s (no snap).
 
 ## Requirements
 
-- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended — SONIC runs three ONNX
-  sessions and is much smoother on GPU). Install the CUDA build that matches your
-  driver (e.g. `onnxruntime-gpu==1.26.0` for a CUDA-12.x driver). Verify with:
+- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended). Verify with:
   ```bash
   python -c "import onnxruntime as ort; print(ort.get_available_providers())"
-  # expect CUDAExecutionProvider in the list for GPU
   ```
 - `mujoco` for simulation (`is_simulation=True`).
-- `pyzmq` only if you use the live SMPL stream (pico headset).
-- The SONIC ONNX models are downloaded automatically from the `nvidia/GEAR-SONIC` Hub repo.
+- The SONIC encoder/decoder ONNX models download automatically from the
+  `nvidia/GEAR-SONIC` Hub repo.
 
-## Running
+## Running a rollout
 
-**Replay an SMPL dataset (motion imitation):**
+Drive the G1 with a 34-D VLA policy (OpenHLM / pi0.5) via `lerobot-rollout`:
 
 ```bash
-lerobot-replay \
-  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
-  --dataset.repo_id=<user>/<smpl_dataset> --dataset.episode=0
+lerobot-rollout \
+  --strategy.type=base \
+  --policy.path=<pi05_openhlm_dir> \
+  --robot.type=unitree_g1 \
+  --robot.controller=SonicWholeBodyController \
+  --robot.is_simulation=true \
+  --robot.publish_hands=true \
+  --task="<language instruction>" \
+  --duration=45 --device=cuda
 ```
 
-**Keyboard teleop** (drives locomotion via the native keyboard teleoperator):
+### Cameras
 
-```bash
-lerobot-teleoperate \
-  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
-  --teleop.type=keyboard
-```
+Image-conditioned policies need camera frames. Two options are available without live
+cameras:
 
-Controls: `WASD` move · `Q`/`E` turn · `1`–`8` mode · `9`/`0` speed · `-`/`=` height ·
-`R` replan · `Space` emergency-stop.
+- **Black frames**: `--robot.empty_cameras='[base, left_wrist, right_wrist]'`.
+- **Replay a recorded episode** as the camera feed:
+  ```bash
+  --robot.replay_camera_parquet=<episode.parquet> \
+  --robot.replay_camera_map='{base: head_image_left, left_wrist: left_wrist_image, right_wrist: right_wrist_image}'
+  ```
 
-**PICO headset teleop — SMPL whole-body** (mode 2, needs PICO Motion Trackers):
+### Hands (Dex3)
 
-```bash
-# 1) publisher (streams rt/smpl from full-body tracking)
-python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50
-# 2) controller
-lerobot-teleoperate \
-  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
-  --teleop.type=pico_headset
-```
+`--robot.publish_hands=true` publishes `rt/dex3/{left,right}/cmd` from the two grip
+scalars (`wb.7.pos` left, `wb.15.pos` right). The scalar is interpolated between
+`hand_open_grip_value` (default 1.0 = open) and `hand_closed_grip_value` (default 0.0 =
+closed) and scaled onto `hand_closed_pose` (7 joints:
+`thumb_0, thumb_1, thumb_2, middle_0, middle_1, index_0, index_1`). Flip the signs in
+`hand_closed_pose` if the fingers curl the wrong way, or raise `hand_kp` for a firmer
+grip.
 
-**PICO headset teleop — 3-point VR** (mode 1, head + 2 controllers only, **no trackers**):
+## Observation state
 
-```bash
-# 1) publisher (head + controllers -> 3-point targets + stick locomotion)
-python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --headset-source devices
-# 2) controller
-lerobot-teleoperate \
-  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
-  --teleop.type=pico_headset --teleop.mode=vr3
-```
-
-3-point controls: left stick move · right stick X turn · right stick Y height ·
-`A`+`B` / `X`+`Y` cycle locomotion mode (walk/run/squat/kneel/…) · hands+head track the
-upper body. **Calibration**: stand in a neutral rest pose and press `A`+`B`+`X`+`Y` — the
-publisher status line flips from `UNCALIBRATED` to `calibrated`. This maps your rest pose
-onto the G1's neutral stance and is required before the hands track well; the SMPL
-(mode 2) path is self-calibrating and needs no such step.
-
-Both require the XRoboToolkit stack — see below.
-
-## PICO headset / XRoboToolkit install
-
-Live full-body teleop needs the **XRoboToolkit** system (a PC Service on your
-workstation + a PICO app on the headset) and its Python binding, `xrobotoolkit_sdk`.
-The full hardware + software walkthrough lives in the SONIC repo:
-[`docs/source/getting_started/vr_teleop_setup.md`](https://nvlabs.github.io/GR00T-WholeBodyControl/getting_started/vr_teleop_setup.html).
-
-Summary:
-
-1. **PC Service** (workstation) — install and run it before connecting the headset.
-   - Ubuntu 22.04 / 24.04 (x86_64): prebuilt `.deb` from the
-     [XRoboToolkit-PC-Service releases](https://github.com/XR-Robotics/XRoboToolkit-PC-Service/releases).
-   - Jetson (aarch64): the arm64 `.deb`.
-   - Windows (x64): the Windows PC Service build.
-2. **PICO app** — install `XRoboToolkit-PICO-*.apk` on the headset (see the guide),
-   enable Developer Mode. For **SMPL whole-body** (mode 2) you also need the PICO Motion
-   Trackers paired/calibrated and "Full body" enabled; for **3-point** (mode 1,
-   `--headset-source devices`) only Head + Controller + Send are required — no trackers.
-3. **`xrobotoolkit_sdk`** — a pybind11/CMake build (not a pip package), from
-   [`XRoboToolkit-PC-Service-Pybind`](https://github.com/XR-Robotics/XRoboToolkit-PC-Service-Pybind):
-   - Linux x86_64: `pip install pybind11 cmake` then `bash setup_ubuntu.sh` (or the
-     SONIC repo's `install_scripts/install_pico.sh`, which builds everything into a
-     `.venv_teleop`).
-   - Jetson aarch64: `bash setup_orin.sh` (builds `libPXREARobotSDK.so` from source).
-   - Windows x64: `pip install pybind11` then `setup_windows.bat` (needs git + an
-     MSVC/CMake toolchain; uses the prebuilt `PXREARobotSDK.dll`/`.lib`).
-4. Connect PICO and workstation to the **same Wi-Fi**, open the XRoboToolkit app, enter
-   the PC IP, and enable Head/Controller/Send (plus Full-body for SMPL mode 2).
-
-### Platform support
-
-| Platform                    | Live headset teleop | Notes                                       |
-| --------------------------- | ------------------- | ------------------------------------------- |
-| Linux x86_64                | ✅                  | Guided `install_pico.sh` (SONIC repo)       |
-| Linux aarch64 (Jetson Orin) | ✅                  | `setup_orin.sh` builds the native lib       |
-| Windows x64                 | ✅ (manual)         | `setup_windows.bat`; no one-shot env script |
-| macOS                       | ❌                  | No PC Service / SDK build for Darwin        |
-
-### No hardware required (any platform, incl. macOS/Windows)
-
-The SMPL pipeline can be exercised without a headset or the SDK — the publisher emits
-`rt/smpl` frames that the controller consumes exactly as it would from the headset:
-
-```bash
-# synthetic motion
-python -m lerobot.teleoperators.pico_headset.pico_publisher --fake
-
-# replay a canned SMPL clip
-python -m lerobot.teleoperators.pico_headset.pico_publisher --motion-file <clip>.npz
-```
-
-## Notes
-
-- SMPL **root motion** into the mode-2 anchor is opt-in (`SonicWholeBodyController(enable_smpl_root=True)`);
-  it stays off by default (untested on hardware). When enabled, the per-frame root quat is
-  spherically smoothed (`root_smoothing_alpha`, default 0.15) before it reaches the anchor,
-  which removes the base-acceleration spikes the raw 30 Hz→50 Hz trajectory used to cause.
-- Direct `rt/smpl` subscription without the pico teleoperator is available via
-  `SonicWholeBodyController(enable_smpl_stream=True, smpl_host=..., smpl_port=...)`.
-- 3-point (mode 1) uses the **headset-yaw frame** as its reference and the `A`+`B`+`X`+`Y`
-  calibration to align to the G1 neutral stance. Calibration maps the operator's rest pose
-  onto the G1's **standing** (`default_angles`) wrist/neck key-frame poses (position **and**
-  orientation) computed by FK — the `default_angles` stand-in for gear_sonic's live
-  measured-q recalibration, since the robot holds `default_angles` at calibration time.
-  Re-aligning the arms only (preserving the neck level) is available via the calibrator's
-  `recalibrate_wrists()`.
-- 3-point **locomotion** from the PICO sticks follows gear_sonic's `PlannerLoop` exactly:
-  a yaw accumulator on the right stick and **mode-dependent speed curves** on the left
-  (slow `0.1+0.5·mag`, run `1.5+3·mag`, walk = planner default). Stick signs replicate
-  gear_sonic's `get_controller_axes` usage (forward `+ly`, strafe `-lx`, turn `-rx`); since
-  the publisher forwards the same raw SDK axes, this is the correct convention by construction.
-- Startup interpolation and the graceful-stop settle are mode-agnostic; set
-  `--robot.graceful_stop_s=0` to restore the old instant zero-torque on disconnect.
+When the whole-body controller is active the robot exposes a 34-D proprio state
+(`wb_state.0.pos … wb_state.33.pos`) in the same OpenHLM layout as the action, which the
+rollout aggregates into `observation.state` for the policy.

src/lerobot/robots/unitree_g1/config_unitree_g1.py

@@ -62,6 +62,24 @@ class UnitreeG1Config(RobotConfig):
     # Socket config for ZMQ bridge
     robot_ip: str = "192.168.123.164"  # default G1 IP
 
+    # Run the locomotion / whole-body controller ONBOARD the robot (policy on the G1
+    # itself, against local DDS at full rate) instead of on the laptop over the ZMQ
+    # socket bridge. In this mode the robot object uses the real Unitree SDK channels
+    # and expects high-level actions (arm targets + joystick axes, or 64-D SONIC
+    # tokens) fed via send_action -- e.g. by run_g1_onboard.py, which receives them
+    # from the laptop over ZMQ. Mutually exclusive with is_simulation.
+    onboard: bool = False
+    # DDS network interface for onboard mode (None = SDK default, matching
+    # run_g1_server.py's ChannelFactoryInitialize(0)).
+    dds_interface: str | None = None
+    # Onboard sub-flags. On a real G1 both are True: the built-in motion services
+    # must be released before we can write lowcmd, and locomotion axes are read from
+    # the physical wireless remote. Against a DDS sim neither applies (no
+    # MotionSwitcher, no physical remote), so set both False so the controller takes
+    # its locomotion axes purely from send_action (ZMQ) input.
+    release_motion_control: bool = True
+    physical_remote: bool = True
+
     # Cameras (ZMQ-based remote cameras)
     cameras: dict[str, CameraConfig] = field(default_factory=dict)
 
@@ -78,11 +96,13 @@ class UnitreeG1Config(RobotConfig):
     # ``hand_closed_pose`` (7 joints: thumb_0/1/2, middle_0/1, index_0/1). Flip signs
     # in ``hand_closed_pose`` if fingers curl the wrong way.
     publish_hands: bool = False
+    # When False, connect() does not start the background controller thread, so a
+    # caller can drive the controller synchronously (one decode per fed action),
+    # reproducing the deploy's single 50Hz control clock for faithful replay.
+    run_controller_thread: bool = True
     hand_open_grip_value: float = 1.0
     hand_closed_grip_value: float = 0.0
-    hand_closed_pose: list[float] = field(
-        default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3]
-    )
+    hand_closed_pose: list[float] = field(default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3])
     hand_kp: float = 1.5
     hand_kd: float = 0.1
 
@@ -94,6 +114,16 @@ class UnitreeG1Config(RobotConfig):
     replay_camera_map: dict[str, str] = field(default_factory=dict)
     replay_camera_loop: bool = True
 
+    # Token-output VLA interface for the SONIC decoder. When True (and the controller
+    # is ``SonicWholeBodyController``), the robot advertises a 64-D latent-token action
+    # space (``motion_token.{i}.pos``) instead of the 34-D whole-body command, and
+    # exposes the last commanded token as a 64-D ``observation.state``
+    # (``motion_token_state.{i}.pos``). This lets ``lerobot-rollout`` drive a policy
+    # that was trained with 64-D SONIC motion tokens as both state and action
+    # (e.g. nepyope/sonic_walk): the decoder consumes the token directly, encoder
+    # bypassed. Ignored unless a SONIC whole-body controller is active.
+    sonic_token_action: bool = False
+
     # Compensates for gravity on the unitree's arms using the arm ik solver
     gravity_compensation: bool = False
 

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

@@ -16,9 +16,13 @@
 
 """Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
 
+from .gr00t_locomotion import GrootLocomotionController
+from .holosoma_locomotion import HolosomaLocomotionController
+from .sonic_whole_body import SonicRuntime, SonicWholeBodyController
+
 __all__ = [
     "GrootLocomotionController",
     "HolosomaLocomotionController",
-    "SonicWholeBodyController",
     "SonicRuntime",
+    "SonicWholeBodyController",
 ]

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

@@ -77,9 +77,15 @@ def load_groot_policies(
         filename="GR00T-WholeBodyControl-Walk.onnx",
     )
 
-    # Load ONNX policies
-    policy_balance = ort.InferenceSession(balance_path)
-    policy_walk = ort.InferenceSession(walk_path)
+    # Load ONNX policies with a capped thread pool. GR00T runs at 50 Hz in a
+    # background thread alongside the (torch) upper-body policy, IK and sim; letting
+    # ORT grab every core starves those and makes the whole rollout stutter. These
+    # are small MLPs, so 1 thread is both enough and lowest-latency.
+    from .sonic_pipeline import make_ort_session_options
+
+    so = make_ort_session_options(intra_op_num_threads=1, inter_op_num_threads=1)
+    policy_balance = ort.InferenceSession(balance_path, sess_options=so)
+    policy_walk = ort.InferenceSession(walk_path, sess_options=so)
 
     logger.info("GR00T policies loaded successfully")
 
@@ -206,6 +212,16 @@ class GrootLocomotionController:
         # Transform action back to target joint positions
         target_dof_pos_15 = GROOT_DEFAULT_ANGLES[:15] + self.groot_action * ACTION_SCALE
 
+        # Waist override: an external upper-body IK can command the 3 waist joints
+        # (indices 12/13/14) via ``kWaist{Yaw,Roll,Pitch}.q`` in the action dict. When
+        # present, we substitute the balance policy's waist target so the torso tracks
+        # the IK while the policy keeps only the legs balanced. Single-publisher stays
+        # intact (this thread still owns joints 0-14).
+        for idx in (G1_29_JointIndex.kWaistYaw, G1_29_JointIndex.kWaistRoll, G1_29_JointIndex.kWaistPitch):
+            key = f"{idx.name}.q"
+            if key in action and action[key] is not None:
+                target_dof_pos_15[idx.value] = float(action[key])
+
         # Build action dict
         action_dict = {}
         for i in range(15):

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

@@ -14,54 +14,30 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""SONIC planner pipeline for the Unitree G1 whole-body controller.
+"""SONIC encoder/decoder pipeline for the Unitree G1 whole-body controller.
 
-This module is a pure-Python/ONNX re-implementation of NVIDIA's SONIC deploy stack
-(mirrors ``g1_deploy_onnx_ref.cpp``). It turns a high-level movement intent
-(walk/run/squat/box/… + speed/height/heading, driven by the joystick) into
-50 Hz joint-position targets for the robot's PD controller.
+Pure-Python/ONNX re-implementation of the reference-tracking half of NVIDIA's SONIC
+deploy stack (mirrors ``g1_deploy_onnx_ref.cpp``). Given a reference motion buffer
+(joint targets + body orientation per frame) it produces 50 Hz joint-position targets
+for the robot's PD controller. The upstream *motion planner* is intentionally absent:
+here the reference is supplied directly by the caller (e.g. a 34-D OpenHLM / pi0.5 VLA
+command per tick, in ``sonic_whole_body.py``).
 
-Data flow (one 50 Hz control tick, orchestrated by ``SonicRuntime`` in
-``sonic_whole_body.py``):
-
-    intent (MovementState) ──► SonicPlanner ──► PlannerController ──► joint targets
-                                   │                    │
-                     (planner ONNX, 30 Hz,       (encoder+decoder ONNX,
-                      async background thread)     runs every tick)
-
-Three cooperating ONNX models:
-  * **planner**  – generates a several-second *reference motion* (body trajectory +
-    joint clip) for the current intent. Slow, so it runs asynchronously in a
-    background thread (``_planner_worker``) and its 30 Hz output is resampled to
-    50 Hz. New motions are cross-faded into the live buffer (``blend_new_motion``).
-  * **encoder**  – compresses the reference window into a 64-D latent ``token``
+Two cooperating ONNX models:
+  * **encoder** – compresses the reference window into a 64-D latent ``token``
     (refreshed every ``ENCODER_UPDATE_EVERY`` ticks).
-  * **decoder**  – every tick, maps the token + recent proprioception history to a
+  * **decoder** – every tick, maps the token + recent proprioception history to a
     residual action that is scaled and added to ``DEFAULT_ANGLES``.
 
-Encoder ``encode_mode`` selects what the reference represents:
-  * ``0`` – locomotion (planner clip drives lower + upper body).
-  * ``1`` – 3-point VR teleop (lower body from planner, arms from VR targets).
-  * ``2`` – SMPL whole-body imitation (720-D SMPL window drives the pose).
-
 Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order)
 and **MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` convert
 between them. Quaternions are scalar-first ``(w, x, y, z)``.
-
-Section map: constants & index tables · PD gains · quaternion helpers · locomotion
-modes · movement state · encoder/decoder · planner motion buffer · async planner
-worker · ``SonicPlanner`` · ``PlannerController`` · joystick input.
 """
 
 from __future__ import annotations
 
 import logging
-import math
-import queue
 import threading
-import time
-from dataclasses import dataclass
-from enum import IntEnum
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -149,19 +125,18 @@ DEFAULT_HEIGHT = 0.788740  # nominal pelvis height (m)
 TOKEN_DIM = 64  # encoder latent size
 ENCODER_UPDATE_EVERY = 5  # refresh the encoder token every N ticks (decoder runs every tick)
 DEBUG_PRINT_EVERY = 100  # ticks between debug prints
-MOTION_LOOK_AHEAD_STEPS = 2  # frames ahead used to seed a replan context (hide planner latency)
-INITIAL_RANDOM_SEED = 1234
-MIN_TOKENS, MAX_TOKENS = 6, 16  # planner prediction-length token range
-K = MAX_TOKENS - MIN_TOKENS + 1
-DEADZONE = 0.05  # joystick dead zone
-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}
 
 
 def _to_mujoco(a):
-    """Reorder a 29-vector from IsaacLab order into MuJoCo/deploy order."""
+    """Apply the ``MUJOCO_TO_ISAACLAB`` gather to a 29-vector (deploy-order reorder).
+
+    NOTE: this returns ``a[MUJOCO_TO_ISAACLAB]``. The ``_mj`` suffixes and the exact
+    permutation direction throughout this module are a fixed convention validated
+    against the deployed SONIC ONNX policy (the encoder/decoder consume vectors in
+    this order). Do not "correct" the table or rename toward the opposite direction
+    without re-validating on hardware — the labels are historical, the ordering is
+    load-bearing.
+    """
     return a[MUJOCO_TO_ISAACLAB]
 
 
@@ -294,139 +269,6 @@ def quat_slerp_batch(q0, q1, t):
     return r / (np.linalg.norm(r, axis=1, keepdims=True) + 1e-12)
 
 
-# ── Locomotion modes ──────────────────────────────────────────────────────────
-
-
-class LocomotionMode(IntEnum):
-    """High-level motion styles understood by the planner (fed as the ``mode`` input)."""
-
-    IDLE = 0
-    SLOW_WALK = 1
-    WALK = 2
-    RUN = 3
-    SQUAT = 4
-    KNEEL_TWO_LEGS = 5
-    KNEEL = 6
-    LYING_FACE_DOWN = 7
-    CRAWLING = 8
-    IDLE_BOXING = 9
-    WALK_BOXING = 10
-    LEFT_PUNCH = 11
-    RIGHT_PUNCH = 12
-    RANDOM_PUNCH = 13
-    ELBOW_CRAWLING = 14
-    LEFT_HOOK = 15
-    RIGHT_HOOK = 16
-    FORWARD_JUMP = 17
-    STEALTH_WALK = 18
-    INJURED_WALK = 19
-    LEDGE_WALKING = 20
-    OBJECT_CARRYING = 21
-    STEALTH_WALK_2 = 22
-    HAPPY_DANCE_WALK = 23
-    ZOMBIE_WALK = 24
-    GUN_WALK = 25
-    SCARE_WALK = 26
-
-
-LM = LocomotionMode
-
-# UI groupings of modes for cycling with the n/p keys; each entry is (label, modes).
-MOTION_SETS = [
-    ("Standing", [LM.SLOW_WALK, LM.WALK, LM.RUN, LM.FORWARD_JUMP, LM.STEALTH_WALK, LM.INJURED_WALK]),
-    ("Squat / Low", [LM.SQUAT, LM.KNEEL_TWO_LEGS, LM.KNEEL, LM.CRAWLING, LM.ELBOW_CRAWLING]),
-    (
-        "Boxing",
-        [
-            LM.IDLE_BOXING,
-            LM.WALK_BOXING,
-            LM.LEFT_PUNCH,
-            LM.RIGHT_PUNCH,
-            LM.RANDOM_PUNCH,
-            LM.LEFT_HOOK,
-            LM.RIGHT_HOOK,
-        ],
-    ),
-    (
-        "Styled Walks",
-        [
-            LM.LEDGE_WALKING,
-            LM.OBJECT_CARRYING,
-            LM.STEALTH_WALK_2,
-            LM.HAPPY_DANCE_WALK,
-            LM.ZOMBIE_WALK,
-            LM.GUN_WALK,
-            LM.SCARE_WALK,
-        ],
-    ),
-]
-
-# Mode classifications used by clamping and replan logic.
-STATIC_MODES = {LM.IDLE, LM.SQUAT, LM.KNEEL_TWO_LEGS, LM.KNEEL, LM.LYING_FACE_DOWN, LM.IDLE_BOXING}
-STANDING_MODES = {
-    LM.IDLE,
-    LM.SLOW_WALK,
-    LM.WALK,
-    LM.RUN,
-    LM.IDLE_BOXING,
-    LM.WALK_BOXING,
-    LM.LEFT_PUNCH,
-    LM.RIGHT_PUNCH,
-    LM.RANDOM_PUNCH,
-    LM.LEFT_HOOK,
-    LM.RIGHT_HOOK,
-    LM.FORWARD_JUMP,
-    LM.STEALTH_WALK,
-    LM.INJURED_WALK,
-    LM.LEDGE_WALKING,
-    LM.OBJECT_CARRYING,
-    LM.STEALTH_WALK_2,
-    LM.HAPPY_DANCE_WALK,
-    LM.ZOMBIE_WALK,
-    LM.GUN_WALK,
-    LM.SCARE_WALK,
-}
-BOXING_MODES = {LM.WALK_BOXING, LM.LEFT_PUNCH, LM.RIGHT_PUNCH, LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK}
-SPEED_RANGES = {
-    LM.SLOW_WALK: (0.2, 0.8),
-    LM.WALK: (0.8, 1.5),
-    LM.RUN: (1.5, 3.0),
-    LM.CRAWLING: (0.4, 1.0),
-    LM.ELBOW_CRAWLING: (0.7, 1.0),
-}
-
-
-def clamp_mode_params(ms):
-    """Clamp ``ms.speed``/``ms.height`` into the valid range for its mode in place.
-
-    ``-1.0`` is a sentinel meaning "use the mode's default" (e.g. standing modes
-    ignore height; static modes ignore speed).
-    """
-    m = LM(ms.mode)
-    ms.height = -1.0 if m in STANDING_MODES else max(0.1, min(0.8, ms.height if ms.height >= 0 else 0.2))
-    if m in STATIC_MODES:
-        ms.speed = -1.0
-    elif m in SPEED_RANGES:
-        lo, hi = SPEED_RANGES[m]
-        ms.speed = max(lo, min(hi, ms.speed if ms.speed >= 0 else lo))
-    elif m in BOXING_MODES:
-        ms.speed = max(0.7, min(1.5, ms.speed if ms.speed >= 0 else 0.7))
-    else:
-        ms.speed = -1.0
-
-
-def replan_interval(mode):
-    """Seconds between automatic replans for the given mode."""
-    m = LM(mode)
-    if m == LM.RUN:
-        return REPLAN_INTERVAL["running"]
-    if m == LM.CRAWLING:
-        return REPLAN_INTERVAL["crawling"]
-    if m in {LM.LEFT_PUNCH, LM.RIGHT_PUNCH, LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK}:
-        return REPLAN_INTERVAL["boxing"]
-    return REPLAN_INTERVAL["default"]
-
-
 def ort_providers(force_cpu: bool = False) -> list[str]:
     """Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed)."""
     avail = ort.get_available_providers()
@@ -435,99 +277,24 @@ def ort_providers(force_cpu: bool = False) -> list[str]:
     return ["CPUExecutionProvider"]
 
 
-def make_ort_session_options():
-    """Build ONNX Runtime SessionOptions (quiet logging, default threading)."""
+def make_ort_session_options(intra_op_num_threads: int | None = None,
+                             inter_op_num_threads: int | None = None):
+    """Build ONNX Runtime SessionOptions (quiet logging).
+
+    Pass thread counts to cap ORT's CPU pool. These tiny MLP policies are latency-
+    bound, not throughput-bound, so letting ORT grab every core just starves the
+    real-time control loop / torch policy / IK solver and causes stutter. 1 intra +
+    1 inter thread is plenty and lowest-latency for a per-step MLP inference.
+    """
     so = ort.SessionOptions()
     so.log_severity_level = 3
+    if intra_op_num_threads is not None:
+        so.intra_op_num_threads = intra_op_num_threads
+    if inter_op_num_threads is not None:
+        so.inter_op_num_threads = inter_op_num_threads
     return so
 
 
-# ── Movement state ────────────────────────────────────────────────────────────
-
-
-@dataclass
-class MovementState:
-    """Mutable high-level intent driven by keyboard/joystick and read by the planner.
-
-    Holds the current locomotion ``mode``, target ``speed``/``height`` (``-1`` =
-    mode default), facing/movement angles, and the ``needs_replan`` flag the control
-    loop watches to decide when to request a fresh motion from the planner.
-    """
-
-    mode: int = LM.SLOW_WALK  # not IDLE — walking modes respond to WASD
-    speed: float = -1.0
-    height: float = -1.0
-    facing_angle: float = 0.0
-    movement_angle: float = 0.0
-    has_movement: bool = False
-    motion_set_idx: int = 0
-    needs_replan: bool = False
-    joy_prev_active: bool = False  # tracks right-stick activity across joystick polls
-
-    @property
-    def movement_direction(self):
-        """Unit XY movement direction (0 vector when not moving)."""
-        if not self.has_movement:
-            return (0.0, 0.0, 0.0)
-        return (math.cos(self.movement_angle), math.sin(self.movement_angle), 0.0)
-
-    @property
-    def facing_direction(self):
-        """Unit XY facing direction."""
-        return (math.cos(self.facing_angle), math.sin(self.facing_angle), 0.0)
-
-    def status_line(self):
-        """Human-readable one-line status for the terminal HUD."""
-        return (
-            f"[{MOTION_SETS[self.motion_set_idx][0]}] mode={self.mode}({LM(self.mode).name}) "
-            f"spd={'default' if self.speed < 0 else f'{self.speed:.1f}'} "
-            f"hgt={'default' if self.height < 0 else f'{self.height:.2f}'} "
-            f"facing={math.degrees(self.facing_angle):.0f}° "
-            f"{'moving' if self.has_movement else 'still'}"
-        )
-
-
-@dataclass
-class MovementSnapshot:
-    """Immutable copy of the intent at the last replan, for change detection."""
-
-    mode: int = 0
-    speed: float = -1.0
-    height: float = -1.0
-    movement: tuple[float, float, float] = (0.0, 0.0, 0.0)
-    facing: tuple[float, float, float] = (1.0, 0.0, 0.0)
-
-
-def snapshot_ms(ms: MovementState) -> MovementSnapshot:
-    """Capture the current movement intent as a comparable snapshot."""
-    md, fd = ms.movement_direction, ms.facing_direction
-    return MovementSnapshot(ms.mode, ms.speed, ms.height, (md[0], md[1], md[2]), (fd[0], fd[1], fd[2]))
-
-
-def should_replan_request(ms: MovementState, last: MovementSnapshot, replan_timer: float, step: int) -> bool:
-    """Decide whether to request a fresh plan this tick.
-
-    Triggers on an explicit ``needs_replan`` flag, any mode/facing/height change, or
-    (for non-static modes) speed/direction changes and periodic timeouts. Mirrors the
-    C++ ``G1Deploy::Planner`` replan triggers (``g1_deploy_onnx_ref.cpp``).
-    """
-    if step <= 0:
-        return False
-    if ms.needs_replan:
-        return True
-    md, fd = ms.movement_direction, ms.facing_direction
-    facing_changed = fd != last.facing
-    height_changed = ms.height != last.height
-    mode_changed = ms.mode != last.mode
-    speed_changed = ms.speed != last.speed
-    dir_changed = md != last.movement
-    is_static = LM(ms.mode) in STATIC_MODES
-    if mode_changed or facing_changed or height_changed:
-        return True
-    time_to_replan = replan_timer >= replan_interval(ms.mode)
-    return not is_static and (speed_changed or dir_changed or (time_to_replan and ms.speed != 0))
-
-
 # ── Encoder / Decoder ─────────────────────────────────────────────────────────
 
 
@@ -570,6 +337,23 @@ class StandingEncoderDecoder:
         self.smpl_root_quat = None
         self.set_zero_reference()
 
+    def reset(self):
+        """Clear the token, 10-frame proprioception history and heading init.
+
+        ``UnitreeG1.reset()`` relies on this so the first decoder outputs of a new
+        episode are not contaminated by the previous episode's state.
+        """
+        self.token = np.zeros(TOKEN_DIM, np.float32)
+        self.last_action_mj = np.zeros(29, np.float32)
+        self.h_q_mj = [np.zeros(29, np.float32)] * 10
+        self.h_dq_mj = [np.zeros(29, np.float32)] * 10
+        self.h_ang = [np.zeros(3, np.float32)] * 10
+        self.h_act_mj = [np.zeros(29, np.float32)] * 10
+        self.h_quat = [np.array([1, 0, 0, 0], np.float32)] * 10
+        self.init_base_quat = np.array([1, 0, 0, 0], np.float32)
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float32)
+        self._heading_init = False
+
     def update_history(self, q, dq, ang, quat):
         """Push the latest proprioception (pos/vel/gyro/orientation) into the 10-frame buffers."""
         quat = quat / (np.linalg.norm(quat) + 1e-8)
@@ -724,314 +508,19 @@ class StandingEncoderDecoder:
         return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
 
 
-# ── Planner motion buffer ─────────────────────────────────────────────────────
-
-
-class PlannerMotion:
-    """Fixed-capacity buffer for a planned motion (joint pos/vel + body pose per frame)."""
-
-    def __init__(self, max_frames=1500):
-        self.timesteps = 0
-        self.joint_positions = np.zeros((max_frames, 29), np.float64)
-        self.joint_velocities = np.zeros((max_frames, 29), np.float64)
-        self.body_positions = np.zeros((max_frames, 3), np.float64)
-        self.body_quaternions = np.zeros((max_frames, 4), np.float64)
-        self.body_quaternions[:, 0] = 1.0
-
-
-# ── Subprocess planner ────────────────────────────────────────────────────────
-
-
-def _resample_30_to_50(qpos, n30):
-    """Resample planner output (30 Hz MuJoCo qpos) to a 50 Hz IsaacLab-order motion.
-
-    Returns a dict with joint positions/velocities (velocities via finite difference)
-    and body position/orientation trajectories at 50 Hz.
-    """
-    t50 = int(np.floor(n30 / 30.0 * 50))
-    f30 = np.arange(t50) / 50.0 * 30.0
-    f0 = np.floor(f30).astype(int)
-    f1 = np.minimum(f0 + 1, n30 - 1)
-    frac, w0 = (f30 - f0).astype(np.float64), None
-    w0 = 1.0 - frac
-    jp = (w0[:, None] * qpos[f0, 7:36] + frac[:, None] * qpos[f1, 7:36])[:, MUJOCO_TO_ISAACLAB]
-    jv = np.zeros_like(jp)
-    if t50 >= 2:
-        jv[: t50 - 1] = (jp[1:] - jp[:-1]) * 50.0
-        jv[-1] = jv[-2]
-    return {
-        "timesteps": t50,
-        "joint_positions": jp,
-        "joint_velocities": jv,
-        "body_positions": w0[:, None] * qpos[f0, :3] + frac[:, None] * qpos[f1, :3],
-        "body_quaternions": quat_slerp_batch(qpos[f0, 3:7], qpos[f1, 3:7], frac),
-    }
-
-
-def _build_planner_inputs(ctx, ms_dict, version, seed):
-    """Build the planner ONNX input dict from a context window + movement intent.
-
-    ``version >= 1`` is the TensorRT-style deploy planner with extra height/target
-    inputs and a token-count mask; ``version 0`` is the minimal input set.
-    """
-    inp = {
-        "context_mujoco_qpos": ctx.astype(np.float32).reshape(1, 4, 36),
-        "target_vel": np.array([ms_dict["speed"]], np.float32),
-        "mode": np.array([ms_dict["mode"]], np.int64),
-        "movement_direction": np.array(ms_dict["movement_direction"], np.float32).reshape(1, 3),
-        "facing_direction": np.array(ms_dict["facing_direction"], np.float32).reshape(1, 3),
-        "random_seed": np.array([seed], np.int64),
-    }
-    if version >= 1:
-        # TensorRT deploy: allow 9–11 prediction tokens only (indices 3–5 for MIN_TOKENS=6).
-        allowed = np.zeros((1, K), np.int64)
-        if K >= 6:
-            allowed[0, 3:6] = 1
-        inp.update(
-            {
-                "height": np.array([ms_dict["height"]], np.float32),
-                "has_specific_target": np.array([[0]], np.int64),
-                "specific_target_positions": np.zeros((1, 4, 3), np.float32),
-                "specific_target_headings": np.zeros((1, 4), np.float32),
-                "allowed_pred_num_tokens": allowed,
-            }
-        )
-    return inp
-
-
-def _planner_worker(path, req_q, res_q, stop_evt, version, seed, use_gpu):
-    """Background thread: consume replan requests, run the planner ONNX, post motions.
-
-    Loads its own ONNX session, then loops pulling ``(ctx, gen_frame, ms_dict)`` off
-    ``req_q``, running inference, resampling to 50 Hz, and putting the newest result
-    on ``res_q`` (dropping stale entries). Runs until ``stop_evt`` is set.
-    """
-    so = make_ort_session_options()
-    providers = ort_providers(force_cpu=not use_gpu)
-    sess = ort.InferenceSession(path, sess_options=so, providers=providers)
-    while not stop_evt.is_set():
-        try:
-            ctx, gf, ms_dict = req_q.get(timeout=0.05)
-        except queue.Empty:  # nosec B112 - idle poll, nothing queued this tick
-            continue
-        try:
-            inp = _build_planner_inputs(ctx, ms_dict, version, seed)
-            t0 = time.time()
-            qpos_out, num_pred = sess.run(None, inp)
-            t_inf = time.time()
-            n = int(num_pred.flat[0])
-            qpos = qpos_out[0, :n]
-            if np.any(np.isnan(qpos)):
-                continue
-            motion = _resample_30_to_50(qpos, n)
-            motion["gen_frame"] = gf
-            logger.debug(
-                "[Planner] inf=%.1fms total=%.1fms frames=%d",
-                1000 * (t_inf - t0),
-                1000 * (time.time() - t0),
-                n,
-            )
-            while not res_q.empty():
-                try:
-                    res_q.get_nowait()
-                except queue.Empty:
-                    break
-            res_q.put(motion)
-        except Exception:
-            logger.exception("[Planner] worker error")
-
-
-# ── SonicPlanner ──────────────────────────────────────────────────────────────
-
-
-class SonicPlanner:
-    """Owns the planner ONNX model and its async background worker.
-
-    Provides the initial motion synchronously (``initialize``), then serves replans
-    off-thread: ``request_replan`` enqueues the current context+intent and
-    ``try_get_new_motion`` non-blockingly returns a freshly planned motion (which the
-    controller cross-fades in). ``version`` selects the planner input schema.
-    """
-
-    def __init__(self, session, planner_path):
-        self.session = session
-        self.planner_path = planner_path
-        self.gen_frame = 0
-        self.random_seed = INITIAL_RANDOM_SEED
-        self.version = 1 if len(session.get_inputs()) >= 11 else 0
-        self.motion_50hz = PlannerMotion()
-        self._snapshot = PlannerMotion()
-        self._req_q = self._res_q = self._stop_evt = self._planner_thread = None
-        self._ctrl = None
-
-    def _build_inputs(self, ctx, ms):
-        return _build_planner_inputs(
-            ctx,
-            {
-                "mode": ms.mode,
-                "speed": ms.speed,
-                "height": ms.height,
-                "movement_direction": list(ms.movement_direction),
-                "facing_direction": list(ms.facing_direction),
-            },
-            self.version,
-            self.random_seed,
-        )
-
-    @staticmethod
-    def build_initial_context(joint_positions):
-        """Build a 4-frame standing context (MuJoCo qpos layout) from a pose."""
-        ctx = np.zeros((4, 36), np.float32)
-        jp_mj = joint_positions.astype(np.float32)[ISAACLAB_TO_MUJOCO]
-        for n in range(4):
-            ctx[n, 2] = DEFAULT_HEIGHT
-            ctx[n, 3] = 1.0
-            ctx[n, 7:36] = jp_mj
-        return ctx
-
-    def _context_from_controller(self, current_frame):
-        """Sample a 4-frame look-ahead context from the controller's live motion buffer.
-
-        The context starts ``MOTION_LOOK_AHEAD_STEPS`` ahead of ``current_frame`` so a
-        replan blends in seamlessly by the time it is ready.
-        """
-        ctrl = self._ctrl
-        gen_frame = current_frame + MOTION_LOOK_AHEAD_STEPS
-        t_arr = gen_frame / 50.0 + np.arange(4) / 30.0
-        f50 = t_arr * 50.0
-        with ctrl.motion_lock:
-            ts = ctrl.motion_timesteps
-            if ts <= 0:
-                return self.build_initial_context(DEFAULT_ANGLES)
-            bp, bq, jp = ctrl.motion_body_pos, ctrl.motion_body_quats, ctrl.motion_joint_positions
-            f0 = np.minimum(np.floor(f50).astype(int), ts - 1)
-            f1 = np.minimum(f0 + 1, ts - 1)
-        frac = f50 - f0
-        w0 = 1.0 - frac
-        ctx = np.zeros((4, 36), np.float32)
-        ctx[:, 0:3] = w0[:, None] * bp[f0] + frac[:, None] * bp[f1]
-        ctx[:, 3:7] = quat_slerp_batch(bq[f0], bq[f1], frac)
-        ij = w0[:, None] * jp[f0] + frac[:, None] * jp[f1]
-        ctx[:, 7:36] = ij[:, ISAACLAB_TO_MUJOCO]
-        self.gen_frame = gen_frame
-        return ctx
-
-    def _load_motion_in_place(self, qpos, n30, target=None):
-        """Resample raw planner qpos to 50 Hz and write it into a ``PlannerMotion`` buffer."""
-        if target is None:
-            target = self.motion_50hz
-        r = _resample_30_to_50(qpos, n30)
-        n = r["timesteps"]
-        target.timesteps = n
-        target.joint_positions[:n] = r["joint_positions"]
-        target.joint_velocities[:n] = r["joint_velocities"]
-        target.body_positions[:n] = r["body_positions"]
-        target.body_quaternions[:n] = r["body_quaternions"]
-        return target
-
-    def initialize(self, joint_positions, ms):
-        """Synchronously run the planner once to produce the first motion buffer."""
-        ctx = self.build_initial_context(joint_positions)
-        qpos_out, num_pred = self.session.run(None, self._build_inputs(ctx, ms))
-        n = int(num_pred.flat[0])
-        qpos = qpos_out[0, :n]
-        if np.any(np.isnan(qpos)):
-            raise RuntimeError("Planner initial output contains NaN")
-        logger.info("[Planner] Init: %d frames @ 30 Hz", n)
-        self._load_motion_in_place(qpos, n)
-        logger.info("[Planner] Resampled to %d frames @ 50 Hz", self.motion_50hz.timesteps)
-        return self.motion_50hz
-
-    def request_replan(self, cursor, ms):
-        """Enqueue a replan for the worker (drops any pending stale request first)."""
-        if self._req_q is None:
-            return
-        ctx = self._context_from_controller(cursor)
-        ms_dict = {
-            "mode": ms.mode,
-            "speed": ms.speed,
-            "height": ms.height,
-            "movement_direction": list(ms.movement_direction),
-            "facing_direction": list(ms.facing_direction),
-        }
-        while not self._req_q.empty():
-            try:
-                self._req_q.get_nowait()
-            except queue.Empty:
-                break
-        self._req_q.put((ctx, self.gen_frame, ms_dict))
-
-    def try_get_new_motion(self):
-        """Non-blocking: return ``(snapshot_motion, gen_frame)`` if a new plan is ready, else None."""
-        if self._res_q is None:
-            return None
-        result = None
-        while not self._res_q.empty():
-            try:
-                result = self._res_q.get_nowait()
-            except queue.Empty:
-                break
-        if result is None:
-            return None
-        n, gf = result["timesteps"], result["gen_frame"]
-        s = self._snapshot
-        s.timesteps = n
-        s.joint_positions[:n] = result["joint_positions"]
-        s.joint_velocities[:n] = result["joint_velocities"]
-        s.body_positions[:n] = result["body_positions"]
-        s.body_quaternions[:n] = result["body_quaternions"]
-        return s, gf
-
-    def start_subprocess(self, controller, use_gpu: bool = False):
-        """Run planner ONNX in a background thread (avoids mp spawn/fork + CUDA/MuJoCo issues)."""
-        self._ctrl = controller
-        self._req_q = queue.Queue()
-        self._res_q = queue.Queue()
-        self._stop_evt = threading.Event()
-        self._planner_thread = threading.Thread(
-            target=_planner_worker,
-            args=(
-                self.planner_path,
-                self._req_q,
-                self._res_q,
-                self._stop_evt,
-                self.version,
-                self.random_seed,
-                use_gpu,
-            ),
-            daemon=True,
-            name="sonic-planner",
-        )
-        self._planner_thread.start()
-        logger.info("[Planner] Background thread started (%s)", "GPU" if use_gpu else "CPU")
-
-    def stop_subprocess(self):
-        """Signal the planner thread to stop and join it."""
-        if self._stop_evt:
-            self._stop_evt.set()
-        if self._planner_thread is not None:
-            self._planner_thread.join(timeout=3.0)
-            logger.info("[Planner] Background thread stopped")
-        self._planner_thread = None
-        self._req_q = self._res_q = self._stop_evt = None
-
-
-# ── PlannerController ─────────────────────────────────────────────────────────
-
-
 class PlannerController(StandingEncoderDecoder):
-    """Encoder/decoder driven by the planner's live, replannable motion buffer.
+    """Encoder/decoder driven by a caller-supplied, rolling motion buffer.
 
-    Extends ``StandingEncoderDecoder`` so the reference comes from a rolling motion
-    (advanced one frame per tick via ``advance_cursor``) instead of a fixed pose.
-    Handles heading re-initialization, cross-fading new plans into the buffer
-    (``blend_new_motion``), and the mode-2 SMPL reference. ``motion_lock`` guards the
-    buffer against the async planner thread.
+    Extends ``StandingEncoderDecoder`` so the reference comes from a motion buffer
+    (a lookahead window with per-frame velocities) instead of a single fixed pose,
+    and handles heading re-initialization on the first frame / after a reset.
+    ``motion_lock`` guards the buffer, which the whole-body controller rewrites each
+    tick from the incoming command. The class name is retained for continuity with
+    the SONIC reference; no motion planner is involved.
     """
 
-    def __init__(self, planner, encoder, decoder):
+    def __init__(self, encoder, decoder):
         super().__init__(encoder, decoder)
-        self.planner = planner
         self.ref_cursor = 0
         self.motion_timesteps = 0
         self.motion_joint_positions = np.zeros((1500, 29), np.float64)
@@ -1046,72 +535,27 @@ class PlannerController(StandingEncoderDecoder):
         self.playing = self.first_motion = False
         self.motion_lock = threading.Lock()
 
-    def load_initial_motion(self, motion):
-        """Copy the planner's first motion into the live buffer and start playback."""
+    def reset(self):
+        """Full reset: clear enc/dec state (super) plus the motion buffer and heading.
+
+        Forces a heading re-init on the next ``step`` so the reference frame is
+        re-latched to the post-reset robot orientation.
+        """
+        super().reset()
         with self.motion_lock:
-            n = motion.timesteps
-            self.motion_timesteps = n
-            self.motion_joint_positions[:n] = motion.joint_positions[:n]
-            self.motion_joint_velocities[:n] = motion.joint_velocities[:n]
-            self.motion_body_quats[:n] = motion.body_quaternions[:n]
-            self.motion_body_pos[:n] = motion.body_positions[:n]
-            self.init_ref_quat = motion.body_quaternions[0].copy()
             self.ref_cursor = 0
-            self.first_motion = True
-            self.playing = True
-            self.delta_heading = 0.0
-
-    def blend_new_motion(self, new_motion, gen_frame):
-        """Blend like C++ CurrentFrameAdvancement: 8-frame cross-fade, then copy tail."""
-        with self.motion_lock:
-            cur = self.ref_cursor
-            new_len = gen_frame - cur + new_motion.timesteps
-            if new_len <= 0:
-                return
-            if self.motion_timesteps == 0:
-                n = new_motion.timesteps
-                self.motion_joint_positions[:n] = new_motion.joint_positions[:n]
-                self.motion_joint_velocities[:n] = new_motion.joint_velocities[:n]
-                self.motion_body_pos[:n] = new_motion.body_positions[:n]
-                self.motion_body_quats[:n] = new_motion.body_quaternions[:n]
-                self.motion_timesteps = n
-                self.ref_cursor = 0
-                self.init_ref_quat = self.motion_body_quats[0].copy()
-                self.first_motion = False
-                return
-
-            blend_start = max(0, gen_frame - cur)
-            blend_end = min(new_len, blend_start + BLEND_FRAMES)
-
-            for f in range(blend_end):
-                f_old = min(f + cur, self.motion_timesteps - 1)
-                f_new = max(0, min(f + cur - gen_frame, new_motion.timesteps - 1))
-                w_new = min(1.0, max(0.0, (f - blend_start) / BLEND_FRAMES))
-                w_old = 1.0 - w_new
-                self.motion_joint_positions[f] = (
-                    w_old * self.motion_joint_positions[f_old] + w_new * new_motion.joint_positions[f_new]
-                )
-                self.motion_joint_velocities[f] = (
-                    w_old * self.motion_joint_velocities[f_old] + w_new * new_motion.joint_velocities[f_new]
-                )
-                self.motion_body_pos[f] = (
-                    w_old * self.motion_body_pos[f_old] + w_new * new_motion.body_positions[f_new]
-                )
-                self.motion_body_quats[f] = quat_slerp(
-                    self.motion_body_quats[f_old], new_motion.body_quaternions[f_new], w_new
-                )
-
-            for f in range(blend_end, new_len):
-                f_new = max(0, min(f + cur - gen_frame, new_motion.timesteps - 1))
-                self.motion_joint_positions[f] = new_motion.joint_positions[f_new]
-                self.motion_joint_velocities[f] = new_motion.joint_velocities[f_new]
-                self.motion_body_pos[f] = new_motion.body_positions[f_new]
-                self.motion_body_quats[f] = new_motion.body_quaternions[f_new].copy()
-
-            self.motion_timesteps = new_len
-            self.first_motion = False
-            self.ref_cursor = 0
-            self.init_ref_quat = self.motion_body_quats[0].copy()
+            self.motion_timesteps = 0
+            self.motion_joint_positions[:] = 0.0
+            self.motion_joint_velocities[:] = 0.0
+            self.motion_body_quats[:] = 0.0
+            self.motion_body_quats[:, 0] = 1.0
+            self.motion_body_pos[:] = 0.0
+        self.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
+        self.heading_init_base_quat = np.array([1, 0, 0, 0], np.float64)
+        self.delta_heading = 0.0
+        self.first_motion = False
+        self.playing = False
+        self.reinit_heading = True
 
     def _heading_apply_delta(self):
         """Heading correction quaternion (init base-vs-ref heading + operator ``delta_heading``)."""
@@ -1224,134 +668,3 @@ class PlannerController(StandingEncoderDecoder):
         with self.motion_lock:
             if self.motion_timesteps > 0:
                 self.ref_cursor = min(self.ref_cursor + 1, self.motion_timesteps - 1)
-
-
-# ── Joystick input ────────────────────────────────────────────────────────────
-
-
-def _parse_wireless(wr):
-    """Parse wireless_remote (bytes or int-array) into (lx, ly, rx, ry)."""
-    import struct as _st
-
-    if not isinstance(wr, (bytes, bytearray)):
-        wr = bytes(wr)
-    if len(wr) < 24:
-        return None
-    lx = _st.unpack("f", wr[4:8])[0]
-    rx = _st.unpack("f", wr[8:12])[0]
-    ry = _st.unpack("f", wr[12:16])[0]
-    ly = _st.unpack("f", wr[20:24])[0]
-    return lx, ly, rx, ry
-
-
-def apply_joystick_axes(lx, ly, rx, ry, ms, controller=None):
-    """Map raw stick axes onto ``MovementState`` (left stick=WASD, right stick X=Q/E,
-    right stick Y=height).
-
-    Shared by the G1 wireless remote (:func:`process_joystick`) and the PICO
-    controller sticks (3-point teleop), so both drive the planner identically. Axes
-    are expected pre-negated to the same convention as the parsed wireless remote:
-    ``ly`` and ``ry`` already flipped, dead zone not yet applied.
-    """
-    # Dead zone + negate both Y axes (bridge already flips them once)
-    lx = 0.0 if abs(lx) < DEADZONE else lx
-    ly = 0.0 if abs(ly) < DEADZONE else -ly
-    rx = 0.0 if abs(rx) < DEADZONE else rx
-    ry = 0.0 if abs(ry) < DEADZONE else -ry
-
-    left_active = abs(lx) > 0 or abs(ly) > 0
-
-    # Left stick → WASD (movement direction relative to facing)
-    if left_active:
-        ms.movement_angle = ms.facing_angle + math.atan2(-lx, -ly)
-        ms.has_movement = True
-        if not ms.joy_prev_active:
-            ms.needs_replan = True
-        ms.joy_prev_active = True
-    elif ms.joy_prev_active and not (abs(rx) > 0 or abs(ry) > 0):
-        ms.joy_prev_active = False
-        ms.has_movement = False
-
-    # Right stick X → Q/E (facing rotation, ~1 rad/s at full deflection)
-    if abs(rx) > 0:
-        delta = -0.02 * rx
-        ms.facing_angle += delta
-        if controller:
-            controller.delta_heading += delta
-
-    # Right stick Y → -/= (height adjustment, ~0.25/s at full deflection)
-    if abs(ry) > 0:
-        step = -0.005 * ry
-        ms.height = max(0.1, min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + step))
-
-
-# gear_sonic PlannerLoop joystick constants (pico_manager_thread_server).
-PICO_YAW_GAIN = 1.5  # rad/s of facing rotation at full right-stick deflection
-PICO_JOY_DEADZONE = 0.15  # gear_sonic JOYSTICK_DEADZONE (distinct from the remote's)
-
-
-def apply_pico_loco_axes(lx, ly, rx, ry, ms, dt=CONTROL_DT):
-    """Map PICO controller sticks onto ``MovementState`` — the gear_sonic 3-point path.
-
-    Ports ``pico_manager_thread_server.PlannerLoop`` (encode_mode 1) rather than the
-    keyboard/remote-parity :func:`apply_joystick_axes`, so speed follows gear_sonic's
-    mode-dependent curves and facing comes from a yaw accumulator:
-
-    - Right stick X -> yaw accumulator (``+= PICO_YAW_GAIN * -rx * dt``), dead-zoned.
-    - Left stick -> magnitude ``mag`` (dead-zone-rescaled to 0..1) driving a
-      *mode-dependent* speed curve (slow ``0.1+0.5·mag``, run ``1.5+3·mag``, walk =
-      default), and a facing-rotated global movement vector -> ``movement_angle``.
-
-    Signs replicate gear_sonic's ``get_controller_axes`` usage exactly (forward = +ly,
-    strafe = -lx, turn = -rx), which resolves the PICO axis-sign question: because the
-    publisher forwards the same raw SDK axes gear_sonic reads, matching its sign usage
-    is by definition the correct mapping. ``ms.speed`` is later re-clamped into each
-    mode's valid range by :func:`clamp_mode_params`.
-    """
-    lx, ly, rx, ry = float(lx), float(ly), float(rx), float(ry)
-
-    # Facing: gear_sonic YawAccumulator — integrate only outside the dead zone.
-    if abs(rx) >= PICO_JOY_DEADZONE:
-        ms.facing_angle += PICO_YAW_GAIN * (-rx) * dt
-
-    raw_mag = min(1.0, math.hypot(lx, ly))
-    if raw_mag < PICO_JOY_DEADZONE:
-        ms.has_movement = False
-        ms.speed = -1.0
-        ms.joy_prev_active = False
-        return
-
-    mag = min(1.0, (raw_mag - PICO_JOY_DEADZONE) / (1.0 - PICO_JOY_DEADZONE))
-    m = LM(ms.mode)
-    if m == LM.SLOW_WALK:
-        ms.speed = 0.1 + 0.5 * mag  # 0.1 .. 0.6
-    elif m == LM.WALK:
-        ms.speed = -1.0  # planner default
-    elif m == LM.RUN:
-        ms.speed = 1.5 + 3.0 * mag  # 1.5 .. 4.5
-    else:
-        ms.speed = mag
-
-    # Facing-rotated movement vector: rotation_facing @ [-lx, ly] * scale.
-    scale = mag / raw_mag
-    mlx, mly = -lx * scale, ly * scale
-    fx, fy = math.cos(ms.facing_angle), math.sin(ms.facing_angle)
-    gx = -fy * mlx + fx * mly
-    gy = fx * mlx + fy * mly
-    if not ms.joy_prev_active:
-        ms.needs_replan = True
-    ms.has_movement = True
-    ms.joy_prev_active = True
-    ms.movement_angle = math.atan2(gy, gx)
-
-
-def process_joystick(obs, ms, controller=None):
-    """Drive ``MovementState`` from the G1 wireless remote in ``obs``."""
-    wr = obs.get("wireless_remote")
-    if wr is None:
-        return
-    parsed = _parse_wireless(wr)
-    if parsed is None:
-        return
-    lx, ly, rx, ry = parsed
-    apply_joystick_axes(lx, ly, rx, ry, ms, controller)

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

@@ -18,54 +18,53 @@
 
 from __future__ import annotations
 
-import logging
-import math
 from collections import deque
+import logging
 from typing import TYPE_CHECKING
 
-import numpy as np
 from huggingface_hub import hf_hub_download
+import numpy as np
 
-from lerobot.teleoperators.pico_headset.smpl_constants import (
-    LOCO_AXES_PREFIX,
-    LOCO_BTN_PREFIX,
-    LOCO_N_AXES,
-    LOCO_N_BTN,
-    ROOT_ACTION_DIM,
-    ROOT_ACTION_PREFIX,
-    SMPL_ACTION_PREFIX,
-    SMPL_OBS_DIM as SMPL_ACTION_DIM,
-    VR3_ORN_DIM,
-    VR3_ORN_PREFIX,
-    VR3_POS_DIM,
-    VR3_POS_PREFIX,
-    WB_ACTION_DIM,
-    wb_action_key,
-)
 from lerobot.utils.import_utils import _onnxruntime_available, require_package
 
-from ..g1_utils import MUJOCO_TO_ISAACLAB, KEYBOARD_KEYS_FIELD, G1_29_JointIndex, lowstate_to_obs
+from ..g1_utils import (
+    MUJOCO_TO_ISAACLAB,
+    WB_ACTION_DIM,
+    G1_29_JointIndex,
+    lowstate_to_obs,
+    wb_action_key,
+)
 from .sonic_pipeline import (
     CONTROL_DT,
-    DEBUG_PRINT_EVERY,
     DEFAULT_ANGLES,
-    DEFAULT_HEIGHT,
     ENCODER_UPDATE_EVERY,
-    LM,
-    MOTION_SETS,
-    MovementState,
+    TOKEN_DIM,
     PlannerController,
-    SonicPlanner,
-    apply_pico_loco_axes,
-    clamp_mode_params,
     compute_kp_kd,
     make_ort_session_options,
     ort_providers,
-    process_joystick,
-    should_replan_request,
-    snapshot_ms,
 )
 
+# Action-feature prefix for the latent-token interface (see _extract_token_from_action).
+TOKEN_ACTION_PREFIX = "motion_token"
+# Proprio-state prefix for the token interface: the robot echoes the last commanded
+# token here so ``lerobot-rollout`` aggregates it into a 64-D ``observation.state``.
+TOKEN_STATE_PREFIX = "motion_token_state"
+
+
+def token_action_key(i: int) -> str:
+    """Action-dict key for the i-th component of the 64-D SONIC latent token.
+
+    The ``.pos`` suffix is required so the value flows through ``lerobot-rollout``,
+    which only routes ``.pos`` scalar features onto the policy action vector.
+    """
+    return f"{TOKEN_ACTION_PREFIX}.{i}.pos"
+
+
+def token_state_key(i: int) -> str:
+    """Observation key for the i-th component of the 64-D SONIC latent token state."""
+    return f"{TOKEN_STATE_PREFIX}.{i}.pos"
+
 if TYPE_CHECKING or _onnxruntime_available:
     import onnxruntime as ort
 else:
@@ -78,61 +77,21 @@ logger = logging.getLogger(__name__)
 # eases in without a snap on the first command.
 INIT_RAMP_S = 3.0
 
-
-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.
-    """
-    # The keys are smpl.0 .. smpl.719; presence of the first element (smpl.0) is the
-    # sentinel that a full SMPL window was sent this tick. If it's absent, there's no
-    # whole-body reference, so bail out.
-    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
-
-
-def _extract_vr3_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None:
-    """Reassemble the 3-point VR targets from ``vr3_pos.{i}`` / ``vr3_orn.{i}`` keys.
-
-    Returns ``(pos (9,), orn (12,))`` for the [l-wrist, r-wrist, neck] keypoints, or
-    None when no VR3 reference was sent this tick. Presence of ``vr3_pos.0`` is the
-    sentinel that a full 3-point frame is available (mirrors the SMPL sentinel).
-    """
-    if not action or f"{VR3_POS_PREFIX}0" not in action:
-        return None
-    pos = np.fromiter(
-        (float(action.get(f"{VR3_POS_PREFIX}{i}", 0.0)) for i in range(VR3_POS_DIM)),
-        dtype=np.float32,
-        count=VR3_POS_DIM,
-    )
-    orn = np.fromiter(
-        (float(action.get(f"{VR3_ORN_PREFIX}{i}", 0.0)) for i in range(VR3_ORN_DIM)),
-        dtype=np.float32,
-        count=VR3_ORN_DIM,
-    )
-    return pos, orn
+# Neutral ("zero pose") SONIC token, held by token_mode until the first real token
+# arrives. Captured from the encoder's own output while the robot stood idle in sim
+# (capture_neutral_token.py): the encoder is an FSQ bottleneck (~5 bit/dim, 15.5 half-
+# width, Div(16)), so its tokens live on the 1/16 grid. We store the integer FSQ codes
+# and rescale by the same 1/16 step, giving an exact on-grid token -- unlike the literal
+# all-zero token, which is off the encoder's learned manifold and decodes to a slightly
+# goofy stance. This one decodes to a stable, natural standing pose.
+_NEUTRAL_TOKEN_CODES = np.array(
+    [-1, 3, 1, -1, 1, -3, 6, 1, 1, 1, -2, -4, -2, 0, -3, -1,
+     2, -1, -3, -5, 3, 1, 1, -4, -1, -1, 1, -7, 0, 1, 2, -2,
+     5, -2, -2, -4, 0, -1, 3, -1, 0, -5, -1, 0, -4, 0, 0, -1,
+     -1, 2, -2, 1, 3, 3, 1, 0, 0, 6, 0, -7, 3, 0, 2, -2],
+    dtype=np.float32,
+)
+NEUTRAL_TOKEN = _NEUTRAL_TOKEN_CODES / 16.0  # FSQ Div(16): integer codes -> on-grid token
 
 
 def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None:
@@ -143,15 +102,40 @@ def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None:
     units. The ``.pos`` suffix lets these flow through ``lerobot-rollout`` as normal
     joint-position action features.
     """
-    if not action or wb_action_key(0) not in action:
+    if not action:
+        return None
+    keys = [wb_action_key(i) for i in range(WB_ACTION_DIM)]
+    # Require the full dense command: a partial action (e.g. only ``wb.0.pos``)
+    # must not be silently zero-filled, which would drive most joints toward 0.
+    if any(key not in action for key in keys):
         return None
     return np.fromiter(
-        (float(action.get(wb_action_key(i), 0.0)) for i in range(WB_ACTION_DIM)),
+        (float(action[key]) for key in keys),
         dtype=np.float32,
         count=WB_ACTION_DIM,
     )
 
 
+def _extract_token_from_action(action: dict | None) -> np.ndarray | None:
+    """Reassemble a dense (64,) latent token from ``motion_token.{i}`` keys, or None.
+
+    This is the token-only replay interface: instead of a joint reference driving the
+    encoder, the caller supplies the 64-D encoder latent directly (e.g. a recorded
+    ``action.motion_token`` column), which the decoder consumes with the encoder
+    bypassed. Requires the full dense token; a partial one is ignored (returns None).
+    """
+    if not action:
+        return None
+    keys = [token_action_key(i) for i in range(TOKEN_DIM)]
+    if any(key not in action for key in keys):
+        return None
+    return np.fromiter(
+        (float(action[key]) for key in keys),
+        dtype=np.float32,
+        count=TOKEN_DIM,
+    )
+
+
 def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     """Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
 
@@ -164,14 +148,15 @@ def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     The 29 joints are first assembled in MuJoCo / Unitree-SDK order
     ([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the
     ``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via
-    ``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) and yaw-rate (33) are not part of the
-    29-DoF SONIC reference.
+    ``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) are not part of the 29-DoF SONIC
+    reference, and yaw-rate (33) is integrated into the heading by the caller (it
+    cannot be represented in this static per-tick anchor).
     """
     ref_mj = np.zeros(29, np.float32)  # MuJoCo / Unitree-SDK grouped order
-    ref_mj[0:6] = wb[16:22]   # left leg
+    ref_mj[0:6] = wb[16:22]  # left leg
     ref_mj[6:12] = wb[22:28]  # right leg
     ref_mj[12:15] = wb[28:31]  # waist
-    ref_mj[15:22] = wb[0:7]   # left arm
+    ref_mj[15:22] = wb[0:7]  # left arm
     ref_mj[22:29] = wb[8:15]  # right arm
     ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32)  # -> IsaacLab order for SONIC
     roll, pitch = float(wb[31]), float(wb[32])
@@ -180,101 +165,51 @@ def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     return ref, anchor
 
 
-def _extract_loco_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None:
-    """Reassemble controller-stick locomotion from ``loco_axes.{i}`` / ``loco_btn.{i}``.
-
-    Returns ``(axes (4,) = [lx, ly, rx, ry], buttons (4,) = [A, B, X, Y])`` or None
-    when no locomotion state was sent this tick (sentinel: ``loco_axes.0``).
-    """
-    if not action or f"{LOCO_AXES_PREFIX}0" not in action:
-        return None
-    axes = np.fromiter(
-        (float(action.get(f"{LOCO_AXES_PREFIX}{i}", 0.0)) for i in range(LOCO_N_AXES)),
-        dtype=np.float32,
-        count=LOCO_N_AXES,
-    )
-    buttons = np.fromiter(
-        (float(action.get(f"{LOCO_BTN_PREFIX}{i}", 0.0)) for i in range(LOCO_N_BTN)),
-        dtype=np.float32,
-        count=LOCO_N_BTN,
-    )
-    return axes, buttons
-
-
 class SonicRuntime:
-    """Shared SONIC control loop state (standalone demo + locomotion controller)."""
+    """Loads the SONIC encoder/decoder ONNX models and owns the controller.
+
+    No motion planner: the reference motion buffer is written directly each tick by
+    :class:`SonicWholeBodyController` from the incoming 34-D whole-body command.
+    """
 
     def __init__(self, force_cpu: bool = False):
         require_package("onnxruntime", extra="unitree_g1")
-        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 = make_ort_session_options()
 
-        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)
 
+        # Report the provider actually bound, not the one requested: ORT silently falls
+        # back to CPU if CUDA can't load (e.g. libcudnn not on LD_LIBRARY_PATH), and a
+        # CPU decoder drifts the closed-loop heading. Warn loudly so it can't hide.
+        self.use_gpu = decoder_sess.get_providers()[0] == "CUDAExecutionProvider"
+        if not force_cpu and not self.use_gpu:
+            print(
+                "[SONIC] WARNING: decoder bound to CPUExecutionProvider (CUDA unavailable). "
+                "Closed-loop replay/control will drift. Ensure libcudnn is on LD_LIBRARY_PATH "
+                "(site-packages/nvidia/*/lib).",
+                flush=True,
+            )
+
         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)
+        self.controller = PlannerController(encoder_sess, decoder_sess)
 
     @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)
+        # Full pipeline reset: clears the encoder token, proprioception history and
+        # heading, and rewinds the motion buffer. reinit_heading is set so the next
+        # step re-latches the reference frame to the current robot orientation.
+        self.controller.reset()
 
     def shutdown(self):
-        self.planner.stop_subprocess()
+        pass
 
 
 class SonicWholeBodyController:
@@ -287,41 +222,12 @@ class SonicWholeBodyController:
     # directly with a 34-D VLA policy.
     wb_action = True
 
-    def __init__(
-        self,
-        force_cpu: bool = False,
-        *,
-        enable_smpl_root: bool = False,
-        root_smoothing_alpha: float = 0.15,
-        enable_smpl_stream: bool = False,
-        smpl_host: str | None = None,
-        smpl_port: int | None = None,
-    ):
+    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
-
-        # When True, the per-frame SMPL root quaternion steers the mode-2 anchor.
-        # Off by default: even with smoothing this changes the anchor/heading and is
-        # untested on hardware, so it stays opt-in. When enabled, the raw per-frame
-        # root quat (from a 30 Hz dataset resampled to a 50 Hz loop) is spherically
-        # smoothed by :meth:`_smooth_root_quat` before it reaches the anchor, which
-        # removes the root-acceleration spikes (NaN QACC at DOF 0) the unsmoothed
-        # trajectory caused. ``root_smoothing_alpha`` in (0, 1] is the per-tick blend
-        # toward the incoming quat (smaller = smoother/laggier, 1 = no smoothing).
-        self.enable_smpl_root = enable_smpl_root
-        self._root_smoothing_alpha = float(np.clip(root_smoothing_alpha, 1e-3, 1.0))
-        self._smoothed_root_quat: np.ndarray | None = None
-
-        # 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.
-        self._prev_keys: set[str] = set()
-        # Edge state for the PICO A+B / X+Y locomotion-mode cycle (3-point teleop).
-        self._prev_loco_mode_pair: tuple[bool, bool] = (False, False)
 
         # Startup blend: ease from the robot's initial pose into the first commanded
         # policy targets over INIT_RAMP_S (captured on the first control tick).
@@ -335,174 +241,22 @@ class SonicWholeBodyController:
         # stream of per-tick whole-body commands, fed to the encoder as a batch.
         self._wb_traj: deque[np.ndarray] = deque(maxlen=50)
         self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50)
+        # Integrated heading (rad) from the whole-body command's yaw-rate (index 33),
+        # forwarded to the pipeline as ``delta_heading`` so turn commands take effect.
+        self._heading = 0.0
 
-        # 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.
-        self._smpl_host = smpl_host
-        self._smpl_port = smpl_port
-        self._smpl_stream = None
-        if enable_smpl_stream:
-            self._init_smpl_stream()
+        # Token-interface state. ``token_mode`` is set True by the robot when the deploy
+        # is token-driven (``UnitreeG1Config.sonic_token_action``): the controller then
+        # holds a stable *neutral* (all-zero) token until the first real token arrives,
+        # and afterwards holds the *last* token received between ticks (the async
+        # controller runs ~50 Hz while a token VLA streams ~30 Hz). This lives here (not
+        # in the entry-point script) so it applies uniformly to run_g1_onboard,
+        # lerobot-rollout and the sim replays. ``token_mode`` stays False for the dense
+        # 34-D whole-body / OpenHLM path, which keeps its own "hold last target" idle.
+        self.token_mode = False
+        self._last_token: np.ndarray | None = None
 
-        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.teleoperators.pico_headset.smpl_stream import (
-            DEFAULT_SMPL_HOST,
-            DEFAULT_SMPL_PORT,
-            SmplStream,
-        )
-
-        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)
-
-    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 _enter_3point(self) -> None:
-        """Switch into 3-point VR upper-body teleop (encode_mode 1).
-
-        The upper body tracks the VR wrist/neck targets while the lower body /
-        locomotion keeps running off the planner (joystick/keyboard-driven).
-        """
-        self.controller.encode_mode = 1
-        self.controller.playing = True
-        self.controller.reinit_heading = True
-        self.ms.needs_replan = True
-        logger.info("SONIC: 3-point VR active -> upper-body tracking + planner locomotion (mode 1)")
-
-    def _exit_wholebody(self) -> None:
-        """Revert to locomotion/standing (encode_mode 0) after a teleop reference 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: teleop reference lost/stale -> reverting to locomotion (standing)")
-
-    def _process_keyboard(self, action: dict | None) -> None:
-        """Translate a native KeyboardTeleop's held-key set into MovementState.
-
-        Mirrors the standalone SONIC demo's keyboard mapping so locomotion (mode 0/1)
-        can be driven with ``--teleop.type=keyboard`` instead of the PICO SMPL stream.
-        Discrete controls act on newly-pressed keys (edge-detected against the previous
-        tick); inherently-continuous controls (facing turn, height, speed) integrate a
-        small per-tick delta while the key is held so they feel smooth at 50 Hz.
-
-        Controls: WASD move, Q/E turn, 1-8 select mode, 9/0 speed down/up,
-        -/= height down/up, R replan, Space emergency-stop -> IDLE.
-        """
-        if action is None:
-            return
-        keys = action.get(KEYBOARD_KEYS_FIELD)
-        if keys is None:
-            return  # No KeyboardTeleop attached; leave joystick/SMPL paths untouched.
-
-        ms, controller = self.ms, self.controller
-        held = {k.lower() if isinstance(k, str) and len(k) == 1 else k for k in keys}
-        prev = self._prev_keys
-        pressed = held - prev  # newly-pressed this tick (edge)
-        self._prev_keys = held
-
-        # ── Discrete: fire once per press ────────────────────────────────────
-        if "space" in pressed:
-            ms.mode = LM.IDLE
-            ms.speed = ms.height = -1.0
-            ms.has_movement = False
-            ms.needs_replan = True
-            controller.playing = False
-            controller.reinit_heading = True
-            logger.info("SONIC keyboard: EMERGENCY STOP -> IDLE")
-        if "r" in pressed:
-            ms.needs_replan = True
-        if "n" in pressed or "p" in pressed:
-            step = 1 if "n" in pressed else -1
-            ms.motion_set_idx = (ms.motion_set_idx + step) % len(MOTION_SETS)
-            logger.info("SONIC keyboard: motion set -> %s", MOTION_SETS[ms.motion_set_idx][0])
-        for digit in ("1", "2", "3", "4", "5", "6", "7", "8"):
-            if digit in pressed:
-                idx = int(digit) - 1
-                modes = MOTION_SETS[ms.motion_set_idx][1]
-                if 0 <= idx < len(modes):
-                    ms.mode = modes[idx]
-                    ms.has_movement = False
-                    ms.needs_replan = True
-                    controller.playing = True
-                    controller.reinit_heading = True
-                    logger.info("SONIC keyboard: mode -> %s", LM(ms.mode).name)
-        # WASD sets the movement direction relative to current facing (press to set,
-        # Space to stop) to match the standalone demo.
-        if "w" in pressed:
-            ms.movement_angle = ms.facing_angle
-        elif "s" in pressed:
-            ms.movement_angle = ms.facing_angle + math.pi
-        elif "a" in pressed:
-            ms.movement_angle = ms.facing_angle + math.pi / 2
-        elif "d" in pressed:
-            ms.movement_angle = ms.facing_angle - math.pi / 2
-        if pressed & {"w", "a", "s", "d"}:
-            ms.has_movement = True
-            ms.needs_replan = True
-
-        # ── Continuous: integrate a small delta while held ───────────────────
-        if "q" in held:
-            ms.facing_angle += 0.02
-            controller.delta_heading += 0.02
-        if "e" in held:
-            ms.facing_angle -= 0.02
-            controller.delta_heading -= 0.02
-        if "0" in held:
-            ms.speed = min(5.0, (ms.speed if ms.speed >= 0 else 1.0) + 0.02)
-        if "9" in held:
-            ms.speed = max(0.0, (ms.speed if ms.speed >= 0 else 1.0) - 0.02)
-        if "=" in held:
-            ms.height = min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + 0.005)
-        if "-" in held:
-            ms.height = max(0.1, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) - 0.005)
-
-    def _process_pico_loco(self, axes: np.ndarray, buttons: np.ndarray) -> None:
-        """Drive locomotion from the PICO controller sticks/buttons (encode_mode 1).
-
-        Mirrors gear_sonic's ``PlannerLoop`` VR-3PT tick: left/right sticks steer
-        movement/facing/speed via :func:`apply_pico_loco_axes` (the faithful gear_sonic
-        yaw-accumulator + mode-dependent speed curves, not the keyboard-parity map), and
-        A+B / X+Y edge-cycle the locomotion mode within the current motion set.
-        """
-        lx, ly, rx, ry = (float(v) for v in axes)
-        apply_pico_loco_axes(lx, ly, rx, ry, self.ms)
-
-        # Mode cycling: step linearly through the LocomotionMode enum (A+B = next,
-        # X+Y = previous), exactly like gear_sonic's PlannerLoop — so the operator can
-        # reach squat/kneel/crawl, not just the modes in one UI motion set.
-        a, b, x, y = (v > 0.5 for v in buttons)
-        ab_now, xy_now = (a and b), (x and y)
-        ab_prev, xy_prev = self._prev_loco_mode_pair
-        mode = int(self.ms.mode)
-        if ab_now and not ab_prev:
-            mode = min(int(LM.INJURED_WALK), mode + 1)
-        elif xy_now and not xy_prev:
-            mode = max(int(LM.IDLE), mode - 1)
-        if mode != int(self.ms.mode):
-            self.ms.mode = LM(mode)
-            self.ms.needs_replan = True
-            self.controller.playing = True
-            logger.info("SONIC 3-point: locomotion mode -> %s", LM(self.ms.mode).name)
-        self._prev_loco_mode_pair = (ab_now, xy_now)
+        logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)")
 
     def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict:
         """Feed a dense 34-D OpenHLM whole-body command as the mode-0 encoder reference.
@@ -517,6 +271,11 @@ class SonicWholeBodyController:
         if c.encode_mode != 0:
             c.encode_mode = 0
             c.reinit_heading = True
+        # Index 33 is a yaw-rate (rad/s): integrate it into a heading offset and hand
+        # it to the pipeline as ``delta_heading`` so commanded turns are tracked rather
+        # than silently dropped (the anchor from _wb34_to_reference only carries r/p).
+        self._heading += float(wb[33]) * CONTROL_DT
+        c.delta_heading = self._heading
         # Capture the heading/anchor reference on the first whole-body tick. The
         # controller only latches ``init_ref_quat`` (and the base heading) inside
         # ``step()`` when ``first_motion or reinit_heading`` — but it already boots in
@@ -564,33 +323,19 @@ class SonicWholeBodyController:
         self._wb_step += 1
         return out
 
-    def _smooth_root_quat(self, root_quat: np.ndarray | None) -> np.ndarray | None:
-        """Spherically smooth the per-frame SMPL root quaternion (mode-2 anchor).
+    def _run_token(self, obs: dict, token: np.ndarray) -> dict:
+        """Decode a supplied 64-D latent token directly (encoder bypassed).
 
-        The reference root trajectory is authored at ~30 Hz and consumed at 50 Hz, so
-        the raw per-tick quat steps unevenly and injects root-acceleration spikes into
-        the anchor. This keeps a persistent estimate and shortest-path nlerp-slerps it
-        toward each incoming (unit) quat by ``root_smoothing_alpha``, yielding a
-        continuous, rate-matched heading. Quaternions are scalar-first (w, x, y, z).
-        Returns ``None`` (leaving the anchor self-driven) for an invalid/zero input.
+        Token-only replay: set the pipeline's cached token to the supplied one and run
+        a decode-only step (``update_encoder=False``). The decoder still closes the loop
+        on live proprioception (history is refreshed inside ``step`` from ``obs``); only
+        the encoder — which would recompute the token from a motion reference — is
+        skipped. Returns the ``<joint>.q`` target dict.
         """
-        if root_quat is None:
-            self._smoothed_root_quat = None
-            return None
-        q = np.asarray(root_quat, np.float64)
-        n = np.linalg.norm(q)
-        if n < 1e-8:
-            return self._smoothed_root_quat
-        q = q / n
-        if self._smoothed_root_quat is None:
-            self._smoothed_root_quat = q
-        else:
-            prev = self._smoothed_root_quat
-            if np.dot(prev, q) < 0.0:  # shortest-path: quats double-cover SO(3)
-                q = -q
-            blended = prev + self._root_smoothing_alpha * (q - prev)
-            self._smoothed_root_quat = blended / (np.linalg.norm(blended) + 1e-12)
-        return self._smoothed_root_quat.astype(np.float32)
+        c = self.controller
+        c.token = np.asarray(token, np.float32)
+        self._wb_step += 1
+        return c.step(obs, update_encoder=False, debug=False)
 
     def _startup_blend(self, obs: dict, out: dict) -> dict:
         """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds,
@@ -623,96 +368,48 @@ class SonicWholeBodyController:
             return {}
         obs = lowstate_to_obs(lowstate)
 
-        # Keyboard teleop (native KeyboardTeleop) drives the same locomotion intent
-        # the joystick does; applied before the SMPL check so whole-body tracking
-        # still takes priority when a headset stream is present.
-        self._process_keyboard(action)
+        # Token-only interface (latent replay / token-output VLA): a dense 64-D
+        # ``motion_token.{i}`` command is decoded directly, bypassing the encoder.
+        # Checked before the joint path so a token action takes precedence.
+        token = _extract_token_from_action(action)
+        if token is not None:
+            self._last_token = token
+        elif self._last_token is None and self.token_mode:
+            # Token-driven deploy, but no token has arrived yet: hold the captured
+            # neutral token (NEUTRAL_TOKEN), which the decoder maps to a stable, natural
+            # standing pose (the encoder's own idle output; see NEUTRAL_TOKEN).
+            self._last_token = NEUTRAL_TOKEN.copy()
+        if self._last_token is not None:
+            # Either a fresh token this tick or the last one received (held between the
+            # ~30 Hz token stream and the ~50 Hz control loop).
+            return self._startup_blend(obs, self._run_token(obs, self._last_token))
 
-        # Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a
-        # 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.
-        # Dense whole-body command (OpenHLM / pi0.5 joint interface) takes priority:
-        # a single 34-D vector drives the mode-0 joint reference directly.
+        # Dense 34-D whole-body command (OpenHLM / pi0.5 joint interface): a single
+        # vector per tick drives the mode-0 encoder reference directly. Until the
+        # policy produces one, hold (no command) so the robot keeps its last target.
         wb = _extract_wb34_from_action(action)
-        if wb is not None:
-            return self._startup_blend(obs, self._run_wholebody34(obs, wb))
-        self._wb_miss = getattr(self, "_wb_miss", 0) + 1
-        if self._wb_miss % 50 == 1:
-            akeys = [k for k in action if isinstance(k, str)]
-            logger.info(
-                "[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s",
-                self._wb_miss,
-                akeys[:8],
-            )
-
-        smpl = _extract_smpl_from_action(action)
-        root_quat = _extract_root_from_action(action)
-        vr3 = _extract_vr3_from_action(action)
-        loco = _extract_loco_from_action(action)
-        if smpl is None and vr3 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)
-            # VR3 is independent of the SMPL window: the controller-state source
-            # (head + controllers only) sends 3-point targets with no SMPL frame.
-            elif self._smpl_stream.has_fresh_vr3:
-                vr3 = (self._smpl_stream.vr3_pos, self._smpl_stream.vr3_orn)
-                if self._smpl_stream.has_fresh_loco:
-                    loco = (self._smpl_stream.loco_axes, self._smpl_stream.loco_buttons)
-
-        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 steers the mode-2 anchor/heading, but only when
-            # explicitly enabled (see enable_smpl_root); the raw per-frame quat is
-            # spherically smoothed first so the 30->50 Hz resample doesn't spike the
-            # anchor. Disabled -> anchor stays self-driven.
-            self.controller.smpl_root_quat = (
-                self._smooth_root_quat(root_quat) if self.enable_smpl_root else None
-            )
-            out = self._runtime.tick(obs, debug=False, use_joystick=False)
-        elif vr3 is not None:
-            # 3-point VR teleop: upper body tracks the wrist/neck targets; the lower
-            # body / locomotion keeps running off the planner, so the joystick (and
-            # keyboard) still steer walking/turning underneath.
-            if self.controller.encode_mode != 1:
-                self._enter_3point()
-            self.controller.vr_3point_local_target = vr3[0]
-            self.controller.vr_3point_local_orn_target = vr3[1]
-            # Replicate the original encode_mode-1 handling: when the PICO controller
-            # sticks are forwarded, drive locomotion from them directly (and skip the
-            # wireless-remote joystick read). Otherwise leave the remote/keyboard path.
-            if loco is not None:
-                self._process_pico_loco(loco[0], loco[1])
-                out = self._runtime.tick(obs, debug=False, use_joystick=False)
-            else:
-                out = self._runtime.tick(obs, debug=False, use_joystick=True)
-        else:
-            # No (or stale) teleop reference: fall back to locomotion so the robot stays balanced.
-            if self.controller.encode_mode != 0:
-                self.controller.smpl_root_quat = None
-                self._smoothed_root_quat = None
-                self._exit_wholebody()
-            out = self._runtime.tick(obs, debug=False)
-
-        # Startup interpolation: blend from the robot's initial pose into the policy's
-        # commanded target over INIT_RAMP_S, regardless of mode.
-        return self._startup_blend(obs, out)
+        if wb is None:
+            self._wb_miss = getattr(self, "_wb_miss", 0) + 1
+            if self._wb_miss % 50 == 1:
+                akeys = [k for k in action if isinstance(k, str)]
+                logger.info(
+                    "[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s",
+                    self._wb_miss,
+                    akeys[:8],
+                )
+            return {}
+        return self._startup_blend(obs, self._run_wholebody34(obs, wb))
 
     def reset(self):
         self._runtime.reset()
         self._init_step = 0  # re-run the startup blend after a reset
         self._start_pose = {}
-        self._smoothed_root_quat = None
         self._wb_step = 0
         self._wb_traj.clear()
         self._wb_quat_traj.clear()
+        self._heading = 0.0
+        # Drop the held token so token_mode re-seeds the neutral token after a reset.
+        self._last_token = None
 
     def shutdown(self):
-        if self._smpl_stream is not None:
-            self._smpl_stream.close()
         self._runtime.shutdown()

src/lerobot/robots/unitree_g1/infer_sonic_g1_onboard.py

@@ -0,0 +1,192 @@
+#!/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.
+
+"""Laptop-side sender for the SONIC whole-body walk policy, onboard deployment.
+
+This is the counterpart to ``run_g1_onboard.py`` (which runs the SONIC decoder on the
+robot). The heavy VLA (``nepyope/sonic_walk``, a pi0.5 token policy) runs here on the
+laptop GPU; only the resulting 64-D latent token is shipped to the robot over ZMQ:
+
+    laptop:  camera frame (ZMQ from robot :5555) + previous token
+             -> pi0.5 -> next 64-D token
+             -> PUSH JSON {motion_token.i.pos: ...} to robot :6004
+    robot:   run_g1_onboard receives the token, SonicWholeBodyController decodes it
+             into whole-body joint commands against local DDS at full rate.
+
+The policy's ``observation.state`` is the token currently being executed, so we close
+the loop by feeding back the *last token we sent* (the decoder holds it until a new one
+arrives). This mirrors what ``lerobot-rollout`` does via the robot's token echo, but
+without a controller / DDS on the laptop.
+
+The policy is pi0.5 with chunk_size=50, so a full diffusion inference runs only about
+once every 50 ticks; ``select_action`` pops one queued token per tick in between.
+
+Run ``run_g1_onboard.py --controller SonicWholeBodyController --sonic-token-action
+--cameras ...`` on the robot first, then this on the laptop:
+
+    python -m lerobot.robots.unitree_g1.infer_sonic_g1_onboard \
+        --policy-path nepyope/sonic_walk --robot-ip 192.168.123.164 \
+        --task "walk back and forth"
+"""
+
+import argparse
+import contextlib
+import json
+import logging
+import signal
+import time
+
+import numpy as np
+import torch
+
+from lerobot.cameras.zmq import ZMQCamera, ZMQCameraConfig
+from lerobot.configs.policies import PreTrainedConfig
+from lerobot.policies.factory import get_policy_class, make_pre_post_processors
+from lerobot.policies.utils import prepare_observation_for_inference
+from lerobot.robots.unitree_g1.controllers.sonic_whole_body import (
+    NEUTRAL_TOKEN,
+    TOKEN_DIM,
+    token_action_key,
+)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
+logger = logging.getLogger("sonic_sender")
+
+ACTION_PORT = 6004  # matches run_g1_onboard.py --action-port
+IMAGE_KEY = "observation.images.ego_view"  # pi05 sonic_walk VISUAL input
+STATE_KEY = "observation.state"
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--policy-path", default="nepyope/sonic_walk", help="Policy repo id or local path")
+    p.add_argument("--robot-ip", default="192.168.123.164", help="Robot IP (camera + action ports)")
+    p.add_argument("--action-port", type=int, default=ACTION_PORT, help="Onboard ZMQ PULL port for actions")
+    p.add_argument("--camera-port", type=int, default=5555, help="Onboard ZMQ camera PUB port")
+    p.add_argument("--camera-name", default="head_camera", help="Camera name served by run_g1_onboard")
+    p.add_argument("--camera-width", type=int, default=640, help="Camera width")
+    p.add_argument("--camera-height", type=int, default=480, help="Camera height")
+    p.add_argument("--task", default="walk back and forth", help="Language prompt for the VLA")
+    p.add_argument("--fps", type=float, default=30.0, help="Token send rate (matches training inference)")
+    p.add_argument("--device", default="cuda", help="Torch device")
+    p.add_argument("--max-ticks", type=int, default=0, help="Stop after N ticks (0 = run forever)")
+    p.add_argument("--dry-run", action="store_true", help="Run inference but do not PUSH tokens to the robot")
+    args = p.parse_args()
+
+    device = torch.device(args.device)
+
+    # --- Policy + processors (normalization stats baked into the checkpoint) ---
+    logger.info("Loading policy from '%s'...", args.policy_path)
+    policy_cfg = PreTrainedConfig.from_pretrained(args.policy_path)
+    policy_cfg.pretrained_path = args.policy_path
+    policy = get_policy_class(policy_cfg.type).from_pretrained(args.policy_path, config=policy_cfg)
+    policy = policy.to(device)
+    policy.eval()
+    policy.reset()
+
+    preprocessor, postprocessor = make_pre_post_processors(
+        policy_cfg=policy_cfg,
+        pretrained_path=args.policy_path,
+        preprocessor_overrides={"device_processor": {"device": str(device)}},
+    )
+    logger.info("Policy loaded (type=%s, device=%s, chunk=%s)", policy_cfg.type, device,
+                getattr(policy_cfg, "chunk_size", "?"))
+
+    # --- Camera (ZMQ from the robot's onboard image server) ---
+    cam = ZMQCamera(
+        ZMQCameraConfig(
+            server_address=args.robot_ip,
+            port=args.camera_port,
+            camera_name=args.camera_name,
+            width=args.camera_width,
+            height=args.camera_height,
+            fps=int(args.fps),
+        )
+    )
+    logger.info("Connecting camera %s@%s:%d ...", args.camera_name, args.robot_ip, args.camera_port)
+    cam.connect()
+
+    # --- Action PUSH socket to the onboard controller ---
+    import zmq
+
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.PUSH)
+    sock.setsockopt(zmq.SNDHWM, 2)
+    sock.setsockopt(zmq.LINGER, 0)
+    sock.connect(f"tcp://{args.robot_ip}:{args.action_port}")
+    logger.info("Sending tokens to tcp://%s:%d (dry_run=%s)", args.robot_ip, args.action_port, args.dry_run)
+
+    stop = {"flag": False}
+    signal.signal(signal.SIGINT, lambda *_: stop.__setitem__("flag", True))
+    signal.signal(signal.SIGTERM, lambda *_: stop.__setitem__("flag", True))
+
+    # observation.state = the token currently executing on the robot (last one we sent);
+    # start at the neutral token the decoder holds before the first send, so the very
+    # first inference sees the true executing token (not zeros).
+    prev_token = NEUTRAL_TOKEN.copy()
+    period = 1.0 / args.fps
+    n = 0
+    t_infer_total = 0.0
+    logger.info("Streaming tokens at %.0f Hz. Ctrl-C to stop.", args.fps)
+    try:
+        while not stop["flag"]:
+            t0 = time.time()
+            try:
+                frame = cam.read()  # HxWxC uint8 RGB
+            except Exception as e:  # noqa: BLE001
+                logger.warning("Camera read failed: %s", e)
+                time.sleep(period)
+                continue
+
+            raw_obs = {
+                IMAGE_KEY: np.ascontiguousarray(frame),
+                STATE_KEY: prev_token.copy(),
+            }
+            with torch.inference_mode():
+                obs = prepare_observation_for_inference(raw_obs, device, args.task, "unitree_g1")
+                obs = preprocessor(obs)
+                action = policy.select_action(obs)
+                action = postprocessor(action)
+            token = action.squeeze(0).to("cpu").numpy().astype(np.float32)
+            prev_token = token
+
+            if not args.dry_run:
+                msg = {token_action_key(i): float(token[i]) for i in range(TOKEN_DIM)}
+                with contextlib.suppress(zmq.Again):
+                    sock.send_string(json.dumps(msg), zmq.NOBLOCK)
+
+            n += 1
+            t_infer_total += time.time() - t0
+            if n % 30 == 0:
+                logger.info(
+                    "tick %d | avg %.1f ms/tick | token[:3]=%s",
+                    n, 1000.0 * t_infer_total / 30.0, np.round(token[:3], 3).tolist(),
+                )
+                t_infer_total = 0.0
+
+            if args.max_ticks and n >= args.max_ticks:
+                break
+            time.sleep(max(0.0, period - (time.time() - t0)))
+    finally:
+        logger.info("Stopping sender after %d ticks.", n)
+        with contextlib.suppress(Exception):
+            cam.disconnect()
+        with contextlib.suppress(Exception):
+            sock.close(linger=0)
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/robots/unitree_g1/run_g1_onboard.py

@@ -0,0 +1,254 @@
+#!/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.
+
+"""Run the G1 locomotion / whole-body controller ONBOARD, driven by high-level actions
+from a laptop.
+
+The controller (GR00T / Holosoma / SONIC whole-body) runs on the robot itself against
+local DDS, at full control rate. The laptop ships only the resulting high-level action
+(arm joint targets + joystick axes + gripper flags, or a 64-D SONIC motion token) as
+JSON over ZMQ. This process applies each action via ``UnitreeG1.send_action`` while the
+onboard controller thread keeps the legs balanced / decodes the token.
+
+This is the real-deploy counterpart to running ``lerobot-rollout`` on the laptop with
+``--robot.is_simulation=false`` (the ZMQ *socket bridge*): there the 50 Hz lowcmd
+crosses the network; here only compact high-level actions do, and the control loop stays
+local to the robot. Pair with a laptop client that produces actions (exo teleop, or a
+policy such as ``nepyope/sonic_walk`` emitting ``motion_token.{i}.pos``).
+
+Besides receiving actions, this process publishes ``observation.state`` (29 joint ``.q``)
+on a ZMQ PUB port so a laptop policy client has proprioception.
+
+Safety: type ``e`` then Enter in this terminal to stop immediately (zero-torque + exit).
+Ctrl-C does the normal graceful shutdown (kp ramp).
+
+Examples (on the robot):
+
+    # GR00T locomotion, arm targets from the laptop:
+    python -m lerobot.robots.unitree_g1.run_g1_onboard --controller GrootLocomotionController
+
+    # SONIC whole-body walk policy: laptop ships 64-D tokens, decoder runs here:
+    python -m lerobot.robots.unitree_g1.run_g1_onboard \
+        --controller SonicWholeBodyController --sonic-token-action \
+        --cameras "head_camera:/dev/v4l/by-path/platform-3610000.usb-usb-0:2.1:1.3-video-index0:640x480"
+"""
+
+import argparse
+import contextlib
+import json
+import logging
+import os
+import signal
+import sys
+import threading
+import time
+
+import numpy as np
+import zmq
+
+from lerobot.cameras.zmq.image_server import ImageServer
+from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
+from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
+from lerobot.robots.unitree_g1.run_g1_server import Gripper, build_gripper, parse_camera_specs
+from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", force=True)
+logger = logging.getLogger("g1_onboard")
+
+ACTION_PORT = 6004
+STATE_PORT = 6005
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--controller", default="GrootLocomotionController", help="Controller class name")
+    p.add_argument("--dds-interface", default=None, help="DDS network interface (default: SDK default)")
+    p.add_argument(
+        "--sim",
+        action="store_true",
+        help="Attach to a DDS MuJoCo sim: skip MotionSwitcher + physical remote, default dds-interface 'lo'.",
+    )
+    p.add_argument(
+        "--sonic-token-action",
+        action="store_true",
+        help="SONIC token interface: actions carry a 64-D motion_token.{i}.pos that the decoder consumes.",
+    )
+    p.add_argument("--action-port", type=int, default=ACTION_PORT, help="ZMQ PULL port for laptop actions")
+    p.add_argument("--state-port", type=int, default=STATE_PORT, help="ZMQ PUB port for observation.state")
+    p.add_argument("--state-fps", type=float, default=30.0, help="observation.state publish rate; <=0 disables")
+    p.add_argument("--gravity-compensation", action="store_true", help="Enable arm gravity compensation")
+    # Gripper control (Damiao over CAN).
+    p.add_argument("--grippers", action="store_true", help="Drive Damiao grippers from action L3/R3 flags")
+    p.add_argument("--gripper-port-left", default="can1", help="CAN interface for LEFT gripper")
+    p.add_argument("--gripper-port-right", default="can0", help="CAN interface for RIGHT gripper")
+    p.add_argument("--gripper-send-id", type=lambda x: int(x, 0), default=0x08, help="Motor send CAN id")
+    p.add_argument("--gripper-recv-id", type=lambda x: int(x, 0), default=0x18, help="Motor recv CAN id")
+    p.add_argument("--gripper-motor-type", default="dm4310", help="Damiao motor type")
+    p.add_argument("--gripper-open-deg", type=float, default=-65.0, help="Gripper OPEN position (deg)")
+    p.add_argument("--gripper-close-deg", type=float, default=0.0, help="Gripper CLOSE position (deg)")
+    p.add_argument("--gripper-kp", type=float, default=15.0, help="MIT position gain (stiffness)")
+    p.add_argument("--gripper-kd", type=float, default=0.5, help="MIT damping gain")
+    p.add_argument("--gripper-no-fd", dest="gripper_fd", action="store_false", help="Classic CAN (non-FD)")
+    p.set_defaults(gripper_fd=True)
+    # Optional camera streaming (ZMQ) so the laptop policy client / viewer can connect.
+    p.add_argument("--cameras", default=None, help="Camera spec 'name:device[:WxH[:FOURCC]]', comma-sep")
+    p.add_argument("--camera-fps", type=int, default=30, help="Camera FPS")
+    p.add_argument("--camera-port", type=int, default=5555, help="Camera ZMQ port")
+    p.add_argument("--camera-width", type=int, default=640, help="Default camera width")
+    p.add_argument("--camera-height", type=int, default=480, help="Default camera height")
+    args = p.parse_args()
+
+    dds_interface = args.dds_interface
+    if args.sim and dds_interface is None:
+        dds_interface = "lo"
+
+    cfg = UnitreeG1Config(
+        is_simulation=False,
+        onboard=True,
+        controller=args.controller,
+        dds_interface=dds_interface,
+        gravity_compensation=args.gravity_compensation,
+        release_motion_control=not args.sim,
+        physical_remote=not args.sim,
+        sonic_token_action=args.sonic_token_action,
+        cameras={},
+    )
+
+    # Optional camera server (background thread; independent of DDS/CAN).
+    camera_server = None
+    if args.cameras:
+        cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
+        camera_server = ImageServer({"fps": args.camera_fps, "cameras": cameras}, port=args.camera_port)
+        threading.Thread(target=camera_server.run, daemon=True).start()
+        cam_summary = ", ".join(f"{name}(dev {c['device_id']})" for name, c in cameras.items())
+        logger.info("Camera server started on :%d: %s", args.camera_port, cam_summary)
+
+    robot = UnitreeG1(cfg)
+    logger.info("Connecting onboard robot (controller=%s, token=%s)...", args.controller, args.sonic_token_action)
+    robot.connect()
+    # Note: with --sonic-token-action the SonicWholeBodyController holds a neutral
+    # (all-zero) token until the first laptop token arrives, then holds the last token
+    # between ticks -- see SonicWholeBodyController.token_mode (set from config).
+
+    grippers: dict[str, Gripper] = {}
+    if args.grippers:
+        for side, port in (("L", args.gripper_port_left), ("R", args.gripper_port_right)):
+            grippers[side] = build_gripper(
+                side, port, args.gripper_send_id, args.gripper_recv_id, args.gripper_motor_type,
+                args.gripper_fd, args.gripper_open_deg, args.gripper_close_deg, args.gripper_kp, args.gripper_kd,
+            )
+        logger.info("Grippers enabled: L3 -> left, R3 -> right")
+
+    ctx = zmq.Context.instance()
+    sock = ctx.socket(zmq.PULL)
+    sock.setsockopt(zmq.CONFLATE, 1)  # only ever act on the freshest command
+    sock.setsockopt(zmq.RCVTIMEO, 200)  # keeps the loop responsive to the stop event
+    sock.bind(f"tcp://0.0.0.0:{args.action_port}")
+    logger.info("Onboard controller live. Waiting for laptop actions on :%d ...", args.action_port)
+    logger.info("Type 'e' then Enter to STOP immediately (or Ctrl-C for graceful shutdown).")
+
+    stop = threading.Event()
+    signal.signal(signal.SIGINT, lambda *_: stop.set())
+    signal.signal(signal.SIGTERM, lambda *_: stop.set())
+
+    def estop_listener() -> None:
+        for line in sys.stdin:
+            if line.strip().lower() == "e":
+                logger.warning("E-STOP ('e'): going passive NOW.")
+                try:
+                    robot._shutdown_event.set()  # stop the controller loop publishing
+                    time.sleep(0.05)
+                    robot._send_zero_torque()  # motors limp; nothing overwrites it now
+                except Exception as e:  # noqa: BLE001
+                    logger.warning("E-stop zero-torque failed: %s", e)
+                os._exit(0)  # immediate hard exit, no slow cleanup
+
+    threading.Thread(target=estop_listener, daemon=True).start()
+
+    # Proprioception feedback: publish observation.state (29 joint .q) so a laptop
+    # inference client can feed it to a policy. DDS stays local; only compact JSON
+    # state crosses the network. (For a token policy the laptop closes the loop on the
+    # token instead, but publishing joint state is harmless and useful for logging.)
+    state_sock = None
+    if args.state_fps > 0:
+        state_sock = ctx.socket(zmq.PUB)
+        state_sock.setsockopt(zmq.SNDHWM, 2)
+        state_sock.setsockopt(zmq.LINGER, 0)
+        state_sock.bind(f"tcp://0.0.0.0:{args.state_port}")
+        logger.info("Publishing observation.state on :%d at %.0f Hz", args.state_port, args.state_fps)
+
+        def publish_state() -> None:
+            period = 1.0 / args.state_fps
+            joint_names = [j.name for j in G1_29_JointIndex]
+            while not stop.is_set():
+                t0 = time.time()
+                obs = robot.get_observation()
+                if obs:
+                    state = {f"{name}.q": float(obs.get(f"{name}.q", 0.0)) for name in joint_names}
+                    with contextlib.suppress(zmq.Again):
+                        state_sock.send_json(state, zmq.NOBLOCK)
+                time.sleep(max(0.0, period - (time.time() - t0)))
+
+        threading.Thread(target=publish_state, daemon=True).start()
+    else:
+        logger.info("observation.state PUB disabled (--state-fps<=0)")
+
+    n = 0
+    try:
+        while not stop.is_set():
+            try:
+                payload = sock.recv()
+            except zmq.Again:
+                continue
+            except zmq.ContextTerminated:
+                break
+
+            try:
+                action = json.loads(payload.decode("utf-8"))
+            except (ValueError, UnicodeDecodeError) as e:
+                logger.warning("Dropping malformed action: %s", e)
+                continue
+
+            robot.send_action(action)
+
+            if grippers:
+                # L3 = remote.button.4 -> left, R3 = remote.button.0 -> right.
+                if "L" in grippers and "remote.button.4" in action:
+                    grippers["L"].apply(bool(action["remote.button.4"]))
+                if "R" in grippers and "remote.button.0" in action:
+                    grippers["R"].apply(bool(action["remote.button.0"]))
+
+            n += 1
+            if n % 60 == 0:
+                axes = {k: round(float(action.get(k, 0.0)), 3) for k in ("remote.lx", "remote.ly", "remote.rx", "remote.ry")}
+                logger.info("Applied %d actions | axes=%s", n, axes)
+    finally:
+        logger.info("Shutting down onboard controller...")
+        stop.set()
+        if state_sock is not None:
+            with contextlib.suppress(Exception):
+                state_sock.close(linger=0)
+        if camera_server is not None:
+            with contextlib.suppress(Exception):
+                camera_server.stop()
+        for g in grippers.values():
+            with contextlib.suppress(Exception):
+                g.bus.disconnect()
+        robot.disconnect()
+
+
+if __name__ == "__main__":
+    main()

src/lerobot/robots/unitree_g1/run_g1_server.py

@@ -28,9 +28,11 @@ import argparse
 import base64
 import contextlib
 import json
+import re
 import threading
 import time
-from typing import Any
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
 
 import zmq
 from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
@@ -41,6 +43,9 @@ from unitree_sdk2py.utils.crc import CRC
 
 from lerobot.cameras.zmq.image_server import ImageServer
 
+if TYPE_CHECKING:
+    from lerobot.motors.damiao.damiao import DamiaoMotorsBus
+
 # DDS topic names follow Unitree SDK naming conventions
 # ruff: noqa: N816
 kTopicLowCommand_Debug = "rt/lowcmd"  # action to robot
@@ -51,6 +56,105 @@ LOWSTATE_PORT = 6001
 NUM_MOTORS = 35
 
 
+@dataclass
+class Gripper:
+    """A single Damiao gripper that only writes to CAN when the open/close state changes."""
+
+    name: str
+    bus: "DamiaoMotorsBus"
+    open_deg: float
+    close_deg: float
+    _last_cmd: str | None = None  # "open" | "close"
+
+    def apply(self, want_close: bool) -> None:
+        want = "close" if want_close else "open"
+        if want == self._last_cmd:
+            return
+        target = self.close_deg if want_close else self.open_deg
+        self.bus.write("Goal_Position", "gripper", target)
+        self._last_cmd = want
+        print(f"[gripper] {self.name} -> {want.upper()} ({target:.1f} deg)")
+
+
+def build_gripper(
+    name: str,
+    port: str,
+    send_id: int,
+    recv_id: int,
+    motor_type: str,
+    use_can_fd: bool,
+    open_deg: float,
+    close_deg: float,
+    kp: float,
+    kd: float,
+) -> Gripper:
+    from lerobot.motors.damiao.damiao import DamiaoMotorsBus
+    from lerobot.motors.motors_bus import Motor, MotorNormMode
+
+    motors = {
+        "gripper": Motor(
+            id=send_id,
+            model=motor_type,
+            norm_mode=MotorNormMode.DEGREES,
+            motor_type_str=motor_type,
+            recv_id=recv_id,
+        )
+    }
+    bus = DamiaoMotorsBus(port=port, motors=motors, use_can_fd=use_can_fd)
+    print(f"Connecting {name} gripper on {port} (fd={use_can_fd})...")
+    bus.connect(handshake=True)
+    bus.write("Kp", "gripper", kp)
+    bus.write("Kd", "gripper", kd)
+    bus.write("Goal_Position", "gripper", open_deg)  # start open
+    print(f"  {name}: connected, torque enabled, opened.")
+    return Gripper(name, bus, open_deg, close_deg, _last_cmd="open")
+
+
+def parse_camera_specs(spec: str, default_width: int, default_height: int) -> dict[str, dict]:
+    """Parse a multi-camera spec string into an ImageServer ``cameras`` dict.
+
+    Format: comma-separated ``name:device[:WxH[:FOURCC]]`` entries, e.g.
+    ``head_camera:6,left_wrist:0``. ``device`` may be an integer index or an explicit
+    device path (e.g. ``/dev/video6``), including stable ``by-path`` names like
+    ``/dev/v4l/by-path/platform-...:2.1:1.3-video-index0`` which survive USB
+    re-enumeration (unlike bare ``/dev/videoN`` indices). Because a by-path name
+    itself contains colons, the optional ``WxH`` and ``FOURCC`` are parsed from the
+    *right* so the device-path colons are preserved.
+    """
+    wh_re = re.compile(r"\d+x\d+", re.IGNORECASE)
+    fourcc_re = re.compile(r"[A-Za-z0-9]{4}")
+
+    cameras: dict[str, dict] = {}
+    for entry in spec.split(","):
+        entry = entry.strip()
+        if not entry:
+            continue
+        if ":" not in entry:
+            raise ValueError(f"Invalid camera spec '{entry}', expected 'name:device[:WxH[:FOURCC]]'")
+        name, rest = entry.split(":", 1)
+        name = name.strip()
+        tokens = [t.strip() for t in rest.split(":")]
+
+        fourcc = None
+        if len(tokens) >= 3 and wh_re.fullmatch(tokens[-2]) and fourcc_re.fullmatch(tokens[-1]):
+            fourcc = tokens.pop().upper()
+        width, height = default_width, default_height
+        if len(tokens) >= 2 and wh_re.fullmatch(tokens[-1]):
+            w, h = tokens.pop().lower().split("x")
+            width, height = int(w), int(h)
+
+        raw_id = ":".join(tokens).strip()
+        if not raw_id:
+            raise ValueError(f"Invalid camera spec '{entry}', missing device")
+        device_id: int | str = int(raw_id) if raw_id.lstrip("-").isdigit() else raw_id
+        if name in cameras:
+            raise ValueError(f"Duplicate camera name '{name}' in --cameras")
+        cameras[name] = {"device_id": device_id, "shape": [height, width], "fourcc": fourcc}
+    if not cameras:
+        raise ValueError("No cameras parsed from --cameras spec")
+    return cameras
+
+
 def lowstate_to_dict(msg: hg_LowState) -> dict[str, Any]:
     """Convert LowState SDK message to a JSON-serializable dictionary."""
     motor_states = []
@@ -155,7 +259,11 @@ def main() -> None:
     """Main entry point for the robot server bridge."""
     parser = argparse.ArgumentParser(description="DDS-to-ZMQ bridge server for Unitree G1")
     parser.add_argument("--camera", action="store_true", help="Also launch camera server")
-    parser.add_argument("--camera-device", type=int, default=4, help="Camera device ID (default: 4)")
+    parser.add_argument("--camera-device", default="4",
+                        help="Camera device: index or /dev/video path or by-path name (default: 4)")
+    parser.add_argument("--cameras", default=None,
+                        help="Multi-camera spec 'name:device[:WxH[:FOURCC]]', comma-separated. Overrides "
+                             "--camera-device; device may be a by-path name to survive USB re-enumeration.")
     parser.add_argument("--camera-fps", type=int, default=30, help="Camera FPS (default: 30)")
     parser.add_argument("--camera-width", type=int, default=640, help="Camera width (default: 640)")
     parser.add_argument("--camera-height", type=int, default=480, help="Camera height (default: 480)")
@@ -164,20 +272,20 @@ def main() -> None:
 
     # Optionally start camera server in background thread
     camera_thread = None
-    if args.camera:
-        camera_config = {
-            "fps": args.camera_fps,
-            "cameras": {
-                "head_camera": {
-                    "device_id": args.camera_device,
-                    "shape": [args.camera_height, args.camera_width],
-                }
-            },
-        }
+    if args.camera or args.cameras:
+        if args.cameras:
+            cameras = parse_camera_specs(args.cameras, args.camera_width, args.camera_height)
+        else:
+            # Single camera; accept an int index or a device/by-path string.
+            dev = args.camera_device
+            dev = int(dev) if str(dev).lstrip("-").isdigit() else dev
+            cameras = {"head_camera": {"device_id": dev, "shape": [args.camera_height, args.camera_width]}}
+        camera_config = {"fps": args.camera_fps, "cameras": cameras}
         camera_server = ImageServer(camera_config, port=args.camera_port)
         camera_thread = threading.Thread(target=camera_server.run, daemon=True)
         camera_thread.start()
-        print(f"Camera server started on port {args.camera_port} (device {args.camera_device})")
+        cam_summary = ", ".join(f"{n}(dev {c['device_id']})" for n, c in cameras.items())
+        print(f"Camera server started on port {args.camera_port}: {cam_summary}")
 
     # initialize DDS
     ChannelFactoryInitialize(0)

src/lerobot/robots/unitree_g1/unitree_g1.py

@@ -87,6 +87,14 @@ class LocomotionController(Protocol):
 kTopicLowCommand_Debug = "rt/lowcmd"
 kTopicLowState = "rt/lowstate"
 
+# Wireless-remote button byte layout, mapped to the positional button indices the
+# locomotion controllers expect. Used in onboard mode to read the physical Unitree
+# remote from lowstate (mirrors the exo teleoperator's RemoteController).
+_REMOTE_BUTTON_MAP: list[str] = [
+    "RB", "LB", "start", "back", "RT", "LT", "", "",
+    "A", "B", "X", "Y", "up", "right", "down", "left",
+]
+
 
 @dataclass
 class MotorState:
@@ -130,8 +138,10 @@ class UnitreeG1(Robot):
         # Initialize cameras config (ZMQ-based) - actual connection in connect()
         self._cameras = make_cameras_from_configs(config.cameras)
 
-        # Import channel classes based on mode
-        if config.is_simulation:
+        # Import channel classes based on mode. Simulation and onboard both talk to a
+        # real (local) DDS via the Unitree SDK; only the laptop-side bridge client uses
+        # the ZMQ socket shim.
+        if config.is_simulation or config.onboard:
             self._ChannelFactoryInitialize = _SDKChannelFactoryInitialize
             self._ChannelPublisher = _SDKChannelPublisher
             self._ChannelSubscriber = _SDKChannelSubscriber
@@ -159,50 +169,100 @@ class UnitreeG1(Robot):
         # Lower-body controller loaded dynamically
         self.controller: LocomotionController | None = make_locomotion_controller(config.controller)
 
+        # Token-driven deploy: let a SONIC controller hold a neutral token until the
+        # first real one arrives, then hold the last token between control ticks.
+        if config.sonic_token_action and hasattr(self.controller, "token_mode"):
+            self.controller.token_mode = True
+
         # Controller thread state
         self._controller_thread = None
+        # When set, the controller loop stops publishing low commands so reset() can
+        # drive the joints directly without two publishers fighting (single-publisher).
+        self._controller_paused = threading.Event()
         self._controller_action_lock = threading.Lock()
         self.controller_input = default_remote_input()
         self.controller_output = {}
 
-        # Replay-camera state (decoded frames per robot camera name + play cursor).
-        self._replay_frames: dict[str, list[np.ndarray]] = {}
+        # Onboard-only: parser for the physical Unitree wireless remote (read straight
+        # from local lowstate so joystick locomotion works without a laptop round-trip).
+        self._joystick = None
+
+        # Replay-camera state: keep the encoded (raw) cells per camera and decode
+        # frames lazily as the play cursor advances, with a small frame cache, so we
+        # don't materialize gigabytes of decoded RGB at construction time.
+        self._replay_raw: dict[str, list] = {}
+        self._replay_cache: dict[tuple[str, int], np.ndarray] = {}
+        self._replay_cache_cap = 8
         self._replay_len = 0
         self._replay_idx = 0
         if config.replay_camera_parquet and config.replay_camera_map:
             self._load_replay_frames()
 
+        # Token-mode state: last 64-D SONIC latent token commanded by the policy,
+        # echoed back as ``observation.state`` so a token-output VLA closes the loop
+        # on its own previous token (see ``sonic_token_action``). Seeded to zeros;
+        # the controller's startup blend eases joints in regardless.
+        self._last_token: np.ndarray | None = None
+        if config.sonic_token_action:
+            from .controllers.sonic_whole_body import TOKEN_DIM
+
+            self._last_token = np.zeros(TOKEN_DIM, dtype=np.float32)
+
     def _load_replay_frames(self) -> None:
-        """Decode recorded episode frames from a parquet into per-camera image lists."""
-        import io
-
+        """Load only the mapped parquet columns (encoded frames); decode on demand."""
         import pyarrow.parquet as pq
-        from PIL import Image
 
-        table = pq.read_table(self.config.replay_camera_parquet)
-        cols = {col: table.column(col).to_pylist() for col in self.config.replay_camera_map.values()}
+        cols_needed = list(dict.fromkeys(self.config.replay_camera_map.values()))
+        table = pq.read_table(self.config.replay_camera_parquet, columns=cols_needed)
         self._replay_len = table.num_rows
-
-        def decode(cell) -> np.ndarray:
-            data = cell["bytes"] if isinstance(cell, dict) else cell
-            return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
-
-        for cam_name, column in self.config.replay_camera_map.items():
-            self._replay_frames[cam_name] = [decode(c) for c in cols[column]]
+        self._replay_raw = {
+            cam_name: table.column(column).to_pylist()
+            for cam_name, column in self.config.replay_camera_map.items()
+        }
         logger.info(
-            "Loaded %d replay frames for cameras %s from %s",
+            "Loaded %d replay frames (lazy-decode) for cameras %s from %s",
             self._replay_len,
             list(self.config.replay_camera_map),
             self.config.replay_camera_parquet,
         )
 
+    def _decode_replay_cell(self, cell) -> np.ndarray:
+        import io
+
+        from PIL import Image
+
+        data = cell["bytes"] if isinstance(cell, dict) else cell
+        return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
+
+    def _replay_frame(self, cam_name: str, idx: int) -> np.ndarray:
+        """Decode (and briefly cache) a single replay frame for a camera."""
+        key = (cam_name, idx)
+        cached = self._replay_cache.get(key)
+        if cached is not None:
+            return cached
+        frame = self._decode_replay_cell(self._replay_raw[cam_name][idx])
+        if len(self._replay_cache) >= self._replay_cache_cap:
+            self._replay_cache.pop(next(iter(self._replay_cache)))
+        self._replay_cache[key] = frame
+        return frame
+
     def _subscribe_lowstate(self):  # polls robot state @ 250Hz
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
             # Step simulation if in simulation mode
             if self.config.is_simulation and self.sim_env is not None:
-                self.sim_env.step()
+                try:
+                    self.sim_env.step()
+                except ValueError as e:
+                    # Startup race: the sim thread can step once before reset() has
+                    # written a valid base pose, giving a zero-norm pelvis quaternion
+                    # (scipy>=1.11 raises instead of normalizing). Skip and retry so
+                    # the thread survives instead of dying and freezing the sim.
+                    if "zero norm" not in str(e).lower():
+                        raise
+                    time.sleep(self.control_dt)
+                    continue
 
             msg = self.lowstate_subscriber.Read()
             if msg is not None:
@@ -278,28 +338,46 @@ class UnitreeG1(Robot):
         (OpenHLM / pi0.5). These ``.pos`` scalars are aggregated by the rollout
         pipeline into a single 34-D ``observation.state`` for the policy.
         """
+        if self.config.sonic_token_action:
+            return {}
         if not getattr(self.controller, "wb_action", False):
             return {}
         from .g1_utils import WB_ACTION_DIM
 
         return {f"wb_state.{i}.pos": float for i in range(WB_ACTION_DIM)}
 
+    @property
+    def _token_state_ft(self) -> dict[str, type]:
+        """64-D SONIC latent-token proprio state (``motion_token_state.{i}.pos``).
+
+        Exposed only in ``sonic_token_action`` mode; aggregated by the rollout into a
+        64-D ``observation.state`` (the last token the policy commanded).
+        """
+        if not self.config.sonic_token_action:
+            return {}
+        from .controllers.sonic_whole_body import TOKEN_DIM, token_state_key
+
+        return {token_state_key(i): float for i in range(TOKEN_DIM)}
+
     @property
     def _empty_cameras_ft(self) -> dict[str, tuple]:
         """Synthetic zero-image cameras (see ``UnitreeG1Config.empty_cameras``)."""
         h, w = self.config.empty_camera_hw
-        return {name: (h, w, 3) for name in self.config.empty_cameras}
+        return dict.fromkeys(self.config.empty_cameras, (h, w, 3))
 
     @property
     def _replay_cameras_ft(self) -> dict[str, tuple]:
-        """Replay cameras, shaped from their first decoded frame."""
-        return {name: frames[0].shape for name, frames in self._replay_frames.items() if frames}
+        """Replay cameras, shaped from their first (lazily decoded) frame."""
+        if not self._replay_len:
+            return {}
+        return {name: self._replay_frame(name, 0).shape for name in self._replay_raw}
 
     @cached_property
     def observation_features(self) -> dict[str, type | tuple]:
         return {
             **self._motors_ft,
             **self._wb_state_ft,
+            **self._token_state_ft,
             **self._empty_cameras_ft,
             **self._replay_cameras_ft,
             **self._cameras_ft,
@@ -310,6 +388,14 @@ class UnitreeG1(Robot):
         if self.controller is None:
             return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
 
+        # Token-output VLA (SONIC decoder): advertise a 64-D latent-token action space
+        # (``motion_token.{i}.pos``) so ``lerobot-rollout`` maps a 64-D policy output
+        # straight onto the decoder, bypassing the encoder.
+        if self.config.sonic_token_action:
+            from .controllers.sonic_whole_body import TOKEN_DIM, token_action_key
+
+            return {token_action_key(i): float for i in range(TOKEN_DIM)}
+
         # Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
         # 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
         # so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
@@ -333,6 +419,11 @@ class UnitreeG1(Robot):
         while not self._shutdown_event.is_set():
             start_time = time.time()
 
+            # Paused during reset() so the reset routine is the sole low-cmd publisher.
+            if self._controller_paused.is_set():
+                time.sleep(control_dt)
+                continue
+
             with self._lowstate_lock:
                 lowstate = self._lowstate
 
@@ -349,6 +440,13 @@ class UnitreeG1(Robot):
                 with self._controller_action_lock:
                     controller_input = dict(self.controller_input)
 
+                # Onboard: the physical Unitree remote (in local lowstate) takes
+                # priority for locomotion when active; otherwise laptop/ZMQ axes stand.
+                if self.config.onboard:
+                    wl = self._wireless_remote_input(lowstate)
+                    if wl is not None:
+                        controller_input.update(wl)
+
                 # Run controller step
                 controller_action = self.controller.run_step(controller_input, lowstate)
 
@@ -371,15 +469,105 @@ class UnitreeG1(Robot):
     def configure(self) -> None:
         pass
 
+    def _wireless_remote_input(self, lowstate) -> dict | None:
+        """Parse the physical Unitree remote from lowstate into controller inputs.
+
+        Onboard only. Returns None when the remote is idle so the laptop-provided
+        (ZMQ) axes keep control; otherwise the physical remote takes priority.
+        """
+        js = self._joystick
+        if js is None:
+            return None
+        wr = getattr(lowstate, "wireless_remote", None)
+        if not wr or len(wr) < 24:
+            return None
+        try:
+            js.extract(wr)
+        except Exception:  # noqa: BLE001
+            return None
+
+        axes = {
+            "remote.lx": float(js.lx.data),
+            "remote.ly": float(js.ly.data),
+            "remote.rx": float(js.rx.data),
+            "remote.ry": float(js.ry.data),
+        }
+        active = any(abs(v) > 1e-2 for v in axes.values())
+        out = dict(axes)
+        for i, name in enumerate(_REMOTE_BUTTON_MAP):
+            if name:
+                val = float(getattr(js, name).data)
+                out[f"remote.button.{i}"] = val
+                if val:
+                    active = True
+        return out if active else None
+
+    def _release_motion_control(self) -> None:
+        """Release the robot's built-in motion services so we can send raw lowcmd.
+
+        Onboard-only. Mirrors run_g1_server.py: on the real robot the factory
+        locomotion/hand services must relinquish control before our controller can
+        write to ``rt/lowcmd``, otherwise commands are ignored or fought.
+        """
+        from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient
+
+        msc = MotionSwitcherClient()
+        msc.SetTimeout(5.0)
+        msc.Init()
+        _, result = msc.CheckMode()
+        while result is not None and "name" in result and result["name"]:
+            logger.info("[UnitreeG1] Releasing built-in mode '%s'...", result["name"])
+            msc.ReleaseMode()
+            _, result = msc.CheckMode()
+            time.sleep(1.0)
+
     def connect(self, calibrate: bool = True) -> None:  # connect to DDS
         # Initialize DDS channel and simulation environment
         if self.config.is_simulation:
-            from lerobot.envs import make_env
+            from lerobot.envs.utils import (
+                _download_hub_file,
+                _import_hub_module,
+                _normalize_hub_result,
+            )
 
             self._ChannelFactoryInitialize(0, "lo")
-            self._env_wrapper = make_env("lerobot/unitree-g1-mujoco", trust_remote_code=True)
+            # Call the hub env's make_env directly so we can disable the offscreen
+            # head_camera renderer. We drive image-conditioned policies from recorded
+            # frames (see replay_camera_parquet / external obs), never the sim's own
+            # camera, so building a MuJoCo offscreen GL context is pure liability: it
+            # crashes with "Failed to make the EGL context current" when GLFW/SDL
+            # already own a context, killing the sim thread and hanging on
+            # "Waiting for robot state...". publish_images=False -> no renderer.
+            repo_id, _, local_file, _ = _download_hub_file(
+                "lerobot/unitree-g1-mujoco", True, None
+            )
+            hub_mod = _import_hub_module(local_file, repo_id)
+            raw = hub_mod.make_env(n_envs=1, use_async_envs=False, publish_images=False, cameras=[])
+            self._env_wrapper = _normalize_hub_result(raw)
             # Extract the actual gym env from the dict structure
             self.sim_env = self._env_wrapper["hub_env"][0].envs[0]
+        elif self.config.onboard:
+            # Real robot, controller running onboard against local DDS. Initialize the
+            # real SDK channel factory on the robot's DDS interface and take low-level
+            # control from the built-in services before we start writing lowcmd.
+            if self.config.dds_interface:
+                self._ChannelFactoryInitialize(0, self.config.dds_interface)
+            else:
+                self._ChannelFactoryInitialize(0)
+            # Real robot: hand low-level control over from the built-in services.
+            # A DDS sim has no MotionSwitcher, so this is skipped there.
+            if self.config.release_motion_control:
+                self._release_motion_control()
+            # Real robot: read the physical wireless remote from lowstate for
+            # locomotion. A sim has no physical remote, so leave _joystick=None and
+            # let send_action (ZMQ) drive the locomotion axes instead.
+            if self.config.physical_remote:
+                from unitree_sdk2py.utils.joystick import Joystick
+
+                self._joystick = Joystick()
+                for axis in (self._joystick.lx, self._joystick.ly, self._joystick.rx, self._joystick.ry):
+                    axis.smooth = 1.0
+                    axis.deadzone = 0.0
         else:
             self._ChannelFactoryInitialize(0, config=self.config)
 
@@ -442,12 +630,16 @@ class UnitreeG1(Robot):
             self.msg.motor_cmd[joint].kd = self.kd[joint.value]
             self.msg.motor_cmd[joint].q = lowstate.motor_state[joint.value].q
 
-        # Start controller thread if enabled
-        if self.controller is not None:
+        # Start controller thread if enabled. Skipped when run_controller_thread is
+        # False so a caller can step the controller synchronously (faithful replay).
+        if self.controller is not None and self.config.run_controller_thread:
             self._controller_thread = threading.Thread(target=self._controller_loop, daemon=True)
             self._controller_thread.start()
             fps = int(1.0 / self.controller.control_dt)
             logger.info(f"Controller thread started ({fps}Hz)")
+        elif self.controller is not None:
+            logger.info("Controller thread disabled (run_controller_thread=False); "
+                        "caller must drive controller.run_step synchronously.")
 
     def _send_zero_torque(self) -> None:
         """Send a zero-gain command to make joints passive before shutting down."""
@@ -495,14 +687,23 @@ class UnitreeG1(Robot):
     def disconnect(self):
         # Stop the controller loop first so it isn't fighting the shutdown ramp.
         self._shutdown_event.set()
+        controller_stopped = True
         if self._controller_thread is not None:
-            self._controller_thread.join(timeout=2.0)
+            # Wait long enough for any in-flight inference tick to finish and the loop
+            # to observe the shutdown flag, so no stray low command is published while
+            # the ramp runs (the shutdown routine must be the single publisher).
+            self._controller_thread.join(timeout=5.0)
             if self._controller_thread.is_alive():
-                logger.warning("Controller thread did not stop cleanly")
+                controller_stopped = False
+                logger.error(
+                    "Controller thread did not stop; skipping graceful ramp to avoid "
+                    "concurrent low commands (fail-safe: joints keep last command until exit)"
+                )
 
         # Soft, damped settle instead of an instant limp (real robot only; the
-        # subscribe thread is still alive here to supply the current pose).
-        if not self.config.is_simulation:
+        # subscribe thread is still alive here to supply the current pose). Only ramp
+        # once the controller thread has definitely exited.
+        if not self.config.is_simulation and controller_stopped:
             self._graceful_stop()
 
         if self.controller is not None and hasattr(self.controller, "shutdown"):
@@ -551,7 +752,15 @@ class UnitreeG1(Robot):
         # Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
         # state as ``wb_state.{i}.pos`` so the rollout aggregates it into
         # ``observation.state`` for the policy.
-        if getattr(self.controller, "wb_action", False):
+        if self.config.sonic_token_action:
+            # Token mode: echo the last commanded latent token as observation.state
+            # so a token-output VLA closes the loop on its own previous token.
+            from .controllers.sonic_whole_body import token_state_key
+
+            token = self._last_token if self._last_token is not None else []
+            for i, v in enumerate(token):
+                obs[token_state_key(i)] = float(v)
+        elif getattr(self.controller, "wb_action", False):
             wb_state = obs_to_wb34_state(obs)
             for i, v in enumerate(wb_state):
                 obs[f"wb_state.{i}.pos"] = float(v)
@@ -569,8 +778,8 @@ class UnitreeG1(Robot):
             idx = self._replay_idx
             if idx >= self._replay_len:
                 idx = self._replay_len - 1 if not self.config.replay_camera_loop else idx % self._replay_len
-            for name, frames in self._replay_frames.items():
-                obs[name] = frames[idx]
+            for name in self._replay_raw:
+                obs[name] = self._replay_frame(name, idx)
             self._replay_idx += 1
 
         # Cameras - read images from ZMQ cameras
@@ -585,6 +794,12 @@ class UnitreeG1(Robot):
     def send_action(self, action: RobotAction) -> RobotAction:
         action_to_publish = action
         if self.controller is not None:
+            if self.config.sonic_token_action:
+                from .controllers.sonic_whole_body import _extract_token_from_action
+
+                token = _extract_token_from_action(action)
+                if token is not None:
+                    self._last_token = token
             self._update_controller_action(action)
             if self.config.publish_hands and getattr(self.controller, "wb_action", False):
                 self._publish_hand_cmds(action)
@@ -709,43 +924,64 @@ class UnitreeG1(Robot):
         if default_positions is None:
             default_positions = np.array(self.config.default_positions, dtype=np.float32)
 
-        if self.config.is_simulation and self.sim_env is not None:
-            self.sim_env.reset()
-            self.publish_lowcmd(
-                {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
-            )
-        else:
-            total_time = 3.0
-            num_steps = int(total_time / control_dt)
+        # Full-body controllers (SONIC / OpenHLM) own the whole 29-DoF command and
+        # ignore ``<joint>.q`` in send_action(), so reset() must publish the default
+        # pose directly. Pause the background controller first so the two aren't both
+        # writing low commands while the robot moves to the default pose.
+        full_body = getattr(self.controller, "full_body", False)
+        paused = False
+        if full_body and self._controller_thread is not None:
+            self._controller_paused.set()
+            paused = True
+            time.sleep(control_dt)  # let any in-flight controller tick settle
 
-            # get current state
-            obs = self.get_observation()
+        try:
+            if self.config.is_simulation and self.sim_env is not None:
+                self.sim_env.reset()
+                self.publish_lowcmd(
+                    {f"{motor.name}.q": float(default_positions[motor.value]) for motor in G1_29_JointIndex}
+                )
+            else:
+                total_time = 3.0
+                num_steps = int(total_time / control_dt)
 
-            # record current positions
-            init_dof_pos = np.zeros(29, dtype=np.float32)
-            for motor in G1_29_JointIndex:
-                init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
+                # get current state
+                obs = self.get_observation()
 
-            # Interpolate to default position
-            for step in range(num_steps):
-                start_time = time.time()
-
-                alpha = step / num_steps
-                action_dict = {}
+                # record current positions
+                init_dof_pos = np.zeros(29, dtype=np.float32)
                 for motor in G1_29_JointIndex:
-                    target_pos = default_positions[motor.value]
-                    interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
-                    action_dict[f"{motor.name}.q"] = float(interp_pos)
+                    init_dof_pos[motor.value] = obs[f"{motor.name}.q"]
 
-                self.send_action(action_dict)
+                # Interpolate to default position
+                for step in range(num_steps):
+                    start_time = time.time()
 
-                # Maintain constant control rate
-                elapsed = time.time() - start_time
-                sleep_time = max(0, control_dt - elapsed)
-                time.sleep(sleep_time)
+                    alpha = step / num_steps
+                    action_dict = {}
+                    for motor in G1_29_JointIndex:
+                        target_pos = default_positions[motor.value]
+                        interp_pos = init_dof_pos[motor.value] * (1 - alpha) + target_pos * alpha
+                        action_dict[f"{motor.name}.q"] = float(interp_pos)
 
-        # Reset controller internal state (gait phase, obs history, etc.)
-        if self.controller is not None and hasattr(self.controller, "reset"):
-            self.controller.reset()
+                    # Full-body controllers no-op in send_action(); publish the pose
+                    # directly (arm-only controllers keep the send_action() path).
+                    if full_body:
+                        self.publish_lowcmd(action_dict)
+                    else:
+                        self.send_action(action_dict)
+
+                    # Maintain constant control rate
+                    elapsed = time.time() - start_time
+                    sleep_time = max(0, control_dt - elapsed)
+                    time.sleep(sleep_time)
+
+            # Reset controller internal state (gait phase, obs history, etc.) before
+            # resuming so its buffers reflect the post-reset pose.
+            if self.controller is not None and hasattr(self.controller, "reset"):
+                self.controller.reset()
+        finally:
+            if paused:
+                self._controller_paused.clear()
 
         logger.info("Reached default position")

src/lerobot/scripts/lerobot_teleoperate.py

@@ -116,7 +116,6 @@ 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

@@ -1,20 +0,0 @@
-#!/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

@@ -1,41 +0,0 @@
-#!/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."""
-    mode: str = "smpl"
-    """Teleop reference to emit: ``"smpl"`` for whole-body imitation (SONIC
-    encode_mode 2) or ``"vr3"`` for sparse 3-point upper-body teleop (encode_mode 1,
-    lower body driven by the joystick/keyboard planner)."""

src/lerobot/teleoperators/pico_headset/pico_headset.py

@@ -1,148 +0,0 @@
-#!/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.types import RobotAction
-
-from ..teleoperator import Teleoperator
-from .config_pico_headset import PicoHeadsetConfig
-from .smpl_constants import (
-    LOCO_AXES_PREFIX,
-    LOCO_BTN_PREFIX,
-    LOCO_N_AXES,
-    LOCO_N_BTN,
-    ROOT_ACTION_DIM,
-    ROOT_ACTION_PREFIX,
-    SMPL_ACTION_PREFIX,
-    SMPL_OBS_DIM,
-    VR3_ORN_DIM,
-    VR3_ORN_PREFIX,
-    VR3_POS_DIM,
-    VR3_POS_PREFIX,
-)
-from .smpl_stream import SmplStream
-
-logger = logging.getLogger(__name__)
-
-
-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:
-        if self.config.mode == "vr3":
-            feats = {f"{VR3_POS_PREFIX}{i}": float for i in range(VR3_POS_DIM)}
-            feats.update({f"{VR3_ORN_PREFIX}{i}": float for i in range(VR3_ORN_DIM)})
-            # Controller-stick locomotion travels alongside the VR targets.
-            feats.update({f"{LOCO_AXES_PREFIX}{i}": float for i in range(LOCO_N_AXES)})
-            feats.update({f"{LOCO_BTN_PREFIX}{i}": float for i in range(LOCO_N_BTN)})
-            return feats
-        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 a reference 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 self.config.mode == "vr3":
-            # Sparse 3-point upper-body teleop (encode_mode 1). Gated on fresh vr3_*
-            # frames only (independent of the SMPL window), so the controller-state
-            # source (head + controllers, no body tracking) works. Emit nothing
-            # otherwise and stay in locomotion.
-            if not self._stream.has_fresh_vr3:
-                return {}
-            action = {f"{VR3_POS_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.vr3_pos)}
-            action.update({f"{VR3_ORN_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.vr3_orn)})
-            # Forward controller-stick locomotion when present, so the planner can
-            # steer walking/turning under the upper-body tracking (encode_mode 1).
-            if self._stream.has_fresh_loco:
-                action.update(
-                    {f"{LOCO_AXES_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.loco_axes)}
-                )
-                action.update(
-                    {f"{LOCO_BTN_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.loco_buttons)}
-                )
-            return action
-        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

@@ -1,333 +0,0 @@
-#!/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.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
-(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,
-    ThreePointCalibrator,
-    canonicalize_smpl_joints,
-    compute_3point,
-    compute_3point_from_devices,
-    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,
-    vr3_pos: np.ndarray | None = None,
-    vr3_orn: np.ndarray | None = None,
-    loco_axes: np.ndarray | None = None,
-    loco_buttons: np.ndarray | None = None,
-) -> bytes:
-    """Build the rt/smpl JSON message (single frame, topic embedded in payload).
-
-    Carries the SMPL whole-body window (``smpl_joints_local`` + ``root_*``) and,
-    when available, the sparse 3-point VR targets (``vr3_pos`` (9,), ``vr3_orn`` (12,))
-    so a single stream can drive either SONIC ``encode_mode`` 1 or 2.
-    """
-    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()
-    if vr3_pos is not None:
-        data["vr3_pos"] = np.asarray(vr3_pos, np.float32).reshape(-1).tolist()
-    if vr3_orn is not None:
-        data["vr3_orn"] = np.asarray(vr3_orn, np.float32).reshape(-1).tolist()
-    if loco_axes is not None:
-        data["loco_axes"] = np.asarray(loco_axes, np.float32).reshape(-1).tolist()
-    if loco_buttons is not None:
-        data["loco_buttons"] = np.asarray(loco_buttons, 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")
-    p.add_argument(
-        "--headset-source",
-        choices=["body", "devices"],
-        default="body",
-        help=(
-            "Live headset keypoint source: 'body' uses full-body tracking "
-            "(get_body_joints_pose, needs PICO Motion Trackers) and drives both SMPL "
-            "(encode_mode 2) and 3-point; 'devices' uses head + 2 controllers only "
-            "(get_headset_pose + get_*_controller_pose, no trackers) and emits 3-point "
-            "(encode_mode 1) exclusively."
-        ),
-    )
-    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 f"headset:{args.headset_source}")
-    )
-    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
-    last_tracked = -1
-    t0 = time.time()
-
-    # 3-point operator calibration (device source only): map the operator's neutral
-    # rest pose onto the G1's neutral stance. Trigger a (re)capture with the A+B+X+Y
-    # controller combo, mirroring gear_sonic's ThreePointPose.calibrate_now.
-    # Device source: the "neck" is the headset (pitches when looking down), so keep the
-    # wrist targets in the yaw-local world frame rather than de-rotating by head tilt.
-    calibrator = (
-        ThreePointCalibrator(neck_relative_wrists=False)
-        if (xrt is not None and args.headset_source == "devices")
-        else None
-    )
-    calib_combo_last = False
-    if calibrator is not None:
-        print(
-            "[pico_publisher] 3-point calibration: stand in a neutral rest pose and press "
-            "A+B+X+Y on the controllers to (re)calibrate."
-        )
-    try:
-        while True:
-            loop_start = time.time()
-            vr3_pos = vr3_orn = None
-            loco_axes = loco_buttons = None
-            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()
-            elif not args.fake and args.headset_source == "devices":
-                # Controller-state 3-point path: head + 2 controllers only, no PICO
-                # Motion Trackers / body tracking. Emits encode_mode-1 targets only;
-                # the SMPL whole-body window is left as a zero placeholder.
-                head = np.asarray(xrt.get_headset_pose(), np.float32)
-                lc = np.asarray(xrt.get_left_controller_pose(), np.float32)
-                rc = np.asarray(xrt.get_right_controller_pose(), np.float32)
-                stamp_ns = int(xrt.get_time_stamp_ns())
-                if head.shape != (7,) or lc.shape != (7,) or rc.shape != (7,):
-                    time.sleep(0.005)
-                    continue
-                last_tracked = int(sum(np.linalg.norm(d[3:7]) > 1e-6 for d in (head, lc, rc)))
-                # Empty SMPL window: this source drives encode_mode 1 only, so the
-                # consumer must not mistake it for a (zero) whole-body reference.
-                joints = np.zeros((0, 3), np.float32)
-                root_quat = None
-                root_transl = None
-                vr3_pos, vr3_orn = compute_3point_from_devices(head, lc, rc)
-                # Edge-triggered (re)calibration on the A+B+X+Y combo.
-                combo_now = bool(
-                    xrt.get_A_button() and xrt.get_B_button() and xrt.get_X_button() and xrt.get_Y_button()
-                )
-                if combo_now and not calib_combo_last:
-                    calibrator.capture(vr3_pos, vr3_orn)
-                    print("\n[pico_publisher] 3-point calibration captured (neutral pose).")
-                calib_combo_last = combo_now
-                vr3_pos, vr3_orn = calibrator.apply(vr3_pos, vr3_orn)
-                # Controller-stick locomotion (encode_mode 1, replicated): left/right
-                # sticks + A/B/X/Y. The all-four combo is reserved for calibration, so
-                # suppress the button pairs on that frame to avoid a spurious mode cycle.
-                la = np.asarray(xrt.get_left_axis(), np.float32).reshape(-1)
-                ra = np.asarray(xrt.get_right_axis(), np.float32).reshape(-1)
-                loco_axes = np.array([la[0], la[1], ra[0], ra[1]], np.float32)
-                btn = (
-                    np.zeros(4, np.float32)
-                    if combo_now
-                    else np.array(
-                        [
-                            float(xrt.get_A_button()),
-                            float(xrt.get_B_button()),
-                            float(xrt.get_X_button()),
-                            float(xrt.get_Y_button()),
-                        ],
-                        np.float32,
-                    )
-                )
-                loco_buttons = btn
-            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
-                # How many joints are actually being tracked (non-zero-norm quat).
-                # If this stays near 0, the headset isn't streaming body data (e.g.
-                # "Full body"/"Send" not enabled, trackers uncalibrated, or a test
-                # device) and the reference will be static regardless of your motion.
-                # In that case, use --headset-source devices for head+controllers only.
-                quat_norms = np.linalg.norm(body_poses[:, 3:7], axis=1)
-                last_tracked = int(np.count_nonzero(quat_norms > 1e-6))
-                out = fk.compute(body_poses)
-                joints = out["smpl_joints_local"]
-                root_quat = out["root_quat"]
-                root_transl = out["root_transl"]
-                # Also emit the sparse 3-point VR targets so the same stream can
-                # drive encode_mode 1 (3-point teleop) without a second producer.
-                vr3_pos, vr3_orn = compute_3point(body_poses)
-
-            sock.send(
-                pack_message(
-                    joints,
-                    frame_index,
-                    stamp_ns,
-                    root_quat=root_quat,
-                    root_transl=root_transl,
-                    vr3_pos=vr3_pos,
-                    vr3_orn=vr3_orn,
-                    loco_axes=loco_axes,
-                    loco_buttons=loco_buttons,
-                )
-            )
-            frame_index += 1
-            if frame_index % int(max(1, args.fps)) == 0:
-                denom = 3 if (not args.fake and clip is None and args.headset_source == "devices") else 24
-                unit = "devices" if denom == 3 else "joints"
-                extra = f" | tracked {last_tracked}/{denom} {unit}" if last_tracked >= 0 else ""
-                if calibrator is not None:
-                    extra += " | calibrated" if calibrator.is_calibrated else " | UNCALIBRATED"
-                print(f"[pico_publisher] sent {frame_index} frames{extra}", 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_constants.py

@@ -1,80 +0,0 @@
-#!/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.
-
-"""Single source of truth for the SONIC SMPL whole-body action protocol.
-
-These constants define the wire format shared between the PICO teleoperator
-(producer, ``pico_headset.py``), the offline motion->dataset converter
-(``smpl_to_dataset.py``), the live stream (``smpl_stream.py``), and the Unitree G1
-``SonicWholeBodyController`` (consumer). Keeping them here avoids the producer and
-consumer silently drifting apart.
-"""
-
-# SMPL encoder window geometry (matches ``smpl_joints_10frame_step1``).
-WINDOW = 10  # frames per encoder window
-N_JOINTS = 24  # SMPL joints per frame
-JOINT_DIM = 3  # xyz per joint
-SMPL_OBS_DIM = WINDOW * N_JOINTS * JOINT_DIM  # 720
-
-# Flat action-dict keys carrying the reference through the standard lerobot action
-# pipeline as scalar floats: ``smpl.0 .. smpl.719`` and ``root.0 .. root.3``.
-SMPL_ACTION_PREFIX = "smpl."
-ROOT_ACTION_PREFIX = "root."
-ROOT_ACTION_DIM = 4  # per-frame SMPL root orientation (wxyz)
-
-# Full per-frame action vector: 720 joint window + 4 root quaternion = 724.
-ACTION_DIM = SMPL_OBS_DIM + ROOT_ACTION_DIM
-
-# ── 3-point VR teleop protocol (SONIC encode_mode 1) ─────────────────────────
-# An alternative, sparse upper-body reference: 3 root-relative keypoints
-# (left wrist, right wrist, neck), each a position + orientation. The lower body /
-# locomotion is driven by the planner (joystick/keyboard), not by these targets.
-VR3_N_POINTS = 3  # left wrist, right wrist, neck
-VR3_POS_DIM = VR3_N_POINTS * 3  # 9  (3 x xyz)
-VR3_ORN_DIM = VR3_N_POINTS * 4  # 12 (3 x wxyz)
-
-# Flat action-dict keys: ``vr3_pos.0 .. vr3_pos.8`` and ``vr3_orn.0 .. vr3_orn.11``.
-VR3_POS_PREFIX = "vr3_pos."
-VR3_ORN_PREFIX = "vr3_orn."
-
-# ── Controller-stick locomotion (SONIC encode_mode 1, replicated exactly) ────
-# In the original 3-point teleop the same tick that sends the VR targets also drives
-# locomotion from the PICO controller sticks/buttons (left stick -> move, right stick
-# -> yaw, A+B / X+Y -> cycle locomotion mode). We forward that raw controller state so
-# the consumer's planner can steer walking/turning underneath the upper-body tracking.
-LOCO_N_AXES = 4  # [left_x, left_y, right_x, right_y]
-LOCO_N_BTN = 4  # [A, B, X, Y]
-LOCO_AXES_PREFIX = "loco_axes."
-LOCO_BTN_PREFIX = "loco_btn."
-
-# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM-style) ─────
-# A single 34-D whole-body command per tick, in the pi0.5 / OpenHLM action layout:
-#   [L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3),
-#    root roll/pitch + yaw-rate(3)]
-# The 29 joint targets (arms/legs/waist, grippers excluded) become the mode-0
-# encoder joint reference and root roll/pitch become the anchor orientation.
-#
-# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix is required so
-# these behave like ordinary joint-position action features: ``lerobot-rollout``
-# only routes ``*.pos`` keys from ``robot.action_features`` into the policy<->robot
-# action mapping, letting a 34-D VLA (OpenHLM / pi0.5) drive the robot directly.
-WB_ACTION_PREFIX = "wb."
-WB_ACTION_DIM = 34
-
-
-def wb_action_key(i: int) -> str:
-    """Action-dict key for the ``i``-th whole-body command scalar (``wb.{i}.pos``)."""
-    return f"{WB_ACTION_PREFIX}{i}.pos"

src/lerobot/teleoperators/pico_headset/smpl_fk.py

@@ -1,571 +0,0 @@
-#!/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
-(SMPL-X rest-pose joints + kinematic tree), hardcoded below as ``_SKELETON_J`` /
-``_SKELETON_PARENTS`` so no external asset download is required.
-
-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 dataclasses import dataclass, field
-from pathlib import Path
-
-import numpy as np
-from scipy.spatial.transform import Rotation as R  # noqa: N817
-
-from .smpl_constants import VR3_N_POINTS, VR3_ORN_DIM, VR3_POS_DIM
-
-# 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])])
-
-# SMPL-X rest-pose skeleton (55 joints), vendored inline to avoid an external
-# asset. ``_SKELETON_PARENTS[i]`` is joint ``i``'s parent (-1 = root); ``_SKELETON_J``
-# holds the (55, 3) rest-pose joint positions.
-_SKELETON_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,
-        15,
-        15,
-        15,
-        20,
-        25,
-        26,
-        20,
-        28,
-        29,
-        20,
-        31,
-        32,
-        20,
-        34,
-        35,
-        20,
-        37,
-        38,
-        21,
-        40,
-        41,
-        21,
-        43,
-        44,
-        21,
-        46,
-        47,
-        21,
-        49,
-        50,
-        21,
-        52,
-        53,
-    ],
-    dtype=np.int64,
-)
-
-_SKELETON_J = np.array(
-    [
-        [0.0031232605688273907, -0.3514074683189392, 0.012036550790071487],
-        [0.06131265312433243, -0.4441709518432617, -0.013964635320007801],
-        [-0.06014421582221985, -0.4553154706954956, -0.009213820099830627],
-        [0.00036056205863133073, -0.2415168583393097, -0.015581080690026283],
-        [0.11600811034440994, -0.8229243755340576, -0.02336069941520691],
-        [-0.10435417294502258, -0.8176955580711365, -0.026037702336907387],
-        [0.009808260947465897, -0.10966360569000244, -0.02152106538414955],
-        [0.07255466282367706, -1.2259838581085205, -0.05523664504289627],
-        [-0.08893736451864243, -1.2284233570098877, -0.046229973435401917],
-        [-0.0015221529174596071, -0.057428449392318726, 0.006925832014530897],
-        [0.11981196701526642, -1.283981204032898, 0.06297968327999115],
-        [-0.12774977087974548, -1.2867517471313477, 0.07281902432441711],
-        [-0.01368661131709814, 0.10773860663175583, -0.024689510464668274],
-        [0.04484200477600098, 0.027515273541212082, -0.0002946509048342705],
-        [-0.04921707883477211, 0.026910223066806793, -0.006474069785326719],
-        [0.011096873320639133, 0.2681904137134552, -0.003952245227992535],
-        [0.16408103704452515, 0.08524329960346222, -0.015755590051412582],
-        [-0.15179482102394104, 0.08043467253446579, -0.019142597913742065],
-        [0.4182038903236389, 0.01309278141707182, -0.058214444667100906],
-        [-0.4229443669319153, 0.04394219070672989, -0.04560968279838562],
-        [0.6701906323432922, 0.03631401062011719, -0.06068652495741844],
-        [-0.6722118258476257, 0.03940964490175247, -0.06093486770987511],
-        [-0.004667762666940689, 0.2676706910133362, -0.009591402485966682],
-        [0.03159928321838379, 0.31083211302757263, 0.062195174396038055],
-        [-0.031599875539541245, 0.3108319342136383, 0.0621943436563015],
-        [0.7720924615859985, 0.02762586995959282, -0.04133538901805878],
-        [0.8040408492088318, 0.02984413132071495, -0.03832494467496872],
-        [0.8265857696533203, 0.027494050562381744, -0.03827037289738655],
-        [0.7795881628990173, 0.029986342415213585, -0.06466733664274216],
-        [0.8101993799209595, 0.030793743208050728, -0.0686899945139885],
-        [0.8337147831916809, 0.028784994035959244, -0.07280422002077103],
-        [0.7542374730110168, 0.02177468128502369, -0.10443613678216934],
-        [0.7697082161903381, 0.020643100142478943, -0.11643557250499725],
-        [0.7852417230606079, 0.018978532403707504, -0.12765252590179443],
-        [0.767634928226471, 0.02704637683928013, -0.08803117275238037],
-        [0.7956817150115967, 0.028531836345791817, -0.09329714626073837],
-        [0.8185034990310669, 0.02707355096936226, -0.1003914326429367],
-        [0.7108263969421387, 0.01833728514611721, -0.03507564589381218],
-        [0.7278420925140381, 0.01931309886276722, -0.010097505524754524],
-        [0.7483652234077454, 0.01415354385972023, 0.005425570998340845],
-        [-0.7720924019813538, 0.027626780793070793, -0.041334930807352066],
-        [-0.8040405511856079, 0.029844673350453377, -0.03832409530878067],
-        [-0.8265854716300964, 0.027495287358760834, -0.03826868534088135],
-        [-0.7795882225036621, 0.029987698420882225, -0.0646686926484108],
-        [-0.8101993799209595, 0.030795171856880188, -0.06869153678417206],
-        [-0.8337149024009705, 0.02878585271537304, -0.07280556112527847],
-        [-0.7542385458946228, 0.021775206550955772, -0.10443780571222305],
-        [-0.7697089910507202, 0.02064313367009163, -0.11643654853105545],
-        [-0.7852423787117004, 0.01897839829325676, -0.1276528388261795],
-        [-0.7676352858543396, 0.02704770304262638, -0.08803359419107437],
-        [-0.7956817150115967, 0.028532907366752625, -0.093299500644207],
-        [-0.818503737449646, 0.027074117213487625, -0.10039224475622177],
-        [-0.7108249664306641, 0.018335221335291862, -0.035073522478342056],
-        [-0.7278403043746948, 0.019311318174004555, -0.01009594276547432],
-        [-0.7483659386634827, 0.014154116623103619, 0.005425604991614819],
-    ],
-    dtype=np.float32,
-)
-
-
-# ── fixed frame corrections (shared by FK + canonicalization) ────────────────
-# Both mirror the SONIC deploy transform. ``_YTOZ_UP`` maps SMPL's Y-up world to the
-# robot's Z-up (a 90 deg rotation about X); ``_SMPL_BASE`` is SMPL's rest-pose base
-# orientation, conjugated out during canonicalization.
-_YTOZ_UP = R.from_euler("x", 90, degrees=True)
-_SMPL_BASE = R.from_quat([0.5, 0.5, 0.5, 0.5])  # scalar-last; symmetric so wxyz==xyzw
-
-
-# ── 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.
-    """
-    root = _YTOZ_UP * R.from_rotvec(root_aa)
-    inv = _SMPL_BASE * root.inv()
-    return np.einsum("tij,tkj->tki", 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.
-    """
-    root = (_YTOZ_UP * R.from_rotvec(root_aa)) * _SMPL_BASE.inv()
-    return root.as_quat(scalar_first=True).astype(np.float32)  # wxyz
-
-
-# ── 3-point VR teleop keypoints (SONIC encode_mode 1) ────────────────────────
-# SMPL body-joint indices for the 3 tracked keypoints, plus the root/pelvis (0)
-# used as the reference frame. Mirrors gear_sonic ``_process_3pt_pose``: neck
-# (joint 12) is used rather than head (15) — it is more rigidly coupled to the
-# torso and less noisy than the free-looking head.
-_VR3_JOINTS = (22, 23, 12)  # left wrist, right wrist, neck
-
-# Per-keypoint rotation offsets aligning each SMPL joint frame to the robot
-# convention (root, left wrist, right wrist, neck), ported verbatim from
-# gear_sonic ``pico_manager_thread_server.OFFSETS`` (extrinsic xyz euler, degrees).
-_VR3_OFFSETS = [
-    R.from_euler("xyz", [0, 0, -90], degrees=True),  # root
-    R.from_euler("xyz", [90, 0, 0], degrees=True),  # left wrist
-    R.from_euler("xyz", [-90, 0, 180], degrees=True),  # right wrist
-    R.from_euler("xyz", [0, 0, -90], degrees=True),  # neck
-]
-
-# Unity (X-right, Y-up, Z-forward, left-handed) -> robot (X-forward, Y-left,
-# Z-up, right-handed) axis remap: Unity [x, y, z] -> robot [-x, z, y].
-_UNITY_TO_ROBOT = np.array([[-1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]])
-
-
-def _safe_quat(quats: np.ndarray) -> np.ndarray:
-    """Replace zero-norm quaternions with the scalar-last identity.
-
-    The headset reports ``[0, 0, 0, 0]`` for joints it isn't currently tracking;
-    ``scipy.Rotation.from_quat`` rejects zero-norm quaternions, so we substitute the
-    identity ``[0, 0, 0, 1]`` (no rotation) for those rows to keep FK robust.
-    """
-    quats = np.asarray(quats, np.float64).copy()
-    bad = np.linalg.norm(quats, axis=-1) < 1e-8
-    quats[bad] = (0.0, 0.0, 0.0, 1.0)  # scalar-last identity
-    return quats
-
-
-def compute_3point(body_poses_np: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-    """Extract the SONIC 3-point VR targets from headset body poses.
-
-    Mirrors gear_sonic ``_process_3pt_pose``: transforms the tracked joints from the
-    Unity frame to the robot frame, applies the per-keypoint rotation offsets, then
-    expresses the left-wrist / right-wrist / neck keypoints relative to the root
-    (pelvis) frame. This is the ``encode_mode == 1`` counterpart to :func:`compute`.
-
-    Note: physical wrist/neck position offsets and the operator calibration done in
-    gear_sonic's ``ThreePointPose.apply_calibration`` are not applied here; the raw
-    tracked joint poses are used.
-
-    Args:
-        body_poses_np: (24, 7) rows of ``[x, y, z, qx, qy, qz, qw]`` (scalar-last),
-            in the Unity frame, as returned by ``xrt.get_body_joints_pose()``.
-
-    Returns:
-        (pos, orn):
-          - pos: (9,) float32, root-relative ``[x, y, z]`` for [l-wrist, r-wrist, neck]
-          - orn: (12,) float32, root-relative ``[w, x, y, z]`` for the same order
-    """
-    body = np.asarray(body_poses_np, np.float64)
-    q = _UNITY_TO_ROBOT
-    # Root (index 0) + the 3 tracked keypoints, each transformed to the robot frame
-    # and rotation-offset-corrected.
-    positions = np.zeros((4, 3), np.float64)
-    rotations: list[R] = []
-    quats = _safe_quat(body[:, 3:7])
-    for out_i, j in enumerate((0, *_VR3_JOINTS)):
-        positions[out_i] = q @ body[j, :3]
-        rot = R.from_quat(quats[j]).as_matrix()  # scalar-last input
-        rotations.append(R.from_matrix(q @ rot @ q.T) * _VR3_OFFSETS[out_i])
-
-    root_inv = rotations[0].inv()
-    root_pos = positions[0]
-    pos = np.zeros(VR3_POS_DIM, np.float32)
-    orn = np.zeros(VR3_ORN_DIM, np.float32)
-    for k in range(VR3_N_POINTS):
-        pos[k * 3 : k * 3 + 3] = root_inv.apply(positions[k + 1] - root_pos)
-        orn[k * 4 : k * 4 + 4] = (root_inv * rotations[k + 1]).as_quat(scalar_first=True)  # wxyz
-    return pos, orn
-
-
-# ── 3-point VR teleop from raw device poses (no body trackers) ───────────────
-# PICO Y-up (X-right, Y-up, Z-back) -> robot Z-up world. Ported verbatim from
-# gear_sonic's controller path (``decoupled_wbc`` ``PicoStreamer.R_HEADSET_TO_WORLD``).
-_HEADSET_TO_WORLD = np.array([[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
-
-
-def _device_pose_to_world(pose: np.ndarray) -> tuple[np.ndarray, R]:
-    """Convert a raw (7,) device pose ``[x, y, z, qx, qy, qz, qw]`` (PICO Y-up frame)
-    to a Z-up world ``(position, Rotation)`` pair.
-
-    Handles the all-zero quaternion the SDK emits when a device is momentarily
-    untracked by substituting the identity, matching ``PicoStreamer._process_xr_pose``.
-    """
-    pose = np.asarray(pose, np.float64)
-    xyz = _HEADSET_TO_WORLD @ pose[:3]
-    quat = pose[3:7]  # scalar-last
-    if np.linalg.norm(quat) < 1e-8:
-        quat = np.array([0.0, 0.0, 0.0, 1.0])
-    rot = _HEADSET_TO_WORLD @ R.from_quat(quat).as_matrix() @ _HEADSET_TO_WORLD.T
-    return xyz, R.from_matrix(rot)
-
-
-def compute_3point_from_devices(
-    head_pose: np.ndarray,
-    left_pose: np.ndarray,
-    right_pose: np.ndarray,
-) -> tuple[np.ndarray, np.ndarray]:
-    """Build the SONIC 3-point VR targets from raw head + controller poses.
-
-    This is the controller-state path (no PICO Motion Trackers / body tracking
-    required): the 3 keypoints are the left controller, right controller, and the
-    headset, each expressed relative to the **headset yaw frame** — mirroring
-    gear_sonic's ``decoupled_wbc`` ``PicoStreamer._process_xr_pose`` (Y-up -> Z-up,
-    then de-headed by the headset yaw). The headset stands in for the "neck" point,
-    so its root-relative position is ~0 and its orientation carries pitch/roll.
-
-    Args:
-        head_pose, left_pose, right_pose: (7,) ``[x, y, z, qx, qy, qz, qw]`` device
-            poses (scalar-last) from ``xrt.get_headset_pose()`` /
-            ``xrt.get_left_controller_pose()`` / ``xrt.get_right_controller_pose()``.
-
-    Returns:
-        (pos, orn):
-          - pos: (9,) float32, headset-yaw-relative ``[x, y, z]`` for [l-wrist, r-wrist, head]
-          - orn: (12,) float32, headset-yaw-relative ``[w, x, y, z]`` for the same order
-    """
-    head_pos, head_rot = _device_pose_to_world(head_pose)
-    left_pos, left_rot = _device_pose_to_world(left_pose)
-    right_pos, right_rot = _device_pose_to_world(right_pose)
-
-    # De-head: cancel the headset yaw so targets are expressed in a heading-local frame.
-    head_yaw = head_rot.as_euler("xyz")[2]
-    inv_yaw = R.from_euler("z", -head_yaw)
-
-    points = ((left_pos, left_rot), (right_pos, right_rot), (head_pos, head_rot))
-    pos = np.zeros(VR3_POS_DIM, np.float32)
-    orn = np.zeros(VR3_ORN_DIM, np.float32)
-    for k, (p_pos, p_rot) in enumerate(points):
-        pos[k * 3 : k * 3 + 3] = inv_yaw.apply(p_pos - head_pos)
-        orn[k * 4 : k * 4 + 4] = (inv_yaw * p_rot).as_quat(scalar_first=True)  # wxyz
-    return pos, orn
-
-
-# ── operator calibration for the 3-point targets ────────────────────────────
-# G1 neutral key-frame targets, pelvis-relative [x, y, z] in metres, from MuJoCo FK on
-# g1_29dof at the robot's *standing* configuration (``default_angles`` — the pose the
-# robot actually holds at calibration time), with gear_sonic's local key-frame offsets
-# (``G1_KEY_FRAME_OFFSETS``: wrists +0.18x ∓0.025y, torso +0.35z). These are the poses
-# the operator's rest pose is mapped onto so the handoff starts at the robot's neutral
-# stance. This is the fixed-``default_angles`` stand-in for gear_sonic's
-# ``get_g1_key_frame_poses(q=measured_q)`` (see :meth:`ThreePointCalibrator.capture`):
-# since the robot stands at ``default_angles`` after the startup ramp, its measured q
-# equals this configuration, so these constants are the measured-q targets for the
-# nominal case. (Per-frame *live* measured q would need a reverse controller->publisher
-# feedback channel; not wired.)
-_G1_NEUTRAL_WRIST_POS = np.array([[0.2232, 0.2177, -0.1555], [0.2232, -0.2177, -0.1555]], np.float64)
-# Wrist neutral orientations (scalar-first w, x, y, z) at ``default_angles`` — NOT
-# identity: the wrists are rolled/pitched in the standing pose. Matches gear_sonic
-# using ``g1_lwrist_rot`` / ``g1_rwrist_rot`` from FK (not identity) as the rotation
-# calibration reference.
-_G1_NEUTRAL_WRIST_ROT = [
-    R.from_quat([0.9168, 0.0897, 0.3864, 0.0463], scalar_first=True),  # left
-    R.from_quat([0.9168, -0.0897, 0.3864, -0.0463], scalar_first=True),  # right
-]
-# Neck reconstruction chain (mirrors ThreePointPose._apply_calibration): torso link
-# +0.05 z, then +0.35 along the neck's local Z.
-_NECK_TORSO_OFFSET_Z = 0.05
-_NECK_LINK_LENGTH = 0.35
-
-
-@dataclass
-class ThreePointCalibrator:
-    """Aligns raw 3-point VR targets to the G1's neutral stance.
-
-    Ports gear_sonic ``ThreePointPose._capture_calibration`` / ``_apply_calibration``:
-    on :meth:`capture` (operator holding a neutral rest pose) it records (a) the
-    inverse of the head/neck orientation, used to de-tilt all points to upright, and
-    (b) per-wrist position + orientation offsets that map the corrected rest pose onto
-    the fixed G1 neutral wrist targets. :meth:`apply` then transforms every subsequent
-    frame by those offsets, and reconstructs the head/neck position from the calibrated
-    neck orientation via the torso->neck kinematic chain.
-
-    All quaternions are scalar-first (w, x, y, z), matching :func:`compute_3point`.
-
-    ``neck_relative_wrists`` selects how the wrist targets are framed:
-
-    - ``True`` (body-source, :func:`compute_3point`): wrists are pelvis-relative and
-      the neck sits roughly upright over the pelvis, so de-rotating them by ``neck_inv``
-      correctly expresses them in the neck/torso frame (gear_sonic's behaviour).
-    - ``False`` (device-source, :func:`compute_3point_from_devices`): the "neck" is the
-      **headset**, which pitches down when the operator looks at their hands; the wrists
-      are already yaw-stabilised, so applying the head pitch/roll would rotate "up" hand
-      motion into "forward". The wrist frame is left in the yaw-local world frame (only
-      the neck point itself is still de-tilted).
-    """
-
-    neck_relative_wrists: bool = True
-    _neck_quat_inv: R | None = field(default=None, init=False)
-    _wrist_pos_offset: np.ndarray | None = field(default=None, init=False)
-    _wrist_rot_offset: list[R] = field(default_factory=list, init=False)
-
-    @property
-    def is_calibrated(self) -> bool:
-        return self._neck_quat_inv is not None and self._wrist_pos_offset is not None
-
-    def reset(self) -> None:
-        self._neck_quat_inv = None
-        self._wrist_pos_offset = None
-        self._wrist_rot_offset = []
-
-    def recalibrate_wrists(self) -> None:
-        """Clear only the wrist offsets, preserving the neck calibration.
-
-        Mirrors gear_sonic ``ThreePointPose.reset_with_measured_q``: the next
-        :meth:`capture` recomputes the wrist offsets (against the G1 neutral targets)
-        while keeping the already-captured neck orientation, so re-aligning the arms
-        doesn't force the operator to re-level their head.
-        """
-        self._wrist_pos_offset = None
-        self._wrist_rot_offset = []
-
-    def capture(self, pos: np.ndarray, orn: np.ndarray) -> None:
-        """Capture calibration offsets from a neutral-pose frame.
-
-        Neck calibration is captured once and then preserved across subsequent
-        captures (matching gear_sonic's ``if self._calibration_neck_quat_inv is
-        None``); call :meth:`reset` to clear it or :meth:`recalibrate_wrists` to
-        re-align only the arms.
-
-        Args:
-            pos: (9,) root-relative ``[x, y, z]`` for [l-wrist, r-wrist, head].
-            orn: (12,) root-relative ``[w, x, y, z]`` for the same order.
-        """
-        pos = np.asarray(pos, np.float64).reshape(3, 3)
-        orn = np.asarray(orn, np.float64).reshape(3, 4)
-        if self._neck_quat_inv is None:
-            self._neck_quat_inv = R.from_quat(orn[2], scalar_first=True).inv()
-        # Wrists use the neck frame only in body-source mode; device-source keeps them
-        # in the already yaw-stabilised world frame (see class docstring).
-        wrist_inv = self._neck_quat_inv if self.neck_relative_wrists else R.identity()
-
-        self._wrist_pos_offset = np.zeros((2, 3), np.float64)
-        self._wrist_rot_offset = []
-        for k in range(2):
-            corrected_pos = wrist_inv.apply(pos[k])
-            corrected_rot = wrist_inv * R.from_quat(orn[k], scalar_first=True)
-            self._wrist_pos_offset[k] = corrected_pos - _G1_NEUTRAL_WRIST_POS[k]
-            # rot_offset maps the corrected rest orientation onto the G1 neutral wrist
-            # orientation: calibrated = rot_offset * (neck_inv * current).
-            self._wrist_rot_offset.append(_G1_NEUTRAL_WRIST_ROT[k] * corrected_rot.inv())
-
-    def apply(self, pos: np.ndarray, orn: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
-        """Apply the stored calibration; returns calibrated ``(pos (9,), orn (12,))``.
-
-        A no-op (returns the inputs unchanged) until :meth:`capture` has been called.
-        """
-        if self._neck_quat_inv is None or self._wrist_pos_offset is None:
-            return (
-                np.asarray(pos, np.float32).reshape(-1),
-                np.asarray(orn, np.float32).reshape(-1),
-            )
-        pos = np.asarray(pos, np.float64).reshape(3, 3)
-        orn = np.asarray(orn, np.float64).reshape(3, 4)
-        neck_inv = self._neck_quat_inv
-        wrist_inv = neck_inv if self.neck_relative_wrists else R.identity()
-
-        out_pos = np.zeros((3, 3), np.float64)
-        out_orn = np.zeros((3, 4), np.float64)
-        for k in range(2):  # wrists
-            out_pos[k] = wrist_inv.apply(pos[k]) - self._wrist_pos_offset[k]
-            corrected_rot = wrist_inv * R.from_quat(orn[k], scalar_first=True)
-            out_orn[k] = (self._wrist_rot_offset[k] * corrected_rot).as_quat(scalar_first=True)
-
-        # Head/neck: orientation de-tilted, position from the torso->neck chain.
-        neck_rot = neck_inv * R.from_quat(orn[2], scalar_first=True)
-        out_orn[2] = neck_rot.as_quat(scalar_first=True)
-        neck_z = neck_rot.apply([0.0, 0.0, 1.0])
-        out_pos[2] = np.array([0.0, 0.0, _NECK_TORSO_OFFSET_Z]) + _NECK_LINK_LENGTH * neck_z
-        return out_pos.reshape(-1).astype(np.float32), out_orn.reshape(-1).astype(np.float32)
-
-
-class SmplForwardKinematics:
-    """Rest-skeleton SMPL forward kinematics (no mesh, no torch)."""
-
-    def __init__(self, skeleton_path: str | Path | None = None):
-        if skeleton_path is not None:
-            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
-        else:
-            self.J = _SKELETON_J.astype(np.float64)
-            self.parents = _SKELETON_PARENTS.copy()
-        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(_safe_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 = _YTOZ_UP * R.from_rotvec(global_orient)
-        global_orient_new = root.as_rotvec()
-
-        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 = root * _SMPL_BASE.inv()
-        smpl_joints_local = root.inv().apply(joints)
-
-        return {
-            "smpl_joints_local": smpl_joints_local.astype(np.float32),
-            "root_quat": root.as_quat(scalar_first=True).astype(np.float32),  # wxyz
-            "root_transl": positions[0].astype(np.float32),
-        }

src/lerobot/teleoperators/pico_headset/smpl_stream.py

@@ -1,277 +0,0 @@
-#!/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 logging
-import time
-from collections import deque
-
-import numpy as np
-import zmq
-
-from .smpl_constants import (
-    JOINT_DIM,
-    LOCO_N_AXES,
-    LOCO_N_BTN,
-    N_JOINTS,
-    SMPL_OBS_DIM,
-    VR3_ORN_DIM,
-    VR3_POS_DIM,
-    WINDOW,
-)
-
-logger = logging.getLogger(__name__)
-
-SMPL_TOPIC = "rt/smpl"
-DEFAULT_SMPL_HOST = "127.0.0.1"
-DEFAULT_SMPL_PORT = 5560
-
-
-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)
-        # Latest sparse 3-point VR targets (encode_mode 1), if the producer sends them.
-        self.vr3_pos = np.zeros(VR3_POS_DIM, np.float32)
-        self.vr3_orn = np.tile([1.0, 0.0, 0.0, 0.0], VR3_ORN_DIM // 4).astype(np.float32)
-        self._got_vr3 = False
-        self._last_vr3_t = 0.0
-        # Latest controller-stick locomotion (encode_mode 1): [lx, ly, rx, ry] + [A,B,X,Y].
-        self.loco_axes = np.zeros(LOCO_N_AXES, np.float32)
-        self.loco_buttons = np.zeros(LOCO_N_BTN, np.float32)
-        self._got_loco = False
-        self._last_loco_t = 0.0
-        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 has_vr3(self) -> bool:
-        """True once the producer has sent at least one 3-point VR frame."""
-        return self._got_vr3
-
-    @property
-    def has_fresh_vr3(self) -> bool:
-        """True when a 3-point VR frame arrived within ``stale_after_s``.
-
-        Unlike :attr:`has_data`, this is independent of the SMPL window, so the
-        controller-state source (head + controllers only, empty SMPL) still drives
-        ``encode_mode 1`` without a whole-body reference.
-        """
-        if not self._got_vr3:
-            return False
-        if not self.stale_after_s:
-            return True
-        return (time.time() - self._last_vr3_t) <= self.stale_after_s
-
-    @property
-    def has_fresh_loco(self) -> bool:
-        """True when controller-stick locomotion arrived within ``stale_after_s``."""
-        if not self._got_loco:
-            return False
-        if not self.stale_after_s:
-            return True
-        return (time.time() - self._last_loco_t) <= self.stale_after_s
-
-    @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
-        self._got_vr3 = False
-        self._last_vr3_t = 0.0
-        self._got_loco = False
-        self._last_loco_t = 0.0
-
-    # -- 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", {})
-
-            # Sparse 3-point VR targets (encode_mode 1). Parsed independently of the
-            # SMPL window so the controller-state source (head + controllers, empty
-            # SMPL) is still handled.
-            vp = data.get("vr3_pos")
-            vo = data.get("vr3_orn")
-            if vp is not None and vo is not None and len(vp) == VR3_POS_DIM and len(vo) == VR3_ORN_DIM:
-                self.vr3_pos = np.asarray(vp, np.float32)
-                self.vr3_orn = np.asarray(vo, np.float32)
-                self._got_vr3 = True
-                self._last_vr3_t = time.time()
-
-            # Controller-stick locomotion (encode_mode 1), also independent of SMPL.
-            la = data.get("loco_axes")
-            lb = data.get("loco_buttons")
-            if la is not None and lb is not None and len(la) == LOCO_N_AXES and len(lb) == LOCO_N_BTN:
-                self.loco_axes = np.asarray(la, np.float32)
-                self.loco_buttons = np.asarray(lb, np.float32)
-                self._got_loco = True
-                self._last_loco_t = time.time()
-
-            # SMPL whole-body window (encode_mode 2), optional on this stream.
-            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
-            ):
-                logger.warning(
-                    "[SmplStream] no %s frame for %.2fs (holding last pose)",
-                    SMPL_TOPIC,
-                    now - self._last_recv_t,
-                )
-                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/teleoperators/pico_headset/smpl_to_dataset.py

@@ -1,147 +0,0 @@
-#!/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.
-
-"""Convert SMPL ``.npz`` motion clips into a LeRobotDataset for SONIC replay.
-
-Each dataset frame's ``action`` is the 720-dim ``smpl.*`` window that the
-``pico_headset`` teleoperator emits and ``SonicWholeBodyController`` reassembles
-into ``encode_mode == 2``. So the resulting dataset can be pushed straight
-through ``lerobot-replay`` to drive SONIC whole-body tracking with no headset:
-
-    lerobot-replay \
-        --robot.type=unitree_g1 \
-        --robot.controller=SonicWholeBodyController \
-        --dataset.repo_id=<user>/<clip> --dataset.episode=0
-
-The 10-frame window is built exactly like the live ``SmplStream`` (oldest->newest,
-the first frame repeated to pre-fill), so replayed actions match a live session.
-
-Usage:
-    # One clip -> one-episode dataset
-    python -m lerobot.teleoperators.pico_headset.smpl_to_dataset \
-        --motion-file examples/unitree_g1/motions/walk_forward.npz \
-        --repo-id me/sonic_walk_forward
-
-    # Every clip in a dir -> one episode each
-    python -m lerobot.teleoperators.pico_headset.smpl_to_dataset \
-        --motion-dir examples/unitree_g1/motions --repo-id me/sonic_motions
-"""
-
-from __future__ import annotations
-
-import argparse
-from pathlib import Path
-
-import numpy as np
-
-from lerobot.teleoperators.pico_headset.smpl_constants import (
-    ACTION_DIM,
-    JOINT_DIM,
-    N_JOINTS,
-    ROOT_ACTION_DIM as ROOT_DIM,
-    ROOT_ACTION_PREFIX,
-    SMPL_ACTION_PREFIX,
-    SMPL_OBS_DIM,
-    WINDOW,
-)
-from lerobot.teleoperators.pico_headset.smpl_fk import canonicalize_smpl_joints, root_quats_from_aa
-
-
-def _load_canonical_joints(path: str) -> tuple[np.ndarray, np.ndarray, float]:
-    """Load an SMPL clip -> (canonical (T,24,3) joints, (T,4) root wxyz, fps)."""
-    data = np.load(path)
-    joints = data["smpl_joints"].astype(np.float32)
-    if joints.ndim != 3 or joints.shape[1:] != (N_JOINTS, JOINT_DIM):
-        raise ValueError(f"{path}: expected smpl_joints (T, 24, 3), got {joints.shape}")
-    t = joints.shape[0]
-    if "pose_aa" in data.files:
-        root_aa = data["pose_aa"].astype(np.float32)[:, :3]
-        joints = canonicalize_smpl_joints(joints, root_aa)
-        root_quat = root_quats_from_aa(root_aa)  # (T, 4) wxyz, matches live stream
-    else:
-        # No global orient available: identity root (anchor falls back to standing).
-        root_quat = np.tile(np.array([1.0, 0.0, 0.0, 0.0], np.float32), (t, 1))
-    fps = float(data["fps"]) if "fps" in data.files else 50.0
-    return joints, root_quat, fps
-
-
-def _windows(joints: np.ndarray) -> np.ndarray:
-    """(T, 24, 3) -> (T, 720): rolling 10-frame window, matching SmplStream.
-
-    Window t = frames [t-9 .. t], clamped to 0 at the start (first frame repeated).
-    """
-    t = joints.shape[0]
-    idx = np.clip(np.arange(t)[:, None] + np.arange(-WINDOW + 1, 1)[None, :], 0, t - 1)
-    return joints[idx].reshape(t, -1).astype(np.float32)
-
-
-def _action_features() -> dict:
-    names = [f"{SMPL_ACTION_PREFIX}{i}" for i in range(SMPL_OBS_DIM)]
-    names += [f"{ROOT_ACTION_PREFIX}{i}" for i in range(ROOT_DIM)]
-    return {"action": {"dtype": "float32", "shape": (ACTION_DIM,), "names": names}}
-
-
-def main() -> None:
-    p = argparse.ArgumentParser(description=__doc__)
-    src = p.add_mutually_exclusive_group(required=True)
-    src.add_argument("--motion-file", type=str, help="Single SMPL .npz clip")
-    src.add_argument("--motion-dir", type=str, help="Directory of .npz clips (one episode each)")
-    p.add_argument("--repo-id", required=True, help="Dataset repo id, e.g. user/name")
-    p.add_argument("--root", type=str, default=None, help="Local dataset root (default HF cache)")
-    p.add_argument("--fps", type=int, default=None, help="Override fps (default: clip fps)")
-    p.add_argument("--task", type=str, default="sonic whole-body SMPL replay")
-    args = p.parse_args()
-
-    from lerobot.datasets.lerobot_dataset import LeRobotDataset
-
-    if args.motion_dir:
-        clips = sorted(str(pth) for pth in Path(args.motion_dir).glob("*.npz"))
-        if not clips:
-            raise SystemExit(f"No .npz clips found in {args.motion_dir}")
-    else:
-        clips = [args.motion_file]
-
-    first_joints, first_root, first_fps = _load_canonical_joints(clips[0])
-    fps = args.fps or int(round(first_fps))
-
-    dataset = LeRobotDataset.create(
-        repo_id=args.repo_id,
-        fps=fps,
-        features=_action_features(),
-        root=args.root,
-        robot_type="unitree_g1",
-        use_videos=False,
-    )
-
-    for clip_i, clip in enumerate(clips):
-        if clip_i == 0:
-            joints, root_quat = first_joints, first_root
-        else:
-            joints, root_quat, _ = _load_canonical_joints(clip)
-        windows = _windows(joints)  # (T, 720)
-        # action = [720 joint window | 4 root wxyz] per frame -> (T, 724)
-        actions = np.concatenate([windows, root_quat.astype(np.float32)], axis=1)
-        for a in actions:
-            dataset.add_frame({"action": a, "task": args.task})
-        dataset.save_episode()
-        print(f"[smpl_to_dataset] episode {clip_i}: {Path(clip).name} ({actions.shape[0]} frames)")
-
-    dataset.finalize()
-    print(f"[smpl_to_dataset] wrote {len(clips)} episode(s) to {dataset.root}")
-
-
-if __name__ == "__main__":
-    main()