diff --git a/docs/SONIC_FIDELITY_AUDIT.md b/docs/SONIC_FIDELITY_AUDIT.md new file mode 100644 index 000000000..5548e5062 --- /dev/null +++ b/docs/SONIC_FIDELITY_AUDIT.md @@ -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. diff --git a/docs/SONIC_REPLAY_DEBUGGING.md b/docs/SONIC_REPLAY_DEBUGGING.md new file mode 100644 index 000000000..af6735d7f --- /dev/null +++ b/docs/SONIC_REPLAY_DEBUGGING.md @@ -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. diff --git a/examples/onnx/PR_DESCRIPTION.md b/examples/onnx/PR_DESCRIPTION.md new file mode 100644 index 000000000..ba4f40e68 --- /dev/null +++ b/examples/onnx/PR_DESCRIPTION.md @@ -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. diff --git a/src/lerobot/robots/unitree_g1/README.md b/src/lerobot/robots/unitree_g1/README.md new file mode 100644 index 000000000..a5fb52345 --- /dev/null +++ b/src/lerobot/robots/unitree_g1/README.md @@ -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=`: + +| 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=/ --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 .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=...)`. diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 3bfb378f2..99a4bfd92 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -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 diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index 6272c706f..b99d8bfaf 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -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 self.controller.encode_mode == 2: + if 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] + 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) diff --git a/src/lerobot/teleoperators/g1_sonic_slider/__init__.py b/src/lerobot/teleoperators/g1_sonic_slider/__init__.py new file mode 100644 index 000000000..1367c57b4 --- /dev/null +++ b/src/lerobot/teleoperators/g1_sonic_slider/__init__.py @@ -0,0 +1,20 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .config_g1_sonic_slider import G1SonicSliderTeleopConfig +from .g1_sonic_slider import G1SonicSliderTeleop + +__all__ = ["G1SonicSliderTeleop", "G1SonicSliderTeleopConfig"] diff --git a/src/lerobot/teleoperators/g1_sonic_slider/config_g1_sonic_slider.py b/src/lerobot/teleoperators/g1_sonic_slider/config_g1_sonic_slider.py new file mode 100644 index 000000000..f89e7a391 --- /dev/null +++ b/src/lerobot/teleoperators/g1_sonic_slider/config_g1_sonic_slider.py @@ -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.""" diff --git a/src/lerobot/teleoperators/g1_sonic_slider/g1_sonic_slider.py b/src/lerobot/teleoperators/g1_sonic_slider/g1_sonic_slider.py new file mode 100644 index 000000000..de8c99a01 --- /dev/null +++ b/src/lerobot/teleoperators/g1_sonic_slider/g1_sonic_slider.py @@ -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 diff --git a/src/lerobot/teleoperators/g1_sonic_slider/joint_limits.py b/src/lerobot/teleoperators/g1_sonic_slider/joint_limits.py new file mode 100644 index 000000000..16d293c12 --- /dev/null +++ b/src/lerobot/teleoperators/g1_sonic_slider/joint_limits.py @@ -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] diff --git a/src/lerobot/teleoperators/pico_headset/config_pico_headset.py b/src/lerobot/teleoperators/pico_headset/config_pico_headset.py index c452dd9c1..81269cf53 100644 --- a/src/lerobot/teleoperators/pico_headset/config_pico_headset.py +++ b/src/lerobot/teleoperators/pico_headset/config_pico_headset.py @@ -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).""" diff --git a/src/lerobot/teleoperators/pico_headset/pico_headset.py b/src/lerobot/teleoperators/pico_headset/pico_headset.py index bbf828388..048180e74 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_headset.py +++ b/src/lerobot/teleoperators/pico_headset/pico_headset.py @@ -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 - # 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). + # 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 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 diff --git a/src/lerobot/teleoperators/pico_headset/pico_publisher.py b/src/lerobot/teleoperators/pico_headset/pico_publisher.py index 09f6e147c..88462a047 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_publisher.py +++ b/src/lerobot/teleoperators/pico_headset/pico_publisher.py @@ -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: diff --git a/src/lerobot/teleoperators/pico_headset/smpl_constants.py b/src/lerobot/teleoperators/pico_headset/smpl_constants.py new file mode 100644 index 000000000..5d2169e95 --- /dev/null +++ b/src/lerobot/teleoperators/pico_headset/smpl_constants.py @@ -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." diff --git a/src/lerobot/teleoperators/pico_headset/smpl_fk.py b/src/lerobot/teleoperators/pico_headset/smpl_fk.py index 098e4a256..e60ea2e51 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_fk.py +++ b/src/lerobot/teleoperators/pico_headset/smpl_fk.py @@ -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).""" diff --git a/src/lerobot/teleoperators/pico_headset/smpl_stream.py b/src/lerobot/teleoperators/pico_headset/smpl_stream.py index 908ad37c6..36f0c321e 100644 --- a/src/lerobot/teleoperators/pico_headset/smpl_stream.py +++ b/src/lerobot/teleoperators/pico_headset/smpl_stream.py @@ -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: diff --git a/src/lerobot/teleoperators/pico_headset/smpl_to_dataset.py b/src/lerobot/teleoperators/pico_headset/smpl_to_dataset.py new file mode 100644 index 000000000..cc6411964 --- /dev/null +++ b/src/lerobot/teleoperators/pico_headset/smpl_to_dataset.py @@ -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=/ --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()