diff --git a/examples/unitree_g1/README.md b/examples/unitree_g1/README.md deleted file mode 100644 index 5753d0294..000000000 --- a/examples/unitree_g1/README.md +++ /dev/null @@ -1,111 +0,0 @@ -# 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 deleted file mode 100644 index a1f10d584..000000000 --- a/examples/unitree_g1/dataset_motion.py +++ /dev/null @@ -1,135 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""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/motion_loader.py b/examples/unitree_g1/motion_loader.py deleted file mode 100644 index b92ebd9f1..000000000 --- a/examples/unitree_g1/motion_loader.py +++ /dev/null @@ -1,170 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Load an SMPL motion clip and expose it in SONIC's encoder format. - -SONIC's whole-body tracking mode (``encode_mode == 2``) consumes a flat -720-vector ``smpl_joints_10frame_step1`` = 10 consecutive frames x 24 SMPL -joints x 3 (xyz) at 50 Hz. - -IMPORTANT - frame convention: the encoder expects each frame's joints with the -body's *root orientation removed* (per-frame canonical), exactly like the live -deploy stream's ``smpl_joints_local`` (see ``process_smpl_joints`` in the GEAR -PICO teleop and ``smpl_joints_multi_future_local`` in training). The reference -``smpl_filtered`` clips instead store **world-frame** joints (heading retained), -so feeding them raw makes the robot move but track poorly / never face-forward. -This loader therefore canonicalizes on load using the clip's per-frame root -orientation (``pose_aa[:, :3]``): - - A = Rx(+90deg) * rotvec(pose_aa[:, :3]) # y-up -> z-up root quat - local = base120 * A^-1 * joints # remove root orient - -with ``base120 = quat(0.5,0.5,0.5,0.5)`` (SMPL base rotation). This reproduces -the deployed transform (verified: per-frame hip-heading std -> 0). - -Clip is read from a numpy ``.npz``. Expected keys: - smpl_joints : (T, 24, 3) float32 -- world-frame joint positions, 50 fps - pose_aa : (T, 72) float32 -- SMPL axis-angle (root = [:, :3]) - transl : (T, 3) float32 -- global root translation (optional) - fps : scalar - -Example: - python examples/unitree_g1/motion_loader.py \ - --motion examples/unitree_g1/motions/walk_forward.npz -""" - -import argparse - -import numpy as np - -WINDOW = 10 # frames per encoder window (smpl_joints_10frame_step1) -N_JOINTS = 24 -JOINT_DIM = 3 -SMPL_OBS_DIM = WINDOW * N_JOINTS * JOINT_DIM # 720 - - -def canonicalize_smpl_joints(smpl_joints: np.ndarray, root_aa: np.ndarray) -> np.ndarray: - """Remove per-frame root orientation -> SONIC ``smpl_joints_local`` format. - - Args: - smpl_joints: (T, 24, 3) world-frame (z-up) SMPL joint positions. - root_aa: (T, 3) SMPL global-orient axis-angle (y-up convention). - - Returns: - (T, 24, 3) per-frame root-orientation-removed joints. - """ - from scipy.spatial.transform import Rotation - - rx90 = Rotation.from_euler("x", 90, degrees=True) # smpl_root_ytoz_up - base120 = Rotation.from_quat([0.5, 0.5, 0.5, 0.5]) # remove_smpl_base_rot - a = rx90 * Rotation.from_rotvec(root_aa) # z-up root quat (left-mult) - b_inv = base120 * a.inv() # inv(remove_smpl_base_rot(a)) - return np.einsum("tij,tkj->tki", b_inv.as_matrix(), smpl_joints).astype(np.float32) - - -class SmplMotion: - """A single SMPL clip with SONIC-format windowing.""" - - def __init__(self, path: str, loop: bool = True, canonicalize: bool = True): - data = np.load(path) - smpl_joints = data["smpl_joints"].astype(np.float32) # (T, 24, 3) - self.pose_aa = data["pose_aa"].astype(np.float32) if "pose_aa" in data.files else None - self.transl = data["transl"].astype(np.float32) if "transl" in data.files else None - self.fps = float(data["fps"]) if "fps" in data.files else 50.0 - self.loop = loop - - if smpl_joints.ndim != 3 or smpl_joints.shape[1:] != (N_JOINTS, JOINT_DIM): - raise ValueError(f"Expected smpl_joints (T, {N_JOINTS}, {JOINT_DIM}), got {smpl_joints.shape}") - - # Reference clips store world-frame joints; the encoder wants per-frame - # root-orientation-removed joints. Canonicalize when we have the root pose. - self.canonicalized = False - if canonicalize and self.pose_aa is not None: - smpl_joints = canonicalize_smpl_joints(smpl_joints, self.pose_aa[:, :3]) - self.canonicalized = True - self.smpl_joints = smpl_joints - - self.num_frames = self.smpl_joints.shape[0] - self._cursor = 0 - - def window(self, start: int) -> np.ndarray: - """Return the 720-vector for the 10-frame window beginning at ``start``. - - Frames are laid out oldest->newest, joint-major within a frame: - [f0_j0_xyz, f0_j1_xyz, ..., f9_j23_xyz]. - """ - idx = np.arange(start, start + WINDOW) - idx = np.mod(idx, self.num_frames) if self.loop else np.clip(idx, 0, self.num_frames - 1) - return self.smpl_joints[idx].reshape(-1).astype(np.float32) - - def reset(self): - self._cursor = 0 - - def step(self) -> np.ndarray: - """Advance one frame and return the current 720-vector window.""" - w = self.window(self._cursor) - self._cursor += 1 - if self.loop: - self._cursor %= self.num_frames - return w - - @property - def done(self) -> bool: - return (not self.loop) and (self._cursor + WINDOW >= self.num_frames) - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--motion", required=True, help="Path to motion .npz") - parser.add_argument("--no-loop", action="store_true") - parser.add_argument( - "--no-canon", action="store_true", help="Skip canonicalization (feed raw stored joints)" - ) - args = parser.parse_args() - - m = SmplMotion(args.motion, loop=not args.no_loop, canonicalize=not args.no_canon) - duration = m.num_frames / m.fps - print(f"Loaded '{args.motion}'") - print(f" frames={m.num_frames} fps={m.fps:.1f} duration={duration:.1f}s") - print( - f" smpl_joints={m.smpl_joints.shape} canonicalized={m.canonicalized} " - f"pose_aa={None if m.pose_aa is None else m.pose_aa.shape} " - f"transl={None if m.transl is None else m.transl.shape}" - ) - - # Sanity: after canonicalization the per-frame body heading should be fixed. - j = m.smpl_joints - v = j[:, 2, :2] - j[:, 1, :2] # R_hip - L_hip, horizontal - a = np.arctan2(v[:, 1], v[:, 0]) - rlen = np.clip(np.hypot(np.cos(a).mean(), np.sin(a).mean()), 1e-9, 1.0) - circ_std = np.degrees(np.sqrt(-2 * np.log(rlen))) - print(f" hip-heading circ-std={circ_std:.1f} deg (~0 => orientation removed; large => world-frame)") - - w0 = m.window(0) - print(f" window(0): shape={w0.shape} (expected {SMPL_OBS_DIM}) min={w0.min():.3f} max={w0.max():.3f}") - assert w0.shape == (SMPL_OBS_DIM,), "window must be 720-dim for obs[922:1642]" - - # Simulate a few control ticks. - print(" stepping 5 ticks:") - for t in range(5): - w = m.step() - print(f" t={t} cursor={m._cursor} window_norm={np.linalg.norm(w):.2f}") - - print("OK: motion loads and yields SONIC-format 720-vec windows.") - - -if __name__ == "__main__": - main() diff --git a/examples/unitree_g1/motions/dance_hiphop.npz b/examples/unitree_g1/motions/dance_hiphop.npz deleted file mode 100644 index 7fa5e7bea..000000000 Binary files a/examples/unitree_g1/motions/dance_hiphop.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/dance_jackson.npz b/examples/unitree_g1/motions/dance_jackson.npz deleted file mode 100644 index cad8c9d2f..000000000 Binary files a/examples/unitree_g1/motions/dance_jackson.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/eating_apple.npz b/examples/unitree_g1/motions/eating_apple.npz deleted file mode 100644 index fc47e1e10..000000000 Binary files a/examples/unitree_g1/motions/eating_apple.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/high_jump.npz b/examples/unitree_g1/motions/high_jump.npz deleted file mode 100644 index 84a86a14f..000000000 Binary files a/examples/unitree_g1/motions/high_jump.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/hiphop_ii.npz b/examples/unitree_g1/motions/hiphop_ii.npz deleted file mode 100644 index 40efdecf7..000000000 Binary files a/examples/unitree_g1/motions/hiphop_ii.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/horse_riding.npz b/examples/unitree_g1/motions/horse_riding.npz deleted file mode 100644 index d8fe0c63b..000000000 Binary files a/examples/unitree_g1/motions/horse_riding.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/jump_360.npz b/examples/unitree_g1/motions/jump_360.npz deleted file mode 100644 index caf9af433..000000000 Binary files a/examples/unitree_g1/motions/jump_360.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/mambo_kicks.npz b/examples/unitree_g1/motions/mambo_kicks.npz deleted file mode 100644 index c2c23dbd8..000000000 Binary files a/examples/unitree_g1/motions/mambo_kicks.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/reach_jump.npz b/examples/unitree_g1/motions/reach_jump.npz deleted file mode 100644 index 984bb53e1..000000000 Binary files a/examples/unitree_g1/motions/reach_jump.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/turn_jump_270.npz b/examples/unitree_g1/motions/turn_jump_270.npz deleted file mode 100644 index 881fcb70e..000000000 Binary files a/examples/unitree_g1/motions/turn_jump_270.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/turn_jump_360.npz b/examples/unitree_g1/motions/turn_jump_360.npz deleted file mode 100644 index e9a4ee21d..000000000 Binary files a/examples/unitree_g1/motions/turn_jump_360.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/victory_dance.npz b/examples/unitree_g1/motions/victory_dance.npz deleted file mode 100644 index 89a8aab10..000000000 Binary files a/examples/unitree_g1/motions/victory_dance.npz and /dev/null differ diff --git a/examples/unitree_g1/motions/walk_forward.npz b/examples/unitree_g1/motions/walk_forward.npz deleted file mode 100644 index c64d72dcd..000000000 Binary files a/examples/unitree_g1/motions/walk_forward.npz and /dev/null differ diff --git a/examples/unitree_g1/pkl_to_npz.py b/examples/unitree_g1/pkl_to_npz.py deleted file mode 100644 index 90806eeb1..000000000 --- a/examples/unitree_g1/pkl_to_npz.py +++ /dev/null @@ -1,88 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Convert a GEAR-SONIC / BONES-SEED ``smpl_filtered`` clip (.pkl) to .npz. - -The reference clips are zlib-compressed joblib pickles holding a dict with -``pose_aa`` (T, 72), ``transl`` (T, 3), ``smpl_joints`` (T, 24, 3), ``fps``. -``motion_loader.SmplMotion`` consumes the .npz form so the runtime needs no -joblib dependency. Canonicalization (root-orientation removal) happens at load -time in ``motion_loader``, so this converter just repackages the raw arrays. - -Run this in an environment that has ``joblib`` (e.g. the sonic teleop venv): - - python examples/unitree_g1/pkl_to_npz.py \ - --pkl sample_data/smpl_filtered/walk_forward_amateur_001__A001.pkl \ - --out examples/unitree_g1/motions/walk_forward.npz -""" - -import argparse -from pathlib import Path - -import numpy as np - - -def load_pkl(path: str) -> dict: - try: - import joblib - - return joblib.load(path) - except Exception: - # joblib clips are zlib-compressed pickles; fall back to manual inflate. - import contextlib - import pickle # nosec B403 - loads trusted local SMPL clips authored by the user - import zlib - - with open(path, "rb") as f: - raw = f.read() - with contextlib.suppress(zlib.error): - raw = zlib.decompress(raw) - return pickle.loads(raw) # nosec B301 - local, user-provided motion files only - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--pkl", required=True, help="Input smpl_filtered .pkl") - parser.add_argument("--out", required=True, help="Output .npz path") - args = parser.parse_args() - - d = load_pkl(args.pkl) - if not isinstance(d, dict) or "smpl_joints" not in d: - raise ValueError(f"Unexpected pkl structure; keys={list(d) if isinstance(d, dict) else type(d)}") - - smpl_joints = np.asarray(d["smpl_joints"], np.float32) - if smpl_joints.ndim != 3 or smpl_joints.shape[1:] != (24, 3): - raise ValueError(f"smpl_joints must be (T,24,3), got {smpl_joints.shape}") - - out = {"smpl_joints": smpl_joints, "fps": np.float32(d.get("fps", 50.0))} - if "pose_aa" in d: - out["pose_aa"] = np.asarray(d["pose_aa"], np.float32) - else: - print("[warn] no pose_aa -> loader cannot canonicalize (will feed raw)") - if "transl" in d: - out["transl"] = np.asarray(d["transl"], np.float32) - - Path(args.out).parent.mkdir(parents=True, exist_ok=True) - np.savez_compressed(args.out, **out) - dur = smpl_joints.shape[0] / float(out["fps"]) - print(f"Wrote {args.out}") - print( - f" frames={smpl_joints.shape[0]} fps={float(out['fps']):.1f} duration={dur:.1f}s keys={sorted(out)}" - ) - - -if __name__ == "__main__": - main() diff --git a/examples/unitree_g1/smpl_stream.py b/examples/unitree_g1/smpl_stream.py deleted file mode 100644 index 4e3d637b7..000000000 --- a/examples/unitree_g1/smpl_stream.py +++ /dev/null @@ -1,79 +0,0 @@ -#!/usr/bin/env python - -# Copyright 2026 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""CLI/self-test shim for the live ``rt/smpl`` SONIC reference stream. - -The implementation now lives in the installed package so it can be shared with the -``pico_headset`` teleoperator: -``lerobot.robots.unitree_g1.smpl_stream``. This example re-exports it (so -``from smpl_stream import SmplStream`` keeps working next to ``sonic.py``) and adds a -standalone smoke-test entrypoint. - -Example: - python examples/unitree_g1/smpl_stream.py --smpl-host 127.0.0.1 --smpl-port 5560 -""" - -from __future__ import annotations - -import argparse -import time - -import numpy as np - -from lerobot.robots.unitree_g1.smpl_stream import ( - DEFAULT_SMPL_HOST, - DEFAULT_SMPL_PORT, - SMPL_OBS_DIM, - SMPL_TOPIC, - SmplStream, -) - -__all__ = [ - "DEFAULT_SMPL_HOST", - "DEFAULT_SMPL_PORT", - "SMPL_OBS_DIM", - "SMPL_TOPIC", - "SmplStream", -] - - -def main(): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--smpl-host", default=DEFAULT_SMPL_HOST) - parser.add_argument("--smpl-port", type=int, default=DEFAULT_SMPL_PORT) - parser.add_argument("--ticks", type=int, default=250, help="control ticks to sample") - args = parser.parse_args() - - stream = SmplStream(host=args.smpl_host, port=args.smpl_port) - print(f"Subscribed to {SMPL_TOPIC} @ tcp://{args.smpl_host}:{args.smpl_port}") - print("Start pico_manager_thread_server.py --manager on the publisher host.") - try: - for t in range(args.ticks): - w = stream.step() - assert w.shape == (SMPL_OBS_DIM,), w.shape - if t % 25 == 0: - print( - f" t={t} idx={stream._last_index} window_norm={np.linalg.norm(w):.3f} " - f"first={stream.has_data}" - ) - time.sleep(1.0 / 50.0) - finally: - stream.close() - print("OK: rt/smpl stream yields SONIC-format 720-vec windows.") - - -if __name__ == "__main__": - main() diff --git a/examples/unitree_g1/sonic.py b/examples/unitree_g1/sonic.py deleted file mode 100644 index eae0049fe..000000000 --- a/examples/unitree_g1/sonic.py +++ /dev/null @@ -1,420 +0,0 @@ -#!/usr/bin/env python -""" -SONIC planner with full mode control. - -Keyboard controls: - N / P - next / previous motion set - 1-8 - select mode within current set - WASD - movement direction - Q / E - rotate facing left / right - 9 / 0 - decrease / increase speed - - / = - decrease / increase height - R - force replan - M - toggle SMPL motion playback <-> locomotion (needs --motion-file) - Space - emergency stop -> IDLE - Esc - quit - -Gamepad controls (Unitree wireless controller): - Left stick Y - speed (forward = fast, back = stop) - Left stick X - movement direction (offset from facing) - Right stick X - facing direction (incremental rotation) - Right stick Y - height (up = tall 0.8m, down = low 0.1m) - Buttons - unused (mode selection is keyboard-only) - -For teleop integration use --robot.controller=SonicWholeBodyController instead. -""" - -import argparse -import contextlib -import faulthandler -import gc -import os -import sys -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 - -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, -) -from lerobot.robots.unitree_g1.controllers.sonic_whole_body import SonicRuntime -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( - "--ip", - type=str, - default=None, - help="Robot IP for real hardware (e.g. 192.168.123.164). Omit for simulation.", - ) - parser.add_argument( - "--log-csv", - action="store_true", - help="Write /tmp/sonic_pose_log.csv (disabled by default for teleop perf)", - ) - parser.add_argument( - "--cpu", - action="store_true", - help="Force CPU ONNX Runtime (skip CUDA even if onnxruntime-gpu is installed)", - ) - parser.add_argument( - "--headless", action="store_true", help="Ignored for sim (stock UnitreeG1 uses hub MuJoCo defaults)" - ) - parser.add_argument( - "--gamepad", - action="store_true", - help="Read Unitree wireless gamepad in sim (default: keyboard-only in sim)", - ) - parser.add_argument( - "--keyboard-only", action="store_true", help="Ignore wireless gamepad (terminal keyboard only)" - ) - parser.add_argument( - "--motion-file", - type=str, - default=None, - help="Play an SMPL motion clip (.npz) via SONIC whole-body mode " - "(encode_mode=2) instead of locomotion planning.", - ) - parser.add_argument( - "--no-loop", action="store_true", help="With --motion-file, play once instead of looping" - ) - parser.add_argument( - "--smpl-stream", - action="store_true", - help="Use the live rt/smpl headset stream as the reference motion " - "(SONIC whole-body, encode_mode=2), instead of a --motion-file clip.", - ) - parser.add_argument( - "--smpl-host", - type=str, - default=DEFAULT_SMPL_HOST, - help=f"Host publishing rt/smpl (default: {DEFAULT_SMPL_HOST})", - ) - parser.add_argument( - "--smpl-port", - type=int, - 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() - - 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. - faulthandler.enable() - with contextlib.suppress(Exception): - sys.stdout.reconfigure(line_buffering=True) - - print("=" * 60) - print("SONIC planner - full mode control") - print(" N/P cycle sets | 1-8 select mode | WASD move") - print(" Q/E rotate | 9/0 speed | -/= height") - print(" R replan | Space IDLE | Esc quit") - if args.ip: - print(f" Robot IP: {args.ip}") - else: - print(" Mode: simulation") - print("=" * 60 + "\n") - - cfg = UnitreeG1Config(controller=None) # full-body SONIC; standalone loop owns publish - if args.ip: - cfg.is_simulation = False - cfg.robot_ip = args.ip - else: - cfg.is_simulation = True - if args.headless: - print("[Note] --headless ignored: sim uses stock UnitreeG1 + hub env") - robot = UnitreeG1(cfg) - robot.connect() - kp, kd = compute_kp_kd() - robot.kp = kp.copy() - robot.kd = kd.copy() - - runtime = SonicRuntime(force_cpu=args.cpu) - 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) - controller.smpl_motion = motion # lets 'M' key toggle streaming - controller.encode_mode = 2 # start in SONIC whole-body SMPL imitation - print(f"\n[Motion] live SMPL stream (rt/smpl @ {args.smpl_host}:{args.smpl_port})") - print(" Source: pico_manager_thread_server.py --manager on the publisher host.") - print(" encode_mode=2. Press 'M' to toggle SMPL stream <-> locomotion at runtime.") - elif args.motion_file: - motion = SmplMotion(args.motion_file, loop=not args.no_loop) - controller.smpl_motion = motion # lets 'M' key toggle playback - controller.encode_mode = 2 # start in SONIC whole-body SMPL imitation - dur = motion.num_frames / motion.fps - print(f"\n[Motion] SMPL whole-body playback: {args.motion_file}") - print( - f" frames={motion.num_frames} fps={motion.fps:.1f} " - f"duration={dur:.1f}s loop={not args.no_loop} encode_mode=2" - ) - print(" Press 'M' to toggle SMPL playback <-> locomotion at runtime.") - - runtime.controller.print_input_diagnostics() - - print(f"\nStarting: {MOTION_SETS[0][0]} (default mode: {LM(ms.mode).name})") - [print(f" {i + 1}: {m.name}") for i, m in enumerate(MOTION_SETS[0][1])] - print( - "\n[Ready] Click THIS terminal, then W/A/S/D to move. 1-6 change mode, 9/0 speed, Esc quit.\n", - flush=True, - ) - - # Sim hub publishes wireless_remote bytes that can fight terminal WASD. - base_joystick = not args.keyboard_only and (args.gamepad or args.ip is not None) - - with RawKeyboard() as kb: - try: - gc.disable() - gc_timer = 0.0 - robot.reset(CONTROL_DT, DEFAULT_ANGLES) - time.sleep(1.0) - - last_status = time.time() - 2.1 - loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], [] - slow_n = blend_n = 0 - stall_src = "" - did_blend = False - t_start = time.time() - - log_path = os.path.join(tempfile.gettempdir(), "sonic_pose_log.csv") - jnames = [m.name for m in G1_29_JointIndex] - log_ctx = open(log_path, "w") if args.log_csv else None # noqa: SIM115 - if log_ctx: - log_ctx.write( - "t,step,cursor,ts,blend,mode," - + ",".join(f"q{i}" for i in range(29)) - + "," - + ",".join(f"ref{i}" for i in range(29)) - + "," - + ",".join(f"act{i}" for i in range(29)) - + ",delta_max,action_norm,token_norm\n" - ) - - try: - while not robot._shutdown_event.is_set(): - t0 = time.time() - if drain_keyboard(kb, ms, controller): - break - - obs = robot.get_observation() - t_obs = time.time() - obs_t.append(1000 * (t_obs - t0)) - if not obs: - runtime.tick({}, use_joystick=False) - 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: - controller.smpl_joints_10frame_step1 = motion.step() - if motion.done: - print("\n[Motion] clip finished") - break - - step_before = runtime.step - t_step = time.time() - action = runtime.tick(obs, use_joystick=base_joystick and not motion_active) - step_ms = 1000 * (time.time() - t_step) - do_enc = step_before % 5 == 0 - (enc_t if do_enc else dec_t).append(step_ms) - - t_act = time.time() - robot.send_action(action) - act_t.append(1000 * (time.time() - t_act)) - - if log_ctx and runtime.step % 5 == 0: - t_rel = time.time() - t_start - q_r = np.array([obs.get(f"{n}.q", 0) for n in jnames]) - a_v = np.array([action.get(f"{n}.q", 0) for n in jnames]) - cur, ts = controller.ref_cursor, controller.motion_timesteps - q_ref = ( - controller.motion_joint_positions[min(cur, ts - 1)] if ts > 0 else np.zeros(29) - ) - log_ctx.write( - f"{t_rel:.4f},{runtime.step},{cur},{ts},{int(did_blend)},{ms.mode}," - + ",".join(f"{v:.6f}" for v in q_r) - + "," - + ",".join(f"{v:.6f}" for v in q_ref) - + "," - + ",".join(f"{v:.6f}" for v in a_v) - + "," - + f"{np.max(np.abs(a_v - q_r)):.6f}," - f"{np.linalg.norm(a_v):.6f}," - f"{np.linalg.norm(controller.token):.6f}\n" - ) - did_blend = False - - now = time.time() - loop_ms = 1000 * (now - t0) - if loop_ms > 50: - stall_src = ( - f"[STALL] {loop_ms:.0f}ms: " - f"obs={obs_t[-1]:.0f} step={step_ms:.0f} act={act_t[-1]:.0f}" - ) - if loop_ms > CONTROL_DT * 1500: - slow_n += 1 - - if now - last_status > 2.0: - - def _avg(lst): - return sum(lst) / len(lst) if lst else 0 - - hz = 1000 / _avg(loop_t) if _avg(loop_t) else 0 - print( - f"\r {ms.status_line()} step={runtime.step} " - f"ref={controller.ref_cursor}/{controller.motion_timesteps} " - f"loop={_avg(loop_t):.1f}ms(max={max(loop_t, default=0):.1f}) hz={hz:.0f} " - f"enc={_avg(enc_t):.1f} dec={_avg(dec_t):.1f} obs={_avg(obs_t):.1f} " - f"slow={slow_n} blends={blend_n}", - end="", - flush=True, - ) - if stall_src: - print(f"\n {stall_src}") - stall_src = "" - last_status = now - loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], [] - slow_n = blend_n = 0 - - gc_timer += CONTROL_DT - if gc_timer >= 10.0: - gc.collect() - gc_timer = 0.0 - loop_t.append(loop_ms) - time.sleep(max(0.0, CONTROL_DT - (time.time() - t0))) - finally: - if log_ctx: - log_ctx.close() - - except KeyboardInterrupt: - pass - finally: - gc.enable() - if args.log_csv: - print(f"\n[Log] Saved to {log_path}") - if motion is not None and hasattr(motion, "close"): - motion.close() - runtime.shutdown() - print("\nStopping...") - if robot.is_connected: - robot.disconnect() - print("Done.") - - -if __name__ == "__main__": - main()