diff --git a/examples/unitree_g1/README.md b/examples/unitree_g1/README.md new file mode 100644 index 000000000..5753d0294 --- /dev/null +++ b/examples/unitree_g1/README.md @@ -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.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.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 +``` diff --git a/examples/unitree_g1/dataset_motion.py b/examples/unitree_g1/dataset_motion.py new file mode 100644 index 000000000..a1f10d584 --- /dev/null +++ b/examples/unitree_g1/dataset_motion.py @@ -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() diff --git a/examples/unitree_g1/sonic.py b/examples/unitree_g1/sonic.py index bd00a1c41..eae0049fe 100644 --- a/examples/unitree_g1/sonic.py +++ b/examples/unitree_g1/sonic.py @@ -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: diff --git a/pyproject.toml b/pyproject.toml index a00deae02..de028e41a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -374,7 +374,11 @@ torch = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }] torchvision = [{ index = "pytorch-cu128", marker = "sys_platform == 'linux'" }] [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] where = ["src"] diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py index 96217e6af..0340f043b 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_pipeline.py @@ -30,7 +30,10 @@ from enum import IntEnum import numpy as np 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 ──────────────────────────────────────────────────────────────── @@ -174,12 +177,6 @@ def _to_mujoco(a): 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) ENCODER_STANDING_REF = DEFAULT_ANGLES.copy() @@ -214,28 +211,6 @@ def compute_kp_kd(): _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 ──────────────────────────────────────────────────────── @@ -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): w, x, y, z = q return np.array( @@ -631,7 +600,7 @@ class StandingEncoderDecoder: obs[off : off + sz] = h[f] off += sz for q in reversed(self.h_quat): - obs[off : off + 3] = gravity_dir(q) + obs[off : off + 3] = get_gravity_orientation(q) off += 3 assert off == 994, f"Decoder obs mismatch: {off}" return obs @@ -689,7 +658,9 @@ class StandingEncoderDecoder: print( 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() 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}") @@ -1122,7 +1093,6 @@ class PlannerController(StandingEncoderDecoder): # ── Keyboard ────────────────────────────────────────────────────────────────── - class RawKeyboard: def __init__(self): self.fd = sys.stdin.fileno() diff --git a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py index 9b813a486..0b4597eaf 100644 --- a/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py +++ b/src/lerobot/robots/unitree_g1/controllers/sonic_whole_body.py @@ -37,10 +37,10 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import ( _snapshot_ms, clamp_mode_params, compute_kp_kd, - lowstate_to_obs, process_joystick, should_replan_request, ) +from lerobot.robots.unitree_g1.g1_utils import lowstate_to_obs logger = logging.getLogger(__name__) @@ -185,25 +185,49 @@ class SonicWholeBodyController: 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 _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: if lowstate is None: return {} obs = lowstate_to_obs(lowstate) # 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) if smpl is None and self._smpl_stream is not None: 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 if smpl is not None: # 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 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) def reset(self): diff --git a/src/lerobot/robots/unitree_g1/g1_utils.py b/src/lerobot/robots/unitree_g1/g1_utils.py index 9ed868bd5..0ea267d76 100644 --- a/src/lerobot/robots/unitree_g1/g1_utils.py +++ b/src/lerobot/robots/unitree_g1/g1_utils.py @@ -63,6 +63,47 @@ class G1_29_JointArmIndex(IntEnum): 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 + ``.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): """Instantiate a locomotion controller by class name. Returns None if name is None.""" if name is None: diff --git a/src/lerobot/robots/unitree_g1/smpl_stream.py b/src/lerobot/robots/unitree_g1/smpl_stream.py index 33fe16c7d..c5dca56bc 100644 --- a/src/lerobot/robots/unitree_g1/smpl_stream.py +++ b/src/lerobot/robots/unitree_g1/smpl_stream.py @@ -112,6 +112,24 @@ class SmplStream: """True once at least one real frame has been received.""" 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): self._buf.clear() self._got_first = False diff --git a/src/lerobot/robots/unitree_g1/unitree_g1.py b/src/lerobot/robots/unitree_g1/unitree_g1.py index f579251f1..a26c5c732 100644 --- a/src/lerobot/robots/unitree_g1/unitree_g1.py +++ b/src/lerobot/robots/unitree_g1/unitree_g1.py @@ -38,6 +38,7 @@ from .g1_utils import ( G1_29_JointArmIndex, G1_29_JointIndex, default_remote_input, + lowstate_to_obs, make_locomotion_controller, ) @@ -428,44 +429,8 @@ class UnitreeG1(Robot): if lowstate is None: return {} - obs = {} - - # 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 + # Motors + IMU + wireless remote (shared lowstate -> obs mapping) + obs = lowstate_to_obs(lowstate) # Cameras - read images from ZMQ cameras for cam_name, cam in self._cameras.items(): diff --git a/src/lerobot/teleoperators/pico_headset/assets/smpl_skeleton.npz b/src/lerobot/teleoperators/pico_headset/assets/smpl_skeleton.npz new file mode 100644 index 000000000..5715dec8b Binary files /dev/null and b/src/lerobot/teleoperators/pico_headset/assets/smpl_skeleton.npz differ diff --git a/src/lerobot/teleoperators/pico_headset/pico_headset.py b/src/lerobot/teleoperators/pico_headset/pico_headset.py index 68eed4a7d..108b29c03 100644 --- a/src/lerobot/teleoperators/pico_headset/pico_headset.py +++ b/src/lerobot/teleoperators/pico_headset/pico_headset.py @@ -91,9 +91,11 @@ class PicoHeadset(Teleoperator): if self._stream is None: raise RuntimeError(f"{self} is not connected") window = self._stream.step() - # Hold back until the headset is actually streaming, so the controller does - # not switch to whole-body tracking of an all-zero (collapsed) pose. - if not self._stream.has_data: + # Emit SMPL only while the headset is actively streaming: hold back before + # the first frame (so the controller doesn't track an all-zero collapsed + # pose) and once the stream goes stale (so the controller falls back to a + # safe standing/locomotion mode instead of freezing on the last pose). + if not self._stream.has_data or self._stream.is_stale: return {} return {f"{SMPL_ACTION_PREFIX}{i}": float(v) for i, v in enumerate(window)} diff --git a/src/lerobot/teleoperators/pico_headset/pico_publisher.py b/src/lerobot/teleoperators/pico_headset/pico_publisher.py new file mode 100644 index 000000000..e9472e779 --- /dev/null +++ b/src/lerobot/teleoperators/pico_headset/pico_publisher.py @@ -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() diff --git a/src/lerobot/teleoperators/pico_headset/smpl_fk.py b/src/lerobot/teleoperators/pico_headset/smpl_fk.py new file mode 100644 index 000000000..7e665e414 --- /dev/null +++ b/src/lerobot/teleoperators/pico_headset/smpl_fk.py @@ -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), + }