test 3-point teleop

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-07-23 01:41:54 +00:00 · 2026-07-15 18:16:46 +02:00
parent f6a845c30c
commit 9c54665a76
17 changed files with 1379 additions and 13 deletions
@@ -0,0 +1,103 @@
 # SONIC fidelity audit — Python/ONNX port vs C++ deploy reference
 Compares the lerobot Python/ONNX SONIC port against the original C++ deploy stack
 (`gear_sonic_deploy/src/g1/g1_deploy_onnx_ref/src/g1_deploy_onnx_ref.cpp` and headers).
 **Verdict:** the port is algorithmically faithful in the core control math (the parts
 that determine stability and pose tracking). The gaps are concentrated in (a) update
 cadence, (b) the safety layer, and (c) two modes/paths present in C++ but never wired up.
 No silent math bug was found; divergences are deliberate or missing-feature.
 ## Genuinely faithful (verified, no action needed)
 - **Action production** — both do `q_target = DEFAULT_ANGLES + policy_out[isaaclab→mujoco] * ACTION_SCALE`
  (residual-to-default, not residual-to-previous). C++ `g1_deploy_onnx_ref.cpp:3119-3127`,
  Python `sonic_pipeline.py:708-724`.
 - **Joint-order remap** — `ISAACLAB_TO_MUJOCO` / `MUJOCO_TO_ISAACLAB` identical arrays.
 - **Gains** — same `Kp = armature·ω²`, `Kd = 2ζ·armature·ω`, ω = 10·2π, ζ = 2, and the same
  2× set `{4,5,10,11,13,14}` (ankles + waist roll/pitch).
 - **History normalization** — robot `q` subtracts defaults, velocities raw,
  gravity = `quat_rotate(conj(base), [0,0,-1])`, oldest→newest ordering.
 - **6D anchor rotation** — verified element-by-element: C++ takes the first two *columns*
  of the rotation matrix flattened row-wise (`.cpp:677-683`); Python `quat_to_6d`
  (`sonic_pipeline.py:227-240`) produces the identical 6 values. Only the Python docstring
  wording ("rows") is misleading — the math is correct.
 - **Planner** — replan intervals (RUN 0.1 / CRAWL 0.2 / boxing 1.0 / default 1.0 s),
  8-frame slerp crossfade blend, 30→50 Hz linear+slerp resample, `MOTION_LOOK_AHEAD = 2`,
  4-frame 30 Hz context. All match.
 - **SLERP / heading / FK conventions** — wxyz, shortest-path slerp with 0.9995 linear
  fallback, `calc_heading` yaw extraction.
 ## Divergences that reduce fidelity (ranked)
 ### 1. Encoder cadence: 10 Hz vs 50 Hz (biggest)
 C++ runs the encoder every control tick — `GatherTokenState → Encode()` is unconditional
 inside the 50 Hz `Control()` loop (`.cpp:1644-1684`, no N-step gate). The port recomputes
 the token only every 5 ticks (`ENCODER_UPDATE_EVERY = 5`, `sonic_pipeline.py:150`), so the
 latent is up to 80 ms stale and the decoder consumes a held token for 4 of every 5 ticks.
 Likely a perf shortcut. For full faithfulness set `ENCODER_UPDATE_EVERY = 1` (cheap on GPU).
 ### 2. SMPL root anchor disabled by default
 C++ always feeds the reference root orientation into the anchor/heading. The port sets
 `smpl_root_quat = None` unless `enable_smpl_root=True` (`sonic_whole_body.py:366`), because
 the raw 30 Hz root caused QACC spikes. Faithful fix: slerp-resample the root 30→50 Hz like
 the joints, then re-enable by default. Until then, mode-2 heading steering isn't faithful.
 See `SONIC_REPLAY_DEBUGGING.md`.
 ### 3. VR 3-point teleop (`encode_mode = 1`) — now wired (was inert)
 Originally the encoder layout for mode 1 existed but nothing set `encode_mode=1` or
 filled `vr_3point_local_target` / `vr_3point_local_orn_target`. **Now implemented
 end-to-end:**
 - Producer `pico_publisher.py` computes the 3 keypoints via `smpl_fk.compute_3point`
  (ported from gear_sonic `_process_3pt_pose`) and adds `vr3_pos` (9) / `vr3_orn` (12)
  to the `rt/smpl` message.
 - `SmplStream` parses them (`has_vr3`); `PicoHeadset(mode="vr3")` emits `vr3_pos.*` /
  `vr3_orn.*` action keys.
 - `SonicWholeBodyController` extracts them, switches to `encode_mode=1`, fills the
  controller targets, and drives locomotion from the joystick/keyboard planner
  (`use_joystick=True`); `PlannerController.build_encoder_obs` gained a mode-1 branch
  (lower body per-frame step 5 + VR targets + anchor).
 Still not ported: `vr_3point_compliance`, and the operator calibration
 (`ThreePointPose.apply_calibration`) / physical wrist offsets — the raw tracked joint
 poses are used, so hand-tracking may need calibration tuning.
 ### 4. Safety layer largely absent
 - **Joint-velocity kill switch** at >35 rad/s (`.cpp:2829-2832`) — missing.
 - **E-stop damping**: C++ e-stop commands `kp=0, kd=8` (active damping, `.cpp:2708-2714`).
  The port's Space/e-stop just sets `playing=False` + `LM.IDLE` and stops the cursor — it
  does not switch to a damped hold. Less safe.
 - **Motor-temperature** monitor (90 °C / 85 °C hysteresis) — missing.
 - **Stale-/late-state watchdogs** (500 ms fail, 50 ms warn, 200 ms token timeout) — not in
  the SONIC layer (partly covered by `UnitreeG1`, not equivalently).
 - **Per-tick delta clamp** — C++ has none, so the reverted `MAX_DELTA_PER_STEP` was correctly
  removed; that part is faithful.
 ### 5. Idle readaptation blend missing
 At planner IDLE at a motion end, C++ runs a double-threshold blend (0.98/0.02 toward
 robot-current or original target; thresholds 0.10/0.05/0.045 rad; `.cpp:3303-3361`). The
 port just holds. Minor; only matters at motion-end idle.
 ### 6. Input/streaming paths not ported
 - **ZMQ Protocol V1 joint streaming** (`encode_mode=0` from a live joint stream via
  `StreamedMotionMerger`) — not implemented; the port's streaming path is SMPL-only.
 - **External token injection** (tokens over ZMQ/ROS2 bypassing the encoder) — not supported;
  the port always encodes locally. Fine for standalone.
 - **Gamepad** — C++ has a full gamepad map (EMA smooth 0.3, deadzone 0.05); the port has
  joystick byte-parsing + keyboard, no gamepad manager. Functionally close.
 ### 7. Minor input-feel differences
 Delta-heading is continuous (±0.02 rad/tick) in the port vs discrete steps (±π/6, ±π/12) in
 C++; speed/height increments differ slightly. Behavioral feel only, not correctness.
 ## Recommendation (priority order)
 1. `ENCODER_UPDATE_EVERY = 1` (or a param defaulting to 1) — closes the biggest gap for
   near-zero cost on GPU.
 2. Rate-match the SMPL root 30→50 Hz (slerp) and re-enable `enable_smpl_root` by default.
 3. Add the safety envelope: joint-velocity kill (35 rad/s) and a proper damped e-stop
   (`kp=0, kd≈8`). Real hardware-safety consequences.
 4. Wire `encode_mode=1` from the pico headset (3-point targets), or document as out of scope.
 5. Fix the `quat_to_6d` docstring wording ("rows" → "first two columns flattened row-wise").
 Items 4–7 are feature-completeness; 1–3 are what to do for faithful behavior on the robot.
@@ -0,0 +1,97 @@
 # SONIC replay instability — root cause & prevention
 This documents the multi-day debugging of "SONIC motion replay is unstable / jitters /
 lags / dies on the floor", so we don't chase the same ghosts again.
 ## TL;DR
 There were **two independent problems**, and they masked each other:
 1. **Wrong conda environment (the "lag"/jitter).** The debugging env `lerobot_sonic`
   had a **CUDA-13** stack that the machine's GPU driver cannot run, so ONNX Runtime
   silently fell back to CPU and oversubscribed threads. The known-good env
   `lerobot312` has a **CUDA-12** stack matching the driver, so the encoder/decoder/
   planner run on the GPU (~12–20 ms planner inference) and the control loop holds
   ~48–50 Hz.
 2. **SMPL root-motion feeding (the NaN/`unstable` crash).** Passing the per-frame SMPL
   root quaternion into the mode-2 anchor produced a root-acceleration spike
   (`Nan, Inf or huge value in QACC at DOF 0`) mid-episode. Disabling it gives clean
   tracking.
 Neither is an algorithmic bug in the ported SONIC pipeline. A lot of earlier "fixes"
 (ORT thread caps, `MAX_DELTA_PER_STEP` clamp, planner-disable toggle, resampling)
 were chasing symptom #1 in the wrong environment and were reverted.
 ## Environment: what "good" looks like
 Run the replay in `lerobot312` (CUDA-12), pointing at the current sonic checkout:
 ```bash
 conda activate lerobot312
 PYTHONPATH=/home/yope/Documents/sonic/lerobot/src \
  lerobot-replay \
  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
  --dataset.repo_id=lerobot/SMPL_samples --dataset.episode=12
 ```
 Known-good versions (`lerobot312`):
 | package         | good (`lerobot312`) | broken (`lerobot_sonic`) |
 |-----------------|---------------------|--------------------------|
 | GPU driver      | CUDA 12.8 (`12080`) | same (unchanged)         |
 | torch           | `2.10.0+cu128`      | `2.11.0+cu130`           |
 | onnxruntime     | `onnxruntime-gpu 1.26.0` | CPU `1.27.0` / cu13 mismatch |
 | cudnn           | cu12 (bundled)      | `nvidia-cudnn-cu13 9.19` |
 | mujoco          | `3.8.1`             | `3.10.0`                 |
 ### How to verify the GPU path is actually live (do this FIRST)
 ```bash
 python -c "import torch; print('cuda', torch.cuda.is_available())"          # must be True
 python -c "import onnxruntime as ort; print(ort.get_available_providers())" # must list CUDAExecutionProvider
 ```
 If `torch.cuda.is_available()` is `False` or `CUDAExecutionProvider` is missing, STOP —
 you are in the wrong/broken env. Do not "optimize" anything else until this passes.
 ### Why CUDA 13 was fatal here
 The GPU driver supports up to **CUDA 12.8**. A CUDA-13 build of torch/onnxruntime
 cannot initialize on it:
 - `torch.cuda.is_available()` returns `False` (silent CPU fallback).
 - `onnxruntime-gpu` (a CUDA-12 build) can't find a matching cuDNN because only the
  CUDA-13 cuDNN is installed → `CUDNN failure 1001: CUDNN_STATUS_NOT_INITIALIZED`.
 Installing `onnxruntime` and `onnxruntime-gpu` **together** also breaks: they share the
 `onnxruntime` namespace and whichever installs last clobbers the other's shared libs.
 Keep only `onnxruntime-gpu` in a GPU env.
 ## Root motion: the NaN/unstable crash
 Symptom:
 ```
 WARNING: Nan, Inf or huge value in QACC at DOF 0. The simulation is unstable. Time = 4.196.
 ```
 `DOF 0` is the floating base/root. Feeding the per-frame SMPL root quaternion
 (`root.*` action keys) into `controller.smpl_root_quat` injected a discontinuity in the
 reference root trajectory (frame-to-frame jump and/or 30 Hz→50 Hz timing mismatch) that
 the tracker converted into an exploding base acceleration.
 Current mitigation (in `sonic_whole_body.py`, `run_step`): the per-frame root quaternion
 is ignored (`self.controller.smpl_root_quat = None`) so the anchor stays self-driven.
 Result: clean tracking, no NaN.
 Proper fix (follow-up, not yet done): smooth/slerp-filter the reference root trajectory
 (or resample to the control rate) before feeding it to the anchor, then re-enable.
 ## Prevention checklist
 - **Always confirm the env before debugging behavior.** Run the two verification
  commands above. Most of the "instability" was environment, not code.
 - **Pin the GPU stack** to match the driver (CUDA 12.8): `torch ==2.10.0+cu128`,
  `onnxruntime-gpu ==1.26.0`, `mujoco ==3.8.1`. Do not let `lerobot_sonic` drift to a
  CUDA-13 stack.
 - **Never install `onnxruntime` and `onnxruntime-gpu` side by side.**
 - **Don't add band-aid clamps/thread-caps/resampling to hide a CPU-fallback**; fix the
  env instead. Those changes were reverted.
 - **Root trajectory must be continuous / rate-matched** before it drives the anchor.
@@ -0,0 +1,79 @@
 # feat: ONNX inference support (ACT)
 ## Summary
 This PR introduces a first, end-to-end path for **ONNX-based policy inference** in LeRobot, currently scoped to the **ACT** policy. The goal is to standardize how we export and run policies through ONNX Runtime so that the same workflow can later cover other policies (including the Unitree G1 whole-body / locomotion policies) and so policies can run on **edge devices** without a full PyTorch stack.
 > ⚠️ **Scope:** today this only works for **ACT**. ACT is the natural starting point because inference is a single deterministic forward pass (ResNet backbone + transformer enc/dec + action head) with a zeros VAE latent — no denoising loop, no KV cache. Other architectures (e.g. PI0.5) need more work before they can be exported the same way.
 ## Motivation
 - **Standardize ONNX inference** across LeRobot policies behind one export + run convention, instead of one-off conversion scripts.
 - **Run on edge devices**: ONNX Runtime has a much smaller footprint than PyTorch and ships CPU / CUDA / TensorRT / mobile execution providers, which is what we want for deploying policies (incl. Unitree G1 policies) on-robot.
 - Keep normalization and control logic in Python (the LeRobot processor pipeline + action queue), and export **only the neural network** as a portable graph.
 ## What's included
 All new files live under `examples/onnx/` (no changes to `src/lerobot/...`):
 - **`export_act.py`** — exports `policy.model` to ONNX as a pure function `(state, images) -> action_chunk`, then runs a numerical parity check (PyTorch vs ONNX Runtime).
 - **`eval_act_onnx.py`** — evaluates ACT in sim with either the PyTorch or the ONNX backend. It swaps **only** `policy.model` with an ONNX Runtime session (wrapped as an `nn.Module`), so processors, action queue and the gym env are identical and any delta is attributable to the backend alone.
 - **`convert_legacy_checkpoint.py`** — helper for older hub checkpoints that bake normalization into weights and lack `policy_preprocessor.json` / `policy_postprocessor.json`.
 ## Design notes
 - Only the network is exported. At inference, ACT's `predict_action_chunk` is effectively `self.model(batch)[0]` with a zeros latent, so the graph is deterministic in `(state, images)`.
 - **Normalization stays outside ONNX**, in the LeRobot processor pipeline. The ONNX graph consumes already-normalized inputs and emits normalized actions.
 - torch 2.9+ defaults to the dynamo exporter (requires `onnxscript`); the exporter uses the legacy TorchScript path (`dynamo=False`) since ACT's graph is fixed-shape.
 ## Results
 **Numerical parity** (PyTorch vs ONNX Runtime):
 ```
 max_abs_diff = 1.073e-06   mean_abs_diff = 1.790e-07   -> PASS
 ```
 **In-sim eval**, `AlohaTransferCube-v0`, identical seed:
 | backend | n_episodes | pc_success |
 |---------|-----------:|-----------:|
 | torch   | 10         | 70.0%      |
 | onnx    | 10         | 70.0%      |
 Identical success rate; sub-1e-6 per-step parity. (Run on CPU here; both backends behave the same on CUDA.)
 ## How to run
 ```bash
 export PYTHONPATH=src
 # export once (also runs the parity check)
 python examples/onnx/export_act.py \
  --policy-path=lerobot/act_aloha_sim_transfer_cube_human \
  --output=outputs/onnx/act_transfer_cube.onnx
 # compare backends in sim
 python examples/onnx/eval_act_onnx.py \
  --policy-path=lerobot/act_aloha_sim_transfer_cube_human \
  --task=AlohaTransferCube-v0 \
  --backend=torch --n-episodes=50 --batch-size=10 --device=cuda
 python examples/onnx/eval_act_onnx.py \
  --policy-path=lerobot/act_aloha_sim_transfer_cube_human \
  --task=AlohaTransferCube-v0 \
  --onnx=outputs/onnx/act_transfer_cube.onnx \
  --backend=onnx --n-episodes=50 --batch-size=10 --device=cuda
 ```
 ## Follow-ups (out of scope for this PR)
 - Generalize the export convention beyond ACT (PI0.5 denoising loop + KV cache, diffusion policies, etc.).
 - Cover the **Unitree G1** policies so they can be deployed via ONNX Runtime on-robot.
 - Provide an edge-device runner / packaging story (CPU / TensorRT / mobile execution providers) and a latency benchmark.
 ## Test plan
 - [x] ONNX export succeeds for ACT and passes the parity check (`max_abs_diff < 1e-3`).
 - [x] In-sim eval matches the PyTorch backend at the same seed.
 - [ ] Full 50-episode eval on CUDA (torch vs onnx) reproduces the baseline success rate.
@@ -0,0 +1,113 @@
 # Unitree G1 — SONIC 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).
 ## Controllers
 Selected with `--robot.controller=<ClassName>`:
 | Controller | Purpose |
 |---|---|
 | `SonicWholeBodyController` | SONIC whole-body: locomotion, keyboard, and SMPL imitation (mode 2) |
 | `GrootLocomotionController` | GR00T locomotion policy |
 | `HolosomaLocomotionController` | Holosoma locomotion policy |
 ## 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:
  ```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.
 ## Running
 **Replay an SMPL dataset (motion imitation):**
 ```bash
 lerobot-replay \
  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
  --dataset.repo_id=<user>/<smpl_dataset> --dataset.episode=0
 ```
 **Keyboard teleop** (drives locomotion via the native keyboard teleoperator):
 ```bash
 lerobot-teleoperate \
  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
  --teleop.type=keyboard
 ```
 Controls: `WASD` move · `Q`/`E` turn · `1`–`8` mode · `9`/`0` speed · `-`/`=` height ·
 `R` replan · `Space` emergency-stop.
 **PICO headset teleop** (live SMPL whole-body):
 ```bash
 lerobot-teleoperate \
  --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
  --teleop.type=pico_headset
 ```
 This requires 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, pair/calibrate the ankle motion 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/Full-body/Send.
 ### 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 is off by default because an unsmoothed 30 Hz→50 Hz root trajectory can spike the
  base acceleration.
 - Direct `rt/smpl` subscription without the pico teleoperator is available via
  `SonicWholeBodyController(enable_smpl_stream=True, smpl_host=..., smpl_port=...)`.
@@ -1153,6 +1153,22 @@ class PlannerController(StandingEncoderDecoder):
                    obs[1642 + 6 * f : 1642 + 6 * (f + 1)] = anchor
                    obs[1702 + 6 * f : 1702 + 6 * (f + 1)] = wrist
                return obs
            if self.encode_mode == 1:
                # 3-point VR teleop: the upper body tracks the VR wrist/neck targets
                # while the planner reference supplies the lower body + anchor. Lower
                # body is per-frame (step 5) like mode 0; the VR targets are current.
                rf = min(self.ref_cursor, self.motion_timesteps - 1)
                obs[595:601] = self._anchor_6d(self.h_quat[0], self.motion_body_quats[rf].astype(np.float32))
                for f in range(10):
                    tf = min(
                        self.ref_cursor + f * 5 if self.playing else self.ref_cursor,
                        self.motion_timesteps - 1,
                    )
                    ref_lower = self.motion_joint_positions[tf].astype(np.float32)[LOWER_BODY_IL]
                    obs[661 + 12 * f : 661 + 12 * (f + 1)] = ref_lower
                obs[901:910] = self.vr_3point_local_target
                obs[910:922] = self.vr_3point_local_orn_target
                return obs
            for f in range(10):
                tf = min(
                    self.ref_cursor + f * 5 if self.playing else self.ref_cursor, self.motion_timesteps - 1
@@ -30,6 +30,10 @@ from lerobot.teleoperators.pico_headset.smpl_constants import (
    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,
 )
 from lerobot.utils.import_utils import _onnxruntime_available, require_package
@@ -96,6 +100,28 @@ def _extract_root_from_action(action: dict | None) -> np.ndarray | 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
 class SonicRuntime:
    """Shared SONIC control loop state (standalone demo + locomotion controller)."""
@@ -240,8 +266,20 @@ class SonicWholeBodyController:
        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 SMPL is lost.
+        """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.
@@ -250,7 +288,7 @@ class SonicWholeBodyController:
        self.controller.playing = True
        self.controller.reinit_heading = True
        self.ms.needs_replan = True
-        logger.warning("SONIC: SMPL stream lost/stale -> reverting to locomotion (standing)")
+        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.
@@ -348,11 +386,14 @@ class SonicWholeBodyController:
        # robot doesn't stay frozen tracking the last pose.
        smpl = _extract_smpl_from_action(action)
        root_quat = _extract_root_from_action(action)
-        if smpl is None and self._smpl_stream is not None:
+        vr3 = _extract_vr3_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)
                if self._smpl_stream.has_vr3:
                    vr3 = (self._smpl_stream.vr3_pos, self._smpl_stream.vr3_orn)
        if smpl is not None:
            # Full-body whole-body tracking: SMPL drives the reference, not joystick.
@@ -366,8 +407,18 @@ class SonicWholeBodyController:
            self.controller.smpl_root_quat = root_quat if self.enable_smpl_root else None
            return self._runtime.tick(obs, debug=False, use_joystick=False)
-        # No (or stale) SMPL: fall back to locomotion so the robot stays balanced.
+        if vr3 is not None:
-        if self.controller.encode_mode == 2:
+            # 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]
            return self._runtime.tick(obs, debug=False, use_joystick=True)
        # 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._exit_wholebody()
        return self._runtime.tick(obs, debug=False)
@@ -0,0 +1,20 @@
 #!/usr/bin/env python
 # Copyright 2026 The HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from .config_g1_sonic_slider import G1SonicSliderTeleopConfig
 from .g1_sonic_slider import G1SonicSliderTeleop
 __all__ = ["G1SonicSliderTeleop", "G1SonicSliderTeleopConfig"]
@@ -0,0 +1,35 @@
 #!/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("g1_sonic_slider")
@dataclass
 class G1SonicSliderTeleopConfig(TeleoperatorConfig):
    """Pygame sliders for 29-DOF G1 poses (SONIC encoder mode 0 reference)."""
    window_width: int = 780
    window_height: int = 720
    slider_width: int = 200
    row_height: int = 22
    scroll_step: int = 40
    foot_panel_width: int = 248
    use_leg_ik: bool = True
    foot_xyz_margin: tuple[float, float, float] = (0.22, 0.18, 0.18)
    """Per-axis slider half-range (m) around standing foot FK position in pelvis frame."""
@@ -0,0 +1,458 @@
 #!/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.
 """Pygame SONIC test UI: foot xyz (leg IK) + waist/arm joint sliders."""
 from __future__ import annotations
 import logging
 import multiprocessing as mp
 import queue
 from functools import cached_property
 import numpy as np
 from lerobot.robots.unitree_g1.controllers.sonic_pipeline import DEFAULT_ANGLES
 from lerobot.utils.import_utils import require_package
 from ..teleoperator import Teleoperator
 from .config_g1_sonic_slider import G1SonicSliderTeleopConfig
 from .joint_limits import JOINT_HI, JOINT_LO, JOINT_NAMES
 logger = logging.getLogger(__name__)
 NUM_JOINTS = 29
 LEG_JOINT_COUNT = 12
 UPPER_BODY_INDICES = list(range(LEG_JOINT_COUNT, NUM_JOINTS))
 NUM_FOOT_SLIDERS = 6
 FOOT_LABELS = ("L foot X", "L foot Y", "L foot Z", "R foot X", "R foot Y", "R foot Z")
 # Pelvis-frame standing foot centers (m) if Pinocchio FK is unavailable at startup.
 _FALLBACK_LEFT_FOOT = np.array([0.02, 0.12, -0.76], dtype=np.float32)
 _FALLBACK_RIGHT_FOOT = np.array([0.02, -0.12, -0.76], dtype=np.float32)
 HEADER_H = 56
 LABEL_W = 148
 MARGIN = 10
 KNOB_W = 10
 def _leg_ik_process(target_q: mp.Queue, result_q: mp.Queue, stop_evt) -> None:
    """Child process: build the leg IK once, then solve for the latest foot target.
    The IPOPT/CasADi solve holds the GIL, so it must run in a separate *process*
    (not a thread) to keep the teleop UI loop responsive.
    """
    import numpy as _np
    from lerobot.robots.unitree_g1.controllers.sonic_pipeline import DEFAULT_ANGLES as _DEFAULTS
    from lerobot.robots.unitree_g1.g1_kinematics import G1_29_LegIK
    try:
        ik = G1_29_LegIK()
        q_legs = _DEFAULTS[:LEG_JOINT_COUNT].astype(_np.float64)
        ik.cache_default_orientation(q_legs)
        left_pos, right_pos = ik.foot_positions(q_legs)
    except Exception as e:  # noqa: BLE001
        result_q.put(("error", str(e)))
        return
    current = q_legs.copy()
    result_q.put(("ready", _np.concatenate([left_pos, right_pos]).astype(_np.float64)))
    while not stop_evt.is_set():
        try:
            target = target_q.get(timeout=0.1)
        except queue.Empty:
            continue
        if target is None:
            break
        # Drain to the most recent target so we never solve stale slider positions.
        while True:
            try:
                newer = target_q.get_nowait()
            except queue.Empty:
                break
            if newer is None:
                target = None
                break
            target = newer
        if target is None:
            break
        target = _np.asarray(target, dtype=_np.float64)
        # Fast damped-least-squares IK: sub-ms per step, warm-started from the last
        # solution so the legs track the sliders in real time.
        leg_q = ik.solve_ik_dls(target[:3], target[3:], current_leg_q_g1=current)
        current = _np.asarray(leg_q, dtype=_np.float64)
        result_q.put(("q", current.copy()))
 class G1SonicSliderTeleop(Teleoperator):
    """Foot xyz + waist/arm sliders feeding SONIC encoder mode 0."""
    config_class = G1SonicSliderTeleopConfig
    name = "g1_sonic_slider"
    def __init__(self, config: G1SonicSliderTeleopConfig):
        super().__init__(config)
        self.config = config
        self._values = DEFAULT_ANGLES.astype(np.float32).copy()
        self._foot_xyz = np.zeros(6, dtype=np.float32)
        self._foot_lo = np.full(6, -0.5, dtype=np.float32)
        self._foot_hi = np.full(6, 0.5, dtype=np.float32)
        self._scroll_y = 0
        self._drag_joint: int | None = None
        self._drag_foot: int | None = None
        self._connected = False
        self._pygame = None
        self._screen = None
        self._font = None
        self._small_font = None
        self._clock = None
        # Leg IK runs in a separate process: the IPOPT/CasADi solve holds the GIL,
        # so a thread would still stall the UI loop. We publish the latest foot
        # target and read back the newest solution over queues, never blocking.
        self._ik_proc: mp.process.BaseProcess | None = None
        self._ik_target_q: mp.Queue | None = None
        self._ik_result_q: mp.Queue | None = None
        self._ik_stop_evt = None
        self._ik_ready = False
        self._standing_foot_xyz: np.ndarray | None = None
        self._leg_ik_ok = False
        self._leg_ik_error: str | None = None
        self._foot_divider_x = config.foot_panel_width
    @cached_property
    def action_features(self) -> dict[str, type]:
        return {f"{name}.q": float for name in JOINT_NAMES}
    @cached_property
    def feedback_features(self) -> dict[str, type]:
        return {}
    @property
    def is_connected(self) -> bool:
        return self._connected
    @property
    def is_calibrated(self) -> bool:
        return True
    def _set_foot_limits_from_positions(self, left_pos: np.ndarray, right_pos: np.ndarray) -> None:
        """Slider range = FK foot position at standing pose ± foot_xyz_margin (meters, pelvis frame)."""
        margin = np.array(self.config.foot_xyz_margin, dtype=np.float32)
        for i, pos in enumerate((left_pos, right_pos)):
            base = i * 3
            self._foot_xyz[base : base + 3] = pos.astype(np.float32)
            self._foot_lo[base : base + 3] = pos.astype(np.float32) - margin
            self._foot_hi[base : base + 3] = pos.astype(np.float32) + margin
    def _set_fallback_foot_limits(self) -> None:
        self._set_foot_limits_from_positions(_FALLBACK_LEFT_FOOT, _FALLBACK_RIGHT_FOOT)
    def _init_leg_ik(self) -> None:
        if not self.config.use_leg_ik:
            return
        # Fallback limits until the solver process reports standing FK foot positions.
        self._set_fallback_foot_limits()
        try:
            # Metadata name for the Pinocchio distribution is "pin" (conda-forge/PyPI),
            # not "pinocchio", which is only the importable module name.
            require_package("pin", extra="unitree_g1", import_name="pinocchio")
            require_package("casadi", extra="unitree_g1", import_name="casadi")
        except Exception as e:
            self._leg_ik_ok = False
            self._leg_ik_error = str(e)
            logger.warning("Leg IK unavailable (%s); foot sliders shown but legs use joint values", e)
            return
        # Use "spawn": the parent already holds CUDA/pygame/threads and forking that
        # state into the solver is unsafe.
        ctx = mp.get_context("spawn")
        self._ik_target_q = ctx.Queue(maxsize=2)
        self._ik_result_q = ctx.Queue()
        self._ik_stop_evt = ctx.Event()
        self._ik_proc = ctx.Process(
            target=_leg_ik_process,
            args=(self._ik_target_q, self._ik_result_q, self._ik_stop_evt),
            name="g1-leg-ik",
            daemon=True,
        )
        self._ik_proc.start()
        self._leg_ik_ok = True
        self._leg_ik_error = None
        logger.info("Leg IK solver process starting (foot limits update once standing FK is ready)...")
    def _publish_ik_target(self) -> None:
        if self._ik_target_q is None:
            return
        # Keep only the newest target; drop if the child is momentarily behind.
        try:
            self._ik_target_q.put_nowait(self._foot_xyz.copy())
        except queue.Full:
            pass
    def _pump_ik_results(self) -> None:
        if self._ik_result_q is None:
            return
        while True:
            try:
                kind, payload = self._ik_result_q.get_nowait()
            except queue.Empty:
                break
            if kind == "q":
                self._values[:LEG_JOINT_COUNT] = np.asarray(payload, dtype=np.float32)
            elif kind == "ready":
                payload = np.asarray(payload, dtype=np.float32)
                self._standing_foot_xyz = payload.copy()
                self._set_foot_limits_from_positions(payload[:3], payload[3:])
                self._ik_ready = True
                logger.info(
                    "Leg IK ready — foot slider limits from standing FK ± %s m (pelvis frame)",
                    self.config.foot_xyz_margin,
                )
            elif kind == "error":
                self._leg_ik_ok = False
                self._leg_ik_error = str(payload)
                logger.warning(
                    "Leg IK unavailable (%s); foot sliders shown but legs use joint values", payload
                )
    def _stop_ik_process(self) -> None:
        if self._ik_proc is None:
            return
        try:
            if self._ik_stop_evt is not None:
                self._ik_stop_evt.set()
            if self._ik_target_q is not None:
                try:
                    self._ik_target_q.put_nowait(None)
                except queue.Full:
                    pass
            self._ik_proc.join(timeout=2.0)
            if self._ik_proc.is_alive():
                self._ik_proc.terminate()
        finally:
            self._ik_proc = None
    def connect(self, calibrate: bool = True) -> None:
        require_package("pygame", extra="pygame-dep", import_name="pygame")
        import pygame
        self._foot_divider_x = self.config.foot_panel_width
        self._pygame = pygame
        pygame.init()
        pygame.display.set_caption("G1 SONIC — foot IK + upper-body sliders")
        self._screen = pygame.display.set_mode((self.config.window_width, self.config.window_height))
        self._font = pygame.font.SysFont("dejavusans", 15)
        self._small_font = pygame.font.SysFont("dejavusans", 12)
        self._clock = pygame.time.Clock()
        self._init_leg_ik()
        self._connected = True
        logger.info("G1 sonic slider UI ready (R=reset, wheel=scroll, Esc=quit)")
    def configure(self) -> None:
        pass
    def calibrate(self) -> None:
        pass
    def _reset_pose(self) -> None:
        self._values[:] = DEFAULT_ANGLES
        if self._leg_ik_ok and self._standing_foot_xyz is not None:
            self._foot_xyz[:] = self._standing_foot_xyz
            self._publish_ik_target()
    def _foot_row_rect(self, foot_idx: int) -> tuple[int, int, int, int]:
        y = HEADER_H + foot_idx * self.config.row_height
        track_x = MARGIN + 88
        track_w = self.config.foot_panel_width - track_x - MARGIN
        return track_x, y, track_w, self.config.row_height
    def _joint_row_rect(self, ui_idx: int) -> tuple[int, int, int, int]:
        joint_idx = UPPER_BODY_INDICES[ui_idx] if self._leg_ik_ok else ui_idx
        row = ui_idx
        y = HEADER_H + row * self.config.row_height - self._scroll_y
        track_x = self._foot_divider_x + MARGIN + LABEL_W
        track_w = self.config.slider_width
        return track_x, y, track_w, self.config.row_height, joint_idx
    def _value_from_track(self, lo: float, hi: float, mouse_x: int, track_x: int, track_w: int) -> float:
        t = (mouse_x - track_x) / max(track_w, 1)
        t = float(np.clip(t, 0.0, 1.0))
        return lo + t * (hi - lo)
    def _knob_x(self, val: float, lo: float, hi: float, track_x: int, track_w: int) -> int:
        span = hi - lo if hi > lo else 1.0
        t = (val - lo) / span
        return int(track_x + t * track_w)
    def _handle_events(self) -> bool:
        pygame = self._pygame
        num_joint_rows = len(UPPER_BODY_INDICES) if self._leg_ik_ok else NUM_JOINTS
        max_scroll = max(0, num_joint_rows * self.config.row_height - self.config.window_height + HEADER_H)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return False
                if event.key == pygame.K_r:
                    self._reset_pose()
            if event.type == pygame.MOUSEWHEEL:
                self._scroll_y = int(np.clip(self._scroll_y - event.y * self.config.scroll_step, 0, max_scroll))
            if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                mx, my = event.pos
                if self.config.use_leg_ik and mx < self._foot_divider_x:
                    for fi in range(NUM_FOOT_SLIDERS):
                        rx, ry, rw, rh = self._foot_row_rect(fi)
                        if ry + 4 <= my <= ry + rh - 4 and rx <= mx <= rx + rw:
                            self._drag_foot = fi
                            self._foot_xyz[fi] = self._value_from_track(
                                float(self._foot_lo[fi]),
                                float(self._foot_hi[fi]),
                                mx,
                                rx,
                                rw,
                            )
                            break
                else:
                    for ui in range(num_joint_rows):
                        rx, ry, rw, rh, ji = self._joint_row_rect(ui)
                        if ry + rh < HEADER_H or ry > self.config.window_height:
                            continue
                        if ry + 4 <= my <= ry + rh - 4 and rx <= mx <= rx + rw:
                            self._drag_joint = ji
                            self._values[ji] = self._value_from_track(
                                float(JOINT_LO[ji]), float(JOINT_HI[ji]), mx, rx, rw
                            )
                            break
            if event.type == pygame.MOUSEBUTTONUP and event.button == 1:
                self._drag_foot = None
                self._drag_joint = None
            if event.type == pygame.MOUSEMOTION:
                mx, my = event.pos
                if self._drag_foot is not None:
                    rx, _, rw, _ = self._foot_row_rect(self._drag_foot)
                    self._foot_xyz[self._drag_foot] = self._value_from_track(
                        float(self._foot_lo[self._drag_foot]),
                        float(self._foot_hi[self._drag_foot]),
                        mx,
                        rx,
                        rw,
                    )
                elif self._drag_joint is not None:
                    for ui in range(num_joint_rows):
                        rx, ry, rw, rh, ji = self._joint_row_rect(ui)
                        if ji == self._drag_joint:
                            self._values[ji] = self._value_from_track(
                                float(JOINT_LO[ji]), float(JOINT_HI[ji]), mx, rx, rw
                            )
                            break
        return True
    def _draw_foot_panel(self) -> None:
        pygame = self._pygame
        screen = self._screen
        panel_title = self._small_font.render("Foot IK (pelvis)", True, (180, 200, 255))
        screen.blit(panel_title, (MARGIN, 38))
        if self._leg_ik_error:
            err = self._leg_ik_error if len(self._leg_ik_error) < 42 else self._leg_ik_error[:39] + "..."
            screen.blit(self._small_font.render(f"IK off: {err}", True, (255, 120, 120)), (MARGIN, 50))
        pygame.draw.line(
            screen,
            (60, 60, 70),
            (self._foot_divider_x - 1, HEADER_H - 4),
            (self._foot_divider_x - 1, self.config.window_height),
            1,
        )
        for fi, label in enumerate(FOOT_LABELS):
            rx, ry, rw, rh = self._foot_row_rect(fi)
            lo, hi = float(self._foot_lo[fi]), float(self._foot_hi[fi])
            txt = self._small_font.render(label, True, (190, 210, 230))
            screen.blit(txt, (MARGIN, ry + 4))
            track_y = ry + rh // 2 - 2
            pygame.draw.rect(screen, (45, 55, 70), (rx, track_y, rw, 4), border_radius=2)
            kx = self._knob_x(float(self._foot_xyz[fi]), lo, hi, rx, rw)
            pygame.draw.rect(screen, (70, 140, 220), (rx, track_y, max(0, kx - rx), 4), border_radius=2)
            pygame.draw.rect(screen, (200, 225, 255), (kx - KNOB_W // 2, track_y - 5, KNOB_W, 14), border_radius=3)
            val_txt = self._small_font.render(f"{self._foot_xyz[fi]:+.3f}", True, (150, 200, 220))
            screen.blit(val_txt, (rx + rw + 4, ry + 4))
    def _draw_joint_panel(self) -> None:
        pygame = self._pygame
        screen = self._screen
        num_joint_rows = len(UPPER_BODY_INDICES) if self._leg_ik_ok else NUM_JOINTS
        joint_title = "Waist + arms" if self._leg_ik_ok else "All joints"
        screen.blit(
            self._small_font.render(joint_title, True, (180, 180, 190)),
            (self._foot_divider_x + MARGIN, 38),
        )
        for ui in range(num_joint_rows):
            rx, ry, rw, rh, ji = self._joint_row_rect(ui)
            if ry + rh < HEADER_H or ry > self.config.window_height:
                continue
            short = JOINT_NAMES[ji].removeprefix("k")
            label = self._small_font.render(f"{ji:02d} {short}", True, (200, 200, 210))
            screen.blit(label, (self._foot_divider_x + MARGIN, ry + 4))
            track_y = ry + rh // 2 - 2
            lo, hi = float(JOINT_LO[ji]), float(JOINT_HI[ji])
            pygame.draw.rect(screen, (55, 55, 65), (rx, track_y, rw, 4), border_radius=2)
            kx = self._knob_x(float(self._values[ji]), lo, hi, rx, rw)
            pygame.draw.rect(screen, (80, 160, 255), (rx, track_y, max(0, kx - rx), 4), border_radius=2)
            pygame.draw.rect(screen, (220, 235, 255), (kx - KNOB_W // 2, track_y - 5, KNOB_W, 14), border_radius=3)
            val_txt = self._small_font.render(f"{self._values[ji]:+.3f}", True, (170, 220, 170))
            screen.blit(val_txt, (rx + rw + 8, ry + 4))
    def _draw(self) -> None:
        pygame = self._pygame
        screen = self._screen
        screen.fill((28, 28, 32))
        title = self._font.render("G1 reference → SONIC encoder mode 0", True, (230, 230, 235))
        hint = self._small_font.render("Foot xyz (left) · waist/arms (right) · R reset · Esc quit", True, (150, 150, 160))
        screen.blit(title, (MARGIN, 10))
        screen.blit(hint, (MARGIN, 28))
        if self.config.use_leg_ik:
            self._draw_foot_panel()
        self._draw_joint_panel()
        pygame.display.flip()
    def get_action(self) -> dict[str, float]:
        if not self._connected:
            return {f"{name}.q": float(self._values[i]) for i, name in enumerate(JOINT_NAMES)}
        if not self._handle_events():
            raise KeyboardInterrupt("G1 sonic slider window closed")
        if self._leg_ik_ok:
            # Read the newest solution from the solver process and hand it the latest
            # foot target — never block the teleop loop on the IPOPT solve.
            self._pump_ik_results()
            self._publish_ik_target()
        self._draw()
        self._clock.tick(60)
        return {f"{name}.q": float(self._values[i]) for i, name in enumerate(JOINT_NAMES)}
    def send_feedback(self, feedback: dict) -> None:
        del feedback
    def disconnect(self) -> None:
        self._stop_ik_process()
        if self._pygame is not None:
            self._pygame.quit()
        self._connected = False
        self._screen = None
@@ -0,0 +1,64 @@
 #!/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.
 """G1 29-DOF joint limits (rad) in MuJoCo / G1_29_JointIndex order — from g1_29dof.xml.
 SONIC encoder mode 0 expects Isaac Lab order; the whole-body controller remaps on ingest.
 """
 import numpy as np
 from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
 # (low, high) per joint index 0..28
 _G1_LIMITS = np.array(
    [
        (-2.5307, 2.8798),  # kLeftHipPitch
        (-0.5236, 2.9671),  # kLeftHipRoll
        (-2.7576, 2.7576),  # kLeftHipYaw
        (-0.087267, 2.8798),  # kLeftKnee
        (-0.87267, 0.5236),  # kLeftAnklePitch
        (-0.2618, 0.2618),  # kLeftAnkleRoll
        (-2.5307, 2.8798),  # kRightHipPitch
        (-2.9671, 0.5236),  # kRightHipRoll
        (-2.7576, 2.7576),  # kRightHipYaw
        (-0.087267, 2.8798),  # kRightKnee
        (-0.87267, 0.5236),  # kRightAnklePitch
        (-0.2618, 0.2618),  # kRightAnkleRoll
        (-2.618, 2.618),  # kWaistYaw
        (-0.52, 0.52),  # kWaistRoll
        (-0.52, 0.52),  # kWaistPitch
        (-3.0892, 2.6704),  # kLeftShoulderPitch
        (-1.5882, 2.2515),  # kLeftShoulderRoll
        (-2.618, 2.618),  # kLeftShoulderYaw
        (-1.0472, 2.0944),  # kLeftElbow
        (-1.97222, 1.97222),  # kLeftWristRoll
        (-1.61443, 1.61443),  # kLeftWristPitch
        (-1.61443, 1.61443),  # kLeftWristYaw
        (-3.0892, 2.6704),  # kRightShoulderPitch
        (-2.2515, 1.5882),  # kRightShoulderRoll
        (-2.618, 2.618),  # kRightShoulderYaw
        (-1.0472, 2.0944),  # kRightElbow
        (-1.97222, 1.97222),  # kRightWristRoll
        (-1.61443, 1.61443),  # kRightWristPitch
        (-1.61443, 1.61443),  # kRightWristYaw
    ],
    dtype=np.float32,
 )
 JOINT_NAMES = [m.name for m in G1_29_JointIndex]
 JOINT_LO = _G1_LIMITS[:, 0]
 JOINT_HI = _G1_LIMITS[:, 1]
@@ -35,3 +35,7 @@ class PicoHeadsetConfig(TeleoperatorConfig):
    """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)."""
@@ -23,7 +23,16 @@ from lerobot.types import RobotAction
 from ..teleoperator import Teleoperator
 from .config_pico_headset import PicoHeadsetConfig
-from .smpl_constants import ROOT_ACTION_DIM, ROOT_ACTION_PREFIX, SMPL_ACTION_PREFIX, SMPL_OBS_DIM
+from .smpl_constants import (
    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__)
@@ -48,6 +57,10 @@ class PicoHeadset(Teleoperator):
    @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)})
            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
@@ -89,12 +102,21 @@ class PicoHeadset(Teleoperator):
        if self._stream is None:
            raise RuntimeError(f"{self} is not connected")
        window = self._stream.step()
-        # Emit SMPL only while the headset is actively streaming: hold back before
+        # Emit a reference only while the headset is actively streaming: hold back
-        # the first frame (so the controller doesn't track an all-zero collapsed
+        # before the first frame (so the controller doesn't track an all-zero
-        # pose) and once the stream goes stale (so the controller falls back to a
+        # collapsed pose) and once the stream goes stale (so the controller falls
-        # safe standing/locomotion mode instead of freezing on the last pose).
+        # back to a safe standing/locomotion mode instead of freezing on the last
        # pose).
        if not self._stream.has_data or self._stream.is_stale:
            return {}
        if self.config.mode == "vr3":
            # Sparse 3-point upper-body teleop (encode_mode 1). Needs the producer to
            # be sending vr3_* fields; otherwise emit nothing and stay in locomotion.
            if not self._stream.has_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)})
            return action
        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
@@ -50,6 +50,7 @@ import zmq
 from lerobot.teleoperators.pico_headset.smpl_fk import (
    SmplForwardKinematics,
    canonicalize_smpl_joints,
    compute_3point,
    root_quats_from_aa,
 )
@@ -63,8 +64,15 @@ def pack_message(
    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,
 ) -> bytes:
-    """Build the rt/smpl JSON message (single frame, topic embedded in payload)."""
+    """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),
@@ -74,6 +82,10 @@ def pack_message(
        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()
    return json.dumps({"topic": SMPL_TOPIC, "data": data}).encode("utf-8")
@@ -157,6 +169,7 @@ def main() -> None:
    try:
        while True:
            loop_start = time.time()
            vr3_pos = vr3_orn = None
            if clip is not None:
                n = clip["joints"].shape[0]
                if args.no_loop and frame_index >= n:
@@ -181,9 +194,20 @@ def main() -> None:
                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)
+                pack_message(
                    joints,
                    frame_index,
                    stamp_ns,
                    root_quat=root_quat,
                    root_transl=root_transl,
                    vr3_pos=vr3_pos,
                    vr3_orn=vr3_orn,
                )
            )
            frame_index += 1
            if frame_index % int(max(1, args.fps)) == 0:
@@ -0,0 +1,51 @@
 #!/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."
@@ -35,6 +35,8 @@ 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(
@@ -210,6 +212,70 @@ def root_quats_from_aa(root_aa: np.ndarray) -> np.ndarray:
    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 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] = []
    for out_i, j in enumerate((0, *_VR3_JOINTS)):
        positions[out_i] = q @ body[j, :3]
        rot = R.from_quat(body[j, 3:7]).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
 class SmplForwardKinematics:
    """Rest-skeleton SMPL forward kinematics (no mesh, no torch)."""
@@ -41,7 +41,7 @@ from collections import deque
 import numpy as np
 import zmq
-from .smpl_constants import JOINT_DIM, N_JOINTS, SMPL_OBS_DIM, WINDOW
+from .smpl_constants import JOINT_DIM, N_JOINTS, SMPL_OBS_DIM, VR3_ORN_DIM, VR3_POS_DIM, WINDOW
 logger = logging.getLogger(__name__)
@@ -91,6 +91,10 @@ class SmplStream:
        # 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_index = -1
        self._last_recv_t = 0.0
        self._warned_stale = False
@@ -112,6 +116,11 @@ class SmplStream:
        """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 seconds_since_last(self) -> float:
        """Wall-clock seconds since the last real frame (inf before the first)."""
@@ -133,6 +142,7 @@ class SmplStream:
    def reset(self):
        self._buf.clear()
        self._got_first = False
        self._got_vr3 = False
    # -- core ----------------------------------------------------------------
    def _drain_latest(self) -> np.ndarray | None:
@@ -156,6 +166,12 @@ class SmplStream:
            rt = data.get("root_transl")
            if rt is not None and len(rt) == 3:
                self.root_transl = np.asarray(rt, np.float32)
            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
        return frame
    def step(self) -> np.ndarray:
@@ -0,0 +1,147 @@
 #!/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()