mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-30 13:09:40 +00:00
Compare commits
41 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| fc7a0bc2fd | |||
| 5f6513551c | |||
| 70e157e00f | |||
| 1837be51bf | |||
| bedd56eed9 | |||
| c165e4df68 | |||
| 5e24da483a | |||
| 9c54665a76 | |||
| f6a845c30c | |||
| 45e8336854 | |||
| 5046e2df32 | |||
| 1c88e26c6d | |||
| 69a3edfa33 | |||
| 2492ce2c29 | |||
| c8e75da55f | |||
| 2eae31ea2b | |||
| c997abe739 | |||
| c73579055e | |||
| 4be438161b | |||
| 806d28a883 | |||
| 573b65ff6b | |||
| bc55713e7c | |||
| 4f53c42583 | |||
| bfced3d149 | |||
| 4969813d4e | |||
| 1c87ca31a3 | |||
| 4bcde762cc | |||
| 943ae78cfe | |||
| 3363688f1e | |||
| 0876629e72 | |||
| 305614b8c6 | |||
| 02d3202c4f | |||
| 3b6de2fdf8 | |||
| 744f3667c0 | |||
| fdde436776 | |||
| 5c683c65c6 | |||
| dfbc25c58f | |||
| 804c76bcc2 | |||
| e6afa69be9 | |||
| 31d1439e29 | |||
| 1c118c6359 |
+5
-1
@@ -374,7 +374,11 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
|
|||||||
torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
|
torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }]
|
||||||
|
|
||||||
[tool.setuptools.package-data]
|
[tool.setuptools.package-data]
|
||||||
lerobot = ["envs/*.json", "annotations/steerable_pipeline/prompts/*.txt"]
|
lerobot = [
|
||||||
|
"envs/*.json",
|
||||||
|
"annotations/steerable_pipeline/prompts/*.txt",
|
||||||
|
"teleoperators/pico_headset/assets/*.npz",
|
||||||
|
]
|
||||||
|
|
||||||
[tool.setuptools.packages.find]
|
[tool.setuptools.packages.find]
|
||||||
where = ["src"]
|
where = ["src"]
|
||||||
|
|||||||
@@ -0,0 +1,162 @@
|
|||||||
|
# Unitree G1 — SONIC whole-body control
|
||||||
|
|
||||||
|
This package runs NVIDIA's **SONIC** whole-body controller (and the GR00T/Holosoma
|
||||||
|
locomotion controllers) on the Unitree G1, in MuJoCo simulation or on real hardware.
|
||||||
|
SONIC turns a high-level movement intent — or a streamed **SMPL** whole-body pose — into
|
||||||
|
50 Hz joint-position targets. It is a pure-Python/ONNX reimplementation of the SONIC
|
||||||
|
deploy stack (no `gear_sonic`/torch dependency).
|
||||||
|
|
||||||
|
## Controllers
|
||||||
|
|
||||||
|
Selected with `--robot.controller=<ClassName>`:
|
||||||
|
|
||||||
|
| Controller | Purpose |
|
||||||
|
| ------------------------------ | ------------------------------------------------------------------------------------------ |
|
||||||
|
| `SonicWholeBodyController` | SONIC whole-body: locomotion (mode 0), 3-point VR teleop (mode 1), SMPL imitation (mode 2) |
|
||||||
|
| `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.
|
||||||
|
|
||||||
|
## Requirements
|
||||||
|
|
||||||
|
- `onnxruntime` (CPU) **or** `onnxruntime-gpu` (recommended — SONIC runs three ONNX
|
||||||
|
sessions and is much smoother on GPU). Install the CUDA build that matches your
|
||||||
|
driver (e.g. `onnxruntime-gpu==1.26.0` for a CUDA-12.x driver). Verify with:
|
||||||
|
```bash
|
||||||
|
python -c "import onnxruntime as ort; print(ort.get_available_providers())"
|
||||||
|
# expect CUDAExecutionProvider in the list for GPU
|
||||||
|
```
|
||||||
|
- `mujoco` for simulation (`is_simulation=True`).
|
||||||
|
- `pyzmq` only if you use the live SMPL stream (pico headset).
|
||||||
|
- The SONIC ONNX models are downloaded automatically from the `nvidia/GEAR-SONIC` Hub repo.
|
||||||
|
|
||||||
|
## Running
|
||||||
|
|
||||||
|
**Replay an SMPL dataset (motion imitation):**
|
||||||
|
|
||||||
|
```bash
|
||||||
|
lerobot-replay \
|
||||||
|
--robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
|
||||||
|
--dataset.repo_id=<user>/<smpl_dataset> --dataset.episode=0
|
||||||
|
```
|
||||||
|
|
||||||
|
**Keyboard teleop** (drives locomotion via the native keyboard teleoperator):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
lerobot-teleoperate \
|
||||||
|
--robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
|
||||||
|
--teleop.type=keyboard
|
||||||
|
```
|
||||||
|
|
||||||
|
Controls: `WASD` move · `Q`/`E` turn · `1`–`8` mode · `9`/`0` speed · `-`/`=` height ·
|
||||||
|
`R` replan · `Space` emergency-stop.
|
||||||
|
|
||||||
|
**PICO headset teleop — SMPL whole-body** (mode 2, needs PICO Motion Trackers):
|
||||||
|
|
||||||
|
```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
|
||||||
|
```
|
||||||
|
|
||||||
|
**PICO headset teleop — 3-point VR** (mode 1, head + 2 controllers only, **no trackers**):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# 1) publisher (head + controllers -> 3-point targets + stick locomotion)
|
||||||
|
python -m lerobot.teleoperators.pico_headset.pico_publisher --fps 50 --headset-source devices
|
||||||
|
# 2) controller
|
||||||
|
lerobot-teleoperate \
|
||||||
|
--robot.type=unitree_g1 --robot.controller=SonicWholeBodyController \
|
||||||
|
--teleop.type=pico_headset --teleop.mode=vr3
|
||||||
|
```
|
||||||
|
|
||||||
|
3-point controls: left stick move · right stick X turn · right stick Y height ·
|
||||||
|
`A`+`B` / `X`+`Y` cycle locomotion mode (walk/run/squat/kneel/…) · hands+head track the
|
||||||
|
upper body. **Calibration**: stand in a neutral rest pose and press `A`+`B`+`X`+`Y` — the
|
||||||
|
publisher status line flips from `UNCALIBRATED` to `calibrated`. This maps your rest pose
|
||||||
|
onto the G1's neutral stance and is required before the hands track well; the SMPL
|
||||||
|
(mode 2) path is self-calibrating and needs no such step.
|
||||||
|
|
||||||
|
Both require the XRoboToolkit stack — see below.
|
||||||
|
|
||||||
|
## PICO headset / XRoboToolkit install
|
||||||
|
|
||||||
|
Live full-body teleop needs the **XRoboToolkit** system (a PC Service on your
|
||||||
|
workstation + a PICO app on the headset) and its Python binding, `xrobotoolkit_sdk`.
|
||||||
|
The full hardware + software walkthrough lives in the SONIC repo:
|
||||||
|
[`docs/source/getting_started/vr_teleop_setup.md`](https://nvlabs.github.io/GR00T-WholeBodyControl/getting_started/vr_teleop_setup.html).
|
||||||
|
|
||||||
|
Summary:
|
||||||
|
|
||||||
|
1. **PC Service** (workstation) — install and run it before connecting the headset.
|
||||||
|
- Ubuntu 22.04 / 24.04 (x86_64): prebuilt `.deb` from the
|
||||||
|
[XRoboToolkit-PC-Service releases](https://github.com/XR-Robotics/XRoboToolkit-PC-Service/releases).
|
||||||
|
- Jetson (aarch64): the arm64 `.deb`.
|
||||||
|
- Windows (x64): the Windows PC Service build.
|
||||||
|
2. **PICO app** — install `XRoboToolkit-PICO-*.apk` on the headset (see the guide),
|
||||||
|
enable Developer Mode. For **SMPL whole-body** (mode 2) you also need the PICO Motion
|
||||||
|
Trackers paired/calibrated and "Full body" enabled; for **3-point** (mode 1,
|
||||||
|
`--headset-source devices`) only Head + Controller + Send are required — no trackers.
|
||||||
|
3. **`xrobotoolkit_sdk`** — a pybind11/CMake build (not a pip package), from
|
||||||
|
[`XRoboToolkit-PC-Service-Pybind`](https://github.com/XR-Robotics/XRoboToolkit-PC-Service-Pybind):
|
||||||
|
- Linux x86_64: `pip install pybind11 cmake` then `bash setup_ubuntu.sh` (or the
|
||||||
|
SONIC repo's `install_scripts/install_pico.sh`, which builds everything into a
|
||||||
|
`.venv_teleop`).
|
||||||
|
- Jetson aarch64: `bash setup_orin.sh` (builds `libPXREARobotSDK.so` from source).
|
||||||
|
- Windows x64: `pip install pybind11` then `setup_windows.bat` (needs git + an
|
||||||
|
MSVC/CMake toolchain; uses the prebuilt `PXREARobotSDK.dll`/`.lib`).
|
||||||
|
4. Connect PICO and workstation to the **same Wi-Fi**, open the XRoboToolkit app, enter
|
||||||
|
the PC IP, and enable Head/Controller/Send (plus Full-body for SMPL mode 2).
|
||||||
|
|
||||||
|
### Platform support
|
||||||
|
|
||||||
|
| Platform | Live headset teleop | Notes |
|
||||||
|
| --------------------------- | ------------------- | ------------------------------------------- |
|
||||||
|
| Linux x86_64 | ✅ | Guided `install_pico.sh` (SONIC repo) |
|
||||||
|
| Linux aarch64 (Jetson Orin) | ✅ | `setup_orin.sh` builds the native lib |
|
||||||
|
| Windows x64 | ✅ (manual) | `setup_windows.bat`; no one-shot env script |
|
||||||
|
| macOS | ❌ | No PC Service / SDK build for Darwin |
|
||||||
|
|
||||||
|
### No hardware required (any platform, incl. macOS/Windows)
|
||||||
|
|
||||||
|
The SMPL pipeline can be exercised without a headset or the SDK — the publisher emits
|
||||||
|
`rt/smpl` frames that the controller consumes exactly as it would from the headset:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
# synthetic motion
|
||||||
|
python -m lerobot.teleoperators.pico_headset.pico_publisher --fake
|
||||||
|
|
||||||
|
# replay a canned SMPL clip
|
||||||
|
python -m lerobot.teleoperators.pico_headset.pico_publisher --motion-file <clip>.npz
|
||||||
|
```
|
||||||
|
|
||||||
|
## Notes
|
||||||
|
|
||||||
|
- SMPL **root motion** into the mode-2 anchor is opt-in (`SonicWholeBodyController(enable_smpl_root=True)`);
|
||||||
|
it stays off by default (untested on hardware). When enabled, the per-frame root quat is
|
||||||
|
spherically smoothed (`root_smoothing_alpha`, default 0.15) before it reaches the anchor,
|
||||||
|
which removes the base-acceleration spikes the raw 30 Hz→50 Hz trajectory used to cause.
|
||||||
|
- Direct `rt/smpl` subscription without the pico teleoperator is available via
|
||||||
|
`SonicWholeBodyController(enable_smpl_stream=True, smpl_host=..., smpl_port=...)`.
|
||||||
|
- 3-point (mode 1) uses the **headset-yaw frame** as its reference and the `A`+`B`+`X`+`Y`
|
||||||
|
calibration to align to the G1 neutral stance. Calibration maps the operator's rest pose
|
||||||
|
onto the G1's **standing** (`default_angles`) wrist/neck key-frame poses (position **and**
|
||||||
|
orientation) computed by FK — the `default_angles` stand-in for gear_sonic's live
|
||||||
|
measured-q recalibration, since the robot holds `default_angles` at calibration time.
|
||||||
|
Re-aligning the arms only (preserving the neck level) is available via the calibrator's
|
||||||
|
`recalibrate_wrists()`.
|
||||||
|
- 3-point **locomotion** from the PICO sticks follows gear_sonic's `PlannerLoop` exactly:
|
||||||
|
a yaw accumulator on the right stick and **mode-dependent speed curves** on the left
|
||||||
|
(slow `0.1+0.5·mag`, run `1.5+3·mag`, walk = planner default). Stick signs replicate
|
||||||
|
gear_sonic's `get_controller_axes` usage (forward `+ly`, strafe `-lx`, turn `-rx`); since
|
||||||
|
the publisher forwards the same raw SDK axes, this is the correct convention by construction.
|
||||||
|
- Startup interpolation and the graceful-stop settle are mode-agnostic; set
|
||||||
|
`--robot.graceful_stop_s=0` to restore the old instant zero-torque on disconnect.
|
||||||
@@ -65,9 +65,43 @@ class UnitreeG1Config(RobotConfig):
|
|||||||
# Cameras (ZMQ-based remote cameras)
|
# Cameras (ZMQ-based remote cameras)
|
||||||
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
cameras: dict[str, CameraConfig] = field(default_factory=dict)
|
||||||
|
|
||||||
|
# Synthetic zero-image cameras exposed as ``observation.images.{name}`` (H×W×3
|
||||||
|
# black frames). Lets image-conditioned policies (e.g. pi0.5 / OpenHLM) run in
|
||||||
|
# sim before real cameras are wired. Empty = disabled.
|
||||||
|
empty_cameras: list[str] = field(default_factory=list)
|
||||||
|
empty_camera_hw: tuple[int, int] = (224, 224)
|
||||||
|
|
||||||
|
# Publish Dex3 hand commands (``rt/dex3/{left,right}/cmd``) driven by the OpenHLM
|
||||||
|
# gripper scalars (``wb.7.pos`` left, ``wb.15.pos`` right). Lets the 43-DoF sim
|
||||||
|
# (or a real Dex3-equipped G1) show grasping. The scalar in [0, 1] is remapped to
|
||||||
|
# a curl amount (``hand_open_grip_value`` -> open) and scaled onto
|
||||||
|
# ``hand_closed_pose`` (7 joints: thumb_0/1/2, middle_0/1, index_0/1). Flip signs
|
||||||
|
# in ``hand_closed_pose`` if fingers curl the wrong way.
|
||||||
|
publish_hands: bool = False
|
||||||
|
hand_open_grip_value: float = 1.0
|
||||||
|
hand_closed_grip_value: float = 0.0
|
||||||
|
hand_closed_pose: list[float] = field(
|
||||||
|
default_factory=lambda: [1.0, 0.9, 0.9, 1.3, 1.3, 1.3, 1.3]
|
||||||
|
)
|
||||||
|
hand_kp: float = 1.5
|
||||||
|
hand_kd: float = 0.1
|
||||||
|
|
||||||
|
# Replay recorded camera frames from a LeRobot parquet episode as the camera
|
||||||
|
# feed (e.g. OpenHLM-data episode). Maps a robot camera name to a parquet image
|
||||||
|
# column; frames advance one per observation and loop. Lets a VLA see the real
|
||||||
|
# task video in sim without live cameras. Empty map = disabled.
|
||||||
|
replay_camera_parquet: str | None = None
|
||||||
|
replay_camera_map: dict[str, str] = field(default_factory=dict)
|
||||||
|
replay_camera_loop: bool = True
|
||||||
|
|
||||||
# Compensates for gravity on the unitree's arms using the arm ik solver
|
# Compensates for gravity on the unitree's arms using the arm ik solver
|
||||||
gravity_compensation: bool = False
|
gravity_compensation: bool = False
|
||||||
|
|
||||||
# Lower-body controller class name, e.g. "GrootLocomotionController" or
|
# Locomotion controller class name, e.g. "GrootLocomotionController",
|
||||||
# "HolosomaLocomotionController". None disables it.
|
# "HolosomaLocomotionController", or "SonicWholeBodyController". None disables it.
|
||||||
controller: str | None = None
|
controller: str | None = None
|
||||||
|
|
||||||
|
# On disconnect (e.g. Ctrl-C), seconds to hold the current pose while ramping joint
|
||||||
|
# stiffness (kp) to zero — a soft, damped settle instead of an instant limp /
|
||||||
|
# free-fall. 0 disables it (immediate zero-torque). Real robot only.
|
||||||
|
graceful_stop_s: float = 1.5
|
||||||
|
|||||||
@@ -0,0 +1,24 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
|
||||||
|
|
||||||
|
__all__ = [
|
||||||
|
"GrootLocomotionController",
|
||||||
|
"HolosomaLocomotionController",
|
||||||
|
"SonicWholeBodyController",
|
||||||
|
"SonicRuntime",
|
||||||
|
]
|
||||||
+12
-2
@@ -14,20 +14,29 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import logging
|
import logging
|
||||||
from collections import deque
|
from collections import deque
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import onnxruntime as ort
|
|
||||||
from huggingface_hub import hf_hub_download
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
from .g1_utils import (
|
from lerobot.utils.import_utils import _onnxruntime_available, require_package
|
||||||
|
|
||||||
|
from ..g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
REMOTE_BUTTONS,
|
REMOTE_BUTTONS,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
get_gravity_orientation,
|
get_gravity_orientation,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _onnxruntime_available:
|
||||||
|
import onnxruntime as ort
|
||||||
|
else:
|
||||||
|
ort = None
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
@@ -83,6 +92,7 @@ class GrootLocomotionController:
|
|||||||
control_dt = CONTROL_DT # Expose for unitree_g1.py
|
control_dt = CONTROL_DT # Expose for unitree_g1.py
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
|
require_package("onnxruntime", extra="unitree_g1")
|
||||||
# Load policies
|
# Load policies
|
||||||
self.policy_balance, self.policy_walk = load_groot_policies()
|
self.policy_balance, self.policy_walk = load_groot_policies()
|
||||||
|
|
||||||
+18
-3
@@ -14,21 +14,34 @@
|
|||||||
# See the License for the specific language governing permissions and
|
# See the License for the specific language governing permissions and
|
||||||
# limitations under the License.
|
# limitations under the License.
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
import json
|
import json
|
||||||
import logging
|
import logging
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import onnx
|
|
||||||
import onnxruntime as ort
|
|
||||||
from huggingface_hub import hf_hub_download
|
from huggingface_hub import hf_hub_download
|
||||||
|
|
||||||
from .g1_utils import (
|
from lerobot.utils.import_utils import _onnx_available, _onnxruntime_available, require_package
|
||||||
|
|
||||||
|
from ..g1_utils import (
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
G1_29_JointArmIndex,
|
G1_29_JointArmIndex,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
get_gravity_orientation,
|
get_gravity_orientation,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _onnxruntime_available:
|
||||||
|
import onnxruntime as ort
|
||||||
|
else:
|
||||||
|
ort = None
|
||||||
|
|
||||||
|
if TYPE_CHECKING or _onnx_available:
|
||||||
|
import onnx
|
||||||
|
else:
|
||||||
|
onnx = None
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
DEFAULT_ANGLES = np.zeros(29, dtype=np.float32)
|
DEFAULT_ANGLES = np.zeros(29, dtype=np.float32)
|
||||||
@@ -101,6 +114,8 @@ class HolosomaLocomotionController:
|
|||||||
control_dt = CONTROL_DT # Expose for unitree_g1.py
|
control_dt = CONTROL_DT # Expose for unitree_g1.py
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
|
require_package("onnxruntime", extra="unitree_g1")
|
||||||
|
require_package("onnx", extra="unitree_g1")
|
||||||
# Load policy and gains
|
# Load policy and gains
|
||||||
self.policy, self.kp, self.kd = load_policy()
|
self.policy, self.kp, self.kd = load_policy()
|
||||||
|
|
||||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,718 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""SONIC full-body controller for Unitree G1."""
|
||||||
|
|
||||||
|
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 .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:
|
||||||
|
import onnxruntime as ort
|
||||||
|
else:
|
||||||
|
ort = None
|
||||||
|
|
||||||
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
# Startup blend duration: over the first control ticks, linearly interpolate every joint
|
||||||
|
# from the robot's initial measured pose into the policy's commanded target, so control
|
||||||
|
# eases in without a snap on the first command.
|
||||||
|
INIT_RAMP_S = 3.0
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_smpl_from_action(action: dict | None) -> np.ndarray | None:
|
||||||
|
"""Reassemble a (720,) SMPL window from ``smpl.{i}`` action keys, or None.
|
||||||
|
|
||||||
|
The pico_headset teleoperator emits the whole-body reference as flat floats so
|
||||||
|
it flows unchanged through the standard lerobot action pipeline.
|
||||||
|
"""
|
||||||
|
# The keys are smpl.0 .. smpl.719; presence of the first element (smpl.0) is the
|
||||||
|
# sentinel that a full SMPL window was sent this tick. If it's absent, there's no
|
||||||
|
# whole-body reference, so bail out.
|
||||||
|
if not action or f"{SMPL_ACTION_PREFIX}0" not in action:
|
||||||
|
return None
|
||||||
|
arr = np.fromiter(
|
||||||
|
(float(action.get(f"{SMPL_ACTION_PREFIX}{i}", 0.0)) for i in range(SMPL_ACTION_DIM)),
|
||||||
|
dtype=np.float32,
|
||||||
|
count=SMPL_ACTION_DIM,
|
||||||
|
)
|
||||||
|
return arr
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_root_from_action(action: dict | None) -> np.ndarray | None:
|
||||||
|
"""Reassemble a (4,) SMPL root quaternion (wxyz) from ``root.{i}`` keys, or None."""
|
||||||
|
if not action or f"{ROOT_ACTION_PREFIX}0" not in action:
|
||||||
|
return None
|
||||||
|
q = np.fromiter(
|
||||||
|
(float(action.get(f"{ROOT_ACTION_PREFIX}{i}", 0.0)) for i in range(ROOT_ACTION_DIM)),
|
||||||
|
dtype=np.float32,
|
||||||
|
count=ROOT_ACTION_DIM,
|
||||||
|
)
|
||||||
|
n = float(np.linalg.norm(q))
|
||||||
|
if n < 1e-6:
|
||||||
|
return None
|
||||||
|
return q / n
|
||||||
|
|
||||||
|
|
||||||
|
def _extract_vr3_from_action(action: dict | None) -> tuple[np.ndarray, np.ndarray] | None:
|
||||||
|
"""Reassemble the 3-point VR targets from ``vr3_pos.{i}`` / ``vr3_orn.{i}`` keys.
|
||||||
|
|
||||||
|
Returns ``(pos (9,), orn (12,))`` for the [l-wrist, r-wrist, neck] keypoints, or
|
||||||
|
None when no VR3 reference was sent this tick. Presence of ``vr3_pos.0`` is the
|
||||||
|
sentinel that a full 3-point frame is available (mirrors the SMPL sentinel).
|
||||||
|
"""
|
||||||
|
if not action or f"{VR3_POS_PREFIX}0" not in action:
|
||||||
|
return None
|
||||||
|
pos = np.fromiter(
|
||||||
|
(float(action.get(f"{VR3_POS_PREFIX}{i}", 0.0)) for i in range(VR3_POS_DIM)),
|
||||||
|
dtype=np.float32,
|
||||||
|
count=VR3_POS_DIM,
|
||||||
|
)
|
||||||
|
orn = np.fromiter(
|
||||||
|
(float(action.get(f"{VR3_ORN_PREFIX}{i}", 0.0)) for i in range(VR3_ORN_DIM)),
|
||||||
|
dtype=np.float32,
|
||||||
|
count=VR3_ORN_DIM,
|
||||||
|
)
|
||||||
|
return pos, orn
|
||||||
|
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
This is the OpenHLM / pi0.5 joint-based interface: one 34-D vector per tick
|
||||||
|
(sentinel: presence of ``wb.0.pos``) carrying absolute joint targets in real
|
||||||
|
units. The ``.pos`` suffix lets these flow through ``lerobot-rollout`` as normal
|
||||||
|
joint-position action features.
|
||||||
|
"""
|
||||||
|
if not action or wb_action_key(0) not in action:
|
||||||
|
return None
|
||||||
|
return np.fromiter(
|
||||||
|
(float(action.get(wb_action_key(i), 0.0)) for i in range(WB_ACTION_DIM)),
|
||||||
|
dtype=np.float32,
|
||||||
|
count=WB_ACTION_DIM,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _wb34_to_reference(wb: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||||||
|
"""Map a 34-D OpenHLM whole-body command to a SONIC mode-0 reference.
|
||||||
|
|
||||||
|
Returns ``(ref29, anchor_quat)`` where ``ref29`` is the 29 joint targets in
|
||||||
|
IsaacLab order (what SONIC's ``motion_joint_positions`` expects) and
|
||||||
|
``anchor_quat`` (wxyz) encodes the root roll/pitch (yaw=0).
|
||||||
|
|
||||||
|
OpenHLM layout : [L-arm 0:7, L-grip 7, R-arm 8:15, R-grip 15,
|
||||||
|
L-leg 16:22, R-leg 22:28, waist 28:31, root rp+yaw 31:34]
|
||||||
|
The 29 joints are first assembled in MuJoCo / Unitree-SDK order
|
||||||
|
([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29] — the
|
||||||
|
``G1_29_JointIndex`` grouping OpenHLM uses), then permuted to IsaacLab order via
|
||||||
|
``MUJOCO_TO_ISAACLAB``. Grippers (7, 15) and yaw-rate (33) are not part of the
|
||||||
|
29-DoF SONIC reference.
|
||||||
|
"""
|
||||||
|
ref_mj = np.zeros(29, np.float32) # MuJoCo / Unitree-SDK grouped order
|
||||||
|
ref_mj[0:6] = wb[16:22] # left leg
|
||||||
|
ref_mj[6:12] = wb[22:28] # right leg
|
||||||
|
ref_mj[12:15] = wb[28:31] # waist
|
||||||
|
ref_mj[15:22] = wb[0:7] # left arm
|
||||||
|
ref_mj[22:29] = wb[8:15] # right arm
|
||||||
|
ref = ref_mj[MUJOCO_TO_ISAACLAB].astype(np.float32) # -> IsaacLab order for SONIC
|
||||||
|
roll, pitch = float(wb[31]), float(wb[32])
|
||||||
|
cr, sr, cp, sp = np.cos(roll / 2), np.sin(roll / 2), np.cos(pitch / 2), np.sin(pitch / 2)
|
||||||
|
anchor = np.array([cr * cp, sr * cp, cr * sp, sr * sp], np.float32) # Rx(roll)·Ry(pitch)
|
||||||
|
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)."""
|
||||||
|
|
||||||
|
def __init__(self, force_cpu: bool = False):
|
||||||
|
require_package("onnxruntime", extra="unitree_g1")
|
||||||
|
planner_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="planner_sonic.onnx")
|
||||||
|
encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx")
|
||||||
|
decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
|
||||||
|
|
||||||
|
providers = ort_providers(force_cpu=force_cpu)
|
||||||
|
self.use_gpu = providers[0] == "CUDAExecutionProvider"
|
||||||
|
so = make_ort_session_options()
|
||||||
|
|
||||||
|
planner_sess = ort.InferenceSession(planner_path, sess_options=so, providers=providers)
|
||||||
|
encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
|
||||||
|
decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
|
||||||
|
|
||||||
|
self.kp, self.kd = compute_kp_kd()
|
||||||
|
self.ms = MovementState()
|
||||||
|
self.planner = SonicPlanner(planner_sess, planner_path)
|
||||||
|
self.controller = PlannerController(self.planner, encoder_sess, decoder_sess)
|
||||||
|
|
||||||
|
motion = self.planner.initialize(DEFAULT_ANGLES, self.ms)
|
||||||
|
self.controller.load_initial_motion(motion)
|
||||||
|
self.planner.start_subprocess(self.controller, use_gpu=self.use_gpu)
|
||||||
|
|
||||||
|
self.step = 0
|
||||||
|
self.replan_timer = 0.0
|
||||||
|
self.last_ms = snapshot_ms(self.ms)
|
||||||
|
|
||||||
|
@property
|
||||||
|
def pipeline(self):
|
||||||
|
return self.controller
|
||||||
|
|
||||||
|
def tick(self, obs: dict, *, debug: bool | None = None, use_joystick: bool = True) -> dict:
|
||||||
|
if not obs:
|
||||||
|
self.step += 1
|
||||||
|
return {}
|
||||||
|
|
||||||
|
if use_joystick:
|
||||||
|
process_joystick(obs, self.ms, self.controller)
|
||||||
|
clamp_mode_params(self.ms)
|
||||||
|
|
||||||
|
if self.step > 0:
|
||||||
|
self.replan_timer += CONTROL_DT
|
||||||
|
if should_replan_request(self.ms, self.last_ms, self.replan_timer, self.step):
|
||||||
|
self.planner.request_replan(self.controller.ref_cursor, self.ms)
|
||||||
|
self.replan_timer = 0.0
|
||||||
|
self.ms.needs_replan = False
|
||||||
|
self.last_ms = snapshot_ms(self.ms)
|
||||||
|
|
||||||
|
do_enc = self.step % ENCODER_UPDATE_EVERY == 0
|
||||||
|
if debug is None:
|
||||||
|
debug = self.step % DEBUG_PRINT_EVERY == 0
|
||||||
|
action = self.controller.step(obs, update_encoder=do_enc, debug=debug)
|
||||||
|
|
||||||
|
result = self.planner.try_get_new_motion()
|
||||||
|
if result:
|
||||||
|
self.controller.blend_new_motion(*result)
|
||||||
|
|
||||||
|
self.controller.advance_cursor()
|
||||||
|
self.step += 1
|
||||||
|
return action
|
||||||
|
|
||||||
|
def reset(self):
|
||||||
|
self.ms = MovementState()
|
||||||
|
self.controller.reinit_heading = True
|
||||||
|
self.controller.playing = True
|
||||||
|
self.step = 0
|
||||||
|
self.replan_timer = 0.0
|
||||||
|
self.last_ms = snapshot_ms(self.ms)
|
||||||
|
|
||||||
|
def shutdown(self):
|
||||||
|
self.planner.stop_subprocess()
|
||||||
|
|
||||||
|
|
||||||
|
class SonicWholeBodyController:
|
||||||
|
"""Full-body SONIC controller for UnitreeG1's background controller thread."""
|
||||||
|
|
||||||
|
control_dt = CONTROL_DT
|
||||||
|
full_body = True
|
||||||
|
# Advertise a dense 34-D whole-body action space (OpenHLM / pi0.5) so the robot
|
||||||
|
# exposes ``wb.{i}.pos`` action features and ``lerobot-rollout`` can drive it
|
||||||
|
# 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,
|
||||||
|
):
|
||||||
|
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).
|
||||||
|
self._init_ramp_steps = max(1, round(INIT_RAMP_S / CONTROL_DT))
|
||||||
|
self._init_step = 0
|
||||||
|
self._start_pose: dict[str, float] = {}
|
||||||
|
|
||||||
|
# Tick counter for the dense whole-body (OpenHLM, mode-0) path's encoder cadence.
|
||||||
|
self._wb_step = 0
|
||||||
|
# Rolling 50-frame reference trajectory (ref29 + anchor quat) built from the
|
||||||
|
# stream of per-tick whole-body commands, fed to the encoder as a batch.
|
||||||
|
self._wb_traj: deque[np.ndarray] = deque(maxlen=50)
|
||||||
|
self._wb_quat_traj: deque[np.ndarray] = deque(maxlen=50)
|
||||||
|
|
||||||
|
# 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)
|
||||||
|
|
||||||
|
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.
|
||||||
|
|
||||||
|
The 29 joint targets are held across the encoder lookahead window (zero
|
||||||
|
velocity) and the root roll/pitch set the anchor orientation, then the
|
||||||
|
encoder/decoder run directly (planner bypassed). One command per tick, so the
|
||||||
|
VLA's commanded pose is what SONIC tracks.
|
||||||
|
"""
|
||||||
|
ref, anchor = _wb34_to_reference(wb)
|
||||||
|
c = self.controller
|
||||||
|
if c.encode_mode != 0:
|
||||||
|
c.encode_mode = 0
|
||||||
|
c.reinit_heading = True
|
||||||
|
# Capture the heading/anchor reference on the first whole-body tick. The
|
||||||
|
# controller only latches ``init_ref_quat`` (and the base heading) inside
|
||||||
|
# ``step()`` when ``first_motion or reinit_heading`` — but it already boots in
|
||||||
|
# mode 0, so the mode-switch guard above misses the very first command and the
|
||||||
|
# anchor would stay identity. This mirrors the GEAR reference, which seeds
|
||||||
|
# ``init_ref_quat`` from the first anchor. Must run before the buffers below so
|
||||||
|
# ``step()`` latches ``motion_body_quats[0]`` = this tick's anchor.
|
||||||
|
if self._wb_step == 0:
|
||||||
|
c.reinit_heading = True
|
||||||
|
|
||||||
|
# Accumulate the per-tick commands into a rolling 50-frame reference
|
||||||
|
# trajectory so the encoder's 10-frame, step-5 lookahead sees an actual
|
||||||
|
# motion sequence (with velocities) instead of one repeated pose. 50 frames
|
||||||
|
# == chunk horizon == 10 lookahead frames × step 5.
|
||||||
|
self._wb_traj.append(ref)
|
||||||
|
self._wb_quat_traj.append(anchor)
|
||||||
|
traj = np.asarray(self._wb_traj, np.float32) # (L, 29), oldest -> newest
|
||||||
|
quats = np.asarray(self._wb_quat_traj, np.float32) # (L, 4)
|
||||||
|
n = len(traj)
|
||||||
|
# Per-frame velocities from finite differences (rad/s at the control rate).
|
||||||
|
vel = np.zeros_like(traj)
|
||||||
|
if n > 1:
|
||||||
|
vel[1:] = (traj[1:] - traj[:-1]) / CONTROL_DT
|
||||||
|
vel[0] = vel[1]
|
||||||
|
with c.motion_lock:
|
||||||
|
c.motion_joint_positions[:n] = traj
|
||||||
|
c.motion_joint_velocities[:n] = vel
|
||||||
|
c.motion_body_quats[:n] = quats
|
||||||
|
c.motion_body_pos[:n] = 0.0
|
||||||
|
c.motion_timesteps = n
|
||||||
|
c.ref_cursor = 0
|
||||||
|
c.playing = True
|
||||||
|
do_enc = self._wb_step % ENCODER_UPDATE_EVERY == 0
|
||||||
|
out = c.step(obs, update_encoder=do_enc, debug=False)
|
||||||
|
if self._wb_step % 25 == 0:
|
||||||
|
tgt = np.array([out[f"{m.name}.q"] for m in G1_29_JointIndex], np.float32)
|
||||||
|
logger.info(
|
||||||
|
"[WB34] step=%d |ref|mean=%.3f |target|mean=%.3f target_std=%.3f init_ref_quat=%s",
|
||||||
|
self._wb_step,
|
||||||
|
float(np.abs(ref).mean()),
|
||||||
|
float(np.abs(tgt).mean()),
|
||||||
|
float(tgt.std()),
|
||||||
|
np.round(c.init_ref_quat, 3).tolist(),
|
||||||
|
)
|
||||||
|
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
|
||||||
|
live commanded target, so the handoff has no snap.
|
||||||
|
|
||||||
|
``out`` is the policy's ``<joint>.q`` target dict for this tick; the blend ratio
|
||||||
|
climbs 0->1 over the ramp, after which the raw policy target passes through.
|
||||||
|
"""
|
||||||
|
if self._init_step >= self._init_ramp_steps or not out:
|
||||||
|
return out
|
||||||
|
if self._init_step == 0:
|
||||||
|
# Capture the robot's actual pose as the interpolation start point.
|
||||||
|
self._start_pose = {
|
||||||
|
f"{m.name}.q": float(obs.get(f"{m.name}.q", DEFAULT_ANGLES[m.value]))
|
||||||
|
for m in G1_29_JointIndex
|
||||||
|
}
|
||||||
|
self._init_step += 1
|
||||||
|
ratio = min(1.0, self._init_step / self._init_ramp_steps)
|
||||||
|
blended = {
|
||||||
|
k: self._start_pose.get(k, float(tgt)) * (1.0 - ratio) + float(tgt) * ratio
|
||||||
|
for k, tgt in out.items()
|
||||||
|
}
|
||||||
|
if self._init_step >= self._init_ramp_steps:
|
||||||
|
logger.info("SONIC startup blend complete -> full policy control")
|
||||||
|
return blended
|
||||||
|
|
||||||
|
def run_step(self, action: dict, lowstate) -> dict:
|
||||||
|
if lowstate is None:
|
||||||
|
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.
|
||||||
|
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)
|
||||||
|
|
||||||
|
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()
|
||||||
@@ -23,10 +23,102 @@ import numpy as np
|
|||||||
|
|
||||||
NUM_MOTORS = 29
|
NUM_MOTORS = 29
|
||||||
|
|
||||||
|
# Joint-order permutations between the two 29-DoF layouts used across the G1 stack:
|
||||||
|
# IsaacLab (policy/training order) and MuJoCo (deploy order). ``a[ISAACLAB_TO_MUJOCO]``
|
||||||
|
# reorders an IsaacLab-ordered vector into MuJoCo order, and vice-versa.
|
||||||
|
ISAACLAB_TO_MUJOCO = np.array(
|
||||||
|
[
|
||||||
|
0,
|
||||||
|
3,
|
||||||
|
6,
|
||||||
|
9,
|
||||||
|
13,
|
||||||
|
17,
|
||||||
|
1,
|
||||||
|
4,
|
||||||
|
7,
|
||||||
|
10,
|
||||||
|
14,
|
||||||
|
18,
|
||||||
|
2,
|
||||||
|
5,
|
||||||
|
8,
|
||||||
|
11,
|
||||||
|
15,
|
||||||
|
19,
|
||||||
|
21,
|
||||||
|
23,
|
||||||
|
25,
|
||||||
|
27,
|
||||||
|
12,
|
||||||
|
16,
|
||||||
|
20,
|
||||||
|
22,
|
||||||
|
24,
|
||||||
|
26,
|
||||||
|
28,
|
||||||
|
],
|
||||||
|
dtype=np.int32,
|
||||||
|
)
|
||||||
|
MUJOCO_TO_ISAACLAB = np.array(
|
||||||
|
[
|
||||||
|
0,
|
||||||
|
6,
|
||||||
|
12,
|
||||||
|
1,
|
||||||
|
7,
|
||||||
|
13,
|
||||||
|
2,
|
||||||
|
8,
|
||||||
|
14,
|
||||||
|
3,
|
||||||
|
9,
|
||||||
|
15,
|
||||||
|
22,
|
||||||
|
4,
|
||||||
|
10,
|
||||||
|
16,
|
||||||
|
23,
|
||||||
|
5,
|
||||||
|
11,
|
||||||
|
17,
|
||||||
|
24,
|
||||||
|
18,
|
||||||
|
25,
|
||||||
|
19,
|
||||||
|
26,
|
||||||
|
20,
|
||||||
|
27,
|
||||||
|
21,
|
||||||
|
28,
|
||||||
|
],
|
||||||
|
dtype=np.int32,
|
||||||
|
)
|
||||||
|
|
||||||
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
|
REMOTE_AXES = ("remote.lx", "remote.ly", "remote.rx", "remote.ry")
|
||||||
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
|
REMOTE_BUTTONS = tuple(f"remote.button.{i}" for i in range(16))
|
||||||
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
REMOTE_KEYS = REMOTE_AXES + REMOTE_BUTTONS
|
||||||
|
|
||||||
|
# Reserved action-dict field used to forward the set of currently-pressed keyboard
|
||||||
|
# keys from a KeyboardTeleop through the standard action pipeline to the SONIC
|
||||||
|
# whole-body controller (see SonicWholeBodyController._process_keyboard).
|
||||||
|
KEYBOARD_KEYS_FIELD = "keyboard.keys"
|
||||||
|
|
||||||
|
# ── Dense whole-body joint reference (SONIC encode_mode 0, OpenHLM / pi0.5) ──────
|
||||||
|
# A single 34-D whole-body command per tick, in the 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)]
|
||||||
|
# Fed as flat scalars ``wb.0.pos .. wb.33.pos``. The ``.pos`` suffix makes these
|
||||||
|
# behave like ordinary joint-position action features so ``lerobot-rollout`` routes
|
||||||
|
# them straight from a 34-D VLA (OpenHLM / pi0.5) onto the robot.
|
||||||
|
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"
|
||||||
|
|
||||||
|
|
||||||
def default_remote_input() -> dict[str, float]:
|
def default_remote_input() -> dict[str, float]:
|
||||||
"""Return a zeroed-out remote input dict (axes + buttons)."""
|
"""Return a zeroed-out remote input dict (axes + buttons)."""
|
||||||
@@ -63,13 +155,92 @@ class G1_29_JointArmIndex(IntEnum):
|
|||||||
kRightWristYaw = 28
|
kRightWristYaw = 28
|
||||||
|
|
||||||
|
|
||||||
|
def lowstate_to_obs(lowstate) -> dict:
|
||||||
|
"""Build a robot observation dict from a Unitree lowstate.
|
||||||
|
|
||||||
|
Shared by ``UnitreeG1.get_observation`` and the SONIC pipeline so the
|
||||||
|
lowstate -> obs mapping lives in exactly one place. Keys match the
|
||||||
|
``<joint>.q``/``imu.*`` schema consumed across the controllers.
|
||||||
|
"""
|
||||||
|
obs: dict = {}
|
||||||
|
|
||||||
|
for motor in G1_29_JointIndex:
|
||||||
|
idx = motor.value
|
||||||
|
obs[f"{motor.name}.q"] = lowstate.motor_state[idx].q
|
||||||
|
obs[f"{motor.name}.dq"] = lowstate.motor_state[idx].dq
|
||||||
|
obs[f"{motor.name}.tau"] = lowstate.motor_state[idx].tau_est
|
||||||
|
|
||||||
|
imu = lowstate.imu_state
|
||||||
|
if imu.gyroscope:
|
||||||
|
obs["imu.gyro.x"] = imu.gyroscope[0]
|
||||||
|
obs["imu.gyro.y"] = imu.gyroscope[1]
|
||||||
|
obs["imu.gyro.z"] = imu.gyroscope[2]
|
||||||
|
if imu.accelerometer:
|
||||||
|
obs["imu.accel.x"] = imu.accelerometer[0]
|
||||||
|
obs["imu.accel.y"] = imu.accelerometer[1]
|
||||||
|
obs["imu.accel.z"] = imu.accelerometer[2]
|
||||||
|
if imu.quaternion:
|
||||||
|
obs["imu.quat.w"] = imu.quaternion[0]
|
||||||
|
obs["imu.quat.x"] = imu.quaternion[1]
|
||||||
|
obs["imu.quat.y"] = imu.quaternion[2]
|
||||||
|
obs["imu.quat.z"] = imu.quaternion[3]
|
||||||
|
if imu.rpy:
|
||||||
|
obs["imu.rpy.roll"] = imu.rpy[0]
|
||||||
|
obs["imu.rpy.pitch"] = imu.rpy[1]
|
||||||
|
obs["imu.rpy.yaw"] = imu.rpy[2]
|
||||||
|
|
||||||
|
wr = getattr(lowstate, "wireless_remote", None)
|
||||||
|
if wr:
|
||||||
|
obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
|
||||||
|
|
||||||
|
return obs
|
||||||
|
|
||||||
|
|
||||||
|
def obs_to_wb34_state(obs: dict) -> np.ndarray:
|
||||||
|
"""Build the 34-D OpenHLM / pi0.5 proprio state from a G1 observation dict.
|
||||||
|
|
||||||
|
Mirrors the whole-body *action* layout so the policy sees state and action in
|
||||||
|
the same coordinates::
|
||||||
|
|
||||||
|
[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)]
|
||||||
|
|
||||||
|
Joint positions come from the ``<joint>.q`` obs keys, which are already in
|
||||||
|
MuJoCo / Unitree-SDK order — the same body-part grouping OpenHLM uses
|
||||||
|
([L-leg 0:6, R-leg 6:12, waist 12:15, L-arm 15:22, R-arm 22:29]) — so they are
|
||||||
|
regrouped directly (no IsaacLab permutation). The G1 has no grippers in its
|
||||||
|
29-DoF body, so both gripper slots are 0. Root roll/pitch are the IMU RPY and
|
||||||
|
the last slot is the IMU yaw rate (gyro z).
|
||||||
|
"""
|
||||||
|
q_mj = np.array(
|
||||||
|
[float(obs.get(f"{m.name}.q", 0.0)) for m in G1_29_JointIndex],
|
||||||
|
dtype=np.float32,
|
||||||
|
)
|
||||||
|
lleg, rleg, waist = q_mj[0:6], q_mj[6:12], q_mj[12:15]
|
||||||
|
larm, rarm = q_mj[15:22], q_mj[22:29]
|
||||||
|
|
||||||
|
state = np.zeros(34, dtype=np.float32)
|
||||||
|
state[0:7] = larm
|
||||||
|
# state[7] left gripper — none on 29-DoF G1
|
||||||
|
state[8:15] = rarm
|
||||||
|
# state[15] right gripper — none on 29-DoF G1
|
||||||
|
state[16:22] = lleg
|
||||||
|
state[22:28] = rleg
|
||||||
|
state[28:31] = waist
|
||||||
|
state[31] = float(obs.get("imu.rpy.roll", 0.0))
|
||||||
|
state[32] = float(obs.get("imu.rpy.pitch", 0.0))
|
||||||
|
state[33] = float(obs.get("imu.gyro.z", 0.0))
|
||||||
|
return state
|
||||||
|
|
||||||
|
|
||||||
def make_locomotion_controller(name: str | None):
|
def make_locomotion_controller(name: str | None):
|
||||||
"""Instantiate a locomotion controller by class name. Returns None if name is None."""
|
"""Instantiate a locomotion controller by class name. Returns None if name is None."""
|
||||||
if name is None:
|
if name is None:
|
||||||
return None
|
return None
|
||||||
controllers = {
|
controllers = {
|
||||||
"GrootLocomotionController": "lerobot.robots.unitree_g1.gr00t_locomotion",
|
"GrootLocomotionController": "lerobot.robots.unitree_g1.controllers.gr00t_locomotion",
|
||||||
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.holosoma_locomotion",
|
"HolosomaLocomotionController": "lerobot.robots.unitree_g1.controllers.holosoma_locomotion",
|
||||||
|
"SonicWholeBodyController": "lerobot.robots.unitree_g1.controllers.sonic_whole_body",
|
||||||
}
|
}
|
||||||
module_path = controllers.get(name)
|
module_path = controllers.get(name)
|
||||||
if module_path is None:
|
if module_path is None:
|
||||||
|
|||||||
@@ -33,12 +33,14 @@ from ..robot import Robot
|
|||||||
from .config_unitree_g1 import UnitreeG1Config
|
from .config_unitree_g1 import UnitreeG1Config
|
||||||
from .g1_kinematics import G1_29_ArmIK
|
from .g1_kinematics import G1_29_ArmIK
|
||||||
from .g1_utils import (
|
from .g1_utils import (
|
||||||
|
KEYBOARD_KEYS_FIELD,
|
||||||
REMOTE_AXES,
|
REMOTE_AXES,
|
||||||
REMOTE_KEYS,
|
|
||||||
G1_29_JointArmIndex,
|
G1_29_JointArmIndex,
|
||||||
G1_29_JointIndex,
|
G1_29_JointIndex,
|
||||||
default_remote_input,
|
default_remote_input,
|
||||||
|
lowstate_to_obs,
|
||||||
make_locomotion_controller,
|
make_locomotion_controller,
|
||||||
|
obs_to_wb34_state,
|
||||||
)
|
)
|
||||||
|
|
||||||
if TYPE_CHECKING or _unitree_sdk_available:
|
if TYPE_CHECKING or _unitree_sdk_available:
|
||||||
@@ -47,8 +49,12 @@ if TYPE_CHECKING or _unitree_sdk_available:
|
|||||||
ChannelPublisher as _SDKChannelPublisher,
|
ChannelPublisher as _SDKChannelPublisher,
|
||||||
ChannelSubscriber as _SDKChannelSubscriber,
|
ChannelSubscriber as _SDKChannelSubscriber,
|
||||||
)
|
)
|
||||||
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
|
from unitree_sdk2py.idl.default import (
|
||||||
|
unitree_hg_msg_dds__HandCmd_ as hg_HandCmd_default,
|
||||||
|
unitree_hg_msg_dds__LowCmd_,
|
||||||
|
)
|
||||||
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
|
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import (
|
||||||
|
HandCmd_ as hg_HandCmd,
|
||||||
LowCmd_ as hg_LowCmd,
|
LowCmd_ as hg_LowCmd,
|
||||||
LowState_ as hg_LowState,
|
LowState_ as hg_LowState,
|
||||||
)
|
)
|
||||||
@@ -58,6 +64,8 @@ else:
|
|||||||
_SDKChannelPublisher = None
|
_SDKChannelPublisher = None
|
||||||
_SDKChannelSubscriber = None
|
_SDKChannelSubscriber = None
|
||||||
unitree_hg_msg_dds__LowCmd_ = None
|
unitree_hg_msg_dds__LowCmd_ = None
|
||||||
|
hg_HandCmd_default = None
|
||||||
|
hg_HandCmd = None
|
||||||
hg_LowCmd = None
|
hg_LowCmd = None
|
||||||
hg_LowState = None
|
hg_LowState = None
|
||||||
CRC = None
|
CRC = None
|
||||||
@@ -157,6 +165,37 @@ class UnitreeG1(Robot):
|
|||||||
self.controller_input = default_remote_input()
|
self.controller_input = default_remote_input()
|
||||||
self.controller_output = {}
|
self.controller_output = {}
|
||||||
|
|
||||||
|
# Replay-camera state (decoded frames per robot camera name + play cursor).
|
||||||
|
self._replay_frames: dict[str, list[np.ndarray]] = {}
|
||||||
|
self._replay_len = 0
|
||||||
|
self._replay_idx = 0
|
||||||
|
if config.replay_camera_parquet and config.replay_camera_map:
|
||||||
|
self._load_replay_frames()
|
||||||
|
|
||||||
|
def _load_replay_frames(self) -> None:
|
||||||
|
"""Decode recorded episode frames from a parquet into per-camera image lists."""
|
||||||
|
import io
|
||||||
|
|
||||||
|
import pyarrow.parquet as pq
|
||||||
|
from PIL import Image
|
||||||
|
|
||||||
|
table = pq.read_table(self.config.replay_camera_parquet)
|
||||||
|
cols = {col: table.column(col).to_pylist() for col in self.config.replay_camera_map.values()}
|
||||||
|
self._replay_len = table.num_rows
|
||||||
|
|
||||||
|
def decode(cell) -> np.ndarray:
|
||||||
|
data = cell["bytes"] if isinstance(cell, dict) else cell
|
||||||
|
return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"), dtype=np.uint8)
|
||||||
|
|
||||||
|
for cam_name, column in self.config.replay_camera_map.items():
|
||||||
|
self._replay_frames[cam_name] = [decode(c) for c in cols[column]]
|
||||||
|
logger.info(
|
||||||
|
"Loaded %d replay frames for cameras %s from %s",
|
||||||
|
self._replay_len,
|
||||||
|
list(self.config.replay_camera_map),
|
||||||
|
self.config.replay_camera_parquet,
|
||||||
|
)
|
||||||
|
|
||||||
def _subscribe_lowstate(self): # polls robot state @ 250Hz
|
def _subscribe_lowstate(self): # polls robot state @ 250Hz
|
||||||
while not self._shutdown_event.is_set():
|
while not self._shutdown_event.is_set():
|
||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
@@ -231,15 +270,54 @@ class UnitreeG1(Robot):
|
|||||||
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
features[f"{cam}_depth"] = (cfg.height, cfg.width, 1)
|
||||||
return features
|
return features
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _wb_state_ft(self) -> dict[str, type]:
|
||||||
|
"""34-D whole-body proprio state (``wb_state.{i}.pos``) for dense controllers.
|
||||||
|
|
||||||
|
Exposed only when the controller consumes a dense whole-body command
|
||||||
|
(OpenHLM / pi0.5). These ``.pos`` scalars are aggregated by the rollout
|
||||||
|
pipeline into a single 34-D ``observation.state`` for the policy.
|
||||||
|
"""
|
||||||
|
if not getattr(self.controller, "wb_action", False):
|
||||||
|
return {}
|
||||||
|
from .g1_utils import WB_ACTION_DIM
|
||||||
|
|
||||||
|
return {f"wb_state.{i}.pos": float for i in range(WB_ACTION_DIM)}
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _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}
|
||||||
|
|
||||||
|
@property
|
||||||
|
def _replay_cameras_ft(self) -> dict[str, tuple]:
|
||||||
|
"""Replay cameras, shaped from their first decoded frame."""
|
||||||
|
return {name: frames[0].shape for name, frames in self._replay_frames.items() if frames}
|
||||||
|
|
||||||
@cached_property
|
@cached_property
|
||||||
def observation_features(self) -> dict[str, type | tuple]:
|
def observation_features(self) -> dict[str, type | tuple]:
|
||||||
return {**self._motors_ft, **self._cameras_ft}
|
return {
|
||||||
|
**self._motors_ft,
|
||||||
|
**self._wb_state_ft,
|
||||||
|
**self._empty_cameras_ft,
|
||||||
|
**self._replay_cameras_ft,
|
||||||
|
**self._cameras_ft,
|
||||||
|
}
|
||||||
|
|
||||||
@cached_property
|
@cached_property
|
||||||
def action_features(self) -> dict[str, type]:
|
def action_features(self) -> dict[str, type]:
|
||||||
if self.controller is None:
|
if self.controller is None:
|
||||||
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
|
return {f"{G1_29_JointIndex(motor).name}.q": float for motor in G1_29_JointIndex}
|
||||||
|
|
||||||
|
# Dense whole-body controllers (SONIC / OpenHLM, pi0.5) consume a single
|
||||||
|
# 34-D command per tick. Expose it as ``wb.{i}.pos`` joint-position features
|
||||||
|
# so ``lerobot-rollout`` maps a 34-D policy output straight onto the robot.
|
||||||
|
if getattr(self.controller, "wb_action", False):
|
||||||
|
from .g1_utils import WB_ACTION_DIM, wb_action_key
|
||||||
|
|
||||||
|
return {wb_action_key(i): float for i in range(WB_ACTION_DIM)}
|
||||||
|
|
||||||
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
|
arm_features = {f"{G1_29_JointArmIndex(motor).name}.q": float for motor in G1_29_JointArmIndex}
|
||||||
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
remote_features = dict.fromkeys(REMOTE_AXES, float)
|
||||||
return {**arm_features, **remote_features}
|
return {**arm_features, **remote_features}
|
||||||
@@ -311,6 +389,17 @@ class UnitreeG1(Robot):
|
|||||||
self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState)
|
self.lowstate_subscriber = self._ChannelSubscriber(kTopicLowState, hg_LowState)
|
||||||
self.lowstate_subscriber.Init()
|
self.lowstate_subscriber.Init()
|
||||||
|
|
||||||
|
# Dex3 hand command publishers (grasping). Driven by the OpenHLM grip scalars.
|
||||||
|
self._hand_publishers = {}
|
||||||
|
if self.config.publish_hands:
|
||||||
|
self._left_hand_cmd = hg_HandCmd_default()
|
||||||
|
self._right_hand_cmd = hg_HandCmd_default()
|
||||||
|
self._hand_publishers["left"] = self._ChannelPublisher("rt/dex3/left/cmd", hg_HandCmd)
|
||||||
|
self._hand_publishers["right"] = self._ChannelPublisher("rt/dex3/right/cmd", hg_HandCmd)
|
||||||
|
for pub in self._hand_publishers.values():
|
||||||
|
pub.Init()
|
||||||
|
logger.info("Dex3 hand command publishers initialized (rt/dex3/{left,right}/cmd)")
|
||||||
|
|
||||||
# Start subscribe thread to read robot state
|
# Start subscribe thread to read robot state
|
||||||
self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
|
self.subscribe_thread = threading.Thread(target=self._subscribe_lowstate)
|
||||||
self.subscribe_thread.start()
|
self.subscribe_thread.start()
|
||||||
@@ -343,6 +432,9 @@ class UnitreeG1(Robot):
|
|||||||
|
|
||||||
self.kp = np.array(self.config.kp, dtype=np.float32)
|
self.kp = np.array(self.config.kp, dtype=np.float32)
|
||||||
self.kd = np.array(self.config.kd, dtype=np.float32)
|
self.kd = np.array(self.config.kd, dtype=np.float32)
|
||||||
|
if self.controller is not None and hasattr(self.controller, "kp"):
|
||||||
|
self.kp = np.array(self.controller.kp, dtype=np.float32)
|
||||||
|
self.kd = np.array(self.controller.kd, dtype=np.float32)
|
||||||
|
|
||||||
for joint in G1_29_JointIndex:
|
for joint in G1_29_JointIndex:
|
||||||
self.msg.motor_cmd[joint].mode = 1
|
self.msg.motor_cmd[joint].mode = 1
|
||||||
@@ -371,13 +463,50 @@ class UnitreeG1(Robot):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.warning(f"Failed to send zero-torque on disconnect: {e}")
|
logger.warning(f"Failed to send zero-torque on disconnect: {e}")
|
||||||
|
|
||||||
def disconnect(self):
|
def _graceful_stop(self) -> None:
|
||||||
# Put robot in passive mode before stopping threads
|
"""Soft shutdown: hold the current pose and ramp joint stiffness (kp) to zero
|
||||||
if not self.config.is_simulation:
|
over ``graceful_stop_s`` while keeping damping (kd), then go passive.
|
||||||
self._send_zero_torque()
|
|
||||||
|
|
||||||
# Signal thread to stop and unblock any waits
|
Prevents the robot from collapsing the instant control ends (a bare
|
||||||
|
zero-torque command is kp=kd=0 ≈ free-fall). Must run after the controller
|
||||||
|
loop has stopped so the two aren't publishing at once.
|
||||||
|
"""
|
||||||
|
if self.config.graceful_stop_s <= 0:
|
||||||
|
self._send_zero_torque()
|
||||||
|
return
|
||||||
|
with self._lowstate_lock:
|
||||||
|
lowstate = self._lowstate
|
||||||
|
if lowstate is None:
|
||||||
|
self._send_zero_torque()
|
||||||
|
return
|
||||||
|
q_hold = {f"{motor.name}.q": lowstate.motor_state[motor.value].q for motor in G1_29_JointIndex}
|
||||||
|
kp = np.array(self.kp, dtype=np.float32)
|
||||||
|
kd = np.array(self.kd, dtype=np.float32)
|
||||||
|
zeros = np.zeros(29, dtype=np.float32)
|
||||||
|
dt = self.controller.control_dt if self.controller is not None else self.config.control_dt
|
||||||
|
steps = max(1, int(self.config.graceful_stop_s / dt))
|
||||||
|
logger.info("Graceful stop: damping down over %.1fs", self.config.graceful_stop_s)
|
||||||
|
for i in range(steps):
|
||||||
|
ratio = (i + 1) / steps
|
||||||
|
self.publish_lowcmd(q_hold, kp=kp * (1.0 - ratio), kd=kd, tau=zeros)
|
||||||
|
time.sleep(dt)
|
||||||
|
self._send_zero_torque()
|
||||||
|
|
||||||
|
def disconnect(self):
|
||||||
|
# Stop the controller loop first so it isn't fighting the shutdown ramp.
|
||||||
self._shutdown_event.set()
|
self._shutdown_event.set()
|
||||||
|
if self._controller_thread is not None:
|
||||||
|
self._controller_thread.join(timeout=2.0)
|
||||||
|
if self._controller_thread.is_alive():
|
||||||
|
logger.warning("Controller thread did not stop cleanly")
|
||||||
|
|
||||||
|
# Soft, damped settle instead of an instant limp (real robot only; the
|
||||||
|
# subscribe thread is still alive here to supply the current pose).
|
||||||
|
if not self.config.is_simulation:
|
||||||
|
self._graceful_stop()
|
||||||
|
|
||||||
|
if self.controller is not None and hasattr(self.controller, "shutdown"):
|
||||||
|
self.controller.shutdown()
|
||||||
|
|
||||||
# Wait for subscribe thread to finish
|
# Wait for subscribe thread to finish
|
||||||
if self.subscribe_thread is not None:
|
if self.subscribe_thread is not None:
|
||||||
@@ -385,12 +514,6 @@ class UnitreeG1(Robot):
|
|||||||
if self.subscribe_thread.is_alive():
|
if self.subscribe_thread.is_alive():
|
||||||
logger.warning("Subscribe thread did not stop cleanly")
|
logger.warning("Subscribe thread did not stop cleanly")
|
||||||
|
|
||||||
# Wait for controller thread to finish
|
|
||||||
if self._controller_thread is not None:
|
|
||||||
self._controller_thread.join(timeout=2.0)
|
|
||||||
if self._controller_thread.is_alive():
|
|
||||||
logger.warning("Controller thread did not stop cleanly")
|
|
||||||
|
|
||||||
# Close simulation environment
|
# Close simulation environment
|
||||||
if self.config.is_simulation and self.sim_env is not None:
|
if self.config.is_simulation and self.sim_env is not None:
|
||||||
try:
|
try:
|
||||||
@@ -422,44 +545,33 @@ class UnitreeG1(Robot):
|
|||||||
if lowstate is None:
|
if lowstate is None:
|
||||||
return {}
|
return {}
|
||||||
|
|
||||||
obs = {}
|
# Motors + IMU + wireless remote (shared lowstate -> obs mapping)
|
||||||
|
obs = lowstate_to_obs(lowstate)
|
||||||
|
|
||||||
# Motors - q, dq, tau for all joints
|
# Dense whole-body controllers (OpenHLM / pi0.5): expose the 34-D proprio
|
||||||
for motor in G1_29_JointIndex:
|
# state as ``wb_state.{i}.pos`` so the rollout aggregates it into
|
||||||
name = motor.name
|
# ``observation.state`` for the policy.
|
||||||
idx = motor.value
|
if getattr(self.controller, "wb_action", False):
|
||||||
obs[f"{name}.q"] = lowstate.motor_state[idx].q
|
wb_state = obs_to_wb34_state(obs)
|
||||||
obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
|
for i, v in enumerate(wb_state):
|
||||||
obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
|
obs[f"wb_state.{i}.pos"] = float(v)
|
||||||
|
|
||||||
# IMU - gyroscope
|
# Synthetic empty cameras: black frames so image-conditioned policies run
|
||||||
if lowstate.imu_state.gyroscope:
|
# before real cameras are wired.
|
||||||
obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
|
if self.config.empty_cameras:
|
||||||
obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
|
h, w = self.config.empty_camera_hw
|
||||||
obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
|
black = np.zeros((h, w, 3), dtype=np.uint8)
|
||||||
|
for name in self.config.empty_cameras:
|
||||||
|
obs[name] = black
|
||||||
|
|
||||||
# IMU - accelerometer
|
# Replay cameras: serve the current recorded frame per camera, then advance.
|
||||||
if lowstate.imu_state.accelerometer:
|
if self._replay_len:
|
||||||
obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
|
idx = self._replay_idx
|
||||||
obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
|
if idx >= self._replay_len:
|
||||||
obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
|
idx = self._replay_len - 1 if not self.config.replay_camera_loop else idx % self._replay_len
|
||||||
|
for name, frames in self._replay_frames.items():
|
||||||
# IMU - quaternion
|
obs[name] = frames[idx]
|
||||||
if lowstate.imu_state.quaternion:
|
self._replay_idx += 1
|
||||||
obs["imu.quat.w"] = lowstate.imu_state.quaternion[0]
|
|
||||||
obs["imu.quat.x"] = lowstate.imu_state.quaternion[1]
|
|
||||||
obs["imu.quat.y"] = lowstate.imu_state.quaternion[2]
|
|
||||||
obs["imu.quat.z"] = lowstate.imu_state.quaternion[3]
|
|
||||||
|
|
||||||
# IMU - rpy
|
|
||||||
if lowstate.imu_state.rpy:
|
|
||||||
obs["imu.rpy.roll"] = lowstate.imu_state.rpy[0]
|
|
||||||
obs["imu.rpy.pitch"] = lowstate.imu_state.rpy[1]
|
|
||||||
obs["imu.rpy.yaw"] = lowstate.imu_state.rpy[2]
|
|
||||||
|
|
||||||
# Wireless remote (raw bytes for teleoperator)
|
|
||||||
if lowstate.wireless_remote:
|
|
||||||
obs["wireless_remote"] = lowstate.wireless_remote
|
|
||||||
|
|
||||||
# Cameras - read images from ZMQ cameras
|
# Cameras - read images from ZMQ cameras
|
||||||
for cam_name, cam in self._cameras.items():
|
for cam_name, cam in self._cameras.items():
|
||||||
@@ -473,9 +585,13 @@ class UnitreeG1(Robot):
|
|||||||
def send_action(self, action: RobotAction) -> RobotAction:
|
def send_action(self, action: RobotAction) -> RobotAction:
|
||||||
action_to_publish = action
|
action_to_publish = action
|
||||||
if self.controller is not None:
|
if self.controller is not None:
|
||||||
|
self._update_controller_action(action)
|
||||||
|
if self.config.publish_hands and getattr(self.controller, "wb_action", False):
|
||||||
|
self._publish_hand_cmds(action)
|
||||||
|
if getattr(self.controller, "full_body", False):
|
||||||
|
return action
|
||||||
# Controller thread owns legs/waist. Here we only update joystick inputs
|
# Controller thread owns legs/waist. Here we only update joystick inputs
|
||||||
# and publish arm targets from the teleoperator.
|
# and publish arm targets from the teleoperator.
|
||||||
self._update_controller_action(action)
|
|
||||||
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
|
arm_prefixes = tuple(j.name for j in G1_29_JointArmIndex)
|
||||||
action_to_publish = {
|
action_to_publish = {
|
||||||
key: value
|
key: value
|
||||||
@@ -503,11 +619,67 @@ class UnitreeG1(Robot):
|
|||||||
return action
|
return action
|
||||||
|
|
||||||
def _update_controller_action(self, action: RobotAction) -> None:
|
def _update_controller_action(self, action: RobotAction) -> None:
|
||||||
"""Update controller input state from incoming teleop action."""
|
"""Update controller input state from an incoming teleop action.
|
||||||
|
|
||||||
|
Controller-agnostic: every value-carrying key is forwarded verbatim into
|
||||||
|
``controller_input`` (whole-body ``wb.{i}.pos`` from a 34-D VLA, or whatever a
|
||||||
|
future controller expects), and each controller extracts only the keys it
|
||||||
|
understands. The robot deliberately does not enumerate any controller's key
|
||||||
|
schema here.
|
||||||
|
|
||||||
|
KeyboardTeleop is the one special case: it emits the currently-pressed keys as
|
||||||
|
bare action keys with a ``None`` value (``dict.fromkeys(pressed, None)``), so
|
||||||
|
those are collected into a single held-key set under ``KEYBOARD_KEYS_FIELD``,
|
||||||
|
rebuilt each tick so releases clear. Special keys arrive as pynput objects and
|
||||||
|
are normalised to their name ("space", ...).
|
||||||
|
"""
|
||||||
with self._controller_action_lock:
|
with self._controller_action_lock:
|
||||||
for key in REMOTE_KEYS:
|
self.controller_input[KEYBOARD_KEYS_FIELD] = {
|
||||||
if key in action:
|
(k if isinstance(k, str) else getattr(k, "name", str(k)))
|
||||||
self.controller_input[key] = action[key]
|
for k, value in action.items()
|
||||||
|
if value is None
|
||||||
|
}
|
||||||
|
for key, value in action.items():
|
||||||
|
if isinstance(key, str) and value is not None:
|
||||||
|
self.controller_input[key] = value
|
||||||
|
|
||||||
|
def _publish_hand_cmds(self, action: RobotAction) -> None:
|
||||||
|
"""Drive the Dex3 hands from the OpenHLM grip scalars in a 34-D wb action.
|
||||||
|
|
||||||
|
``wb.7.pos`` is the left grip and ``wb.15.pos`` the right grip. Each scalar in
|
||||||
|
[0, 1] (``hand_open_grip_value`` == fully open) is turned into a curl amount and
|
||||||
|
scaled onto ``hand_closed_pose`` (7 joints), then published as a PD target on
|
||||||
|
``rt/dex3/{left,right}/cmd`` so the fingers close when the policy grips.
|
||||||
|
"""
|
||||||
|
if not self._hand_publishers:
|
||||||
|
return
|
||||||
|
from .g1_utils import wb_action_key
|
||||||
|
|
||||||
|
open_val = float(self.config.hand_open_grip_value)
|
||||||
|
closed_val = float(self.config.hand_closed_grip_value)
|
||||||
|
closed_pose = self.config.hand_closed_pose
|
||||||
|
kp, kd = float(self.config.hand_kp), float(self.config.hand_kd)
|
||||||
|
span = (closed_val - open_val) or 1.0
|
||||||
|
|
||||||
|
def curl_amount(grip: float) -> float:
|
||||||
|
# Fraction of the way from the open scalar to the closed scalar, in [0, 1].
|
||||||
|
return float(min(max((grip - open_val) / span, 0.0), 1.0))
|
||||||
|
|
||||||
|
for side, grip_idx, cmd in (
|
||||||
|
("left", 7, self._left_hand_cmd),
|
||||||
|
("right", 15, self._right_hand_cmd),
|
||||||
|
):
|
||||||
|
grip = action.get(wb_action_key(grip_idx))
|
||||||
|
if grip is None:
|
||||||
|
continue
|
||||||
|
amount = curl_amount(float(grip))
|
||||||
|
for i, closed_q in enumerate(closed_pose):
|
||||||
|
cmd.motor_cmd[i].q = float(closed_q) * amount
|
||||||
|
cmd.motor_cmd[i].dq = 0.0
|
||||||
|
cmd.motor_cmd[i].kp = kp
|
||||||
|
cmd.motor_cmd[i].kd = kd
|
||||||
|
cmd.motor_cmd[i].tau = 0.0
|
||||||
|
self._hand_publishers[side].Write(cmd)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_calibrated(self) -> bool:
|
def is_calibrated(self) -> bool:
|
||||||
|
|||||||
@@ -116,6 +116,7 @@ from lerobot.teleoperators import ( # noqa: F401
|
|||||||
omx_leader,
|
omx_leader,
|
||||||
openarm_leader,
|
openarm_leader,
|
||||||
openarm_mini,
|
openarm_mini,
|
||||||
|
pico_headset,
|
||||||
reachy2_teleoperator,
|
reachy2_teleoperator,
|
||||||
rebot_102_leader,
|
rebot_102_leader,
|
||||||
so_leader,
|
so_leader,
|
||||||
|
|||||||
@@ -0,0 +1,20 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
from .config_pico_headset import PicoHeadsetConfig
|
||||||
|
from .pico_headset import PicoHeadset
|
||||||
|
|
||||||
|
__all__ = ["PicoHeadset", "PicoHeadsetConfig"]
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
#!/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)."""
|
||||||
@@ -0,0 +1,148 @@
|
|||||||
|
#!/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
|
||||||
@@ -0,0 +1,333 @@
|
|||||||
|
#!/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()
|
||||||
@@ -0,0 +1,80 @@
|
|||||||
|
#!/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"
|
||||||
@@ -0,0 +1,571 @@
|
|||||||
|
#!/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),
|
||||||
|
}
|
||||||
@@ -0,0 +1,277 @@
|
|||||||
|
#!/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)
|
||||||
@@ -0,0 +1,147 @@
|
|||||||
|
#!/usr/bin/env python
|
||||||
|
|
||||||
|
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||||
|
#
|
||||||
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||||
|
# you may not use this file except in compliance with the License.
|
||||||
|
# You may obtain a copy of the License at
|
||||||
|
#
|
||||||
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||||||
|
#
|
||||||
|
# Unless required by applicable law or agreed to in writing, software
|
||||||
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||||
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||||
|
# See the License for the specific language governing permissions and
|
||||||
|
# limitations under the License.
|
||||||
|
|
||||||
|
"""Convert SMPL ``.npz`` motion clips into a LeRobotDataset for SONIC replay.
|
||||||
|
|
||||||
|
Each dataset frame's ``action`` is the 720-dim ``smpl.*`` window that the
|
||||||
|
``pico_headset`` teleoperator emits and ``SonicWholeBodyController`` reassembles
|
||||||
|
into ``encode_mode == 2``. So the resulting dataset can be pushed straight
|
||||||
|
through ``lerobot-replay`` to drive SONIC whole-body tracking with no headset:
|
||||||
|
|
||||||
|
lerobot-replay \
|
||||||
|
--robot.type=unitree_g1 \
|
||||||
|
--robot.controller=SonicWholeBodyController \
|
||||||
|
--dataset.repo_id=<user>/<clip> --dataset.episode=0
|
||||||
|
|
||||||
|
The 10-frame window is built exactly like the live ``SmplStream`` (oldest->newest,
|
||||||
|
the first frame repeated to pre-fill), so replayed actions match a live session.
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
# One clip -> one-episode dataset
|
||||||
|
python -m lerobot.teleoperators.pico_headset.smpl_to_dataset \
|
||||||
|
--motion-file examples/unitree_g1/motions/walk_forward.npz \
|
||||||
|
--repo-id me/sonic_walk_forward
|
||||||
|
|
||||||
|
# Every clip in a dir -> one episode each
|
||||||
|
python -m lerobot.teleoperators.pico_headset.smpl_to_dataset \
|
||||||
|
--motion-dir examples/unitree_g1/motions --repo-id me/sonic_motions
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
from lerobot.teleoperators.pico_headset.smpl_constants import (
|
||||||
|
ACTION_DIM,
|
||||||
|
JOINT_DIM,
|
||||||
|
N_JOINTS,
|
||||||
|
ROOT_ACTION_DIM as ROOT_DIM,
|
||||||
|
ROOT_ACTION_PREFIX,
|
||||||
|
SMPL_ACTION_PREFIX,
|
||||||
|
SMPL_OBS_DIM,
|
||||||
|
WINDOW,
|
||||||
|
)
|
||||||
|
from lerobot.teleoperators.pico_headset.smpl_fk import canonicalize_smpl_joints, root_quats_from_aa
|
||||||
|
|
||||||
|
|
||||||
|
def _load_canonical_joints(path: str) -> tuple[np.ndarray, np.ndarray, float]:
|
||||||
|
"""Load an SMPL clip -> (canonical (T,24,3) joints, (T,4) root wxyz, fps)."""
|
||||||
|
data = np.load(path)
|
||||||
|
joints = data["smpl_joints"].astype(np.float32)
|
||||||
|
if joints.ndim != 3 or joints.shape[1:] != (N_JOINTS, JOINT_DIM):
|
||||||
|
raise ValueError(f"{path}: expected smpl_joints (T, 24, 3), got {joints.shape}")
|
||||||
|
t = joints.shape[0]
|
||||||
|
if "pose_aa" in data.files:
|
||||||
|
root_aa = data["pose_aa"].astype(np.float32)[:, :3]
|
||||||
|
joints = canonicalize_smpl_joints(joints, root_aa)
|
||||||
|
root_quat = root_quats_from_aa(root_aa) # (T, 4) wxyz, matches live stream
|
||||||
|
else:
|
||||||
|
# No global orient available: identity root (anchor falls back to standing).
|
||||||
|
root_quat = np.tile(np.array([1.0, 0.0, 0.0, 0.0], np.float32), (t, 1))
|
||||||
|
fps = float(data["fps"]) if "fps" in data.files else 50.0
|
||||||
|
return joints, root_quat, fps
|
||||||
|
|
||||||
|
|
||||||
|
def _windows(joints: np.ndarray) -> np.ndarray:
|
||||||
|
"""(T, 24, 3) -> (T, 720): rolling 10-frame window, matching SmplStream.
|
||||||
|
|
||||||
|
Window t = frames [t-9 .. t], clamped to 0 at the start (first frame repeated).
|
||||||
|
"""
|
||||||
|
t = joints.shape[0]
|
||||||
|
idx = np.clip(np.arange(t)[:, None] + np.arange(-WINDOW + 1, 1)[None, :], 0, t - 1)
|
||||||
|
return joints[idx].reshape(t, -1).astype(np.float32)
|
||||||
|
|
||||||
|
|
||||||
|
def _action_features() -> dict:
|
||||||
|
names = [f"{SMPL_ACTION_PREFIX}{i}" for i in range(SMPL_OBS_DIM)]
|
||||||
|
names += [f"{ROOT_ACTION_PREFIX}{i}" for i in range(ROOT_DIM)]
|
||||||
|
return {"action": {"dtype": "float32", "shape": (ACTION_DIM,), "names": names}}
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> None:
|
||||||
|
p = argparse.ArgumentParser(description=__doc__)
|
||||||
|
src = p.add_mutually_exclusive_group(required=True)
|
||||||
|
src.add_argument("--motion-file", type=str, help="Single SMPL .npz clip")
|
||||||
|
src.add_argument("--motion-dir", type=str, help="Directory of .npz clips (one episode each)")
|
||||||
|
p.add_argument("--repo-id", required=True, help="Dataset repo id, e.g. user/name")
|
||||||
|
p.add_argument("--root", type=str, default=None, help="Local dataset root (default HF cache)")
|
||||||
|
p.add_argument("--fps", type=int, default=None, help="Override fps (default: clip fps)")
|
||||||
|
p.add_argument("--task", type=str, default="sonic whole-body SMPL replay")
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||||
|
|
||||||
|
if args.motion_dir:
|
||||||
|
clips = sorted(str(pth) for pth in Path(args.motion_dir).glob("*.npz"))
|
||||||
|
if not clips:
|
||||||
|
raise SystemExit(f"No .npz clips found in {args.motion_dir}")
|
||||||
|
else:
|
||||||
|
clips = [args.motion_file]
|
||||||
|
|
||||||
|
first_joints, first_root, first_fps = _load_canonical_joints(clips[0])
|
||||||
|
fps = args.fps or int(round(first_fps))
|
||||||
|
|
||||||
|
dataset = LeRobotDataset.create(
|
||||||
|
repo_id=args.repo_id,
|
||||||
|
fps=fps,
|
||||||
|
features=_action_features(),
|
||||||
|
root=args.root,
|
||||||
|
robot_type="unitree_g1",
|
||||||
|
use_videos=False,
|
||||||
|
)
|
||||||
|
|
||||||
|
for clip_i, clip in enumerate(clips):
|
||||||
|
if clip_i == 0:
|
||||||
|
joints, root_quat = first_joints, first_root
|
||||||
|
else:
|
||||||
|
joints, root_quat, _ = _load_canonical_joints(clip)
|
||||||
|
windows = _windows(joints) # (T, 720)
|
||||||
|
# action = [720 joint window | 4 root wxyz] per frame -> (T, 724)
|
||||||
|
actions = np.concatenate([windows, root_quat.astype(np.float32)], axis=1)
|
||||||
|
for a in actions:
|
||||||
|
dataset.add_frame({"action": a, "task": args.task})
|
||||||
|
dataset.save_episode()
|
||||||
|
print(f"[smpl_to_dataset] episode {clip_i}: {Path(clip).name} ({actions.shape[0]} frames)")
|
||||||
|
|
||||||
|
dataset.finalize()
|
||||||
|
print(f"[smpl_to_dataset] wrote {len(clips)} episode(s) to {dataset.root}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -60,8 +60,18 @@ def is_package_available(
|
|||||||
# If the package can't be imported, it's not available
|
# If the package can't be imported, it's not available
|
||||||
package_exists = False
|
package_exists = False
|
||||||
else:
|
else:
|
||||||
# For packages other than "torch", don't attempt the fallback and set as not available
|
# The distribution may be published under a name that differs from the
|
||||||
package_exists = False
|
# import name (e.g. ``onnxruntime`` imports from ``onnxruntime-gpu`` /
|
||||||
|
# ``onnxruntime-silicon``). Resolve the import name to its actual
|
||||||
|
# distribution(s) and read the version from there before giving up.
|
||||||
|
try:
|
||||||
|
dists = importlib.metadata.packages_distributions().get(import_name, [])
|
||||||
|
if dists:
|
||||||
|
package_version = importlib.metadata.version(dists[0])
|
||||||
|
else:
|
||||||
|
package_exists = False
|
||||||
|
except importlib.metadata.PackageNotFoundError:
|
||||||
|
package_exists = False
|
||||||
logging.debug(f"Detected {pkg_name} version: {package_version}")
|
logging.debug(f"Detected {pkg_name} version: {package_version}")
|
||||||
if return_version:
|
if return_version:
|
||||||
return package_exists, package_version
|
return package_exists, package_version
|
||||||
@@ -123,6 +133,8 @@ _pyrealsense2_available = is_package_available("pyrealsense2") or is_package_ava
|
|||||||
"pyrealsense2-macosx", import_name="pyrealsense2"
|
"pyrealsense2-macosx", import_name="pyrealsense2"
|
||||||
)
|
)
|
||||||
_zmq_available = is_package_available("pyzmq", import_name="zmq")
|
_zmq_available = is_package_available("pyzmq", import_name="zmq")
|
||||||
|
_onnxruntime_available = is_package_available("onnxruntime")
|
||||||
|
_onnx_available = is_package_available("onnx")
|
||||||
_hebi_available = is_package_available("hebi-py", import_name="hebi")
|
_hebi_available = is_package_available("hebi-py", import_name="hebi")
|
||||||
_teleop_available = is_package_available("teleop")
|
_teleop_available = is_package_available("teleop")
|
||||||
_placo_available = is_package_available("placo")
|
_placo_available = is_package_available("placo")
|
||||||
|
|||||||
Reference in New Issue
Block a user