remove examples inlcuding npz motion files

This commit is contained in:
Martino Russi
2026-07-10 15:47:21 +02:00
parent 4bcde762cc
commit 1c87ca31a3
19 changed files with 0 additions and 1003 deletions
-111
View File
@@ -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_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
```
-135
View File
@@ -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()
-170
View File
@@ -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()
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
-88
View File
@@ -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()
-79
View File
@@ -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()
-420
View File
@@ -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()