From 4209639f338dd0a42851d420d7fcb6467a0cdf87 Mon Sep 17 00:00:00 2001 From: Martino Russi Date: Mon, 20 Jul 2026 20:12:27 +0200 Subject: [PATCH] refactor(unitree_g1): isolate SONIC encoder/decoder whole-body path Strip everything except the OpenHLM/pi0.5 -> SONIC encoder/decoder rollout path so this branch does exactly that and nothing more: - Remove the SONIC motion planner (planner ONNX + subprocess worker, PlannerMotion, replanning, MovementState/LocomotionMode, joystick) from sonic_pipeline; keep the encoder/decoder and the caller-fed reference buffer (PlannerController) intact. - Slim SonicRuntime to load only the encoder/decoder; SonicWholeBodyController now runs solely the 34-D whole-body command path (drop SMPL/VR3/keyboard teleop). - Delete the pico_headset teleoperator (SONIC's SMPL/VR3 teleop source). - Move WB action constants into g1_utils; repoint imports. GR00T/Holosoma locomotion controllers are left untouched. Co-authored-by: Cursor --- src/lerobot/robots/unitree_g1/README.md | 199 ++--- .../unitree_g1/controllers/sonic_pipeline.py | 781 +----------------- .../controllers/sonic_whole_body.py | 491 +---------- src/lerobot/robots/unitree_g1/unitree_g1.py | 2 +- src/lerobot/scripts/lerobot_teleoperate.py | 1 - .../teleoperators/pico_headset/__init__.py | 20 - .../pico_headset/config_pico_headset.py | 41 - .../pico_headset/pico_headset.py | 148 ---- .../pico_headset/pico_publisher.py | 333 -------- .../pico_headset/smpl_constants.py | 80 -- .../teleoperators/pico_headset/smpl_fk.py | 571 ------------- .../teleoperators/pico_headset/smpl_stream.py | 277 ------- .../pico_headset/smpl_to_dataset.py | 147 ---- 13 files changed, 112 insertions(+), 2979 deletions(-) delete mode 100644 src/lerobot/teleoperators/pico_headset/__init__.py delete mode 100644 src/lerobot/teleoperators/pico_headset/config_pico_headset.py delete mode 100644 src/lerobot/teleoperators/pico_headset/pico_headset.py delete mode 100644 src/lerobot/teleoperators/pico_headset/pico_publisher.py delete mode 100644 src/lerobot/teleoperators/pico_headset/smpl_constants.py delete mode 100644 src/lerobot/teleoperators/pico_headset/smpl_fk.py delete mode 100644 src/lerobot/teleoperators/pico_headset/smpl_stream.py delete mode 100644 src/lerobot/teleoperators/pico_headset/smpl_to_dataset.py diff --git a/src/lerobot/robots/unitree_g1/README.md b/src/lerobot/robots/unitree_g1/README.md index d1568572e..6c2799266 100644 --- a/src/lerobot/robots/unitree_g1/README.md +++ b/src/lerobot/robots/unitree_g1/README.md @@ -1,162 +1,89 @@ -# Unitree G1 — SONIC whole-body control +# Unitree G1 — SONIC encoder/decoder whole-body control -This package runs NVIDIA's **SONIC** whole-body controller (and the GR00T/Holosoma -locomotion controllers) on the Unitree G1, in MuJoCo simulation or on real hardware. -SONIC turns a high-level movement intent — or a streamed **SMPL** whole-body pose — into -50 Hz joint-position targets. It is a pure-Python/ONNX reimplementation of the SONIC -deploy stack (no `gear_sonic`/torch dependency). +This package runs NVIDIA's **SONIC** encoder/decoder on the Unitree G1, in MuJoCo +simulation or on real hardware, driven by a dense **34-D whole-body command** (the +OpenHLM / pi0.5 action layout). It is a pure-Python/ONNX reimplementation of the +reference-tracking half of the SONIC deploy stack (no `gear_sonic`/torch dependency, and +no motion planner): the encoder compresses a reference motion window into a latent token +and the decoder maps that token + proprioception history into 50 Hz joint-position +targets for the robot's PD controller. ## Controllers Selected with `--robot.controller=`: -| Controller | Purpose | -| ------------------------------ | ------------------------------------------------------------------------------------------ | -| `SonicWholeBodyController` | SONIC whole-body: locomotion (mode 0), 3-point VR teleop (mode 1), SMPL imitation (mode 2) | -| `GrootLocomotionController` | GR00T locomotion policy | -| `HolosomaLocomotionController` | Holosoma locomotion policy | +| Controller | Purpose | +| ------------------------------ | ------------------------------------------------------------ | +| `SonicWholeBodyController` | SONIC encoder/decoder driven by a 34-D OpenHLM/pi0.5 command | +| `GrootLocomotionController` | GR00T locomotion policy | +| `HolosomaLocomotionController` | Holosoma locomotion policy | -On startup the controller **interpolates** from the robot's measured pose into the -policy's commanded target over ~3 s (no snap), and on disconnect (Ctrl-C) it performs a -**graceful damped settle** — holding pose while ramping stiffness to zero over -`--robot.graceful_stop_s` (default 1.5 s) instead of going instantly limp. Both apply in -every mode. +The rest of this document covers the SONIC whole-body path. + +Each tick the `SonicWholeBodyController` takes a 34-D command (`wb.0.pos … wb.33.pos`) in the OpenHLM +layout: + +``` +[L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3), + root roll/pitch + yaw-rate(3)] +``` + +The 29 joint targets become the SONIC encode-mode-0 reference (accumulated into a rolling +50-frame trajectory with finite-difference velocities so the encoder's lookahead sees a +real motion sequence), the root roll/pitch set the anchor orientation, and the two grip +scalars can drive the Dex3 hands (see below). On startup the controller **interpolates** +from the robot's measured pose into the policy's commanded target over ~3 s (no snap). ## Requirements -- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended — SONIC runs three ONNX - sessions and is much smoother on GPU). Install the CUDA build that matches your - driver (e.g. `onnxruntime-gpu==1.26.0` for a CUDA-12.x driver). Verify with: +- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended). Verify with: ```bash python -c "import onnxruntime as ort; print(ort.get_available_providers())" - # expect CUDAExecutionProvider in the list for GPU ``` - `mujoco` for simulation (`is_simulation=True`). -- `pyzmq` only if you use the live SMPL stream (pico headset). -- The SONIC ONNX models are downloaded automatically from the `nvidia/GEAR-SONIC` Hub repo. +- The SONIC encoder/decoder ONNX models download automatically from the + `nvidia/GEAR-SONIC` Hub repo. -## Running +## Running a rollout -**Replay an SMPL dataset (motion imitation):** +Drive the G1 with a 34-D VLA policy (OpenHLM / pi0.5) via `lerobot-rollout`: ```bash -lerobot-replay \ - --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \ - --dataset.repo_id=/ --dataset.episode=0 +lerobot-rollout \ + --strategy.type=base \ + --policy.path= \ + --robot.type=unitree_g1 \ + --robot.controller=SonicWholeBodyController \ + --robot.is_simulation=true \ + --robot.publish_hands=true \ + --task="" \ + --duration=45 --device=cuda ``` -**Keyboard teleop** (drives locomotion via the native keyboard teleoperator): +### Cameras -```bash -lerobot-teleoperate \ - --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \ - --teleop.type=keyboard -``` +Image-conditioned policies need camera frames. Two options are available without live +cameras: -Controls: `WASD` move · `Q`/`E` turn · `1`–`8` mode · `9`/`0` speed · `-`/`=` height · -`R` replan · `Space` emergency-stop. +- **Black frames**: `--robot.empty_cameras='[base, left_wrist, right_wrist]'`. +- **Replay a recorded episode** as the camera feed: + ```bash + --robot.replay_camera_parquet= \ + --robot.replay_camera_map='{base: head_image_left, left_wrist: left_wrist_image, right_wrist: right_wrist_image}' + ``` -**PICO headset teleop — SMPL whole-body** (mode 2, needs PICO Motion Trackers): +### Hands (Dex3) -```bash -# 1) publisher (streams rt/smpl from full-body tracking) -python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 -# 2) controller -lerobot-teleoperate \ - --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \ - --teleop.type=pico_headset -``` +`--robot.publish_hands=true` publishes `rt/dex3/{left,right}/cmd` from the two grip +scalars (`wb.7.pos` left, `wb.15.pos` right). The scalar is interpolated between +`hand_open_grip_value` (default 1.0 = open) and `hand_closed_grip_value` (default 0.0 = +closed) and scaled onto `hand_closed_pose` (7 joints: +`thumb_0, thumb_1, thumb_2, middle_0, middle_1, index_0, index_1`). Flip the signs in +`hand_closed_pose` if the fingers curl the wrong way, or raise `hand_kp` for a firmer +grip. -**PICO headset teleop — 3-point VR** (mode 1, head + 2 controllers only, **no trackers**): +## Observation state -```bash -# 1) publisher (head + controllers -> 3-point targets + stick locomotion) -python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --headset-source devices -# 2) controller -lerobot-teleoperate \ - --robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \ - --teleop.type=pico_headset --teleop.mode=vr3 -``` - -3-point controls: left stick move · right stick X turn · right stick Y height · -`A`+`B` / `X`+`Y` cycle locomotion mode (walk/run/squat/kneel/…) · hands+head track the -upper body. **Calibration**: stand in a neutral rest pose and press `A`+`B`+`X`+`Y` — the -publisher status line flips from `UNCALIBRATED` to `calibrated`. This maps your rest pose -onto the G1's neutral stance and is required before the hands track well; the SMPL -(mode 2) path is self-calibrating and needs no such step. - -Both require the XRoboToolkit stack — see below. - -## PICO headset / XRoboToolkit install - -Live full-body teleop needs the **XRoboToolkit** system (a PC Service on your -workstation + a PICO app on the headset) and its Python binding, `xrobotoolkit_sdk`. -The full hardware + software walkthrough lives in the SONIC repo: -[`docs/source/getting_started/vr_teleop_setup.md`](https://nvlabs.github.io/GR00T-WholeBodyControl/getting_started/vr_teleop_setup.html). - -Summary: - -1. **PC Service** (workstation) — install and run it before connecting the headset. - - Ubuntu 22.04 / 24.04 (x86_64): prebuilt `.deb` from the - [XRoboToolkit-PC-Service releases](https://github.com/XR-Robotics/XRoboToolkit-PC-Service/releases). - - Jetson (aarch64): the arm64 `.deb`. - - Windows (x64): the Windows PC Service build. -2. **PICO app** — install `XRoboToolkit-PICO-*.apk` on the headset (see the guide), - enable Developer Mode. For **SMPL whole-body** (mode 2) you also need the PICO Motion - Trackers paired/calibrated and "Full body" enabled; for **3-point** (mode 1, - `--headset-source devices`) only Head + Controller + Send are required — no trackers. -3. **`xrobotoolkit_sdk`** — a pybind11/CMake build (not a pip package), from - [`XRoboToolkit-PC-Service-Pybind`](https://github.com/XR-Robotics/XRoboToolkit-PC-Service-Pybind): - - Linux x86_64: `pip install pybind11 cmake` then `bash setup_ubuntu.sh` (or the - SONIC repo's `install_scripts/install_pico.sh`, which builds everything into a - `.venv_teleop`). - - Jetson aarch64: `bash setup_orin.sh` (builds `libPXREARobotSDK.so` from source). - - Windows x64: `pip install pybind11` then `setup_windows.bat` (needs git + an - MSVC/CMake toolchain; uses the prebuilt `PXREARobotSDK.dll`/`.lib`). -4. Connect PICO and workstation to the **same Wi-Fi**, open the XRoboToolkit app, enter - the PC IP, and enable Head/Controller/Send (plus Full-body for SMPL mode 2). - -### Platform support - -| Platform | Live headset teleop | Notes | -| --------------------------- | ------------------- | ------------------------------------------- | -| Linux x86_64 | ✅ | Guided `install_pico.sh` (SONIC repo) | -| Linux aarch64 (Jetson Orin) | ✅ | `setup_orin.sh` builds the native lib | -| Windows x64 | ✅ (manual) | `setup_windows.bat`; no one-shot env script | -| macOS | ❌ | No PC Service / SDK build for Darwin | - -### No hardware required (any platform, incl. macOS/Windows) - -The SMPL pipeline can be exercised without a headset or the SDK — the publisher emits -`rt/smpl` frames that the controller consumes exactly as it would from the headset: - -```bash -# synthetic motion -python -m lerobot.teleoperators.pico_headset.pico_publisher --fake - -# replay a canned SMPL clip -python -m lerobot.teleoperators.pico_headset.pico_publisher --motion-file .npz -``` - -## Notes - -- SMPL **root motion** into the mode-2 anchor is opt-in (`SonicWholeBodyController(enable_smpl_root=True)`); - it stays off by default (untested on hardware). When enabled, the per-frame root quat is - spherically smoothed (`root_smoothing_alpha`, default 0.15) before it reaches the anchor, - which removes the base-acceleration spikes the raw 30 Hz→50 Hz trajectory used to cause. -- Direct `rt/smpl` subscription without the pico teleoperator is available via - `SonicWholeBodyController(enable_smpl_stream=True, smpl_host=..., smpl_port=...)`. -- 3-point (mode 1) uses the **headset-yaw frame** as its reference and the `A`+`B`+`X`+`Y` - calibration to align to the G1 neutral stance. Calibration maps the operator's rest pose - onto the G1's **standing** (`default_angles`) wrist/neck key-frame poses (position **and** - orientation) computed by FK — the `default_angles` stand-in for gear_sonic's live - measured-q recalibration, since the robot holds `default_angles` at calibration time. - Re-aligning the arms only (preserving the neck level) is available via the calibrator's - `recalibrate_wrists()`. -- 3-point **locomotion** from the PICO sticks follows gear_sonic's `PlannerLoop` exactly: - a yaw accumulator on the right stick and **mode-dependent speed curves** on the left - (slow `0.1+0.5·mag`, run `1.5+3·mag`, walk = planner default). Stick signs replicate - gear_sonic's `get_controller_axes` usage (forward `+ly`, strafe `-lx`, turn `-rx`); since - the publisher forwards the same raw SDK axes, this is the correct convention by construction. -- Startup interpolation and the graceful-stop settle are mode-agnostic; set - `--robot.graceful_stop_s=0` to restore the old instant zero-torque on disconnect. +When the whole-body controller is active the robot exposes a 34-D proprio state +(`wb_state.0.pos … wb_state.33.pos`) in the same OpenHLM layout as the action, which the +rollout aggregates into `observation.state` for the policy. diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 359a8c2a5..8f261074f 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -14,54 +14,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""SONIC planner pipeline for the Unitree G1 whole-body controller. +"""SONIC encoder/decoder pipeline for the Unitree G1 whole-body controller. -This module is a pure-Python/ONNX re-implementation of NVIDIA's SONIC deploy stack -(mirrors ``g1_deploy_onnx_ref.cpp``). It turns a high-level movement intent -(walk/run/squat/box/… + speed/height/heading, driven by the joystick) into -50 Hz joint-position targets for the robot's PD controller. +Pure-Python/ONNX re-implementation of the reference-tracking half of NVIDIA's SONIC +deploy stack (mirrors ``g1_deploy_onnx_ref.cpp``). Given a reference motion buffer +(joint targets + body orientation per frame) it produces 50 Hz joint-position targets +for the robot's PD controller. The upstream *motion planner* is intentionally absent: +here the reference is supplied directly by the caller (e.g. a 34-D OpenHLM / pi0.5 VLA +command per tick, in ``sonic_whole_body.py``). -Data flow (one 50 Hz control tick, orchestrated by ``SonicRuntime`` in -``sonic_whole_body.py``): - - intent (MovementState) ──► SonicPlanner ──► PlannerController ──► joint targets - │ │ - (planner ONNX, 30 Hz, (encoder+decoder ONNX, - async background thread) runs every tick) - -Three cooperating ONNX models: - * **planner** – generates a several-second *reference motion* (body trajectory + - joint clip) for the current intent. Slow, so it runs asynchronously in a - background thread (``_planner_worker``) and its 30 Hz output is resampled to - 50 Hz. New motions are cross-faded into the live buffer (``blend_new_motion``). - * **encoder** – compresses the reference window into a 64-D latent ``token`` +Two cooperating ONNX models: + * **encoder** – compresses the reference window into a 64-D latent ``token`` (refreshed every ``ENCODER_UPDATE_EVERY`` ticks). - * **decoder** – every tick, maps the token + recent proprioception history to a + * **decoder** – every tick, maps the token + recent proprioception history to a residual action that is scaled and added to ``DEFAULT_ANGLES``. -Encoder ``encode_mode`` selects what the reference represents: - * ``0`` – locomotion (planner clip drives lower + upper body). - * ``1`` – 3-point VR teleop (lower body from planner, arms from VR targets). - * ``2`` – SMPL whole-body imitation (720-D SMPL window drives the pose). - Index spaces: joints exist in two orderings — **IsaacLab** (policy/training order) and **MuJoCo** (deploy order). ``ISAACLAB_TO_MUJOCO`` / ``MUJOCO_TO_ISAACLAB`` convert between them. Quaternions are scalar-first ``(w, x, y, z)``. - -Section map: constants & index tables · PD gains · quaternion helpers · locomotion -modes · movement state · encoder/decoder · planner motion buffer · async planner -worker · ``SonicPlanner`` · ``PlannerController`` · joystick input. """ from __future__ import annotations import logging -import math -import queue import threading -import time -from dataclasses import dataclass -from enum import IntEnum from typing import TYPE_CHECKING import numpy as np @@ -149,15 +125,6 @@ DEFAULT_HEIGHT = 0.788740 # nominal pelvis height (m) TOKEN_DIM = 64 # encoder latent size ENCODER_UPDATE_EVERY = 5 # refresh the encoder token every N ticks (decoder runs every tick) DEBUG_PRINT_EVERY = 100 # ticks between debug prints -MOTION_LOOK_AHEAD_STEPS = 2 # frames ahead used to seed a replan context (hide planner latency) -INITIAL_RANDOM_SEED = 1234 -MIN_TOKENS, MAX_TOKENS = 6, 16 # planner prediction-length token range -K = MAX_TOKENS - MIN_TOKENS + 1 -DEADZONE = 0.05 # joystick dead zone -BLEND_FRAMES = 8 # cross-fade length when swapping in a freshly planned motion - -# Seconds between automatic replans, per motion class (faster for dynamic motions). -REPLAN_INTERVAL = {"running": 0.1, "crawling": 0.2, "boxing": 1.0, "default": 1.0} def _to_mujoco(a): @@ -294,139 +261,6 @@ def quat_slerp_batch(q0, q1, t): return r / (np.linalg.norm(r, axis=1, keepdims=True) + 1e-12) -# ── Locomotion modes ────────────────────────────────────────────────────────── - - -class LocomotionMode(IntEnum): - """High-level motion styles understood by the planner (fed as the ``mode`` input).""" - - IDLE = 0 - SLOW_WALK = 1 - WALK = 2 - RUN = 3 - SQUAT = 4 - KNEEL_TWO_LEGS = 5 - KNEEL = 6 - LYING_FACE_DOWN = 7 - CRAWLING = 8 - IDLE_BOXING = 9 - WALK_BOXING = 10 - LEFT_PUNCH = 11 - RIGHT_PUNCH = 12 - RANDOM_PUNCH = 13 - ELBOW_CRAWLING = 14 - LEFT_HOOK = 15 - RIGHT_HOOK = 16 - FORWARD_JUMP = 17 - STEALTH_WALK = 18 - INJURED_WALK = 19 - LEDGE_WALKING = 20 - OBJECT_CARRYING = 21 - STEALTH_WALK_2 = 22 - HAPPY_DANCE_WALK = 23 - ZOMBIE_WALK = 24 - GUN_WALK = 25 - SCARE_WALK = 26 - - -LM = LocomotionMode - -# UI groupings of modes for cycling with the n/p keys; each entry is (label, modes). -MOTION_SETS = [ - ("Standing", [LM.SLOW_WALK, LM.WALK, LM.RUN, LM.FORWARD_JUMP, LM.STEALTH_WALK, LM.INJURED_WALK]), - ("Squat / Low", [LM.SQUAT, LM.KNEEL_TWO_LEGS, LM.KNEEL, LM.CRAWLING, LM.ELBOW_CRAWLING]), - ( - "Boxing", - [ - LM.IDLE_BOXING, - LM.WALK_BOXING, - LM.LEFT_PUNCH, - LM.RIGHT_PUNCH, - LM.RANDOM_PUNCH, - LM.LEFT_HOOK, - LM.RIGHT_HOOK, - ], - ), - ( - "Styled Walks", - [ - LM.LEDGE_WALKING, - LM.OBJECT_CARRYING, - LM.STEALTH_WALK_2, - LM.HAPPY_DANCE_WALK, - LM.ZOMBIE_WALK, - LM.GUN_WALK, - LM.SCARE_WALK, - ], - ), -] - -# Mode classifications used by clamping and replan logic. -STATIC_MODES = {LM.IDLE, LM.SQUAT, LM.KNEEL_TWO_LEGS, LM.KNEEL, LM.LYING_FACE_DOWN, LM.IDLE_BOXING} -STANDING_MODES = { - LM.IDLE, - LM.SLOW_WALK, - LM.WALK, - LM.RUN, - LM.IDLE_BOXING, - LM.WALK_BOXING, - LM.LEFT_PUNCH, - LM.RIGHT_PUNCH, - LM.RANDOM_PUNCH, - LM.LEFT_HOOK, - LM.RIGHT_HOOK, - LM.FORWARD_JUMP, - LM.STEALTH_WALK, - LM.INJURED_WALK, - LM.LEDGE_WALKING, - LM.OBJECT_CARRYING, - LM.STEALTH_WALK_2, - LM.HAPPY_DANCE_WALK, - LM.ZOMBIE_WALK, - LM.GUN_WALK, - LM.SCARE_WALK, -} -BOXING_MODES = {LM.WALK_BOXING, LM.LEFT_PUNCH, LM.RIGHT_PUNCH, LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK} -SPEED_RANGES = { - LM.SLOW_WALK: (0.2, 0.8), - LM.WALK: (0.8, 1.5), - LM.RUN: (1.5, 3.0), - LM.CRAWLING: (0.4, 1.0), - LM.ELBOW_CRAWLING: (0.7, 1.0), -} - - -def clamp_mode_params(ms): - """Clamp ``ms.speed``/``ms.height`` into the valid range for its mode in place. - - ``-1.0`` is a sentinel meaning "use the mode's default" (e.g. standing modes - ignore height; static modes ignore speed). - """ - m = LM(ms.mode) - ms.height = -1.0 if m in STANDING_MODES else max(0.1, min(0.8, ms.height if ms.height >= 0 else 0.2)) - if m in STATIC_MODES: - ms.speed = -1.0 - elif m in SPEED_RANGES: - lo, hi = SPEED_RANGES[m] - ms.speed = max(lo, min(hi, ms.speed if ms.speed >= 0 else lo)) - elif m in BOXING_MODES: - ms.speed = max(0.7, min(1.5, ms.speed if ms.speed >= 0 else 0.7)) - else: - ms.speed = -1.0 - - -def replan_interval(mode): - """Seconds between automatic replans for the given mode.""" - m = LM(mode) - if m == LM.RUN: - return REPLAN_INTERVAL["running"] - if m == LM.CRAWLING: - return REPLAN_INTERVAL["crawling"] - if m in {LM.LEFT_PUNCH, LM.RIGHT_PUNCH, LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK}: - return REPLAN_INTERVAL["boxing"] - return REPLAN_INTERVAL["default"] - - def ort_providers(force_cpu: bool = False) -> list[str]: """Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed).""" avail = ort.get_available_providers() @@ -442,92 +276,6 @@ def make_ort_session_options(): return so -# ── Movement state ──────────────────────────────────────────────────────────── - - -@dataclass -class MovementState: - """Mutable high-level intent driven by keyboard/joystick and read by the planner. - - Holds the current locomotion ``mode``, target ``speed``/``height`` (``-1`` = - mode default), facing/movement angles, and the ``needs_replan`` flag the control - loop watches to decide when to request a fresh motion from the planner. - """ - - mode: int = LM.SLOW_WALK # not IDLE — walking modes respond to WASD - speed: float = -1.0 - height: float = -1.0 - facing_angle: float = 0.0 - movement_angle: float = 0.0 - has_movement: bool = False - motion_set_idx: int = 0 - needs_replan: bool = False - joy_prev_active: bool = False # tracks right-stick activity across joystick polls - - @property - def movement_direction(self): - """Unit XY movement direction (0 vector when not moving).""" - if not self.has_movement: - return (0.0, 0.0, 0.0) - return (math.cos(self.movement_angle), math.sin(self.movement_angle), 0.0) - - @property - def facing_direction(self): - """Unit XY facing direction.""" - return (math.cos(self.facing_angle), math.sin(self.facing_angle), 0.0) - - def status_line(self): - """Human-readable one-line status for the terminal HUD.""" - return ( - f"[{MOTION_SETS[self.motion_set_idx][0]}] mode={self.mode}({LM(self.mode).name}) " - f"spd={'default' if self.speed < 0 else f'{self.speed:.1f}'} " - f"hgt={'default' if self.height < 0 else f'{self.height:.2f}'} " - f"facing={math.degrees(self.facing_angle):.0f}° " - f"{'moving' if self.has_movement else 'still'}" - ) - - -@dataclass -class MovementSnapshot: - """Immutable copy of the intent at the last replan, for change detection.""" - - mode: int = 0 - speed: float = -1.0 - height: float = -1.0 - movement: tuple[float, float, float] = (0.0, 0.0, 0.0) - facing: tuple[float, float, float] = (1.0, 0.0, 0.0) - - -def snapshot_ms(ms: MovementState) -> MovementSnapshot: - """Capture the current movement intent as a comparable snapshot.""" - md, fd = ms.movement_direction, ms.facing_direction - return MovementSnapshot(ms.mode, ms.speed, ms.height, (md[0], md[1], md[2]), (fd[0], fd[1], fd[2])) - - -def should_replan_request(ms: MovementState, last: MovementSnapshot, replan_timer: float, step: int) -> bool: - """Decide whether to request a fresh plan this tick. - - Triggers on an explicit ``needs_replan`` flag, any mode/facing/height change, or - (for non-static modes) speed/direction changes and periodic timeouts. Mirrors the - C++ ``G1Deploy::Planner`` replan triggers (``g1_deploy_onnx_ref.cpp``). - """ - if step <= 0: - return False - if ms.needs_replan: - return True - md, fd = ms.movement_direction, ms.facing_direction - facing_changed = fd != last.facing - height_changed = ms.height != last.height - mode_changed = ms.mode != last.mode - speed_changed = ms.speed != last.speed - dir_changed = md != last.movement - is_static = LM(ms.mode) in STATIC_MODES - if mode_changed or facing_changed or height_changed: - return True - time_to_replan = replan_timer >= replan_interval(ms.mode) - return not is_static and (speed_changed or dir_changed or (time_to_replan and ms.speed != 0)) - - # ── Encoder / Decoder ───────────────────────────────────────────────────────── @@ -724,314 +472,19 @@ class StandingEncoderDecoder: return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex} -# ── Planner motion buffer ───────────────────────────────────────────────────── - - -class PlannerMotion: - """Fixed-capacity buffer for a planned motion (joint pos/vel + body pose per frame).""" - - def __init__(self, max_frames=1500): - self.timesteps = 0 - self.joint_positions = np.zeros((max_frames, 29), np.float64) - self.joint_velocities = np.zeros((max_frames, 29), np.float64) - self.body_positions = np.zeros((max_frames, 3), np.float64) - self.body_quaternions = np.zeros((max_frames, 4), np.float64) - self.body_quaternions[:, 0] = 1.0 - - -# ── Subprocess planner ──────────────────────────────────────────────────────── - - -def _resample_30_to_50(qpos, n30): - """Resample planner output (30 Hz MuJoCo qpos) to a 50 Hz IsaacLab-order motion. - - Returns a dict with joint positions/velocities (velocities via finite difference) - and body position/orientation trajectories at 50 Hz. - """ - t50 = int(np.floor(n30 / 30.0 * 50)) - f30 = np.arange(t50) / 50.0 * 30.0 - f0 = np.floor(f30).astype(int) - f1 = np.minimum(f0 + 1, n30 - 1) - frac, w0 = (f30 - f0).astype(np.float64), None - w0 = 1.0 - frac - jp = (w0[:, None] * qpos[f0, 7:36] + frac[:, None] * qpos[f1, 7:36])[:, MUJOCO_TO_ISAACLAB] - jv = np.zeros_like(jp) - if t50 >= 2: - jv[: t50 - 1] = (jp[1:] - jp[:-1]) * 50.0 - jv[-1] = jv[-2] - return { - "timesteps": t50, - "joint_positions": jp, - "joint_velocities": jv, - "body_positions": w0[:, None] * qpos[f0, :3] + frac[:, None] * qpos[f1, :3], - "body_quaternions": quat_slerp_batch(qpos[f0, 3:7], qpos[f1, 3:7], frac), - } - - -def _build_planner_inputs(ctx, ms_dict, version, seed): - """Build the planner ONNX input dict from a context window + movement intent. - - ``version >= 1`` is the TensorRT-style deploy planner with extra height/target - inputs and a token-count mask; ``version 0`` is the minimal input set. - """ - inp = { - "context_mujoco_qpos": ctx.astype(np.float32).reshape(1, 4, 36), - "target_vel": np.array([ms_dict["speed"]], np.float32), - "mode": np.array([ms_dict["mode"]], np.int64), - "movement_direction": np.array(ms_dict["movement_direction"], np.float32).reshape(1, 3), - "facing_direction": np.array(ms_dict["facing_direction"], np.float32).reshape(1, 3), - "random_seed": np.array([seed], np.int64), - } - if version >= 1: - # TensorRT deploy: allow 9–11 prediction tokens only (indices 3–5 for MIN_TOKENS=6). - allowed = np.zeros((1, K), np.int64) - if K >= 6: - allowed[0, 3:6] = 1 - inp.update( - { - "height": np.array([ms_dict["height"]], np.float32), - "has_specific_target": np.array([[0]], np.int64), - "specific_target_positions": np.zeros((1, 4, 3), np.float32), - "specific_target_headings": np.zeros((1, 4), np.float32), - "allowed_pred_num_tokens": allowed, - } - ) - return inp - - -def _planner_worker(path, req_q, res_q, stop_evt, version, seed, use_gpu): - """Background thread: consume replan requests, run the planner ONNX, post motions. - - Loads its own ONNX session, then loops pulling ``(ctx, gen_frame, ms_dict)`` off - ``req_q``, running inference, resampling to 50 Hz, and putting the newest result - on ``res_q`` (dropping stale entries). Runs until ``stop_evt`` is set. - """ - so = make_ort_session_options() - providers = ort_providers(force_cpu=not use_gpu) - sess = ort.InferenceSession(path, sess_options=so, providers=providers) - while not stop_evt.is_set(): - try: - ctx, gf, ms_dict = req_q.get(timeout=0.05) - except queue.Empty: # nosec B112 - idle poll, nothing queued this tick - continue - try: - inp = _build_planner_inputs(ctx, ms_dict, version, seed) - t0 = time.time() - qpos_out, num_pred = sess.run(None, inp) - t_inf = time.time() - n = int(num_pred.flat[0]) - qpos = qpos_out[0, :n] - if np.any(np.isnan(qpos)): - continue - motion = _resample_30_to_50(qpos, n) - motion["gen_frame"] = gf - logger.debug( - "[Planner] inf=%.1fms total=%.1fms frames=%d", - 1000 * (t_inf - t0), - 1000 * (time.time() - t0), - n, - ) - while not res_q.empty(): - try: - res_q.get_nowait() - except queue.Empty: - break - res_q.put(motion) - except Exception: - logger.exception("[Planner] worker error") - - -# ── SonicPlanner ────────────────────────────────────────────────────────────── - - -class SonicPlanner: - """Owns the planner ONNX model and its async background worker. - - Provides the initial motion synchronously (``initialize``), then serves replans - off-thread: ``request_replan`` enqueues the current context+intent and - ``try_get_new_motion`` non-blockingly returns a freshly planned motion (which the - controller cross-fades in). ``version`` selects the planner input schema. - """ - - def __init__(self, session, planner_path): - self.session = session - self.planner_path = planner_path - self.gen_frame = 0 - self.random_seed = INITIAL_RANDOM_SEED - self.version = 1 if len(session.get_inputs()) >= 11 else 0 - self.motion_50hz = PlannerMotion() - self._snapshot = PlannerMotion() - self._req_q = self._res_q = self._stop_evt = self._planner_thread = None - self._ctrl = None - - def _build_inputs(self, ctx, ms): - return _build_planner_inputs( - ctx, - { - "mode": ms.mode, - "speed": ms.speed, - "height": ms.height, - "movement_direction": list(ms.movement_direction), - "facing_direction": list(ms.facing_direction), - }, - self.version, - self.random_seed, - ) - - @staticmethod - def build_initial_context(joint_positions): - """Build a 4-frame standing context (MuJoCo qpos layout) from a pose.""" - ctx = np.zeros((4, 36), np.float32) - jp_mj = joint_positions.astype(np.float32)[ISAACLAB_TO_MUJOCO] - for n in range(4): - ctx[n, 2] = DEFAULT_HEIGHT - ctx[n, 3] = 1.0 - ctx[n, 7:36] = jp_mj - return ctx - - def _context_from_controller(self, current_frame): - """Sample a 4-frame look-ahead context from the controller's live motion buffer. - - The context starts ``MOTION_LOOK_AHEAD_STEPS`` ahead of ``current_frame`` so a - replan blends in seamlessly by the time it is ready. - """ - ctrl = self._ctrl - gen_frame = current_frame + MOTION_LOOK_AHEAD_STEPS - t_arr = gen_frame / 50.0 + np.arange(4) / 30.0 - f50 = t_arr * 50.0 - with ctrl.motion_lock: - ts = ctrl.motion_timesteps - if ts <= 0: - return self.build_initial_context(DEFAULT_ANGLES) - bp, bq, jp = ctrl.motion_body_pos, ctrl.motion_body_quats, ctrl.motion_joint_positions - f0 = np.minimum(np.floor(f50).astype(int), ts - 1) - f1 = np.minimum(f0 + 1, ts - 1) - frac = f50 - f0 - w0 = 1.0 - frac - ctx = np.zeros((4, 36), np.float32) - ctx[:, 0:3] = w0[:, None] * bp[f0] + frac[:, None] * bp[f1] - ctx[:, 3:7] = quat_slerp_batch(bq[f0], bq[f1], frac) - ij = w0[:, None] * jp[f0] + frac[:, None] * jp[f1] - ctx[:, 7:36] = ij[:, ISAACLAB_TO_MUJOCO] - self.gen_frame = gen_frame - return ctx - - def _load_motion_in_place(self, qpos, n30, target=None): - """Resample raw planner qpos to 50 Hz and write it into a ``PlannerMotion`` buffer.""" - if target is None: - target = self.motion_50hz - r = _resample_30_to_50(qpos, n30) - n = r["timesteps"] - target.timesteps = n - target.joint_positions[:n] = r["joint_positions"] - target.joint_velocities[:n] = r["joint_velocities"] - target.body_positions[:n] = r["body_positions"] - target.body_quaternions[:n] = r["body_quaternions"] - return target - - def initialize(self, joint_positions, ms): - """Synchronously run the planner once to produce the first motion buffer.""" - ctx = self.build_initial_context(joint_positions) - qpos_out, num_pred = self.session.run(None, self._build_inputs(ctx, ms)) - n = int(num_pred.flat[0]) - qpos = qpos_out[0, :n] - if np.any(np.isnan(qpos)): - raise RuntimeError("Planner initial output contains NaN") - logger.info("[Planner] Init: %d frames @ 30 Hz", n) - self._load_motion_in_place(qpos, n) - logger.info("[Planner] Resampled to %d frames @ 50 Hz", self.motion_50hz.timesteps) - return self.motion_50hz - - def request_replan(self, cursor, ms): - """Enqueue a replan for the worker (drops any pending stale request first).""" - if self._req_q is None: - return - ctx = self._context_from_controller(cursor) - ms_dict = { - "mode": ms.mode, - "speed": ms.speed, - "height": ms.height, - "movement_direction": list(ms.movement_direction), - "facing_direction": list(ms.facing_direction), - } - while not self._req_q.empty(): - try: - self._req_q.get_nowait() - except queue.Empty: - break - self._req_q.put((ctx, self.gen_frame, ms_dict)) - - def try_get_new_motion(self): - """Non-blocking: return ``(snapshot_motion, gen_frame)`` if a new plan is ready, else None.""" - if self._res_q is None: - return None - result = None - while not self._res_q.empty(): - try: - result = self._res_q.get_nowait() - except queue.Empty: - break - if result is None: - return None - n, gf = result["timesteps"], result["gen_frame"] - s = self._snapshot - s.timesteps = n - s.joint_positions[:n] = result["joint_positions"] - s.joint_velocities[:n] = result["joint_velocities"] - s.body_positions[:n] = result["body_positions"] - s.body_quaternions[:n] = result["body_quaternions"] - return s, gf - - def start_subprocess(self, controller, use_gpu: bool = False): - """Run planner ONNX in a background thread (avoids mp spawn/fork + CUDA/MuJoCo issues).""" - self._ctrl = controller - self._req_q = queue.Queue() - self._res_q = queue.Queue() - self._stop_evt = threading.Event() - self._planner_thread = threading.Thread( - target=_planner_worker, - args=( - self.planner_path, - self._req_q, - self._res_q, - self._stop_evt, - self.version, - self.random_seed, - use_gpu, - ), - daemon=True, - name="sonic-planner", - ) - self._planner_thread.start() - logger.info("[Planner] Background thread started (%s)", "GPU" if use_gpu else "CPU") - - def stop_subprocess(self): - """Signal the planner thread to stop and join it.""" - if self._stop_evt: - self._stop_evt.set() - if self._planner_thread is not None: - self._planner_thread.join(timeout=3.0) - logger.info("[Planner] Background thread stopped") - self._planner_thread = None - self._req_q = self._res_q = self._stop_evt = None - - -# ── PlannerController ───────────────────────────────────────────────────────── - - class PlannerController(StandingEncoderDecoder): - """Encoder/decoder driven by the planner's live, replannable motion buffer. + """Encoder/decoder driven by a caller-supplied, rolling motion buffer. - Extends ``StandingEncoderDecoder`` so the reference comes from a rolling motion - (advanced one frame per tick via ``advance_cursor``) instead of a fixed pose. - Handles heading re-initialization, cross-fading new plans into the buffer - (``blend_new_motion``), and the mode-2 SMPL reference. ``motion_lock`` guards the - buffer against the async planner thread. + Extends ``StandingEncoderDecoder`` so the reference comes from a motion buffer + (a lookahead window with per-frame velocities) instead of a single fixed pose, + and handles heading re-initialization on the first frame / after a reset. + ``motion_lock`` guards the buffer, which the whole-body controller rewrites each + tick from the incoming command. The class name is retained for continuity with + the SONIC reference; no motion planner is involved. """ - def __init__(self, planner, encoder, decoder): + def __init__(self, encoder, decoder): super().__init__(encoder, decoder) - self.planner = planner self.ref_cursor = 0 self.motion_timesteps = 0 self.motion_joint_positions = np.zeros((1500, 29), np.float64) @@ -1046,73 +499,6 @@ class PlannerController(StandingEncoderDecoder): self.playing = self.first_motion = False self.motion_lock = threading.Lock() - def load_initial_motion(self, motion): - """Copy the planner's first motion into the live buffer and start playback.""" - with self.motion_lock: - n = motion.timesteps - self.motion_timesteps = n - self.motion_joint_positions[:n] = motion.joint_positions[:n] - self.motion_joint_velocities[:n] = motion.joint_velocities[:n] - self.motion_body_quats[:n] = motion.body_quaternions[:n] - self.motion_body_pos[:n] = motion.body_positions[:n] - self.init_ref_quat = motion.body_quaternions[0].copy() - self.ref_cursor = 0 - self.first_motion = True - self.playing = True - self.delta_heading = 0.0 - - def blend_new_motion(self, new_motion, gen_frame): - """Blend like C++ CurrentFrameAdvancement: 8-frame cross-fade, then copy tail.""" - with self.motion_lock: - cur = self.ref_cursor - new_len = gen_frame - cur + new_motion.timesteps - if new_len <= 0: - return - if self.motion_timesteps == 0: - n = new_motion.timesteps - self.motion_joint_positions[:n] = new_motion.joint_positions[:n] - self.motion_joint_velocities[:n] = new_motion.joint_velocities[:n] - self.motion_body_pos[:n] = new_motion.body_positions[:n] - self.motion_body_quats[:n] = new_motion.body_quaternions[:n] - self.motion_timesteps = n - self.ref_cursor = 0 - self.init_ref_quat = self.motion_body_quats[0].copy() - self.first_motion = False - return - - blend_start = max(0, gen_frame - cur) - blend_end = min(new_len, blend_start + BLEND_FRAMES) - - for f in range(blend_end): - f_old = min(f + cur, self.motion_timesteps - 1) - f_new = max(0, min(f + cur - gen_frame, new_motion.timesteps - 1)) - w_new = min(1.0, max(0.0, (f - blend_start) / BLEND_FRAMES)) - w_old = 1.0 - w_new - self.motion_joint_positions[f] = ( - w_old * self.motion_joint_positions[f_old] + w_new * new_motion.joint_positions[f_new] - ) - self.motion_joint_velocities[f] = ( - w_old * self.motion_joint_velocities[f_old] + w_new * new_motion.joint_velocities[f_new] - ) - self.motion_body_pos[f] = ( - w_old * self.motion_body_pos[f_old] + w_new * new_motion.body_positions[f_new] - ) - self.motion_body_quats[f] = quat_slerp( - self.motion_body_quats[f_old], new_motion.body_quaternions[f_new], w_new - ) - - for f in range(blend_end, new_len): - f_new = max(0, min(f + cur - gen_frame, new_motion.timesteps - 1)) - self.motion_joint_positions[f] = new_motion.joint_positions[f_new] - self.motion_joint_velocities[f] = new_motion.joint_velocities[f_new] - self.motion_body_pos[f] = new_motion.body_positions[f_new] - self.motion_body_quats[f] = new_motion.body_quaternions[f_new].copy() - - self.motion_timesteps = new_len - self.first_motion = False - self.ref_cursor = 0 - self.init_ref_quat = self.motion_body_quats[0].copy() - def _heading_apply_delta(self): """Heading correction quaternion (init base-vs-ref heading + operator ``delta_heading``).""" delta = quat_mul( @@ -1224,134 +610,3 @@ class PlannerController(StandingEncoderDecoder): with self.motion_lock: if self.motion_timesteps > 0: self.ref_cursor = min(self.ref_cursor + 1, self.motion_timesteps - 1) - - -# ── Joystick input ──────────────────────────────────────────────────────────── - - -def _parse_wireless(wr): - """Parse wireless_remote (bytes or int-array) into (lx, ly, rx, ry).""" - import struct as _st - - if not isinstance(wr, (bytes, bytearray)): - wr = bytes(wr) - if len(wr) < 24: - return None - lx = _st.unpack("f", wr[4:8])[0] - rx = _st.unpack("f", wr[8:12])[0] - ry = _st.unpack("f", wr[12:16])[0] - ly = _st.unpack("f", wr[20:24])[0] - return lx, ly, rx, ry - - -def apply_joystick_axes(lx, ly, rx, ry, ms, controller=None): - """Map raw stick axes onto ``MovementState`` (left stick=WASD, right stick X=Q/E, - right stick Y=height). - - Shared by the G1 wireless remote (:func:`process_joystick`) and the PICO - controller sticks (3-point teleop), so both drive the planner identically. Axes - are expected pre-negated to the same convention as the parsed wireless remote: - ``ly`` and ``ry`` already flipped, dead zone not yet applied. - """ - # Dead zone + negate both Y axes (bridge already flips them once) - lx = 0.0 if abs(lx) < DEADZONE else lx - ly = 0.0 if abs(ly) < DEADZONE else -ly - rx = 0.0 if abs(rx) < DEADZONE else rx - ry = 0.0 if abs(ry) < DEADZONE else -ry - - left_active = abs(lx) > 0 or abs(ly) > 0 - - # Left stick → WASD (movement direction relative to facing) - if left_active: - ms.movement_angle = ms.facing_angle + math.atan2(-lx, -ly) - ms.has_movement = True - if not ms.joy_prev_active: - ms.needs_replan = True - ms.joy_prev_active = True - elif ms.joy_prev_active and not (abs(rx) > 0 or abs(ry) > 0): - ms.joy_prev_active = False - ms.has_movement = False - - # Right stick X → Q/E (facing rotation, ~1 rad/s at full deflection) - if abs(rx) > 0: - delta = -0.02 * rx - ms.facing_angle += delta - if controller: - controller.delta_heading += delta - - # Right stick Y → -/= (height adjustment, ~0.25/s at full deflection) - if abs(ry) > 0: - step = -0.005 * ry - ms.height = max(0.1, min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + step)) - - -# gear_sonic PlannerLoop joystick constants (pico_manager_thread_server). -PICO_YAW_GAIN = 1.5 # rad/s of facing rotation at full right-stick deflection -PICO_JOY_DEADZONE = 0.15 # gear_sonic JOYSTICK_DEADZONE (distinct from the remote's) - - -def apply_pico_loco_axes(lx, ly, rx, ry, ms, dt=CONTROL_DT): - """Map PICO controller sticks onto ``MovementState`` — the gear_sonic 3-point path. - - Ports ``pico_manager_thread_server.PlannerLoop`` (encode_mode 1) rather than the - keyboard/remote-parity :func:`apply_joystick_axes`, so speed follows gear_sonic's - mode-dependent curves and facing comes from a yaw accumulator: - - - Right stick X -> yaw accumulator (``+= PICO_YAW_GAIN * -rx * dt``), dead-zoned. - - Left stick -> magnitude ``mag`` (dead-zone-rescaled to 0..1) driving a - *mode-dependent* speed curve (slow ``0.1+0.5·mag``, run ``1.5+3·mag``, walk = - default), and a facing-rotated global movement vector -> ``movement_angle``. - - Signs replicate gear_sonic's ``get_controller_axes`` usage exactly (forward = +ly, - strafe = -lx, turn = -rx), which resolves the PICO axis-sign question: because the - publisher forwards the same raw SDK axes gear_sonic reads, matching its sign usage - is by definition the correct mapping. ``ms.speed`` is later re-clamped into each - mode's valid range by :func:`clamp_mode_params`. - """ - lx, ly, rx, ry = float(lx), float(ly), float(rx), float(ry) - - # Facing: gear_sonic YawAccumulator — integrate only outside the dead zone. - if abs(rx) >= PICO_JOY_DEADZONE: - ms.facing_angle += PICO_YAW_GAIN * (-rx) * dt - - raw_mag = min(1.0, math.hypot(lx, ly)) - if raw_mag < PICO_JOY_DEADZONE: - ms.has_movement = False - ms.speed = -1.0 - ms.joy_prev_active = False - return - - mag = min(1.0, (raw_mag - PICO_JOY_DEADZONE) / (1.0 - PICO_JOY_DEADZONE)) - m = LM(ms.mode) - if m == LM.SLOW_WALK: - ms.speed = 0.1 + 0.5 * mag # 0.1 .. 0.6 - elif m == LM.WALK: - ms.speed = -1.0 # planner default - elif m == LM.RUN: - ms.speed = 1.5 + 3.0 * mag # 1.5 .. 4.5 - else: - ms.speed = mag - - # Facing-rotated movement vector: rotation_facing @ [-lx, ly] * scale. - scale = mag / raw_mag - mlx, mly = -lx * scale, ly * scale - fx, fy = math.cos(ms.facing_angle), math.sin(ms.facing_angle) - gx = -fy * mlx + fx * mly - gy = fx * mlx + fy * mly - if not ms.joy_prev_active: - ms.needs_replan = True - ms.has_movement = True - ms.joy_prev_active = True - ms.movement_angle = math.atan2(gy, gx) - - -def process_joystick(obs, ms, controller=None): - """Drive ``MovementState`` from the G1 wireless remote in ``obs``.""" - wr = obs.get("wireless_remote") - if wr is None: - return - parsed = _parse_wireless(wr) - if parsed is None: - return - lx, ly, rx, ry = parsed - apply_joystick_axes(lx, ly, rx, ry, ms, controller) 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 120550504..921e00cb9 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -19,51 +19,29 @@ from __future__ import annotations import logging -import math from collections import deque from typing import TYPE_CHECKING import numpy as np from huggingface_hub import hf_hub_download -from lerobot.teleoperators.pico_headset.smpl_constants import ( - LOCO_AXES_PREFIX, - LOCO_BTN_PREFIX, - LOCO_N_AXES, - LOCO_N_BTN, - ROOT_ACTION_DIM, - ROOT_ACTION_PREFIX, - SMPL_ACTION_PREFIX, - SMPL_OBS_DIM as SMPL_ACTION_DIM, - VR3_ORN_DIM, - VR3_ORN_PREFIX, - VR3_POS_DIM, - VR3_POS_PREFIX, - WB_ACTION_DIM, - wb_action_key, -) from lerobot.utils.import_utils import _onnxruntime_available, require_package -from ..g1_utils import MUJOCO_TO_ISAACLAB, KEYBOARD_KEYS_FIELD, G1_29_JointIndex, lowstate_to_obs +from ..g1_utils import ( + MUJOCO_TO_ISAACLAB, + WB_ACTION_DIM, + G1_29_JointIndex, + lowstate_to_obs, + wb_action_key, +) from .sonic_pipeline import ( CONTROL_DT, - DEBUG_PRINT_EVERY, DEFAULT_ANGLES, - DEFAULT_HEIGHT, ENCODER_UPDATE_EVERY, - LM, - MOTION_SETS, - MovementState, PlannerController, - SonicPlanner, - apply_pico_loco_axes, - clamp_mode_params, compute_kp_kd, make_ort_session_options, ort_providers, - process_joystick, - should_replan_request, - snapshot_ms, ) if TYPE_CHECKING or _onnxruntime_available: @@ -79,62 +57,6 @@ logger = logging.getLogger(__name__) INIT_RAMP_S = 3.0 -def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None: - """Reassemble a (720,) SMPL window from ``smpl.{i}`` action keys, or None. - - The pico_headset teleoperator emits the whole-body reference as flat floats so - it flows unchanged through the standard lerobot action pipeline. - """ - # The keys are smpl.0 .. smpl.719; presence of the first element (smpl.0) is the - # sentinel that a full SMPL window was sent this tick. If it's absent, there's no - # whole-body reference, so bail out. - if not action or f"{SMPL_ACTION_PREFIX}0" not in action: - return None - arr = np.fromiter( - (float(action.get(f"{SMPL_ACTION_PREFIX}{i}", 0.0)) for i in range(SMPL_ACTION_DIM)), - dtype=np.float32, - count=SMPL_ACTION_DIM, - ) - return arr - - -def _extract_root_from_action(action: dict | None) -> np.ndarray | None: - """Reassemble a (4,) SMPL root quaternion (wxyz) from ``root.{i}`` keys, or None.""" - if not action or f"{ROOT_ACTION_PREFIX}0" not in action: - return None - q = np.fromiter( - (float(action.get(f"{ROOT_ACTION_PREFIX}{i}", 0.0)) for i in range(ROOT_ACTION_DIM)), - dtype=np.float32, - count=ROOT_ACTION_DIM, - ) - n = float(np.linalg.norm(q)) - if n < 1e-6: - return None - return q / n - - -def _extract_vr3_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None: - """Reassemble the 3-point VR targets from ``vr3_pos.{i}`` / ``vr3_orn.{i}`` keys. - - Returns ``(pos (9,), orn (12,))`` for the [l-wrist, r-wrist, neck] keypoints, or - None when no VR3 reference was sent this tick. Presence of ``vr3_pos.0`` is the - sentinel that a full 3-point frame is available (mirrors the SMPL sentinel). - """ - if not action or f"{VR3_POS_PREFIX}0" not in action: - return None - pos = np.fromiter( - (float(action.get(f"{VR3_POS_PREFIX}{i}", 0.0)) for i in range(VR3_POS_DIM)), - dtype=np.float32, - count=VR3_POS_DIM, - ) - orn = np.fromiter( - (float(action.get(f"{VR3_ORN_PREFIX}{i}", 0.0)) for i in range(VR3_ORN_DIM)), - dtype=np.float32, - count=VR3_ORN_DIM, - ) - return pos, orn - - def _extract_wb34_from_action(action: dict | None) -> np.ndarray | None: """Reassemble a dense (34,) whole-body command from ``wb.{i}.pos`` keys, or None. @@ -180,33 +102,15 @@ def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]: return ref, anchor -def _extract_loco_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None: - """Reassemble controller-stick locomotion from ``loco_axes.{i}`` / ``loco_btn.{i}``. - - Returns ``(axes (4,) = [lx, ly, rx, ry], buttons (4,) = [A, B, X, Y])`` or None - when no locomotion state was sent this tick (sentinel: ``loco_axes.0``). - """ - if not action or f"{LOCO_AXES_PREFIX}0" not in action: - return None - axes = np.fromiter( - (float(action.get(f"{LOCO_AXES_PREFIX}{i}", 0.0)) for i in range(LOCO_N_AXES)), - dtype=np.float32, - count=LOCO_N_AXES, - ) - buttons = np.fromiter( - (float(action.get(f"{LOCO_BTN_PREFIX}{i}", 0.0)) for i in range(LOCO_N_BTN)), - dtype=np.float32, - count=LOCO_N_BTN, - ) - return axes, buttons - - class SonicRuntime: - """Shared SONIC control loop state (standalone demo + locomotion controller).""" + """Loads the SONIC encoder/decoder ONNX models and owns the controller. + + No motion planner: the reference motion buffer is written directly each tick by + :class:`SonicWholeBodyController` from the incoming 34-D whole-body command. + """ def __init__(self, force_cpu: bool = False): require_package("onnxruntime", extra="unitree_g1") - planner_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="planner_sonic.onnx") encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx") decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx") @@ -214,67 +118,22 @@ class SonicRuntime: self.use_gpu = providers[0] == "CUDAExecutionProvider" so = make_ort_session_options() - planner_sess = ort.InferenceSession(planner_path, sess_options=so, providers=providers) encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers) decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers) self.kp, self.kd = compute_kp_kd() - self.ms = MovementState() - self.planner = SonicPlanner(planner_sess, planner_path) - self.controller = PlannerController(self.planner, encoder_sess, decoder_sess) - - motion = self.planner.initialize(DEFAULT_ANGLES, self.ms) - self.controller.load_initial_motion(motion) - self.planner.start_subprocess(self.controller, use_gpu=self.use_gpu) - - self.step = 0 - self.replan_timer = 0.0 - self.last_ms = snapshot_ms(self.ms) + self.controller = PlannerController(encoder_sess, decoder_sess) @property def pipeline(self): return self.controller - def tick(self, obs: dict, *, debug: bool | None = None, use_joystick: bool = True) -> dict: - if not obs: - self.step += 1 - return {} - - if use_joystick: - process_joystick(obs, self.ms, self.controller) - clamp_mode_params(self.ms) - - if self.step > 0: - self.replan_timer += CONTROL_DT - if should_replan_request(self.ms, self.last_ms, self.replan_timer, self.step): - self.planner.request_replan(self.controller.ref_cursor, self.ms) - self.replan_timer = 0.0 - self.ms.needs_replan = False - self.last_ms = snapshot_ms(self.ms) - - do_enc = self.step % ENCODER_UPDATE_EVERY == 0 - if debug is None: - debug = self.step % DEBUG_PRINT_EVERY == 0 - action = self.controller.step(obs, update_encoder=do_enc, debug=debug) - - result = self.planner.try_get_new_motion() - if result: - self.controller.blend_new_motion(*result) - - self.controller.advance_cursor() - self.step += 1 - return action - def reset(self): - self.ms = MovementState() self.controller.reinit_heading = True self.controller.playing = True - self.step = 0 - self.replan_timer = 0.0 - self.last_ms = snapshot_ms(self.ms) def shutdown(self): - self.planner.stop_subprocess() + pass class SonicWholeBodyController: @@ -287,41 +146,12 @@ class SonicWholeBodyController: # directly with a 34-D VLA policy. wb_action = True - def __init__( - self, - force_cpu: bool = False, - *, - enable_smpl_root: bool = False, - root_smoothing_alpha: float = 0.15, - enable_smpl_stream: bool = False, - smpl_host: str | None = None, - smpl_port: int | None = None, - ): + def __init__(self, force_cpu: bool = False): logger.info("Loading SONIC whole-body controller...") self._runtime = SonicRuntime(force_cpu=force_cpu) self.kp = self._runtime.kp self.kd = self._runtime.kd self.controller = self._runtime.controller - self.ms = self._runtime.ms - - # When True, the per-frame SMPL root quaternion steers the mode-2 anchor. - # Off by default: even with smoothing this changes the anchor/heading and is - # untested on hardware, so it stays opt-in. When enabled, the raw per-frame - # root quat (from a 30 Hz dataset resampled to a 50 Hz loop) is spherically - # smoothed by :meth:`_smooth_root_quat` before it reaches the anchor, which - # removes the root-acceleration spikes (NaN QACC at DOF 0) the unsmoothed - # trajectory caused. ``root_smoothing_alpha`` in (0, 1] is the per-tick blend - # toward the incoming quat (smaller = smoother/laggier, 1 = no smoothing). - self.enable_smpl_root = enable_smpl_root - self._root_smoothing_alpha = float(np.clip(root_smoothing_alpha, 1e-3, 1.0)) - self._smoothed_root_quat: np.ndarray | None = None - - # Tracks the previous keyboard held-key set so discrete controls (mode, - # motion set, replan, e-stop, WASD direction) fire once per physical press - # instead of every 50 Hz tick while the key is held. - self._prev_keys: set[str] = set() - # Edge state for the PICO A+B / X+Y locomotion-mode cycle (3-point teleop). - self._prev_loco_mode_pair: tuple[bool, bool] = (False, False) # Startup blend: ease from the robot's initial pose into the first commanded # policy targets over INIT_RAMP_S (captured on the first control tick). @@ -336,173 +166,7 @@ class SonicWholeBodyController: self._wb_traj: deque[np.ndarray] = deque(maxlen=50) self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50) - # Optional: subscribe directly to the rt/smpl headset stream so full-body - # teleop works with ANY teleoperator (e.g. --teleop.type=unitree_g1 for the - # estop/joystick) before the dedicated pico_headset teleop exists. - self._smpl_host = smpl_host - self._smpl_port = smpl_port - self._smpl_stream = None - if enable_smpl_stream: - self._init_smpl_stream() - - logger.info( - "SONIC ready: %s (default mode: %s, smpl_stream=%s)", - MOTION_SETS[0][0], - LM(self.ms.mode).name, - self._smpl_stream is not None, - ) - - def _init_smpl_stream(self) -> None: - # Lazy import so the zmq dependency is only required when streaming is on. - from lerobot.teleoperators.pico_headset.smpl_stream import ( - DEFAULT_SMPL_HOST, - DEFAULT_SMPL_PORT, - SmplStream, - ) - - host = self._smpl_host or DEFAULT_SMPL_HOST - port = self._smpl_port or DEFAULT_SMPL_PORT - self._smpl_stream = SmplStream(host=host, port=port) - logger.info("SONIC subscribed to rt/smpl @ tcp://%s:%d", host, port) - - def _enter_wholebody(self) -> None: - """Switch into SMPL whole-body tracking (encode_mode 2).""" - self.controller.encode_mode = 2 - self.controller.reinit_heading = True - logger.info("SONIC: SMPL stream active -> whole-body tracking (mode 2)") - - def _enter_3point(self) -> None: - """Switch into 3-point VR upper-body teleop (encode_mode 1). - - The upper body tracks the VR wrist/neck targets while the lower body / - locomotion keeps running off the planner (joystick/keyboard-driven). - """ - self.controller.encode_mode = 1 - self.controller.playing = True - self.controller.reinit_heading = True - self.ms.needs_replan = True - logger.info("SONIC: 3-point VR active -> upper-body tracking + planner locomotion (mode 1)") - - def _exit_wholebody(self) -> None: - """Revert to locomotion/standing (encode_mode 0) after a teleop reference is lost. - - Mirrors the 'M' toggle in sonic.py so the handoff is clean: the robot holds - a standing reference and (if a joystick teleop is attached) can be driven. - """ - self.controller.encode_mode = 0 - self.controller.playing = True - self.controller.reinit_heading = True - self.ms.needs_replan = True - logger.warning("SONIC: teleop reference lost/stale -> reverting to locomotion (standing)") - - def _process_keyboard(self, action: dict | None) -> None: - """Translate a native KeyboardTeleop's held-key set into MovementState. - - Mirrors the standalone SONIC demo's keyboard mapping so locomotion (mode 0/1) - can be driven with ``--teleop.type=keyboard`` instead of the PICO SMPL stream. - Discrete controls act on newly-pressed keys (edge-detected against the previous - tick); inherently-continuous controls (facing turn, height, speed) integrate a - small per-tick delta while the key is held so they feel smooth at 50 Hz. - - Controls: WASD move, Q/E turn, 1-8 select mode, 9/0 speed down/up, - -/= height down/up, R replan, Space emergency-stop -> IDLE. - """ - if action is None: - return - keys = action.get(KEYBOARD_KEYS_FIELD) - if keys is None: - return # No KeyboardTeleop attached; leave joystick/SMPL paths untouched. - - ms, controller = self.ms, self.controller - held = {k.lower() if isinstance(k, str) and len(k) == 1 else k for k in keys} - prev = self._prev_keys - pressed = held - prev # newly-pressed this tick (edge) - self._prev_keys = held - - # ── Discrete: fire once per press ──────────────────────────────────── - if "space" in pressed: - ms.mode = LM.IDLE - ms.speed = ms.height = -1.0 - ms.has_movement = False - ms.needs_replan = True - controller.playing = False - controller.reinit_heading = True - logger.info("SONIC keyboard: EMERGENCY STOP -> IDLE") - if "r" in pressed: - ms.needs_replan = True - if "n" in pressed or "p" in pressed: - step = 1 if "n" in pressed else -1 - ms.motion_set_idx = (ms.motion_set_idx + step) % len(MOTION_SETS) - logger.info("SONIC keyboard: motion set -> %s", MOTION_SETS[ms.motion_set_idx][0]) - for digit in ("1", "2", "3", "4", "5", "6", "7", "8"): - if digit in pressed: - idx = int(digit) - 1 - modes = MOTION_SETS[ms.motion_set_idx][1] - if 0 <= idx < len(modes): - ms.mode = modes[idx] - ms.has_movement = False - ms.needs_replan = True - controller.playing = True - controller.reinit_heading = True - logger.info("SONIC keyboard: mode -> %s", LM(ms.mode).name) - # WASD sets the movement direction relative to current facing (press to set, - # Space to stop) to match the standalone demo. - if "w" in pressed: - ms.movement_angle = ms.facing_angle - elif "s" in pressed: - ms.movement_angle = ms.facing_angle + math.pi - elif "a" in pressed: - ms.movement_angle = ms.facing_angle + math.pi / 2 - elif "d" in pressed: - ms.movement_angle = ms.facing_angle - math.pi / 2 - if pressed & {"w", "a", "s", "d"}: - ms.has_movement = True - ms.needs_replan = True - - # ── Continuous: integrate a small delta while held ─────────────────── - if "q" in held: - ms.facing_angle += 0.02 - controller.delta_heading += 0.02 - if "e" in held: - ms.facing_angle -= 0.02 - controller.delta_heading -= 0.02 - if "0" in held: - ms.speed = min(5.0, (ms.speed if ms.speed >= 0 else 1.0) + 0.02) - if "9" in held: - ms.speed = max(0.0, (ms.speed if ms.speed >= 0 else 1.0) - 0.02) - if "=" in held: - ms.height = min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + 0.005) - if "-" in held: - ms.height = max(0.1, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) - 0.005) - - def _process_pico_loco(self, axes: np.ndarray, buttons: np.ndarray) -> None: - """Drive locomotion from the PICO controller sticks/buttons (encode_mode 1). - - Mirrors gear_sonic's ``PlannerLoop`` VR-3PT tick: left/right sticks steer - movement/facing/speed via :func:`apply_pico_loco_axes` (the faithful gear_sonic - yaw-accumulator + mode-dependent speed curves, not the keyboard-parity map), and - A+B / X+Y edge-cycle the locomotion mode within the current motion set. - """ - lx, ly, rx, ry = (float(v) for v in axes) - apply_pico_loco_axes(lx, ly, rx, ry, self.ms) - - # Mode cycling: step linearly through the LocomotionMode enum (A+B = next, - # X+Y = previous), exactly like gear_sonic's PlannerLoop — so the operator can - # reach squat/kneel/crawl, not just the modes in one UI motion set. - a, b, x, y = (v > 0.5 for v in buttons) - ab_now, xy_now = (a and b), (x and y) - ab_prev, xy_prev = self._prev_loco_mode_pair - mode = int(self.ms.mode) - if ab_now and not ab_prev: - mode = min(int(LM.INJURED_WALK), mode + 1) - elif xy_now and not xy_prev: - mode = max(int(LM.IDLE), mode - 1) - if mode != int(self.ms.mode): - self.ms.mode = LM(mode) - self.ms.needs_replan = True - self.controller.playing = True - logger.info("SONIC 3-point: locomotion mode -> %s", LM(self.ms.mode).name) - self._prev_loco_mode_pair = (ab_now, xy_now) + logger.info("SONIC ready (encoder/decoder, 34-D whole-body command path)") def _run_wholebody34(self, obs: dict, wb: np.ndarray) -> dict: """Feed a dense 34-D OpenHLM whole-body command as the mode-0 encoder reference. @@ -564,34 +228,6 @@ class SonicWholeBodyController: self._wb_step += 1 return out - def _smooth_root_quat(self, root_quat: np.ndarray | None) -> np.ndarray | None: - """Spherically smooth the per-frame SMPL root quaternion (mode-2 anchor). - - The reference root trajectory is authored at ~30 Hz and consumed at 50 Hz, so - the raw per-tick quat steps unevenly and injects root-acceleration spikes into - the anchor. This keeps a persistent estimate and shortest-path nlerp-slerps it - toward each incoming (unit) quat by ``root_smoothing_alpha``, yielding a - continuous, rate-matched heading. Quaternions are scalar-first (w, x, y, z). - Returns ``None`` (leaving the anchor self-driven) for an invalid/zero input. - """ - if root_quat is None: - self._smoothed_root_quat = None - return None - q = np.asarray(root_quat, np.float64) - n = np.linalg.norm(q) - if n < 1e-8: - return self._smoothed_root_quat - q = q / n - if self._smoothed_root_quat is None: - self._smoothed_root_quat = q - else: - prev = self._smoothed_root_quat - if np.dot(prev, q) < 0.0: # shortest-path: quats double-cover SO(3) - q = -q - blended = prev + self._root_smoothing_alpha * (q - prev) - self._smoothed_root_quat = blended / (np.linalg.norm(blended) + 1e-12) - return self._smoothed_root_quat.astype(np.float32) - def _startup_blend(self, obs: dict, out: dict) -> dict: """Ease into policy control at startup: for the first ``INIT_RAMP_S`` seconds, interpolate between the robot's pose captured on the first tick and the policy's @@ -623,96 +259,29 @@ class SonicWholeBodyController: return {} obs = lowstate_to_obs(lowstate) - # Keyboard teleop (native KeyboardTeleop) drives the same locomotion intent - # the joystick does; applied before the SMPL check so whole-body tracking - # still takes priority when a headset stream is present. - self._process_keyboard(action) - - # Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a - # direct rt/smpl subscription when enabled (enable_smpl_stream). A stale - # stream (headset silent past its timeout) is treated as "no SMPL" so the - # robot doesn't stay frozen tracking the last pose. - # Dense whole-body command (OpenHLM / pi0.5 joint interface) takes priority: - # a single 34-D vector drives the mode-0 joint reference directly. + # Dense 34-D whole-body command (OpenHLM / pi0.5 joint interface): a single + # vector per tick drives the mode-0 encoder reference directly. Until the + # policy produces one, hold (no command) so the robot keeps its last target. wb = _extract_wb34_from_action(action) - if wb is not None: - return self._startup_blend(obs, self._run_wholebody34(obs, wb)) - self._wb_miss = getattr(self, "_wb_miss", 0) + 1 - if self._wb_miss % 50 == 1: - akeys = [k for k in action if isinstance(k, str)] - logger.info( - "[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s", - self._wb_miss, - akeys[:8], - ) - - smpl = _extract_smpl_from_action(action) - root_quat = _extract_root_from_action(action) - vr3 = _extract_vr3_from_action(action) - loco = _extract_loco_from_action(action) - if smpl is None and vr3 is None and self._smpl_stream is not None: - window = self._smpl_stream.step() - if self._smpl_stream.has_data and not self._smpl_stream.is_stale: - smpl = window - root_quat = np.asarray(self._smpl_stream.root_quat, np.float32) - # VR3 is independent of the SMPL window: the controller-state source - # (head + controllers only) sends 3-point targets with no SMPL frame. - elif self._smpl_stream.has_fresh_vr3: - vr3 = (self._smpl_stream.vr3_pos, self._smpl_stream.vr3_orn) - if self._smpl_stream.has_fresh_loco: - loco = (self._smpl_stream.loco_axes, self._smpl_stream.loco_buttons) - - if smpl is not None: - # Full-body whole-body tracking: SMPL drives the reference, not joystick. - if self.controller.encode_mode != 2: - self._enter_wholebody() - self.controller.smpl_joints_10frame_step1 = smpl - # Root orientation steers the mode-2 anchor/heading, but only when - # explicitly enabled (see enable_smpl_root); the raw per-frame quat is - # spherically smoothed first so the 30->50 Hz resample doesn't spike the - # anchor. Disabled -> anchor stays self-driven. - self.controller.smpl_root_quat = ( - self._smooth_root_quat(root_quat) if self.enable_smpl_root else None - ) - out = self._runtime.tick(obs, debug=False, use_joystick=False) - elif vr3 is not None: - # 3-point VR teleop: upper body tracks the wrist/neck targets; the lower - # body / locomotion keeps running off the planner, so the joystick (and - # keyboard) still steer walking/turning underneath. - if self.controller.encode_mode != 1: - self._enter_3point() - self.controller.vr_3point_local_target = vr3[0] - self.controller.vr_3point_local_orn_target = vr3[1] - # Replicate the original encode_mode-1 handling: when the PICO controller - # sticks are forwarded, drive locomotion from them directly (and skip the - # wireless-remote joystick read). Otherwise leave the remote/keyboard path. - if loco is not None: - self._process_pico_loco(loco[0], loco[1]) - out = self._runtime.tick(obs, debug=False, use_joystick=False) - else: - out = self._runtime.tick(obs, debug=False, use_joystick=True) - else: - # No (or stale) teleop reference: fall back to locomotion so the robot stays balanced. - if self.controller.encode_mode != 0: - self.controller.smpl_root_quat = None - self._smoothed_root_quat = None - self._exit_wholebody() - out = self._runtime.tick(obs, debug=False) - - # Startup interpolation: blend from the robot's initial pose into the policy's - # commanded target over INIT_RAMP_S, regardless of mode. - return self._startup_blend(obs, out) + if wb is None: + self._wb_miss = getattr(self, "_wb_miss", 0) + 1 + if self._wb_miss % 50 == 1: + akeys = [k for k in action if isinstance(k, str)] + logger.info( + "[WB34] no wb.*.pos in action this tick (miss=%d). action keys sample: %s", + self._wb_miss, + akeys[:8], + ) + return {} + return self._startup_blend(obs, self._run_wholebody34(obs, wb)) def reset(self): self._runtime.reset() self._init_step = 0 # re-run the startup blend after a reset self._start_pose = {} - self._smoothed_root_quat = None self._wb_step = 0 self._wb_traj.clear() self._wb_quat_traj.clear() def shutdown(self): - if self._smpl_stream is not None: - self._smpl_stream.close() self._runtime.shutdown() diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index 5349aec23..9c9caf9cb 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -288,7 +288,7 @@ class UnitreeG1(Robot): def _empty_cameras_ft(self) -> dict[str, tuple]: """Synthetic zero-image cameras (see ``UnitreeG1Config.empty_cameras``).""" h, w = self.config.empty_camera_hw - return {name: (h, w, 3) for name in self.config.empty_cameras} + return dict.fromkeys(self.config.empty_cameras, (h, w, 3)) @property def _replay_cameras_ft(self) -> dict[str, tuple]: diff --git a/src/lerobot/scripts/lerobot_teleoperate.py b/src/lerobot/scripts/lerobot_teleoperate.py index c52b22f73..30f13987e 100644 --- a/src/lerobot/scripts/lerobot_teleoperate.py +++ b/src/lerobot/scripts/lerobot_teleoperate.py @@ -116,7 +116,6 @@ from lerobot.teleoperators import ( # noqa: F401 omx_leader, openarm_leader, openarm_mini, - pico_headset, reachy2_teleoperator, rebot_102_leader, so_leader, diff --git a/src/lerobot/teleoperators/pico_headset/__init__.py b/src/lerobot/teleoperators/pico_headset/__init__.py deleted file mode 100644 index b800c4879..000000000 --- a/src/lerobot/teleoperators/pico_headset/__init__.py +++ /dev/null @@ -1,20 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from .config_pico_headset import PicoHeadsetConfig -from .pico_headset import PicoHeadset - -__all__ = ["PicoHeadset", "PicoHeadsetConfig"] diff --git a/src/lerobot/teleoperators/pico_headset/config_pico_headset.py b/src/lerobot/teleoperators/pico_headset/config_pico_headset.py deleted file mode 100644 index 81269cf53..000000000 --- a/src/lerobot/teleoperators/pico_headset/config_pico_headset.py +++ /dev/null @@ -1,41 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from dataclasses import dataclass - -from ..config import TeleoperatorConfig - - -@TeleoperatorConfig.register_subclass("pico_headset") -@dataclass -class PicoHeadsetConfig(TeleoperatorConfig): - """PICO full-body headset teleop: live SMPL over the rt/smpl ZMQ stream. - - Consumes the ``rt/smpl`` channel published by the GEAR PICO manager - (``gear_sonic/scripts/pico_manager_thread_server.py``) and emits the whole-body - SONIC reference window (``encode_mode == 2``) for SonicWholeBodyController. - """ - - smpl_host: str = "127.0.0.1" - """Host of the pico_manager rt/smpl publisher (the laptop bridging the PICO).""" - smpl_port: int = 5560 - """Port of the rt/smpl publisher.""" - stale_after_s: float = 0.5 - """Warn if no fresh headset frame arrives within this many seconds.""" - mode: str = "smpl" - """Teleop reference to emit: ``"smpl"`` for whole-body imitation (SONIC - encode_mode 2) or ``"vr3"`` for sparse 3-point upper-body teleop (encode_mode 1, - lower body driven by the joystick/keyboard planner).""" diff --git a/src/lerobot/teleoperators/pico_headset/pico_headset.py b/src/lerobot/teleoperators/pico_headset/pico_headset.py deleted file mode 100644 index 3dbba0404..000000000 --- a/src/lerobot/teleoperators/pico_headset/pico_headset.py +++ /dev/null @@ -1,148 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""PICO full-body headset teleoperator (live SMPL -> SONIC whole-body).""" - -import logging -from typing import Any - -from lerobot.types import RobotAction - -from ..teleoperator import Teleoperator -from .config_pico_headset import PicoHeadsetConfig -from .smpl_constants import ( - LOCO_AXES_PREFIX, - LOCO_BTN_PREFIX, - LOCO_N_AXES, - LOCO_N_BTN, - ROOT_ACTION_DIM, - ROOT_ACTION_PREFIX, - SMPL_ACTION_PREFIX, - SMPL_OBS_DIM, - VR3_ORN_DIM, - VR3_ORN_PREFIX, - VR3_POS_DIM, - VR3_POS_PREFIX, -) -from .smpl_stream import SmplStream - -logger = logging.getLogger(__name__) - - -class PicoHeadset(Teleoperator): - """Streams full-body SMPL from a PICO headset as a SONIC whole-body reference. - - Subscribes to the ``rt/smpl`` ZMQ channel and, once real frames are flowing, - emits the 720-element encoder window as ``smpl.{i}`` floats. Before the first - frame arrives it emits no SMPL keys, so the robot stays in safe locomotion mode - rather than tracking a zero pose. - """ - - config_class = PicoHeadsetConfig - name = "pico_headset" - - def __init__(self, config: PicoHeadsetConfig): - super().__init__(config) - self.config = config - self._stream: SmplStream | None = None - - @property - def action_features(self) -> dict: - if self.config.mode == "vr3": - feats = {f"{VR3_POS_PREFIX}{i}": float for i in range(VR3_POS_DIM)} - feats.update({f"{VR3_ORN_PREFIX}{i}": float for i in range(VR3_ORN_DIM)}) - # Controller-stick locomotion travels alongside the VR targets. - feats.update({f"{LOCO_AXES_PREFIX}{i}": float for i in range(LOCO_N_AXES)}) - feats.update({f"{LOCO_BTN_PREFIX}{i}": float for i in range(LOCO_N_BTN)}) - return feats - feats = {f"{SMPL_ACTION_PREFIX}{i}": float for i in range(SMPL_OBS_DIM)} - feats.update({f"{ROOT_ACTION_PREFIX}{i}": float for i in range(ROOT_ACTION_DIM)}) - return feats - - @property - def feedback_features(self) -> dict: - return {} - - @property - def is_connected(self) -> bool: - return self._stream is not None - - def connect(self, calibrate: bool = True) -> None: - if self._stream is not None: - raise RuntimeError(f"{self} already connected") - self._stream = SmplStream( - host=self.config.smpl_host, - port=self.config.smpl_port, - stale_after_s=self.config.stale_after_s, - ) - logger.info( - "PicoHeadset subscribed to rt/smpl @ tcp://%s:%d", - self.config.smpl_host, - self.config.smpl_port, - ) - - @property - def is_calibrated(self) -> bool: - return True - - def calibrate(self) -> None: - # Calibration happens on the headset / pico_manager side, not here. - pass - - def configure(self) -> None: - pass - - def get_action(self) -> RobotAction: - if self._stream is None: - raise RuntimeError(f"{self} is not connected") - window = self._stream.step() - # Emit a reference only while the headset is actively streaming: hold back - # before the first frame (so the controller doesn't track an all-zero - # collapsed pose) and once the stream goes stale (so the controller falls - # back to a safe standing/locomotion mode instead of freezing on the last - # pose). - if self.config.mode == "vr3": - # Sparse 3-point upper-body teleop (encode_mode 1). Gated on fresh vr3_* - # frames only (independent of the SMPL window), so the controller-state - # source (head + controllers, no body tracking) works. Emit nothing - # otherwise and stay in locomotion. - if not self._stream.has_fresh_vr3: - return {} - action = {f"{VR3_POS_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.vr3_pos)} - action.update({f"{VR3_ORN_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.vr3_orn)}) - # Forward controller-stick locomotion when present, so the planner can - # steer walking/turning under the upper-body tracking (encode_mode 1). - if self._stream.has_fresh_loco: - action.update( - {f"{LOCO_AXES_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.loco_axes)} - ) - action.update( - {f"{LOCO_BTN_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.loco_buttons)} - ) - return action - if not self._stream.has_data or self._stream.is_stale: - return {} - action = {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)} - action.update({f"{ROOT_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(self._stream.root_quat)}) - return action - - def send_feedback(self, feedback: dict[str, Any]) -> None: - pass - - def disconnect(self) -> None: - if self._stream is not None: - self._stream.close() - self._stream = None diff --git a/src/lerobot/teleoperators/pico_headset/pico_publisher.py b/src/lerobot/teleoperators/pico_headset/pico_publisher.py deleted file mode 100644 index e4d249ff6..000000000 --- a/src/lerobot/teleoperators/pico_headset/pico_publisher.py +++ /dev/null @@ -1,333 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Standalone ``rt/smpl`` publisher for the PICO headset (no gear_sonic / torch). - -Reads 24 body-joint poses from the XRoboToolkit SDK, runs pure-numpy SMPL forward -kinematics + canonicalization (``smpl_fk.py``), and publishes one canonical -``(24, 3)`` SMPL frame per tick over ZMQ on the ``rt/smpl`` topic — the exact -message ``lerobot.teleoperators.pico_headset.smpl_stream.SmplStream`` consumes. - -This makes the LeRobot side self-contained: the only runtime dependency to drive -SONIC whole-body teleop from the headset is the ``xrobotoolkit_sdk`` Python package -(plus numpy/scipy/pyzmq), not the full ``gear_sonic`` stack. - -Usage: - # Real headset (XRoboToolkit PC Service must be running and connected): - python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --port 5560 - - # No hardware — emit a synthetic waving motion to test the consumer end-to-end: - python -m lerobot.teleoperators.pico_headset.pico_publisher --fake - - # Replay a canned SMPL clip to the robot through the same rt/smpl -> SONIC path: - python -m lerobot.teleoperators.pico_headset.pico_publisher \ - --motion-file examples/unitree_g1/motions/walk_forward.npz -""" - -from __future__ import annotations - -import argparse -import contextlib -import json -import time - -import numpy as np -import zmq - -from lerobot.teleoperators.pico_headset.smpl_fk import ( - SmplForwardKinematics, - ThreePointCalibrator, - canonicalize_smpl_joints, - compute_3point, - compute_3point_from_devices, - root_quats_from_aa, -) - -SMPL_TOPIC = "rt/smpl" -DEFAULT_SMPL_PORT = 5560 - - -def pack_message( - smpl_joints_local: np.ndarray, - frame_index: int, - stamp_ns: int, - root_quat: np.ndarray | None = None, - root_transl: np.ndarray | None = None, - vr3_pos: np.ndarray | None = None, - vr3_orn: np.ndarray | None = None, - loco_axes: np.ndarray | None = None, - loco_buttons: np.ndarray | None = None, -) -> bytes: - """Build the rt/smpl JSON message (single frame, topic embedded in payload). - - Carries the SMPL whole-body window (``smpl_joints_local`` + ``root_*``) and, - when available, the sparse 3-point VR targets (``vr3_pos`` (9,), ``vr3_orn`` (12,)) - so a single stream can drive either SONIC ``encode_mode`` 1 or 2. - """ - data = { - "smpl_joints_local": np.asarray(smpl_joints_local, np.float32).reshape(-1).tolist(), - "frame_index": int(frame_index), - "stamp_ns": int(stamp_ns), - } - if root_quat is not None: - data["root_quat"] = np.asarray(root_quat, np.float32).reshape(-1).tolist() - if root_transl is not None: - data["root_transl"] = np.asarray(root_transl, np.float32).reshape(-1).tolist() - if vr3_pos is not None: - data["vr3_pos"] = np.asarray(vr3_pos, np.float32).reshape(-1).tolist() - if vr3_orn is not None: - data["vr3_orn"] = np.asarray(vr3_orn, np.float32).reshape(-1).tolist() - if loco_axes is not None: - data["loco_axes"] = np.asarray(loco_axes, np.float32).reshape(-1).tolist() - if loco_buttons is not None: - data["loco_buttons"] = np.asarray(loco_buttons, np.float32).reshape(-1).tolist() - return json.dumps({"topic": SMPL_TOPIC, "data": data}).encode("utf-8") - - -def _fake_body_poses(t: float) -> np.ndarray: - """Synthetic (24, 7) body poses: identity rotations + a gently waving right arm.""" - poses = np.zeros((24, 7), np.float32) - poses[:, 6] = 1.0 # unit quaternion (qw = 1), scalar-last - poses[:, 1] = 1.0 # ~1 m pelvis height (positions only matter for root_transl) - # Wave the right shoulder (SMPL body joint 17) about Z. - ang = 0.5 * np.sin(2.0 * np.pi * 0.5 * t) - poses[17, 3:7] = [0.0, 0.0, np.sin(ang / 2), np.cos(ang / 2)] - return poses - - -def _load_motion_clip(path: str) -> dict: - """Load an SMPL ``.npz`` clip and canonicalize it for rt/smpl streaming. - - Expects the same keys as ``motion_loader.SmplMotion``: - smpl_joints (T, 24, 3), pose_aa (T, 72) optional, transl (T, 3) optional. - Returns per-frame joints already in the encoder's root-removed convention, - plus optional per-frame root quat/translation. - """ - data = np.load(path) - joints = data["smpl_joints"].astype(np.float32) - if joints.ndim != 3 or joints.shape[1:] != (24, 3): - raise ValueError(f"Expected smpl_joints (T, 24, 3), got {joints.shape}") - - pose_aa = data["pose_aa"].astype(np.float32) if "pose_aa" in data.files else None - root_quat = None - if pose_aa is not None: - joints = canonicalize_smpl_joints(joints, pose_aa[:, :3]) - root_quat = root_quats_from_aa(pose_aa[:, :3]) - transl = data["transl"].astype(np.float32) if "transl" in data.files else None - return {"joints": joints, "root_quat": root_quat, "transl": transl} - - -def main() -> None: - p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--port", type=int, default=DEFAULT_SMPL_PORT, help="ZMQ PUB port for rt/smpl") - p.add_argument("--fps", type=float, default=50.0, help="Target publish rate (Hz)") - p.add_argument("--skeleton", type=str, default=None, help="Path to smpl_skeleton.npz") - p.add_argument( - "--headset-source", - choices=["body", "devices"], - default="body", - help=( - "Live headset keypoint source: 'body' uses full-body tracking " - "(get_body_joints_pose, needs PICO Motion Trackers) and drives both SMPL " - "(encode_mode 2) and 3-point; 'devices' uses head + 2 controllers only " - "(get_headset_pose + get_*_controller_pose, no trackers) and emits 3-point " - "(encode_mode 1) exclusively." - ), - ) - src = p.add_mutually_exclusive_group() - src.add_argument("--fake", action="store_true", help="Publish synthetic motion (no headset)") - src.add_argument("--motion-file", type=str, default=None, help="Replay an SMPL .npz clip over rt/smpl") - p.add_argument("--no-loop", action="store_true", help="Play a --motion-file once, then stop") - args = p.parse_args() - - clip = _load_motion_clip(args.motion_file) if args.motion_file else None - - # FK is only needed for live/synthetic (24,7) body poses; clips are pre-canonical. - fk = None - if clip is None: - fk = SmplForwardKinematics(args.skeleton) if args.skeleton else SmplForwardKinematics() - - xrt = None - if clip is None and not args.fake: - try: - import xrobotoolkit_sdk as xrt # noqa: PLC0415 - except ImportError as e: - raise SystemExit( - "xrobotoolkit_sdk not available. Install it, or run with --fake / --motion-file " - "to test the pipeline without a headset." - ) from e - xrt.init() - print("[pico_publisher] XRoboToolkit initialized") - - ctx = zmq.Context.instance() - sock = ctx.socket(zmq.PUB) - sock.bind(f"tcp://*:{args.port}") - src_desc = ( - f"motion-file {args.motion_file}" - if clip - else ("fake" if args.fake else f"headset:{args.headset_source}") - ) - print( - f"[pico_publisher] '{SMPL_TOPIC}' bound to tcp://*:{args.port} @ {args.fps:.0f} Hz " - f"[source: {src_desc}]" - ) - if clip is not None: - print(f"[pico_publisher] clip frames={clip['joints'].shape[0]} loop={not args.no_loop}") - - period = 1.0 / max(1.0, args.fps) - frame_index = 0 - last_tracked = -1 - t0 = time.time() - - # 3-point operator calibration (device source only): map the operator's neutral - # rest pose onto the G1's neutral stance. Trigger a (re)capture with the A+B+X+Y - # controller combo, mirroring gear_sonic's ThreePointPose.calibrate_now. - # Device source: the "neck" is the headset (pitches when looking down), so keep the - # wrist targets in the yaw-local world frame rather than de-rotating by head tilt. - calibrator = ( - ThreePointCalibrator(neck_relative_wrists=False) - if (xrt is not None and args.headset_source == "devices") - else None - ) - calib_combo_last = False - if calibrator is not None: - print( - "[pico_publisher] 3-point calibration: stand in a neutral rest pose and press " - "A+B+X+Y on the controllers to (re)calibrate." - ) - try: - while True: - loop_start = time.time() - vr3_pos = vr3_orn = None - loco_axes = loco_buttons = None - if clip is not None: - n = clip["joints"].shape[0] - if args.no_loop and frame_index >= n: - print("\n[pico_publisher] clip finished") - break - i = frame_index % n - joints = clip["joints"][i] - root_quat = None if clip["root_quat"] is None else clip["root_quat"][i] - root_transl = None if clip["transl"] is None else clip["transl"][i] - stamp_ns = time.time_ns() - elif not args.fake and args.headset_source == "devices": - # Controller-state 3-point path: head + 2 controllers only, no PICO - # Motion Trackers / body tracking. Emits encode_mode-1 targets only; - # the SMPL whole-body window is left as a zero placeholder. - head = np.asarray(xrt.get_headset_pose(), np.float32) - lc = np.asarray(xrt.get_left_controller_pose(), np.float32) - rc = np.asarray(xrt.get_right_controller_pose(), np.float32) - stamp_ns = int(xrt.get_time_stamp_ns()) - if head.shape != (7,) or lc.shape != (7,) or rc.shape != (7,): - time.sleep(0.005) - continue - last_tracked = int(sum(np.linalg.norm(d[3:7]) > 1e-6 for d in (head, lc, rc))) - # Empty SMPL window: this source drives encode_mode 1 only, so the - # consumer must not mistake it for a (zero) whole-body reference. - joints = np.zeros((0, 3), np.float32) - root_quat = None - root_transl = None - vr3_pos, vr3_orn = compute_3point_from_devices(head, lc, rc) - # Edge-triggered (re)calibration on the A+B+X+Y combo. - combo_now = bool( - xrt.get_A_button() and xrt.get_B_button() and xrt.get_X_button() and xrt.get_Y_button() - ) - if combo_now and not calib_combo_last: - calibrator.capture(vr3_pos, vr3_orn) - print("\n[pico_publisher] 3-point calibration captured (neutral pose).") - calib_combo_last = combo_now - vr3_pos, vr3_orn = calibrator.apply(vr3_pos, vr3_orn) - # Controller-stick locomotion (encode_mode 1, replicated): left/right - # sticks + A/B/X/Y. The all-four combo is reserved for calibration, so - # suppress the button pairs on that frame to avoid a spurious mode cycle. - la = np.asarray(xrt.get_left_axis(), np.float32).reshape(-1) - ra = np.asarray(xrt.get_right_axis(), np.float32).reshape(-1) - loco_axes = np.array([la[0], la[1], ra[0], ra[1]], np.float32) - btn = ( - np.zeros(4, np.float32) - if combo_now - else np.array( - [ - float(xrt.get_A_button()), - float(xrt.get_B_button()), - float(xrt.get_X_button()), - float(xrt.get_Y_button()), - ], - np.float32, - ) - ) - loco_buttons = btn - else: - if args.fake: - body_poses = _fake_body_poses(loop_start - t0) - stamp_ns = time.time_ns() - else: - body_poses = np.asarray(xrt.get_body_joints_pose(), np.float32) - stamp_ns = int(xrt.get_time_stamp_ns()) - if body_poses.shape != (24, 7): - time.sleep(0.005) - continue - # How many joints are actually being tracked (non-zero-norm quat). - # If this stays near 0, the headset isn't streaming body data (e.g. - # "Full body"/"Send" not enabled, trackers uncalibrated, or a test - # device) and the reference will be static regardless of your motion. - # In that case, use --headset-source devices for head+controllers only. - quat_norms = np.linalg.norm(body_poses[:, 3:7], axis=1) - last_tracked = int(np.count_nonzero(quat_norms > 1e-6)) - out = fk.compute(body_poses) - joints = out["smpl_joints_local"] - root_quat = out["root_quat"] - root_transl = out["root_transl"] - # Also emit the sparse 3-point VR targets so the same stream can - # drive encode_mode 1 (3-point teleop) without a second producer. - vr3_pos, vr3_orn = compute_3point(body_poses) - - sock.send( - pack_message( - joints, - frame_index, - stamp_ns, - root_quat=root_quat, - root_transl=root_transl, - vr3_pos=vr3_pos, - vr3_orn=vr3_orn, - loco_axes=loco_axes, - loco_buttons=loco_buttons, - ) - ) - frame_index += 1 - if frame_index % int(max(1, args.fps)) == 0: - denom = 3 if (not args.fake and clip is None and args.headset_source == "devices") else 24 - unit = "devices" if denom == 3 else "joints" - extra = f" | tracked {last_tracked}/{denom} {unit}" if last_tracked >= 0 else "" - if calibrator is not None: - extra += " | calibrated" if calibrator.is_calibrated else " | UNCALIBRATED" - print(f"[pico_publisher] sent {frame_index} frames{extra}", end="\r") - - dt = time.time() - loop_start - if dt < period: - time.sleep(period - dt) - except KeyboardInterrupt: - print("\n[pico_publisher] stopping") - finally: - sock.close(0) - if xrt is not None: - with contextlib.suppress(Exception): - xrt.close() - - -if __name__ == "__main__": - main() diff --git a/src/lerobot/teleoperators/pico_headset/smpl_constants.py b/src/lerobot/teleoperators/pico_headset/smpl_constants.py deleted file mode 100644 index a59b320df..000000000 --- a/src/lerobot/teleoperators/pico_headset/smpl_constants.py +++ /dev/null @@ -1,80 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Single source of truth for the SONIC SMPL whole-body action protocol. - -These constants define the wire format shared between the PICO teleoperator -(producer, ``pico_headset.py``), the offline motion->dataset converter -(``smpl_to_dataset.py``), the live stream (``smpl_stream.py``), and the Unitree G1 -``SonicWholeBodyController`` (consumer). Keeping them here avoids the producer and -consumer silently drifting apart. -""" - -# SMPL encoder window geometry (matches ``smpl_joints_10frame_step1``). -WINDOW = 10 # frames per encoder window -N_JOINTS = 24 # SMPL joints per frame -JOINT_DIM = 3 # xyz per joint -SMPL_OBS_DIM = WINDOW * N_JOINTS * JOINT_DIM # 720 - -# Flat action-dict keys carrying the reference through the standard lerobot action -# pipeline as scalar floats: ``smpl.0 .. smpl.719`` and ``root.0 .. root.3``. -SMPL_ACTION_PREFIX = "smpl." -ROOT_ACTION_PREFIX = "root." -ROOT_ACTION_DIM = 4 # per-frame SMPL root orientation (wxyz) - -# Full per-frame action vector: 720 joint window + 4 root quaternion = 724. -ACTION_DIM = SMPL_OBS_DIM + ROOT_ACTION_DIM - -# ── 3-point VR teleop protocol (SONIC encode_mode 1) ───────────────────────── -# An alternative, sparse upper-body reference: 3 root-relative keypoints -# (left wrist, right wrist, neck), each a position + orientation. The lower body / -# locomotion is driven by the planner (joystick/keyboard), not by these targets. -VR3_N_POINTS = 3 # left wrist, right wrist, neck -VR3_POS_DIM = VR3_N_POINTS * 3 # 9 (3 x xyz) -VR3_ORN_DIM = VR3_N_POINTS * 4 # 12 (3 x wxyz) - -# Flat action-dict keys: ``vr3_pos.0 .. vr3_pos.8`` and ``vr3_orn.0 .. vr3_orn.11``. -VR3_POS_PREFIX = "vr3_pos." -VR3_ORN_PREFIX = "vr3_orn." - -# ── Controller-stick locomotion (SONIC encode_mode 1, replicated exactly) ──── -# In the original 3-point teleop the same tick that sends the VR targets also drives -# locomotion from the PICO controller sticks/buttons (left stick -> move, right stick -# -> yaw, A+B / X+Y -> cycle locomotion mode). We forward that raw controller state so -# the consumer's planner can steer walking/turning underneath the upper-body tracking. -LOCO_N_AXES = 4 # [left_x, left_y, right_x, right_y] -LOCO_N_BTN = 4 # [A, B, X, Y] -LOCO_AXES_PREFIX = "loco_axes." -LOCO_BTN_PREFIX = "loco_btn." - -# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM-style) ───── -# A single 34-D whole-body command per tick, in the pi0.5 / OpenHLM action layout: -# [L-arm(7), L-grip(1), R-arm(7), R-grip(1), L-leg(6), R-leg(6), waist(3), -# root roll/pitch + yaw-rate(3)] -# The 29 joint targets (arms/legs/waist, grippers excluded) become the mode-0 -# encoder joint reference and root roll/pitch become the anchor orientation. -# -# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix is required so -# these behave like ordinary joint-position action features: ``lerobot-rollout`` -# only routes ``*.pos`` keys from ``robot.action_features`` into the policy<->robot -# action mapping, letting a 34-D VLA (OpenHLM / pi0.5) drive the robot directly. -WB_ACTION_PREFIX = "wb." -WB_ACTION_DIM = 34 - - -def wb_action_key(i: int) -> str: - """Action-dict key for the ``i``-th whole-body command scalar (``wb.{i}.pos``).""" - return f"{WB_ACTION_PREFIX}{i}.pos" diff --git a/src/lerobot/teleoperators/pico_headset/smpl_fk.py b/src/lerobot/teleoperators/pico_headset/smpl_fk.py deleted file mode 100644 index a0791919b..000000000 --- a/src/lerobot/teleoperators/pico_headset/smpl_fk.py +++ /dev/null @@ -1,571 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Standalone SMPL forward kinematics + canonicalization in pure numpy/scipy. - -This mirrors the ``rt/smpl`` producer path in ``gear_sonic`` (``compute_from_body_poses`` --> ``process_smpl_joints`` -> ``compute_human_joints``) without depending on torch, -the SMPL mesh model, or ``gear_sonic``. It only needs a small fixed skeleton table -(SMPL-X rest-pose joints + kinematic tree), hardcoded below as ``_SKELETON_J`` / -``_SKELETON_PARENTS`` so no external asset download is required. - -Given the 24 body-joint poses reported by the XRoboToolkit headset SDK -(``xrt.get_body_joints_pose()`` -> (24, 7) of ``[x, y, z, qx, qy, qz, qw]``), it -produces the root-orientation-removed 24x3 SMPL joints the SONIC encoder expects, -plus the root orientation quaternion and pelvis translation. - -Quaternions are scalar-first (w, x, y, z) unless noted. -""" - -from dataclasses import dataclass, field -from pathlib import Path - -import numpy as np -from scipy.spatial.transform import Rotation as R # noqa: N817 - -from .smpl_constants import VR3_N_POINTS, VR3_ORN_DIM, VR3_POS_DIM - -# 24-joint parent tree used by the headset body-pose stream (SMPL-X body subset). -# Matches PoseStreamer.parent_indices in gear_sonic's pico_manager_thread_server.py. -BODY24_PARENTS = np.array( - [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14, 16, 17, 18, 19, 20, 22], - dtype=np.int64, -) - -# FK output joints: first 22 SMPL body joints + two thumb tips (SMPL-X indices 39, 54). -OUTPUT_JOINT_INDEX = np.concatenate([np.arange(22), np.array([39, 54])]) - -# SMPL-X rest-pose skeleton (55 joints), vendored inline to avoid an external -# asset. ``_SKELETON_PARENTS[i]`` is joint ``i``'s parent (-1 = root); ``_SKELETON_J`` -# holds the (55, 3) rest-pose joint positions. -_SKELETON_PARENTS = np.array( - [ - -1, - 0, - 0, - 0, - 1, - 2, - 3, - 4, - 5, - 6, - 7, - 8, - 9, - 9, - 9, - 12, - 13, - 14, - 16, - 17, - 18, - 19, - 15, - 15, - 15, - 20, - 25, - 26, - 20, - 28, - 29, - 20, - 31, - 32, - 20, - 34, - 35, - 20, - 37, - 38, - 21, - 40, - 41, - 21, - 43, - 44, - 21, - 46, - 47, - 21, - 49, - 50, - 21, - 52, - 53, - ], - dtype=np.int64, -) - -_SKELETON_J = np.array( - [ - [0.0031232605688273907, -0.3514074683189392, 0.012036550790071487], - [0.06131265312433243, -0.4441709518432617, -0.013964635320007801], - [-0.06014421582221985, -0.4553154706954956, -0.009213820099830627], - [0.00036056205863133073, -0.2415168583393097, -0.015581080690026283], - [0.11600811034440994, -0.8229243755340576, -0.02336069941520691], - [-0.10435417294502258, -0.8176955580711365, -0.026037702336907387], - [0.009808260947465897, -0.10966360569000244, -0.02152106538414955], - [0.07255466282367706, -1.2259838581085205, -0.05523664504289627], - [-0.08893736451864243, -1.2284233570098877, -0.046229973435401917], - [-0.0015221529174596071, -0.057428449392318726, 0.006925832014530897], - [0.11981196701526642, -1.283981204032898, 0.06297968327999115], - [-0.12774977087974548, -1.2867517471313477, 0.07281902432441711], - [-0.01368661131709814, 0.10773860663175583, -0.024689510464668274], - [0.04484200477600098, 0.027515273541212082, -0.0002946509048342705], - [-0.04921707883477211, 0.026910223066806793, -0.006474069785326719], - [0.011096873320639133, 0.2681904137134552, -0.003952245227992535], - [0.16408103704452515, 0.08524329960346222, -0.015755590051412582], - [-0.15179482102394104, 0.08043467253446579, -0.019142597913742065], - [0.4182038903236389, 0.01309278141707182, -0.058214444667100906], - [-0.4229443669319153, 0.04394219070672989, -0.04560968279838562], - [0.6701906323432922, 0.03631401062011719, -0.06068652495741844], - [-0.6722118258476257, 0.03940964490175247, -0.06093486770987511], - [-0.004667762666940689, 0.2676706910133362, -0.009591402485966682], - [0.03159928321838379, 0.31083211302757263, 0.062195174396038055], - [-0.031599875539541245, 0.3108319342136383, 0.0621943436563015], - [0.7720924615859985, 0.02762586995959282, -0.04133538901805878], - [0.8040408492088318, 0.02984413132071495, -0.03832494467496872], - [0.8265857696533203, 0.027494050562381744, -0.03827037289738655], - [0.7795881628990173, 0.029986342415213585, -0.06466733664274216], - [0.8101993799209595, 0.030793743208050728, -0.0686899945139885], - [0.8337147831916809, 0.028784994035959244, -0.07280422002077103], - [0.7542374730110168, 0.02177468128502369, -0.10443613678216934], - [0.7697082161903381, 0.020643100142478943, -0.11643557250499725], - [0.7852417230606079, 0.018978532403707504, -0.12765252590179443], - [0.767634928226471, 0.02704637683928013, -0.08803117275238037], - [0.7956817150115967, 0.028531836345791817, -0.09329714626073837], - [0.8185034990310669, 0.02707355096936226, -0.1003914326429367], - [0.7108263969421387, 0.01833728514611721, -0.03507564589381218], - [0.7278420925140381, 0.01931309886276722, -0.010097505524754524], - [0.7483652234077454, 0.01415354385972023, 0.005425570998340845], - [-0.7720924019813538, 0.027626780793070793, -0.041334930807352066], - [-0.8040405511856079, 0.029844673350453377, -0.03832409530878067], - [-0.8265854716300964, 0.027495287358760834, -0.03826868534088135], - [-0.7795882225036621, 0.029987698420882225, -0.0646686926484108], - [-0.8101993799209595, 0.030795171856880188, -0.06869153678417206], - [-0.8337149024009705, 0.02878585271537304, -0.07280556112527847], - [-0.7542385458946228, 0.021775206550955772, -0.10443780571222305], - [-0.7697089910507202, 0.02064313367009163, -0.11643654853105545], - [-0.7852423787117004, 0.01897839829325676, -0.1276528388261795], - [-0.7676352858543396, 0.02704770304262638, -0.08803359419107437], - [-0.7956817150115967, 0.028532907366752625, -0.093299500644207], - [-0.818503737449646, 0.027074117213487625, -0.10039224475622177], - [-0.7108249664306641, 0.018335221335291862, -0.035073522478342056], - [-0.7278403043746948, 0.019311318174004555, -0.01009594276547432], - [-0.7483659386634827, 0.014154116623103619, 0.005425604991614819], - ], - dtype=np.float32, -) - - -# ── fixed frame corrections (shared by FK + canonicalization) ──────────────── -# Both mirror the SONIC deploy transform. ``_YTOZ_UP`` maps SMPL's Y-up world to the -# robot's Z-up (a 90 deg rotation about X); ``_SMPL_BASE`` is SMPL's rest-pose base -# orientation, conjugated out during canonicalization. -_YTOZ_UP = R.from_euler("x", 90, degrees=True) -_SMPL_BASE = R.from_quat([0.5, 0.5, 0.5, 0.5]) # scalar-last; symmetric so wxyz==xyzw - - -# ── forward kinematics ─────────────────────────────────────────────────────── - - -def canonicalize_smpl_joints(smpl_joints: np.ndarray, root_aa: np.ndarray) -> np.ndarray: - """Remove per-frame root orientation -> SONIC ``smpl_joints_local`` format. - - Mirrors the deploy transform (and ``motion_loader.canonicalize_smpl_joints``): - reference clips store world-frame joints, but the encoder wants each frame's - joints with the body root orientation removed. - - Args: - smpl_joints: (T, 24, 3) world-frame (z-up) SMPL joint positions. - root_aa: (T, 3) SMPL global-orient axis-angle (y-up convention). - - Returns: - (T, 24, 3) per-frame root-orientation-removed joints. - """ - root = _YTOZ_UP * R.from_rotvec(root_aa) - inv = _SMPL_BASE * root.inv() - return np.einsum("tij,tkj->tki", inv.as_matrix(), smpl_joints).astype(np.float32) - - -def root_quats_from_aa(root_aa: np.ndarray) -> np.ndarray: - """Per-frame root orientation as (T, 4) wxyz, matching the live ``root_quat``. - - Same convention as the headset stream: ytoz-up then base-rotation removed. - """ - root = (_YTOZ_UP * R.from_rotvec(root_aa)) * _SMPL_BASE.inv() - return root.as_quat(scalar_first=True).astype(np.float32) # wxyz - - -# ── 3-point VR teleop keypoints (SONIC encode_mode 1) ──────────────────────── -# SMPL body-joint indices for the 3 tracked keypoints, plus the root/pelvis (0) -# used as the reference frame. Mirrors gear_sonic ``_process_3pt_pose``: neck -# (joint 12) is used rather than head (15) — it is more rigidly coupled to the -# torso and less noisy than the free-looking head. -_VR3_JOINTS = (22, 23, 12) # left wrist, right wrist, neck - -# Per-keypoint rotation offsets aligning each SMPL joint frame to the robot -# convention (root, left wrist, right wrist, neck), ported verbatim from -# gear_sonic ``pico_manager_thread_server.OFFSETS`` (extrinsic xyz euler, degrees). -_VR3_OFFSETS = [ - R.from_euler("xyz", [0, 0, -90], degrees=True), # root - R.from_euler("xyz", [90, 0, 0], degrees=True), # left wrist - R.from_euler("xyz", [-90, 0, 180], degrees=True), # right wrist - R.from_euler("xyz", [0, 0, -90], degrees=True), # neck -] - -# Unity (X-right, Y-up, Z-forward, left-handed) -> robot (X-forward, Y-left, -# Z-up, right-handed) axis remap: Unity [x, y, z] -> robot [-x, z, y]. -_UNITY_TO_ROBOT = np.array([[-1.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]]) - - -def _safe_quat(quats: np.ndarray) -> np.ndarray: - """Replace zero-norm quaternions with the scalar-last identity. - - The headset reports ``[0, 0, 0, 0]`` for joints it isn't currently tracking; - ``scipy.Rotation.from_quat`` rejects zero-norm quaternions, so we substitute the - identity ``[0, 0, 0, 1]`` (no rotation) for those rows to keep FK robust. - """ - quats = np.asarray(quats, np.float64).copy() - bad = np.linalg.norm(quats, axis=-1) < 1e-8 - quats[bad] = (0.0, 0.0, 0.0, 1.0) # scalar-last identity - return quats - - -def compute_3point(body_poses_np: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Extract the SONIC 3-point VR targets from headset body poses. - - Mirrors gear_sonic ``_process_3pt_pose``: transforms the tracked joints from the - Unity frame to the robot frame, applies the per-keypoint rotation offsets, then - expresses the left-wrist / right-wrist / neck keypoints relative to the root - (pelvis) frame. This is the ``encode_mode == 1`` counterpart to :func:`compute`. - - Note: physical wrist/neck position offsets and the operator calibration done in - gear_sonic's ``ThreePointPose.apply_calibration`` are not applied here; the raw - tracked joint poses are used. - - Args: - body_poses_np: (24, 7) rows of ``[x, y, z, qx, qy, qz, qw]`` (scalar-last), - in the Unity frame, as returned by ``xrt.get_body_joints_pose()``. - - Returns: - (pos, orn): - - pos: (9,) float32, root-relative ``[x, y, z]`` for [l-wrist, r-wrist, neck] - - orn: (12,) float32, root-relative ``[w, x, y, z]`` for the same order - """ - body = np.asarray(body_poses_np, np.float64) - q = _UNITY_TO_ROBOT - # Root (index 0) + the 3 tracked keypoints, each transformed to the robot frame - # and rotation-offset-corrected. - positions = np.zeros((4, 3), np.float64) - rotations: list[R] = [] - quats = _safe_quat(body[:, 3:7]) - for out_i, j in enumerate((0, *_VR3_JOINTS)): - positions[out_i] = q @ body[j, :3] - rot = R.from_quat(quats[j]).as_matrix() # scalar-last input - rotations.append(R.from_matrix(q @ rot @ q.T) * _VR3_OFFSETS[out_i]) - - root_inv = rotations[0].inv() - root_pos = positions[0] - pos = np.zeros(VR3_POS_DIM, np.float32) - orn = np.zeros(VR3_ORN_DIM, np.float32) - for k in range(VR3_N_POINTS): - pos[k * 3 : k * 3 + 3] = root_inv.apply(positions[k + 1] - root_pos) - orn[k * 4 : k * 4 + 4] = (root_inv * rotations[k + 1]).as_quat(scalar_first=True) # wxyz - return pos, orn - - -# ── 3-point VR teleop from raw device poses (no body trackers) ─────────────── -# PICO Y-up (X-right, Y-up, Z-back) -> robot Z-up world. Ported verbatim from -# gear_sonic's controller path (``decoupled_wbc`` ``PicoStreamer.R_HEADSET_TO_WORLD``). -_HEADSET_TO_WORLD = np.array([[0.0, 0.0, -1.0], [-1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) - - -def _device_pose_to_world(pose: np.ndarray) -> tuple[np.ndarray, R]: - """Convert a raw (7,) device pose ``[x, y, z, qx, qy, qz, qw]`` (PICO Y-up frame) - to a Z-up world ``(position, Rotation)`` pair. - - Handles the all-zero quaternion the SDK emits when a device is momentarily - untracked by substituting the identity, matching ``PicoStreamer._process_xr_pose``. - """ - pose = np.asarray(pose, np.float64) - xyz = _HEADSET_TO_WORLD @ pose[:3] - quat = pose[3:7] # scalar-last - if np.linalg.norm(quat) < 1e-8: - quat = np.array([0.0, 0.0, 0.0, 1.0]) - rot = _HEADSET_TO_WORLD @ R.from_quat(quat).as_matrix() @ _HEADSET_TO_WORLD.T - return xyz, R.from_matrix(rot) - - -def compute_3point_from_devices( - head_pose: np.ndarray, - left_pose: np.ndarray, - right_pose: np.ndarray, -) -> tuple[np.ndarray, np.ndarray]: - """Build the SONIC 3-point VR targets from raw head + controller poses. - - This is the controller-state path (no PICO Motion Trackers / body tracking - required): the 3 keypoints are the left controller, right controller, and the - headset, each expressed relative to the **headset yaw frame** — mirroring - gear_sonic's ``decoupled_wbc`` ``PicoStreamer._process_xr_pose`` (Y-up -> Z-up, - then de-headed by the headset yaw). The headset stands in for the "neck" point, - so its root-relative position is ~0 and its orientation carries pitch/roll. - - Args: - head_pose, left_pose, right_pose: (7,) ``[x, y, z, qx, qy, qz, qw]`` device - poses (scalar-last) from ``xrt.get_headset_pose()`` / - ``xrt.get_left_controller_pose()`` / ``xrt.get_right_controller_pose()``. - - Returns: - (pos, orn): - - pos: (9,) float32, headset-yaw-relative ``[x, y, z]`` for [l-wrist, r-wrist, head] - - orn: (12,) float32, headset-yaw-relative ``[w, x, y, z]`` for the same order - """ - head_pos, head_rot = _device_pose_to_world(head_pose) - left_pos, left_rot = _device_pose_to_world(left_pose) - right_pos, right_rot = _device_pose_to_world(right_pose) - - # De-head: cancel the headset yaw so targets are expressed in a heading-local frame. - head_yaw = head_rot.as_euler("xyz")[2] - inv_yaw = R.from_euler("z", -head_yaw) - - points = ((left_pos, left_rot), (right_pos, right_rot), (head_pos, head_rot)) - pos = np.zeros(VR3_POS_DIM, np.float32) - orn = np.zeros(VR3_ORN_DIM, np.float32) - for k, (p_pos, p_rot) in enumerate(points): - pos[k * 3 : k * 3 + 3] = inv_yaw.apply(p_pos - head_pos) - orn[k * 4 : k * 4 + 4] = (inv_yaw * p_rot).as_quat(scalar_first=True) # wxyz - return pos, orn - - -# ── operator calibration for the 3-point targets ──────────────────────────── -# G1 neutral key-frame targets, pelvis-relative [x, y, z] in metres, from MuJoCo FK on -# g1_29dof at the robot's *standing* configuration (``default_angles`` — the pose the -# robot actually holds at calibration time), with gear_sonic's local key-frame offsets -# (``G1_KEY_FRAME_OFFSETS``: wrists +0.18x ∓0.025y, torso +0.35z). These are the poses -# the operator's rest pose is mapped onto so the handoff starts at the robot's neutral -# stance. This is the fixed-``default_angles`` stand-in for gear_sonic's -# ``get_g1_key_frame_poses(q=measured_q)`` (see :meth:`ThreePointCalibrator.capture`): -# since the robot stands at ``default_angles`` after the startup ramp, its measured q -# equals this configuration, so these constants are the measured-q targets for the -# nominal case. (Per-frame *live* measured q would need a reverse controller->publisher -# feedback channel; not wired.) -_G1_NEUTRAL_WRIST_POS = np.array([[0.2232, 0.2177, -0.1555], [0.2232, -0.2177, -0.1555]], np.float64) -# Wrist neutral orientations (scalar-first w, x, y, z) at ``default_angles`` — NOT -# identity: the wrists are rolled/pitched in the standing pose. Matches gear_sonic -# using ``g1_lwrist_rot`` / ``g1_rwrist_rot`` from FK (not identity) as the rotation -# calibration reference. -_G1_NEUTRAL_WRIST_ROT = [ - R.from_quat([0.9168, 0.0897, 0.3864, 0.0463], scalar_first=True), # left - R.from_quat([0.9168, -0.0897, 0.3864, -0.0463], scalar_first=True), # right -] -# Neck reconstruction chain (mirrors ThreePointPose._apply_calibration): torso link -# +0.05 z, then +0.35 along the neck's local Z. -_NECK_TORSO_OFFSET_Z = 0.05 -_NECK_LINK_LENGTH = 0.35 - - -@dataclass -class ThreePointCalibrator: - """Aligns raw 3-point VR targets to the G1's neutral stance. - - Ports gear_sonic ``ThreePointPose._capture_calibration`` / ``_apply_calibration``: - on :meth:`capture` (operator holding a neutral rest pose) it records (a) the - inverse of the head/neck orientation, used to de-tilt all points to upright, and - (b) per-wrist position + orientation offsets that map the corrected rest pose onto - the fixed G1 neutral wrist targets. :meth:`apply` then transforms every subsequent - frame by those offsets, and reconstructs the head/neck position from the calibrated - neck orientation via the torso->neck kinematic chain. - - All quaternions are scalar-first (w, x, y, z), matching :func:`compute_3point`. - - ``neck_relative_wrists`` selects how the wrist targets are framed: - - - ``True`` (body-source, :func:`compute_3point`): wrists are pelvis-relative and - the neck sits roughly upright over the pelvis, so de-rotating them by ``neck_inv`` - correctly expresses them in the neck/torso frame (gear_sonic's behaviour). - - ``False`` (device-source, :func:`compute_3point_from_devices`): the "neck" is the - **headset**, which pitches down when the operator looks at their hands; the wrists - are already yaw-stabilised, so applying the head pitch/roll would rotate "up" hand - motion into "forward". The wrist frame is left in the yaw-local world frame (only - the neck point itself is still de-tilted). - """ - - neck_relative_wrists: bool = True - _neck_quat_inv: R | None = field(default=None, init=False) - _wrist_pos_offset: np.ndarray | None = field(default=None, init=False) - _wrist_rot_offset: list[R] = field(default_factory=list, init=False) - - @property - def is_calibrated(self) -> bool: - return self._neck_quat_inv is not None and self._wrist_pos_offset is not None - - def reset(self) -> None: - self._neck_quat_inv = None - self._wrist_pos_offset = None - self._wrist_rot_offset = [] - - def recalibrate_wrists(self) -> None: - """Clear only the wrist offsets, preserving the neck calibration. - - Mirrors gear_sonic ``ThreePointPose.reset_with_measured_q``: the next - :meth:`capture` recomputes the wrist offsets (against the G1 neutral targets) - while keeping the already-captured neck orientation, so re-aligning the arms - doesn't force the operator to re-level their head. - """ - self._wrist_pos_offset = None - self._wrist_rot_offset = [] - - def capture(self, pos: np.ndarray, orn: np.ndarray) -> None: - """Capture calibration offsets from a neutral-pose frame. - - Neck calibration is captured once and then preserved across subsequent - captures (matching gear_sonic's ``if self._calibration_neck_quat_inv is - None``); call :meth:`reset` to clear it or :meth:`recalibrate_wrists` to - re-align only the arms. - - Args: - pos: (9,) root-relative ``[x, y, z]`` for [l-wrist, r-wrist, head]. - orn: (12,) root-relative ``[w, x, y, z]`` for the same order. - """ - pos = np.asarray(pos, np.float64).reshape(3, 3) - orn = np.asarray(orn, np.float64).reshape(3, 4) - if self._neck_quat_inv is None: - self._neck_quat_inv = R.from_quat(orn[2], scalar_first=True).inv() - # Wrists use the neck frame only in body-source mode; device-source keeps them - # in the already yaw-stabilised world frame (see class docstring). - wrist_inv = self._neck_quat_inv if self.neck_relative_wrists else R.identity() - - self._wrist_pos_offset = np.zeros((2, 3), np.float64) - self._wrist_rot_offset = [] - for k in range(2): - corrected_pos = wrist_inv.apply(pos[k]) - corrected_rot = wrist_inv * R.from_quat(orn[k], scalar_first=True) - self._wrist_pos_offset[k] = corrected_pos - _G1_NEUTRAL_WRIST_POS[k] - # rot_offset maps the corrected rest orientation onto the G1 neutral wrist - # orientation: calibrated = rot_offset * (neck_inv * current). - self._wrist_rot_offset.append(_G1_NEUTRAL_WRIST_ROT[k] * corrected_rot.inv()) - - def apply(self, pos: np.ndarray, orn: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Apply the stored calibration; returns calibrated ``(pos (9,), orn (12,))``. - - A no-op (returns the inputs unchanged) until :meth:`capture` has been called. - """ - if self._neck_quat_inv is None or self._wrist_pos_offset is None: - return ( - np.asarray(pos, np.float32).reshape(-1), - np.asarray(orn, np.float32).reshape(-1), - ) - pos = np.asarray(pos, np.float64).reshape(3, 3) - orn = np.asarray(orn, np.float64).reshape(3, 4) - neck_inv = self._neck_quat_inv - wrist_inv = neck_inv if self.neck_relative_wrists else R.identity() - - out_pos = np.zeros((3, 3), np.float64) - out_orn = np.zeros((3, 4), np.float64) - for k in range(2): # wrists - out_pos[k] = wrist_inv.apply(pos[k]) - self._wrist_pos_offset[k] - corrected_rot = wrist_inv * R.from_quat(orn[k], scalar_first=True) - out_orn[k] = (self._wrist_rot_offset[k] * corrected_rot).as_quat(scalar_first=True) - - # Head/neck: orientation de-tilted, position from the torso->neck chain. - neck_rot = neck_inv * R.from_quat(orn[2], scalar_first=True) - out_orn[2] = neck_rot.as_quat(scalar_first=True) - neck_z = neck_rot.apply([0.0, 0.0, 1.0]) - out_pos[2] = np.array([0.0, 0.0, _NECK_TORSO_OFFSET_Z]) + _NECK_LINK_LENGTH * neck_z - return out_pos.reshape(-1).astype(np.float32), out_orn.reshape(-1).astype(np.float32) - - -class SmplForwardKinematics: - """Rest-skeleton SMPL forward kinematics (no mesh, no torch).""" - - def __init__(self, skeleton_path: str | Path | None = None): - if skeleton_path is not None: - data = np.load(skeleton_path) - self.J = data["J"].astype(np.float64) # (55, 3) rest joint positions - self.parents = data["parents"].astype(np.int64) # (55,) kinematic tree - else: - self.J = _SKELETON_J.astype(np.float64) - self.parents = _SKELETON_PARENTS.copy() - self.n_joints = self.J.shape[0] - - def _fk(self, full_pose_aa: np.ndarray) -> np.ndarray: - """full_pose_aa: (n_joints, 3) axis-angle (joint 0 = global). Returns (24, 3).""" - rot = R.from_rotvec(full_pose_aa).as_matrix() # (n, 3, 3) - rel = self.J.copy() - rel[1:] -= self.J[self.parents[1:]] - - transforms = np.zeros((self.n_joints, 4, 4), np.float64) - transforms[:, :3, :3] = rot - transforms[:, :3, 3] = rel - transforms[:, 3, 3] = 1.0 - - chain = [transforms[0]] - for i in range(1, self.n_joints): - chain.append(chain[self.parents[i]] @ transforms[i]) - joints = np.stack(chain)[:, :3, 3] - return joints[OUTPUT_JOINT_INDEX] - - def compute(self, body_poses_np: np.ndarray) -> dict: - """Convert (24, 7) headset body poses to canonical SMPL joints. - - Args: - body_poses_np: (24, 7) rows of [x, y, z, qx, qy, qz, qw] (scalar-last). - - Returns: - dict with: - - smpl_joints_local: (24, 3) root-orientation-removed joints - - root_quat: (4,) root/torso orientation (w, x, y, z) - - root_transl: (3,) pelvis translation - """ - body_poses_np = np.asarray(body_poses_np, np.float64) - positions = body_poses_np[:, :3] - - # Global joint rotations from the headset (scalar-last), with the SMPL - # +180 deg-about-Y frame fix, converted to per-joint local axis-angle. - global_rots = R.from_quat(_safe_quat(body_poses_np[:, 3:7])) * R.from_euler("y", 180, degrees=True) - gm = global_rots.as_matrix() # (24, 3, 3) - - local_aa = np.zeros((24, 3), np.float64) - for i in range(24): - p = BODY24_PARENTS[i] - m = gm[i] if p == -1 else gm[p].T @ gm[i] - local_aa[i] = R.from_matrix(m).as_rotvec() - - global_orient = local_aa[0] - body_pose = local_aa[1:].reshape(-1)[:63] # 21 body joints - - # Root: Y-up -> Z-up, then run FK with the transformed root. - root = _YTOZ_UP * R.from_rotvec(global_orient) - global_orient_new = root.as_rotvec() - - full_pose = np.concatenate([global_orient_new, body_pose, np.zeros(3 * self.n_joints - 66)]).reshape( - self.n_joints, 3 - ) - joints = self._fk(full_pose) # (24, 3) - - # Canonicalize: strip SMPL base rot and the root orientation. - root = root * _SMPL_BASE.inv() - smpl_joints_local = root.inv().apply(joints) - - return { - "smpl_joints_local": smpl_joints_local.astype(np.float32), - "root_quat": root.as_quat(scalar_first=True).astype(np.float32), # wxyz - "root_transl": positions[0].astype(np.float32), - } diff --git a/src/lerobot/teleoperators/pico_headset/smpl_stream.py b/src/lerobot/teleoperators/pico_headset/smpl_stream.py deleted file mode 100644 index 756e78081..000000000 --- a/src/lerobot/teleoperators/pico_headset/smpl_stream.py +++ /dev/null @@ -1,277 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Live SMPL stream as a SONIC reference motion (drop-in for ``SmplMotion``). - -Instead of reading an ``.npz`` clip, this pulls per-frame SMPL joints live off the -``rt/smpl`` ZMQ channel published by the GEAR PICO teleop -(``gear_sonic/scripts/pico_manager_thread_server.py``). Each message carries one -frame of **canonical** (root-orientation-removed) SMPL local joints ``(24, 3)`` -- -the exact per-frame format ``SmplMotion`` emits -- so this class exposes the same -``step() -> (720,)`` window interface and can be handed to ``sonic.py`` (or the -``pico_headset`` teleoperator) wherever a reference motion is expected -(``encode_mode == 2``). - -Transport mirrors the Unitree SDK socket bridge (``unitree_sdk2_socket.py``): a ZMQ -``SUB`` socket with ``CONFLATE`` (keep only the latest frame) subscribed to the -``rt/smpl`` topic, JSON payloads. -""" - -from __future__ import annotations - -import contextlib -import json -import logging -import time -from collections import deque - -import numpy as np -import zmq - -from .smpl_constants import ( - JOINT_DIM, - LOCO_N_AXES, - LOCO_N_BTN, - N_JOINTS, - SMPL_OBS_DIM, - VR3_ORN_DIM, - VR3_POS_DIM, - WINDOW, -) - -logger = logging.getLogger(__name__) - -SMPL_TOPIC = "rt/smpl" -DEFAULT_SMPL_HOST = "127.0.0.1" -DEFAULT_SMPL_PORT = 5560 - - -class SmplStream: - """Live ``rt/smpl`` consumer with the ``SmplMotion`` interface. - - Args: - host: publisher host (the laptop running pico_manager_thread_server.py). - port: publisher port for the ``rt/smpl`` channel. - fps: nominal source rate, only used for status/reporting. - stale_after_s: log a warning if no fresh frame arrives within this window. - loop: accepted for API parity with ``SmplMotion`` (ignored; a stream never ends). - """ - - def __init__( - self, - host: str = DEFAULT_SMPL_HOST, - port: int = DEFAULT_SMPL_PORT, - fps: float = 50.0, - stale_after_s: float = 0.5, - loop: bool = True, - ): - self.host = host - self.port = port - self.fps = float(fps) - self.loop = loop - self.stale_after_s = stale_after_s - - self._ctx = zmq.Context.instance() - self._sock = self._ctx.socket(zmq.SUB) - self._sock.setsockopt(zmq.CONFLATE, 1) # keep only the most recent frame - self._sock.connect(f"tcp://{host}:{port}") - # Single-frame JSON messages (topic embedded in payload); CONFLATE does not - # support multipart, so subscribe to everything on this dedicated port. - self._sock.setsockopt_string(zmq.SUBSCRIBE, "") - self._poller = zmq.Poller() - self._poller.register(self._sock, zmq.POLLIN) - - # Rolling window, oldest -> newest (matches SmplMotion.window layout). - self._buf: deque[np.ndarray] = deque(maxlen=WINDOW) - self._last_frame = np.zeros((N_JOINTS, JOINT_DIM), np.float32) - # Latest root/torso pose (updated every received frame). - self.root_quat = np.array([1.0, 0.0, 0.0, 0.0], np.float32) # (w, x, y, z) - self.root_transl = np.zeros(3, np.float32) - # Latest sparse 3-point VR targets (encode_mode 1), if the producer sends them. - self.vr3_pos = np.zeros(VR3_POS_DIM, np.float32) - self.vr3_orn = np.tile([1.0, 0.0, 0.0, 0.0], VR3_ORN_DIM // 4).astype(np.float32) - self._got_vr3 = False - self._last_vr3_t = 0.0 - # Latest controller-stick locomotion (encode_mode 1): [lx, ly, rx, ry] + [A,B,X,Y]. - self.loco_axes = np.zeros(LOCO_N_AXES, np.float32) - self.loco_buttons = np.zeros(LOCO_N_BTN, np.float32) - self._got_loco = False - self._last_loco_t = 0.0 - self._last_index = -1 - self._last_recv_t = 0.0 - self._warned_stale = False - self._got_first = False - - # -- SmplMotion-compatible attributes ------------------------------------ - @property - def num_frames(self) -> int: - """Streams are unbounded; report 0 (kept for API parity).""" - return 0 - - @property - def done(self) -> bool: - """A live stream never finishes.""" - return False - - @property - def has_data(self) -> bool: - """True once at least one real frame has been received.""" - return self._got_first - - @property - def has_vr3(self) -> bool: - """True once the producer has sent at least one 3-point VR frame.""" - return self._got_vr3 - - @property - def has_fresh_vr3(self) -> bool: - """True when a 3-point VR frame arrived within ``stale_after_s``. - - Unlike :attr:`has_data`, this is independent of the SMPL window, so the - controller-state source (head + controllers only, empty SMPL) still drives - ``encode_mode 1`` without a whole-body reference. - """ - if not self._got_vr3: - return False - if not self.stale_after_s: - return True - return (time.time() - self._last_vr3_t) <= self.stale_after_s - - @property - def has_fresh_loco(self) -> bool: - """True when controller-stick locomotion arrived within ``stale_after_s``.""" - if not self._got_loco: - return False - if not self.stale_after_s: - return True - return (time.time() - self._last_loco_t) <= self.stale_after_s - - @property - def seconds_since_last(self) -> float: - """Wall-clock seconds since the last real frame (inf before the first).""" - if not self._got_first: - return float("inf") - return time.time() - self._last_recv_t - - @property - def is_stale(self) -> bool: - """True when the stream has gone silent past ``stale_after_s``. - - Consumers use this to stop feeding a frozen pose and let the controller - fall back to a safe standing/locomotion mode. - """ - if not self._got_first or not self.stale_after_s: - return False - return self.seconds_since_last > self.stale_after_s - - def reset(self): - self._buf.clear() - self._got_first = False - self._got_vr3 = False - self._last_vr3_t = 0.0 - self._got_loco = False - self._last_loco_t = 0.0 - - # -- core ---------------------------------------------------------------- - def _drain_latest(self) -> np.ndarray | None: - """Return the newest available (24, 3) frame, or None if nothing new. - - CONFLATE already keeps only the last message, but we poll non-blocking so - the 50 Hz control loop never stalls waiting on the headset. - """ - frame = None - while dict(self._poller.poll(0)).get(self._sock) == zmq.POLLIN: - payload = self._sock.recv() - data = json.loads(payload.decode("utf-8")).get("data", {}) - - # Sparse 3-point VR targets (encode_mode 1). Parsed independently of the - # SMPL window so the controller-state source (head + controllers, empty - # SMPL) is still handled. - vp = data.get("vr3_pos") - vo = data.get("vr3_orn") - if vp is not None and vo is not None and len(vp) == VR3_POS_DIM and len(vo) == VR3_ORN_DIM: - self.vr3_pos = np.asarray(vp, np.float32) - self.vr3_orn = np.asarray(vo, np.float32) - self._got_vr3 = True - self._last_vr3_t = time.time() - - # Controller-stick locomotion (encode_mode 1), also independent of SMPL. - la = data.get("loco_axes") - lb = data.get("loco_buttons") - if la is not None and lb is not None and len(la) == LOCO_N_AXES and len(lb) == LOCO_N_BTN: - self.loco_axes = np.asarray(la, np.float32) - self.loco_buttons = np.asarray(lb, np.float32) - self._got_loco = True - self._last_loco_t = time.time() - - # SMPL whole-body window (encode_mode 2), optional on this stream. - joints = np.asarray(data.get("smpl_joints_local", []), np.float32) - if joints.size != N_JOINTS * JOINT_DIM: - continue - frame = joints.reshape(N_JOINTS, JOINT_DIM) - self._last_index = int(data.get("frame_index", self._last_index + 1)) - rq = data.get("root_quat") - if rq is not None and len(rq) == 4: - self.root_quat = np.asarray(rq, np.float32) - rt = data.get("root_transl") - if rt is not None and len(rt) == 3: - self.root_transl = np.asarray(rt, np.float32) - return frame - - def step(self) -> np.ndarray: - """Advance one control tick, returning the current 720-vec window. - - If no new headset frame arrived this tick we hold the last one, so the - policy keeps tracking the latest pose rather than snapping to zero. - """ - frame = self._drain_latest() - now = time.time() - - if frame is not None: - self._last_frame = frame - self._last_recv_t = now - self._warned_stale = False - if not self._got_first: - # Pre-fill the window so the first send is a full, coherent clip. - self._buf.extend([frame.copy() for _ in range(WINDOW)]) - self._got_first = True - else: - self._buf.append(frame) - elif self._got_first: - # No fresh frame: repeat the most recent to keep the window moving. - self._buf.append(self._last_frame.copy()) - if ( - self.stale_after_s - and not self._warned_stale - and (now - self._last_recv_t) > self.stale_after_s - ): - logger.warning( - "[SmplStream] no %s frame for %.2fs (holding last pose)", - SMPL_TOPIC, - now - self._last_recv_t, - ) - self._warned_stale = True - - if not self._got_first: - return np.zeros(SMPL_OBS_DIM, np.float32) - # Flatten to (720,): frames oldest->newest, joint-major within a frame - # [f0_j0_xyz, f0_j1_xyz, ..., f9_j23_xyz] — matches SmplMotion.window. - return np.concatenate(list(self._buf), dtype=np.float32).reshape(-1) - - def close(self): - with contextlib.suppress(Exception): - self._poller.unregister(self._sock) - self._sock.close(0) diff --git a/src/lerobot/teleoperators/pico_headset/smpl_to_dataset.py b/src/lerobot/teleoperators/pico_headset/smpl_to_dataset.py deleted file mode 100644 index cc6411964..000000000 --- a/src/lerobot/teleoperators/pico_headset/smpl_to_dataset.py +++ /dev/null @@ -1,147 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Convert SMPL ``.npz`` motion clips into a LeRobotDataset for SONIC replay. - -Each dataset frame's ``action`` is the 720-dim ``smpl.*`` window that the -``pico_headset`` teleoperator emits and ``SonicWholeBodyController`` reassembles -into ``encode_mode == 2``. So the resulting dataset can be pushed straight -through ``lerobot-replay`` to drive SONIC whole-body tracking with no headset: - - lerobot-replay \ - --robot.type=unitree_g1 \ - --robot.controller=SonicWholeBodyController \ - --dataset.repo_id=/ --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()