mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-17 06:51:48 +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:
Martino Russi
@@ -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 numpy as np
|
||||
from dataset_motion import DatasetJointMotion
|
||||
from motion_loader import SmplMotion
|
||||
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 (
|
||||
CONTROL_DT,
|
||||
DEFAULT_ANGLES,
|
||||
ENCODER_UPDATE_EVERY,
|
||||
LM,
|
||||
MOTION_SETS,
|
||||
MUJOCO_TO_ISAACLAB,
|
||||
RawKeyboard,
|
||||
compute_kp_kd,
|
||||
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
|
||||
|
||||
|
||||
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():
|
||||
parser = argparse.ArgumentParser(description="SONIC planner with keyboard + gamepad control")
|
||||
parser.add_argument(
|
||||
@@ -109,10 +155,27 @@ def main():
|
||||
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()
|
||||
|
||||
if args.smpl_stream and args.motion_file:
|
||||
parser.error("--smpl-stream and --motion-file are mutually exclusive")
|
||||
exclusive = [bool(args.smpl_stream), bool(args.motion_file), bool(args.replay_dataset)]
|
||||
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
|
||||
# avoid losing buffered diagnostics if the process dies mid-loop.
|
||||
@@ -149,6 +212,22 @@ def main():
|
||||
controller = runtime.controller
|
||||
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
|
||||
if args.smpl_stream:
|
||||
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)))
|
||||
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.
|
||||
motion_active = motion is not None and controller.encode_mode == 2
|
||||
if motion_active:
|
||||
|
||||
Reference in New Issue
Block a user