mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-24 10:16:09 +00:00
feat(unitree_g1): standalone PICO SMPL publisher + dedup/replay fixes
Add a self-contained rt/smpl publisher in the pico_headset teleoperator (pico_publisher.py + numpy SMPL FK in smpl_fk.py + vendored skeleton table) so headset whole-body teleop no longer depends on gear_sonic/torch; only xrobotoolkit_sdk is needed at the headset. Also: share lowstate_to_obs/get_gravity_orientation via g1_utils (dedup sonic_pipeline and UnitreeG1.get_observation), and fix dataset-replay joint ordering (Unitree -> IsaacLab) for sonic.py --replay-dataset. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,111 @@
|
|||||||
|
# Unitree G1 — SONIC whole-body teleop (PICO headset)
|
||||||
|
|
||||||
|
Drive the G1 with the **SONIC** policy from live **full-body SMPL** streamed off a PICO
|
||||||
|
headset — running entirely through `lerobot-teleoperate` (no C++ deploy stack).
|
||||||
|
|
||||||
|
## Architecture
|
||||||
|
|
||||||
|
Two processes talk over one ZMQ channel (`rt/smpl`, TCP port `5560`):
|
||||||
|
|
||||||
|
```
|
||||||
|
PICO headset ──► XRoboToolkit PC Service ──► pico_manager_thread_server.py ──(rt/smpl)──► lerobot-teleoperate
|
||||||
|
(gear_sonic env, publisher) (lerobot env, subscriber)
|
||||||
|
└─ SonicWholeBodyController (encode_mode=2)
|
||||||
|
```
|
||||||
|
|
||||||
|
- **Publisher** (`gear_sonic/scripts/pico_manager_thread_server.py`): reads PICO body
|
||||||
|
tracking, converts to canonical SMPL joints, publishes each frame on `rt/smpl`.
|
||||||
|
Lives in `gear_sonic` because it needs the XRoboToolkit SDK.
|
||||||
|
- **Subscriber** (this repo): the `pico_headset` teleoperator (or the
|
||||||
|
`SONIC_SMPL_STREAM` fallback in `SonicWholeBodyController`) consumes `rt/smpl` and
|
||||||
|
feeds the SONIC encoder's 10-frame (720-vec) whole-body reference window.
|
||||||
|
|
||||||
|
The subscriber never launches the publisher — you start it yourself. Until real
|
||||||
|
frames arrive, the robot stays in safe locomotion mode; it switches to whole-body
|
||||||
|
tracking automatically once frames flow.
|
||||||
|
|
||||||
|
## Install
|
||||||
|
|
||||||
|
**lerobot side (subscriber):**
|
||||||
|
```bash
|
||||||
|
# in your lerobot env
|
||||||
|
pip install -e ".[unitree_g1]" # provides pyzmq, unitree_sdk2py, etc.
|
||||||
|
```
|
||||||
|
|
||||||
|
**publisher side (gear_sonic):** see the repo root `docs/TELEOP_QUICKSTART.md`
|
||||||
|
("One-time setup"). In short: install the gear_sonic teleop env, the
|
||||||
|
[XRoboToolkit PC Service](https://github.com/XR-Robotics/XRoboToolkit-PC-Service/releases),
|
||||||
|
and the [PICO APK](https://github.com/XR-Robotics/XRoboToolkit-Unity-Client/releases).
|
||||||
|
On the PICO app set: PC Service IP = laptop Wi-Fi IP, Motion Tracker = **Full body**,
|
||||||
|
Data/Control = **Send**.
|
||||||
|
|
||||||
|
> The publisher and subscriber can run on the same laptop (use `127.0.0.1`) or on
|
||||||
|
> separate machines (point `--teleop.smpl_host` at the publisher's IP).
|
||||||
|
|
||||||
|
## Run
|
||||||
|
|
||||||
|
**1. Start the publisher** (gear_sonic env). Any `--manager` run publishes `rt/smpl`
|
||||||
|
on port 5560; the `--vis_*` flags only add the calibration/preview windows:
|
||||||
|
```bash
|
||||||
|
cd ~/Documents/sonic
|
||||||
|
python gear_sonic/scripts/pico_manager_thread_server.py --manager --vis_vr3pt --vis_smpl
|
||||||
|
# look for: [Manager] ZMQ 'rt/smpl' socket bound to port 5560
|
||||||
|
```
|
||||||
|
|
||||||
|
**2. Start teleoperate** (lerobot env).
|
||||||
|
|
||||||
|
Simulation:
|
||||||
|
```bash
|
||||||
|
lerobot-teleoperate \
|
||||||
|
--robot.type=unitree_g1 --robot.is_simulation=true \
|
||||||
|
--robot.controller=SonicWholeBodyController \
|
||||||
|
--teleop.type=pico_headset --teleop.smpl_host=127.0.0.1 --teleop.smpl_port=5560 \
|
||||||
|
--fps=50
|
||||||
|
```
|
||||||
|
|
||||||
|
Real robot:
|
||||||
|
```bash
|
||||||
|
lerobot-teleoperate \
|
||||||
|
--robot.type=unitree_g1 --robot.is_simulation=false --robot.robot_ip=<ROBOT_IP> \
|
||||||
|
--robot.controller=SonicWholeBodyController \
|
||||||
|
--teleop.type=pico_headset --teleop.smpl_host=127.0.0.1 --teleop.smpl_port=5560 \
|
||||||
|
--fps=50
|
||||||
|
```
|
||||||
|
|
||||||
|
> Skip `--display_data=true` with `pico_headset`: it would print all 720 `smpl.*`
|
||||||
|
> values every tick.
|
||||||
|
|
||||||
|
### Fallback: no teleoperator
|
||||||
|
|
||||||
|
To test the whole-body controller before wiring a teleop (e.g. to keep the
|
||||||
|
`unitree_g1` remote for e-stop), let the controller subscribe to `rt/smpl` directly:
|
||||||
|
```bash
|
||||||
|
SONIC_SMPL_STREAM=1 \
|
||||||
|
lerobot-teleoperate \
|
||||||
|
--robot.type=unitree_g1 --robot.is_simulation=false --robot.robot_ip=<ROBOT_IP> \
|
||||||
|
--robot.controller=SonicWholeBodyController \
|
||||||
|
--teleop.type=unitree_g1 --fps=50
|
||||||
|
# override endpoint with SONIC_SMPL_HOST / SONIC_SMPL_PORT
|
||||||
|
```
|
||||||
|
|
||||||
|
## Safety (real robot)
|
||||||
|
|
||||||
|
- The `pico_headset` teleop is **SMPL-only** — it does not pass a software e-stop.
|
||||||
|
Keep the **hardware remote** in hand.
|
||||||
|
- Start in a neutral standing pose: the robot flips to whole-body tracking the
|
||||||
|
instant real frames arrive.
|
||||||
|
- If frames drop, the controller **holds the last pose** (it does not snap to zero),
|
||||||
|
but it won't auto-return to locomotion — exit via the hardware remote.
|
||||||
|
- Clear a ~3 m safety zone; move slowly at first.
|
||||||
|
|
||||||
|
## Standalone (no teleoperate)
|
||||||
|
|
||||||
|
`sonic.py` can consume the same stream directly for quick tests:
|
||||||
|
```bash
|
||||||
|
python examples/unitree_g1/sonic.py --smpl-stream --smpl-host 127.0.0.1 --smpl-port 5560
|
||||||
|
```
|
||||||
|
|
||||||
|
`smpl_stream.py` is a self-test that just prints window stats:
|
||||||
|
```bash
|
||||||
|
python examples/unitree_g1/smpl_stream.py --smpl-host 127.0.0.1 --smpl-port 5560
|
||||||
|
```
|
||||||
@@ -0,0 +1,135 @@
|
|||||||
|
#!/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.
|
||||||
|
|
||||||
|
"""Load a 29-DoF joint trajectory from a LeRobot dataset episode for SONIC mode 0.
|
||||||
|
|
||||||
|
SONIC's locomotion/tracking mode (``encode_mode == 0``) references the robot in
|
||||||
|
**29-DoF joint space** (see ``build_encoder_obs`` -> ``motion_joint_positions``).
|
||||||
|
Humanoid teleop datasets like ``BitRobot/HIW-500-lerobot`` store exactly that under
|
||||||
|
``observation.state`` (29 joints, same G1 index order as ``G1_29_JointIndex``), so
|
||||||
|
we can feed a recorded episode straight in as the reference and let SONIC try to
|
||||||
|
track it.
|
||||||
|
|
||||||
|
Note the dataset's ``action`` feature is a 23-dim whole-body command (pivot
|
||||||
|
velocities + EE poses), *not* joint targets -- so we deliberately read
|
||||||
|
``observation.state`` (the measured 29-DoF joints), not ``action``.
|
||||||
|
|
||||||
|
The dataset runs at 30 fps; SONIC ticks at 50 Hz and consumes one reference frame
|
||||||
|
per tick, so we resample to 50 fps to preserve real-time speed.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
python examples/unitree_g1/dataset_motion.py \
|
||||||
|
--repo-id BitRobot/HIW-500-lerobot --episode 0
|
||||||
|
"""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
|
||||||
|
STATE_KEY = "observation.state"
|
||||||
|
N_JOINTS = 29
|
||||||
|
SONIC_FPS = 50.0
|
||||||
|
|
||||||
|
|
||||||
|
def _resample(traj: np.ndarray, src_fps: float, dst_fps: float) -> np.ndarray:
|
||||||
|
"""Linearly resample a (T, D) trajectory from src_fps to dst_fps."""
|
||||||
|
if abs(src_fps - dst_fps) < 1e-6:
|
||||||
|
return traj.astype(np.float32)
|
||||||
|
t_in = np.arange(traj.shape[0]) / src_fps
|
||||||
|
dur = t_in[-1] if traj.shape[0] > 1 else 0.0
|
||||||
|
t_out = np.arange(0.0, dur + 1e-9, 1.0 / dst_fps)
|
||||||
|
out = np.empty((t_out.shape[0], traj.shape[1]), np.float32)
|
||||||
|
for j in range(traj.shape[1]):
|
||||||
|
out[:, j] = np.interp(t_out, t_in, traj[:, j])
|
||||||
|
return out
|
||||||
|
|
||||||
|
|
||||||
|
class DatasetJointMotion:
|
||||||
|
"""A recorded 29-DoF joint episode, resampled to SONIC's 50 Hz tick.
|
||||||
|
|
||||||
|
Attributes:
|
||||||
|
joints: (T, 29) float32 reference joint positions at ``fps``.
|
||||||
|
velocities: (T, 29) float32 finite-difference joint velocities.
|
||||||
|
fps: output rate (50 Hz).
|
||||||
|
src_fps: original dataset rate.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
repo_id: str,
|
||||||
|
episode: int = 0,
|
||||||
|
max_frames: int | None = None,
|
||||||
|
root: str | None = None,
|
||||||
|
revision: str = "main",
|
||||||
|
):
|
||||||
|
# Imported lazily so the heavy datasets stack is only pulled in on demand.
|
||||||
|
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||||
|
|
||||||
|
# Pin the branch (default "main"): many community datasets aren't tagged with a
|
||||||
|
# LeRobot codebase_version, and the version-resolution path crashes on them.
|
||||||
|
# A non-PEP440 revision like "main" skips that resolution entirely.
|
||||||
|
ds = LeRobotDataset(
|
||||||
|
repo_id,
|
||||||
|
root=root,
|
||||||
|
episodes=[episode],
|
||||||
|
revision=revision,
|
||||||
|
download_videos=False, # we only need observation.state, skip ~TB of video
|
||||||
|
)
|
||||||
|
self.src_fps = float(ds.fps)
|
||||||
|
|
||||||
|
# Read the joint column straight from the underlying table. Going through
|
||||||
|
# ds[i] would trigger video decoding (the dataset has camera features) and
|
||||||
|
# fail because we intentionally skipped the mp4 download.
|
||||||
|
raw = np.asarray(ds.hf_dataset[STATE_KEY], np.float32) # (T_src, 29)
|
||||||
|
if raw.ndim != 2 or raw.shape[0] == 0:
|
||||||
|
raise ValueError(f"Episode {episode} of {repo_id} has no usable {STATE_KEY}")
|
||||||
|
if raw.shape[1] != N_JOINTS:
|
||||||
|
raise ValueError(f"{STATE_KEY} must be (T, {N_JOINTS}), got {raw.shape}")
|
||||||
|
|
||||||
|
self.joints = _resample(raw, self.src_fps, SONIC_FPS)
|
||||||
|
if max_frames is not None:
|
||||||
|
self.joints = self.joints[:max_frames]
|
||||||
|
self.fps = SONIC_FPS
|
||||||
|
|
||||||
|
# Finite-difference velocities (rad/s) at the resampled rate.
|
||||||
|
self.velocities = np.gradient(self.joints, axis=0).astype(np.float32) * self.fps
|
||||||
|
self.num_frames = self.joints.shape[0]
|
||||||
|
self.repo_id = repo_id
|
||||||
|
self.episode = episode
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(description=__doc__)
|
||||||
|
parser.add_argument("--repo-id", default="BitRobot/HIW-500-lerobot")
|
||||||
|
parser.add_argument("--episode", type=int, default=0)
|
||||||
|
parser.add_argument("--max-frames", type=int, default=None)
|
||||||
|
parser.add_argument("--revision", default="main", help="Repo branch/tag (default: main)")
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
m = DatasetJointMotion(
|
||||||
|
args.repo_id, episode=args.episode, max_frames=args.max_frames, revision=args.revision
|
||||||
|
)
|
||||||
|
dur = m.num_frames / m.fps
|
||||||
|
print(f"Loaded {args.repo_id} episode {args.episode}")
|
||||||
|
print(f" src_fps={m.src_fps:.1f} -> {m.fps:.1f} frames={m.num_frames} duration={dur:.1f}s")
|
||||||
|
print(f" joints={m.joints.shape} range=[{m.joints.min():.3f}, {m.joints.max():.3f}]")
|
||||||
|
print(f" |velocity| max={np.abs(m.velocities).max():.3f} rad/s")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
@@ -34,6 +34,7 @@ import tempfile
|
|||||||
import time
|
import time
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
from dataset_motion import DatasetJointMotion
|
||||||
from motion_loader import SmplMotion
|
from motion_loader import SmplMotion
|
||||||
from smpl_stream import DEFAULT_SMPL_HOST, DEFAULT_SMPL_PORT, SmplStream
|
from smpl_stream import DEFAULT_SMPL_HOST, DEFAULT_SMPL_PORT, SmplStream
|
||||||
|
|
||||||
@@ -41,8 +42,10 @@ from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
|
|||||||
from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
|
from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
|
||||||
CONTROL_DT,
|
CONTROL_DT,
|
||||||
DEFAULT_ANGLES,
|
DEFAULT_ANGLES,
|
||||||
|
ENCODER_UPDATE_EVERY,
|
||||||
LM,
|
LM,
|
||||||
MOTION_SETS,
|
MOTION_SETS,
|
||||||
|
MUJOCO_TO_ISAACLAB,
|
||||||
RawKeyboard,
|
RawKeyboard,
|
||||||
compute_kp_kd,
|
compute_kp_kd,
|
||||||
drain_keyboard,
|
drain_keyboard,
|
||||||
@@ -52,6 +55,49 @@ from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
|
|||||||
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
|
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
|
||||||
|
|
||||||
|
|
||||||
|
def _load_joint_trajectory(controller, joints: np.ndarray, velocities: np.ndarray) -> None:
|
||||||
|
"""Load a (T, 29) joint reference into the controller for encode_mode=0 tracking.
|
||||||
|
|
||||||
|
The dataset provides joints in Unitree/G1_29_JointIndex order, but the SONIC
|
||||||
|
encoder reference (motion_joint_positions) is in IsaacLab order. Reorder here.
|
||||||
|
"""
|
||||||
|
joints = np.asarray(joints)[:, MUJOCO_TO_ISAACLAB]
|
||||||
|
velocities = np.asarray(velocities)[:, MUJOCO_TO_ISAACLAB]
|
||||||
|
t = joints.shape[0]
|
||||||
|
with controller.motion_lock:
|
||||||
|
cap = controller.motion_joint_positions.shape[0]
|
||||||
|
if t > cap:
|
||||||
|
controller.motion_joint_positions = np.zeros((t, 29), np.float64)
|
||||||
|
controller.motion_joint_velocities = np.zeros((t, 29), np.float64)
|
||||||
|
controller.motion_body_quats = np.zeros((t, 4), np.float64)
|
||||||
|
controller.motion_body_quats[:, 0] = 1.0
|
||||||
|
controller.motion_body_pos = np.zeros((t, 3), np.float64)
|
||||||
|
controller.motion_joint_positions[:t] = joints
|
||||||
|
controller.motion_joint_velocities[:t] = velocities
|
||||||
|
controller.motion_body_quats[:t, 0] = 1.0
|
||||||
|
controller.motion_body_quats[:t, 1:] = 0.0
|
||||||
|
controller.motion_body_pos[:t] = 0.0
|
||||||
|
controller.motion_timesteps = t
|
||||||
|
controller.ref_cursor = 0
|
||||||
|
controller.init_ref_quat = np.array([1, 0, 0, 0], np.float64)
|
||||||
|
controller.encode_mode = 0
|
||||||
|
controller.playing = True
|
||||||
|
controller.first_motion = True # triggers heading init on first obs
|
||||||
|
controller.reinit_heading = True
|
||||||
|
|
||||||
|
|
||||||
|
def _tick_replay(runtime, obs: dict) -> dict:
|
||||||
|
"""One control tick for dataset replay: encode/decode + advance, no planner."""
|
||||||
|
if not obs:
|
||||||
|
runtime.step += 1
|
||||||
|
return {}
|
||||||
|
do_enc = runtime.step % ENCODER_UPDATE_EVERY == 0
|
||||||
|
action = runtime.controller.step(obs, update_encoder=do_enc, debug=False)
|
||||||
|
runtime.controller.advance_cursor()
|
||||||
|
runtime.step += 1
|
||||||
|
return action
|
||||||
|
|
||||||
|
|
||||||
def main():
|
def main():
|
||||||
parser = argparse.ArgumentParser(description="SONIC planner with keyboard + gamepad control")
|
parser = argparse.ArgumentParser(description="SONIC planner with keyboard + gamepad control")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
@@ -109,10 +155,27 @@ def main():
|
|||||||
default=DEFAULT_SMPL_PORT,
|
default=DEFAULT_SMPL_PORT,
|
||||||
help=f"Port for the rt/smpl stream (default: {DEFAULT_SMPL_PORT})",
|
help=f"Port for the rt/smpl stream (default: {DEFAULT_SMPL_PORT})",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--replay-dataset",
|
||||||
|
type=str,
|
||||||
|
default=None,
|
||||||
|
help="Replay a LeRobot dataset episode's 29-DoF observation.state as a SONIC "
|
||||||
|
"encode_mode=0 joint reference (e.g. BitRobot/HIW-500-lerobot).",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--episode", type=int, default=0, help="Episode index for --replay-dataset (default: 0)"
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--replay-frames",
|
||||||
|
type=int,
|
||||||
|
default=None,
|
||||||
|
help="Cap the number of replayed frames (default: whole episode)",
|
||||||
|
)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
if args.smpl_stream and args.motion_file:
|
exclusive = [bool(args.smpl_stream), bool(args.motion_file), bool(args.replay_dataset)]
|
||||||
parser.error("--smpl-stream and --motion-file are mutually exclusive")
|
if sum(exclusive) > 1:
|
||||||
|
parser.error("--smpl-stream, --motion-file and --replay-dataset are mutually exclusive")
|
||||||
|
|
||||||
# Surface native crashes (onnxruntime / mujoco) with a real traceback, and
|
# Surface native crashes (onnxruntime / mujoco) with a real traceback, and
|
||||||
# avoid losing buffered diagnostics if the process dies mid-loop.
|
# avoid losing buffered diagnostics if the process dies mid-loop.
|
||||||
@@ -149,6 +212,22 @@ def main():
|
|||||||
controller = runtime.controller
|
controller = runtime.controller
|
||||||
ms = runtime.ms
|
ms = runtime.ms
|
||||||
|
|
||||||
|
# --replay-dataset drives SONIC mode 0: load a recorded 29-DoF joint trajectory
|
||||||
|
# into the controller's reference buffers and let the policy try to track it,
|
||||||
|
# bypassing the locomotion planner (which would otherwise overwrite the ref).
|
||||||
|
replay = None
|
||||||
|
if args.replay_dataset:
|
||||||
|
replay = DatasetJointMotion(
|
||||||
|
args.replay_dataset, episode=args.episode, max_frames=args.replay_frames
|
||||||
|
)
|
||||||
|
_load_joint_trajectory(controller, replay.joints, replay.velocities)
|
||||||
|
dur = replay.num_frames / replay.fps
|
||||||
|
print(f"\n[Replay] {args.replay_dataset} episode {args.episode} -> SONIC mode 0")
|
||||||
|
print(
|
||||||
|
f" frames={replay.num_frames} fps={replay.fps:.0f} duration={dur:.1f}s "
|
||||||
|
f"(src {replay.src_fps:.0f} fps, encode_mode=0, planner bypassed)"
|
||||||
|
)
|
||||||
|
|
||||||
motion = None
|
motion = None
|
||||||
if args.smpl_stream:
|
if args.smpl_stream:
|
||||||
motion = SmplStream(host=args.smpl_host, port=args.smpl_port)
|
motion = SmplStream(host=args.smpl_host, port=args.smpl_port)
|
||||||
@@ -223,6 +302,22 @@ def main():
|
|||||||
time.sleep(max(0.0, CONTROL_DT - (time.time() - t0)))
|
time.sleep(max(0.0, CONTROL_DT - (time.time() - t0)))
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
# Dataset replay: SONIC tracks the recorded 29-DoF joint clip.
|
||||||
|
if replay is not None:
|
||||||
|
step_before = runtime.step
|
||||||
|
t_step = time.time()
|
||||||
|
action = _tick_replay(runtime, obs)
|
||||||
|
step_ms = 1000 * (time.time() - t_step)
|
||||||
|
(enc_t if step_before % 5 == 0 else dec_t).append(step_ms)
|
||||||
|
robot.send_action(action)
|
||||||
|
if controller.ref_cursor >= controller.motion_timesteps - 1:
|
||||||
|
print("\n[Replay] episode finished")
|
||||||
|
break
|
||||||
|
now = time.time()
|
||||||
|
loop_t.append(1000 * (now - t0))
|
||||||
|
time.sleep(max(0.0, CONTROL_DT - (now - t0)))
|
||||||
|
continue
|
||||||
|
|
||||||
# SMPL playback only while in whole-body mode; 'M' toggles it.
|
# SMPL playback only while in whole-body mode; 'M' toggles it.
|
||||||
motion_active = motion is not None and controller.encode_mode == 2
|
motion_active = motion is not None and controller.encode_mode == 2
|
||||||
if motion_active:
|
if motion_active:
|
||||||
|
|||||||
+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"]
|
||||||
|
|||||||
@@ -30,7 +30,10 @@ from enum import IntEnum
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
import onnxruntime as ort
|
import onnxruntime as ort
|
||||||
|
|
||||||
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
|
from lerobot.robots.unitree_g1.g1_utils import (
|
||||||
|
G1_29_JointIndex,
|
||||||
|
get_gravity_orientation,
|
||||||
|
)
|
||||||
|
|
||||||
# ── Constants ────────────────────────────────────────────────────────────────
|
# ── Constants ────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
@@ -174,12 +177,6 @@ def _to_mujoco(a):
|
|||||||
return a[MUJOCO_TO_ISAACLAB]
|
return a[MUJOCO_TO_ISAACLAB]
|
||||||
|
|
||||||
|
|
||||||
def _to_runtime(a):
|
|
||||||
r = np.zeros(29, np.float32)
|
|
||||||
r[MUJOCO_TO_ISAACLAB] = a
|
|
||||||
return r
|
|
||||||
|
|
||||||
|
|
||||||
DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
|
DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
|
||||||
ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
|
ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
|
||||||
|
|
||||||
@@ -214,28 +211,6 @@ def compute_kp_kd():
|
|||||||
_kp_kd = compute_kp_kd # backward-compatible alias
|
_kp_kd = compute_kp_kd # backward-compatible alias
|
||||||
|
|
||||||
|
|
||||||
def lowstate_to_obs(lowstate) -> dict:
|
|
||||||
"""Build a robot observation dict from Unitree lowstate."""
|
|
||||||
obs: dict = {}
|
|
||||||
for motor in G1_29_JointIndex:
|
|
||||||
idx = motor.value
|
|
||||||
obs[f"{motor.name}.q"] = float(lowstate.motor_state[idx].q)
|
|
||||||
obs[f"{motor.name}.dq"] = float(lowstate.motor_state[idx].dq)
|
|
||||||
quat = lowstate.imu_state.quaternion
|
|
||||||
obs["imu.quat.w"] = float(quat[0])
|
|
||||||
obs["imu.quat.x"] = float(quat[1])
|
|
||||||
obs["imu.quat.y"] = float(quat[2])
|
|
||||||
obs["imu.quat.z"] = float(quat[3])
|
|
||||||
gyro = lowstate.imu_state.gyroscope
|
|
||||||
obs["imu.gyro.x"] = float(gyro[0])
|
|
||||||
obs["imu.gyro.y"] = float(gyro[1])
|
|
||||||
obs["imu.gyro.z"] = float(gyro[2])
|
|
||||||
wr = getattr(lowstate, "wireless_remote", None)
|
|
||||||
if wr is not None:
|
|
||||||
obs["wireless_remote"] = bytes(wr) if not isinstance(wr, (bytes, bytearray)) else wr
|
|
||||||
return obs
|
|
||||||
|
|
||||||
|
|
||||||
# ── Quaternion helpers ────────────────────────────────────────────────────────
|
# ── Quaternion helpers ────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
@@ -257,12 +232,6 @@ def quat_mul(q1, q2):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
def gravity_dir(q):
|
|
||||||
q = q / (np.linalg.norm(q) + 1e-8)
|
|
||||||
qv = np.array([0, 0, 0, -1], dtype=np.float32)
|
|
||||||
return quat_mul(quat_mul(quat_conj(q), qv), q)[1:]
|
|
||||||
|
|
||||||
|
|
||||||
def quat_to_6d(q):
|
def quat_to_6d(q):
|
||||||
w, x, y, z = q
|
w, x, y, z = q
|
||||||
return np.array(
|
return np.array(
|
||||||
@@ -631,7 +600,7 @@ class StandingEncoderDecoder:
|
|||||||
obs[off : off + sz] = h[f]
|
obs[off : off + sz] = h[f]
|
||||||
off += sz
|
off += sz
|
||||||
for q in reversed(self.h_quat):
|
for q in reversed(self.h_quat):
|
||||||
obs[off : off + 3] = gravity_dir(q)
|
obs[off : off + 3] = get_gravity_orientation(q)
|
||||||
off += 3
|
off += 3
|
||||||
assert off == 994, f"Decoder obs mismatch: {off}"
|
assert off == 994, f"Decoder obs mismatch: {off}"
|
||||||
return obs
|
return obs
|
||||||
@@ -689,7 +658,9 @@ class StandingEncoderDecoder:
|
|||||||
print(
|
print(
|
||||||
f" anchor_6d(identity): {self._anchor_6d(np.array([1, 0, 0, 0], np.float32), np.array([1, 0, 0, 0], np.float32))}"
|
f" anchor_6d(identity): {self._anchor_6d(np.array([1, 0, 0, 0], np.float32), np.array([1, 0, 0, 0], np.float32))}"
|
||||||
)
|
)
|
||||||
print(f" gravity(identity): {gravity_dir(np.array([1, 0, 0, 0], np.float32))} (expect [0,0,-1])")
|
print(
|
||||||
|
f" gravity(identity): {get_gravity_orientation(np.array([1, 0, 0, 0], np.float32))} (expect [0,0,-1])"
|
||||||
|
)
|
||||||
dec0 = self.build_decoder_obs()
|
dec0 = self.build_decoder_obs()
|
||||||
print(f" decoder q-delta max: {np.max(np.abs(dec0[94:384])):.6f}")
|
print(f" decoder q-delta max: {np.max(np.abs(dec0[94:384])):.6f}")
|
||||||
print(f" decoder dq max: {np.max(np.abs(dec0[384:674])):.6f}")
|
print(f" decoder dq max: {np.max(np.abs(dec0[384:674])):.6f}")
|
||||||
@@ -1122,7 +1093,6 @@ class PlannerController(StandingEncoderDecoder):
|
|||||||
|
|
||||||
# ── Keyboard ──────────────────────────────────────────────────────────────────
|
# ── Keyboard ──────────────────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
class RawKeyboard:
|
class RawKeyboard:
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.fd = sys.stdin.fileno()
|
self.fd = sys.stdin.fileno()
|
||||||
|
|||||||
@@ -37,10 +37,10 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
|
|||||||
_snapshot_ms,
|
_snapshot_ms,
|
||||||
clamp_mode_params,
|
clamp_mode_params,
|
||||||
compute_kp_kd,
|
compute_kp_kd,
|
||||||
lowstate_to_obs,
|
|
||||||
process_joystick,
|
process_joystick,
|
||||||
should_replan_request,
|
should_replan_request,
|
||||||
)
|
)
|
||||||
|
from lerobot.robots.unitree_g1.g1_utils import lowstate_to_obs
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -185,25 +185,49 @@ class SonicWholeBodyController:
|
|||||||
self._smpl_stream = SmplStream(host=host, port=port)
|
self._smpl_stream = SmplStream(host=host, port=port)
|
||||||
logger.info("SONIC subscribed to rt/smpl @ tcp://%s:%d", host, 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 _exit_wholebody(self) -> None:
|
||||||
|
"""Revert to locomotion/standing (encode_mode 0) after SMPL 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: SMPL stream lost/stale -> reverting to locomotion (standing)")
|
||||||
|
|
||||||
def run_step(self, action: dict, lowstate) -> dict:
|
def run_step(self, action: dict, lowstate) -> dict:
|
||||||
if lowstate is None:
|
if lowstate is None:
|
||||||
return {}
|
return {}
|
||||||
obs = lowstate_to_obs(lowstate)
|
obs = lowstate_to_obs(lowstate)
|
||||||
|
|
||||||
# Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a
|
# Prefer SMPL delivered via the teleop action (pico_headset). Fall back to a
|
||||||
# direct rt/smpl subscription when SONIC_SMPL_STREAM is enabled.
|
# direct rt/smpl subscription when SONIC_SMPL_STREAM is enabled. A stale
|
||||||
|
# stream (headset silent past its timeout) is treated as "no SMPL" so the
|
||||||
|
# robot doesn't stay frozen tracking the last pose.
|
||||||
smpl = _extract_smpl_from_action(action)
|
smpl = _extract_smpl_from_action(action)
|
||||||
if smpl is None and self._smpl_stream is not None:
|
if smpl is None and self._smpl_stream is not None:
|
||||||
window = self._smpl_stream.step()
|
window = self._smpl_stream.step()
|
||||||
if self._smpl_stream.has_data:
|
if self._smpl_stream.has_data and not self._smpl_stream.is_stale:
|
||||||
smpl = window
|
smpl = window
|
||||||
|
|
||||||
if smpl is not None:
|
if smpl is not None:
|
||||||
# Full-body whole-body tracking: SMPL drives the reference, not joystick.
|
# Full-body whole-body tracking: SMPL drives the reference, not joystick.
|
||||||
self.controller.encode_mode = 2
|
if self.controller.encode_mode != 2:
|
||||||
|
self._enter_wholebody()
|
||||||
self.controller.smpl_joints_10frame_step1 = smpl
|
self.controller.smpl_joints_10frame_step1 = smpl
|
||||||
return self._runtime.tick(obs, debug=False, use_joystick=False)
|
return self._runtime.tick(obs, debug=False, use_joystick=False)
|
||||||
|
|
||||||
|
# No (or stale) SMPL: fall back to locomotion so the robot stays balanced.
|
||||||
|
if self.controller.encode_mode == 2:
|
||||||
|
self._exit_wholebody()
|
||||||
return self._runtime.tick(obs, debug=False)
|
return self._runtime.tick(obs, debug=False)
|
||||||
|
|
||||||
def reset(self):
|
def reset(self):
|
||||||
|
|||||||
@@ -63,6 +63,47 @@ 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 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:
|
||||||
|
|||||||
@@ -112,6 +112,24 @@ class SmplStream:
|
|||||||
"""True once at least one real frame has been received."""
|
"""True once at least one real frame has been received."""
|
||||||
return self._got_first
|
return self._got_first
|
||||||
|
|
||||||
|
@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):
|
def reset(self):
|
||||||
self._buf.clear()
|
self._buf.clear()
|
||||||
self._got_first = False
|
self._got_first = False
|
||||||
|
|||||||
@@ -38,6 +38,7 @@ from .g1_utils import (
|
|||||||
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,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -428,44 +429,8 @@ 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
|
|
||||||
for motor in G1_29_JointIndex:
|
|
||||||
name = motor.name
|
|
||||||
idx = motor.value
|
|
||||||
obs[f"{name}.q"] = lowstate.motor_state[idx].q
|
|
||||||
obs[f"{name}.dq"] = lowstate.motor_state[idx].dq
|
|
||||||
obs[f"{name}.tau"] = lowstate.motor_state[idx].tau_est
|
|
||||||
|
|
||||||
# IMU - gyroscope
|
|
||||||
if lowstate.imu_state.gyroscope:
|
|
||||||
obs["imu.gyro.x"] = lowstate.imu_state.gyroscope[0]
|
|
||||||
obs["imu.gyro.y"] = lowstate.imu_state.gyroscope[1]
|
|
||||||
obs["imu.gyro.z"] = lowstate.imu_state.gyroscope[2]
|
|
||||||
|
|
||||||
# IMU - accelerometer
|
|
||||||
if lowstate.imu_state.accelerometer:
|
|
||||||
obs["imu.accel.x"] = lowstate.imu_state.accelerometer[0]
|
|
||||||
obs["imu.accel.y"] = lowstate.imu_state.accelerometer[1]
|
|
||||||
obs["imu.accel.z"] = lowstate.imu_state.accelerometer[2]
|
|
||||||
|
|
||||||
# IMU - quaternion
|
|
||||||
if lowstate.imu_state.quaternion:
|
|
||||||
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():
|
||||||
|
|||||||
Binary file not shown.
@@ -91,9 +91,11 @@ class PicoHeadset(Teleoperator):
|
|||||||
if self._stream is None:
|
if self._stream is None:
|
||||||
raise RuntimeError(f"{self} is not connected")
|
raise RuntimeError(f"{self} is not connected")
|
||||||
window = self._stream.step()
|
window = self._stream.step()
|
||||||
# Hold back until the headset is actually streaming, so the controller does
|
# Emit SMPL only while the headset is actively streaming: hold back before
|
||||||
# not switch to whole-body tracking of an all-zero (collapsed) pose.
|
# the first frame (so the controller doesn't track an all-zero collapsed
|
||||||
if not self._stream.has_data:
|
# pose) and once the stream goes stale (so the controller falls back to a
|
||||||
|
# safe standing/locomotion mode instead of freezing on the last pose).
|
||||||
|
if not self._stream.has_data or self._stream.is_stale:
|
||||||
return {}
|
return {}
|
||||||
return {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)}
|
return {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)}
|
||||||
|
|
||||||
|
|||||||
@@ -0,0 +1,153 @@
|
|||||||
|
#!/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.robots.unitree_g1.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
|
||||||
|
"""
|
||||||
|
|
||||||
|
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
|
||||||
|
|
||||||
|
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,
|
||||||
|
) -> bytes:
|
||||||
|
"""Build the rt/smpl JSON message (single frame, topic embedded in payload)."""
|
||||||
|
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()
|
||||||
|
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 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("--fake", action="store_true", help="Publish synthetic motion (no headset)")
|
||||||
|
args = p.parse_args()
|
||||||
|
|
||||||
|
fk = SmplForwardKinematics(args.skeleton) if args.skeleton else SmplForwardKinematics()
|
||||||
|
|
||||||
|
xrt = None
|
||||||
|
if 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 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}")
|
||||||
|
print(f"[pico_publisher] '{SMPL_TOPIC}' bound to tcp://*:{args.port} @ {args.fps:.0f} Hz")
|
||||||
|
|
||||||
|
period = 1.0 / max(1.0, args.fps)
|
||||||
|
frame_index = 0
|
||||||
|
t0 = time.time()
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
loop_start = time.time()
|
||||||
|
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
|
||||||
|
|
||||||
|
out = fk.compute(body_poses)
|
||||||
|
sock.send(
|
||||||
|
pack_message(
|
||||||
|
out["smpl_joints_local"],
|
||||||
|
frame_index,
|
||||||
|
stamp_ns,
|
||||||
|
root_quat=out["root_quat"],
|
||||||
|
root_transl=out["root_transl"],
|
||||||
|
)
|
||||||
|
)
|
||||||
|
frame_index += 1
|
||||||
|
if frame_index % int(max(1, args.fps)) == 0:
|
||||||
|
print(f"[pico_publisher] sent {frame_index} frames", 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,194 @@
|
|||||||
|
#!/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
|
||||||
|
(rest-pose joints + kinematic tree), vendored in ``assets/smpl_skeleton.npz``.
|
||||||
|
|
||||||
|
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 pathlib import Path
|
||||||
|
|
||||||
|
import numpy as np
|
||||||
|
from scipy.spatial.transform import Rotation as R
|
||||||
|
|
||||||
|
# 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])])
|
||||||
|
|
||||||
|
_SKELETON_PATH = Path(__file__).parent / "assets" / "smpl_skeleton.npz"
|
||||||
|
|
||||||
|
|
||||||
|
# ── quaternion helpers (scalar-first w, x, y, z) ─────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
def aa_to_quat(aa: np.ndarray) -> np.ndarray:
|
||||||
|
aa = np.asarray(aa, np.float64)
|
||||||
|
angle = np.linalg.norm(aa, axis=-1, keepdims=True)
|
||||||
|
small = angle < 1e-8
|
||||||
|
safe = np.where(small, 1.0, angle)
|
||||||
|
axis = np.where(small, 0.0, aa / safe)
|
||||||
|
half = angle * 0.5
|
||||||
|
return np.concatenate([np.cos(half), axis * np.sin(half)], axis=-1)
|
||||||
|
|
||||||
|
|
||||||
|
def quat_to_aa(q: np.ndarray) -> np.ndarray:
|
||||||
|
q = np.asarray(q, np.float64)
|
||||||
|
w = q[..., 0:1]
|
||||||
|
xyz = q[..., 1:]
|
||||||
|
n = np.linalg.norm(xyz, axis=-1, keepdims=True)
|
||||||
|
angle = 2.0 * np.arctan2(n, w)
|
||||||
|
small = n < 1e-8
|
||||||
|
safe = np.where(small, 1.0, n)
|
||||||
|
axis = np.where(small, 0.0, xyz / safe)
|
||||||
|
return axis * angle
|
||||||
|
|
||||||
|
|
||||||
|
def quat_mul(a: np.ndarray, b: np.ndarray) -> np.ndarray:
|
||||||
|
aw, ax, ay, az = a[..., 0], a[..., 1], a[..., 2], a[..., 3]
|
||||||
|
bw, bx, by, bz = b[..., 0], b[..., 1], b[..., 2], b[..., 3]
|
||||||
|
return np.stack(
|
||||||
|
[
|
||||||
|
aw * bw - ax * bx - ay * by - az * bz,
|
||||||
|
aw * bx + ax * bw + ay * bz - az * by,
|
||||||
|
aw * by - ax * bz + ay * bw + az * bx,
|
||||||
|
aw * bz + ax * by - ay * bx + az * bw,
|
||||||
|
],
|
||||||
|
axis=-1,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def quat_conj(q: np.ndarray) -> np.ndarray:
|
||||||
|
return np.concatenate([q[..., 0:1], -q[..., 1:]], axis=-1)
|
||||||
|
|
||||||
|
|
||||||
|
def quat_inv(q: np.ndarray) -> np.ndarray:
|
||||||
|
c = quat_conj(q)
|
||||||
|
return c / (np.linalg.norm(c, axis=-1, keepdims=True) + 1e-12)
|
||||||
|
|
||||||
|
|
||||||
|
def quat_apply(q: np.ndarray, v: np.ndarray) -> np.ndarray:
|
||||||
|
xyz = q[..., 1:]
|
||||||
|
w = q[..., 0:1]
|
||||||
|
t = 2.0 * np.cross(xyz, v)
|
||||||
|
return v + w * t + np.cross(xyz, t)
|
||||||
|
|
||||||
|
|
||||||
|
def smpl_root_ytoz_up(root_quat: np.ndarray) -> np.ndarray:
|
||||||
|
"""Rotate the root quaternion 90 deg about X to map SMPL Y-up to robot Z-up."""
|
||||||
|
base = aa_to_quat(np.array([np.pi / 2.0, 0.0, 0.0]))
|
||||||
|
return quat_mul(base, root_quat)
|
||||||
|
|
||||||
|
|
||||||
|
def remove_smpl_base_rot(root_quat: np.ndarray) -> np.ndarray:
|
||||||
|
"""Conjugate out SMPL's default rest orientation ([0.5, 0.5, 0.5, 0.5])."""
|
||||||
|
base_conj = quat_conj(np.array([0.5, 0.5, 0.5, 0.5]))
|
||||||
|
return quat_mul(root_quat, base_conj)
|
||||||
|
|
||||||
|
|
||||||
|
# ── forward kinematics ───────────────────────────────────────────────────────
|
||||||
|
|
||||||
|
|
||||||
|
class SmplForwardKinematics:
|
||||||
|
"""Rest-skeleton SMPL forward kinematics (no mesh, no torch)."""
|
||||||
|
|
||||||
|
def __init__(self, skeleton_path: str | Path = _SKELETON_PATH):
|
||||||
|
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
|
||||||
|
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(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_quat = aa_to_quat(global_orient)
|
||||||
|
root_quat = smpl_root_ytoz_up(root_quat)
|
||||||
|
global_orient_new = quat_to_aa(root_quat)
|
||||||
|
|
||||||
|
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_quat = remove_smpl_base_rot(root_quat)
|
||||||
|
inv = np.broadcast_to(quat_inv(root_quat), (joints.shape[0], 4))
|
||||||
|
smpl_joints_local = quat_apply(inv, joints)
|
||||||
|
|
||||||
|
return {
|
||||||
|
"smpl_joints_local": smpl_joints_local.astype(np.float32),
|
||||||
|
"root_quat": root_quat.astype(np.float32),
|
||||||
|
"root_transl": positions[0].astype(np.float32),
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user