fix(unitree_g1): satisfy ruff lint/format and address review comments

This commit is contained in:
Martino Russi
2026-06-25 14:36:05 +02:00
parent 804c76bcc2
commit dfbc25c58f
6 changed files with 814 additions and 405 deletions
+16 -20
View File
@@ -66,11 +66,11 @@ def canonicalize_smpl_joints(smpl_joints: np.ndarray, root_aa: np.ndarray) -> np
Returns: Returns:
(T, 24, 3) per-frame root-orientation-removed joints. (T, 24, 3) per-frame root-orientation-removed joints.
""" """
from scipy.spatial.transform import Rotation as R from scipy.spatial.transform import Rotation
rx90 = R.from_euler("x", 90, degrees=True) # smpl_root_ytoz_up rx90 = Rotation.from_euler("x", 90, degrees=True) # smpl_root_ytoz_up
base120 = R.from_quat([0.5, 0.5, 0.5, 0.5]) # remove_smpl_base_rot base120 = Rotation.from_quat([0.5, 0.5, 0.5, 0.5]) # remove_smpl_base_rot
a = rx90 * R.from_rotvec(root_aa) # z-up root quat (left-mult) a = rx90 * Rotation.from_rotvec(root_aa) # z-up root quat (left-mult)
b_inv = base120 * a.inv() # inv(remove_smpl_base_rot(a)) 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) return np.einsum("tij,tkj->tki", b_inv.as_matrix(), smpl_joints).astype(np.float32)
@@ -87,9 +87,7 @@ class SmplMotion:
self.loop = loop self.loop = loop
if smpl_joints.ndim != 3 or smpl_joints.shape[1:] != (N_JOINTS, JOINT_DIM): if smpl_joints.ndim != 3 or smpl_joints.shape[1:] != (N_JOINTS, JOINT_DIM):
raise ValueError( raise ValueError(f"Expected smpl_joints (T, {N_JOINTS}, {JOINT_DIM}), got {smpl_joints.shape}")
f"Expected smpl_joints (T, {N_JOINTS}, {JOINT_DIM}), got {smpl_joints.shape}"
)
# Reference clips store world-frame joints; the encoder wants per-frame # Reference clips store world-frame joints; the encoder wants per-frame
# root-orientation-removed joints. Canonicalize when we have the root pose. # root-orientation-removed joints. Canonicalize when we have the root pose.
@@ -109,10 +107,7 @@ class SmplMotion:
[f0_j0_xyz, f0_j1_xyz, ..., f9_j23_xyz]. [f0_j0_xyz, f0_j1_xyz, ..., f9_j23_xyz].
""" """
idx = np.arange(start, start + WINDOW) idx = np.arange(start, start + WINDOW)
if self.loop: idx = np.mod(idx, self.num_frames) if self.loop else np.clip(idx, 0, self.num_frames - 1)
idx = np.mod(idx, self.num_frames)
else:
idx = np.clip(idx, 0, self.num_frames - 1)
return self.smpl_joints[idx].reshape(-1).astype(np.float32) return self.smpl_joints[idx].reshape(-1).astype(np.float32)
def reset(self): def reset(self):
@@ -135,30 +130,31 @@ def main():
parser = argparse.ArgumentParser(description=__doc__) parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--motion", required=True, help="Path to motion .npz") parser.add_argument("--motion", required=True, help="Path to motion .npz")
parser.add_argument("--no-loop", action="store_true") parser.add_argument("--no-loop", action="store_true")
parser.add_argument("--no-canon", action="store_true", parser.add_argument(
help="Skip canonicalization (feed raw stored joints)") "--no-canon", action="store_true", help="Skip canonicalization (feed raw stored joints)"
)
args = parser.parse_args() args = parser.parse_args()
m = SmplMotion(args.motion, loop=not args.no_loop, canonicalize=not args.no_canon) m = SmplMotion(args.motion, loop=not args.no_loop, canonicalize=not args.no_canon)
duration = m.num_frames / m.fps duration = m.num_frames / m.fps
print(f"Loaded '{args.motion}'") print(f"Loaded '{args.motion}'")
print(f" frames={m.num_frames} fps={m.fps:.1f} duration={duration:.1f}s") 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} " 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"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}") f"transl={None if m.transl is None else m.transl.shape}"
)
# Sanity: after canonicalization the per-frame body heading should be fixed. # Sanity: after canonicalization the per-frame body heading should be fixed.
j = m.smpl_joints j = m.smpl_joints
v = (j[:, 2, :2] - j[:, 1, :2]) # R_hip - L_hip, horizontal v = j[:, 2, :2] - j[:, 1, :2] # R_hip - L_hip, horizontal
a = np.arctan2(v[:, 1], v[:, 0]) a = np.arctan2(v[:, 1], v[:, 0])
rlen = np.clip(np.hypot(np.cos(a).mean(), np.sin(a).mean()), 1e-9, 1.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))) circ_std = np.degrees(np.sqrt(-2 * np.log(rlen)))
print(f" hip-heading circ-std={circ_std:.1f} deg " print(f" hip-heading circ-std={circ_std:.1f} deg (~0 => orientation removed; large => world-frame)")
f"(~0 => orientation removed; large => world-frame)")
w0 = m.window(0) w0 = m.window(0)
print(f" window(0): shape={w0.shape} (expected {SMPL_OBS_DIM}) " print(f" window(0): shape={w0.shape} (expected {SMPL_OBS_DIM}) min={w0.min():.3f} max={w0.max():.3f}")
f"min={w0.min():.3f} max={w0.max():.3f}")
assert w0.shape == (SMPL_OBS_DIM,), "window must be 720-dim for obs[922:1642]" assert w0.shape == (SMPL_OBS_DIM,), "window must be 720-dim for obs[922:1642]"
# Simulate a few control ticks. # Simulate a few control ticks.
+5 -5
View File
@@ -42,15 +42,14 @@ def load_pkl(path: str) -> dict:
return joblib.load(path) return joblib.load(path)
except Exception: except Exception:
# joblib clips are zlib-compressed pickles; fall back to manual inflate. # joblib clips are zlib-compressed pickles; fall back to manual inflate.
import contextlib
import pickle import pickle
import zlib import zlib
with open(path, "rb") as f: with open(path, "rb") as f:
raw = f.read() raw = f.read()
try: with contextlib.suppress(zlib.error):
raw = zlib.decompress(raw) raw = zlib.decompress(raw)
except zlib.error:
pass
return pickle.loads(raw) return pickle.loads(raw)
@@ -80,8 +79,9 @@ def main():
np.savez_compressed(args.out, **out) np.savez_compressed(args.out, **out)
dur = smpl_joints.shape[0] / float(out["fps"]) dur = smpl_joints.shape[0] / float(out["fps"])
print(f"Wrote {args.out}") print(f"Wrote {args.out}")
print(f" frames={smpl_joints.shape[0]} fps={float(out['fps']):.1f} duration={dur:.1f}s " print(
f"keys={sorted(out)}") f" frames={smpl_joints.shape[0]} fps={float(out['fps']):.1f} duration={dur:.1f}s keys={sorted(out)}"
)
if __name__ == "__main__": if __name__ == "__main__":
+86 -49
View File
@@ -25,15 +25,16 @@ For teleop integration use --robot.controller=SonicWholeBodyController instead.
""" """
import argparse import argparse
import contextlib
import faulthandler import faulthandler
import gc import gc
import sys import sys
import time import time
import numpy as np import numpy as np
from motion_loader import SmplMotion
from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
from lerobot.robots.unitree_g1.controllers.sonic_whole_body import SonicRuntime
from lerobot.robots.unitree_g1.controllers.sonic_pipeline import ( from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
CONTROL_DT, CONTROL_DT,
DEFAULT_ANGLES, DEFAULT_ANGLES,
@@ -43,41 +44,57 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
compute_kp_kd, compute_kp_kd,
drain_keyboard, 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.g1_utils import G1_29_JointIndex
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1 from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
from motion_loader import SmplMotion
def main(): def main():
parser = argparse.ArgumentParser(description="SONIC planner with keyboard + gamepad control") parser = argparse.ArgumentParser(description="SONIC planner with keyboard + gamepad control")
parser.add_argument("--ip", type=str, default=None, parser.add_argument(
help="Robot IP for real hardware (e.g. 192.168.123.164). " "--ip",
"Omit for simulation.") type=str,
parser.add_argument("--log-csv", action="store_true", default=None,
help="Write /tmp/sonic_pose_log.csv (disabled by default for teleop perf)") help="Robot IP for real hardware (e.g. 192.168.123.164). Omit for simulation.",
parser.add_argument("--cpu", action="store_true", )
help="Force CPU ONNX Runtime (skip CUDA even if onnxruntime-gpu is installed)") parser.add_argument(
parser.add_argument("--headless", action="store_true", "--log-csv",
help="Ignored for sim (stock UnitreeG1 uses hub MuJoCo defaults)") action="store_true",
parser.add_argument("--gamepad", action="store_true", help="Write /tmp/sonic_pose_log.csv (disabled by default for teleop perf)",
help="Read Unitree wireless gamepad in sim (default: keyboard-only in sim)") )
parser.add_argument("--keyboard-only", action="store_true", parser.add_argument(
help="Ignore wireless gamepad (terminal keyboard only)") "--cpu",
parser.add_argument("--motion-file", type=str, default=None, 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 " help="Play an SMPL motion clip (.npz) via SONIC whole-body mode "
"(encode_mode=2) instead of locomotion planning.") "(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(
"--no-loop", action="store_true", help="With --motion-file, play once instead of looping"
)
args = parser.parse_args() args = parser.parse_args()
# Surface native crashes (onnxruntime / mujoco) with a real traceback, and # Surface native crashes (onnxruntime / mujoco) with a real traceback, and
# avoid losing buffered diagnostics if the process dies mid-loop. # avoid losing buffered diagnostics if the process dies mid-loop.
faulthandler.enable() faulthandler.enable()
try: with contextlib.suppress(Exception):
sys.stdout.reconfigure(line_buffering=True) sys.stdout.reconfigure(line_buffering=True)
except Exception:
pass
print("=" * 60) print("=" * 60)
print("SONIC planner - full mode control") print("SONIC planner - full mode control")
@@ -115,16 +132,20 @@ def main():
controller.encode_mode = 2 # start in SONIC whole-body SMPL imitation controller.encode_mode = 2 # start in SONIC whole-body SMPL imitation
dur = motion.num_frames / motion.fps dur = motion.num_frames / motion.fps
print(f"\n[Motion] SMPL whole-body playback: {args.motion_file}") print(f"\n[Motion] SMPL whole-body playback: {args.motion_file}")
print(f" frames={motion.num_frames} fps={motion.fps:.1f} " print(
f"duration={dur:.1f}s loop={not args.no_loop} encode_mode=2") 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.") print(" Press 'M' to toggle SMPL playback <-> locomotion at runtime.")
runtime.controller.print_input_diagnostics() runtime.controller.print_input_diagnostics()
print(f"\nStarting: {MOTION_SETS[0][0]} (default mode: {LM(ms.mode).name})") 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(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. " print(
"1-6 change mode, 9/0 speed, Esc quit.\n", flush=True) "\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. # 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) base_joystick = not args.keyboard_only and (args.gamepad or args.ip is not None)
@@ -137,22 +158,25 @@ def main():
time.sleep(1.0) time.sleep(1.0)
last_status = time.time() - 2.1 last_status = time.time() - 2.1
loop_t = enc_t = dec_t = obs_t = act_t = [] loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], []
slow_n = blend_n = 0 slow_n = blend_n = 0
stall_src = "" stall_src = ""
did_blend = False did_blend = False
prev_end = time.time()
t_start = time.time() t_start = time.time()
log_path = "/tmp/sonic_pose_log.csv" log_path = "/tmp/sonic_pose_log.csv"
jnames = [m.name for m in G1_29_JointIndex] jnames = [m.name for m in G1_29_JointIndex]
log_ctx = open(log_path, "w") if args.log_csv else None log_ctx = open(log_path, "w") if args.log_csv else None # noqa: SIM115
if log_ctx: if log_ctx:
log_ctx.write("t,step,cursor,ts,blend,mode," + log_ctx.write(
",".join(f"q{i}" for i in range(29)) + "," + "t,step,cursor,ts,blend,mode,"
",".join(f"ref{i}" for i in range(29)) + "," + + ",".join(f"q{i}" for i in range(29))
",".join(f"act{i}" for i in range(29)) + + ","
",delta_max,action_norm,token_norm\n") + ",".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: try:
while not robot._shutdown_event.is_set(): while not robot._shutdown_event.is_set():
@@ -192,42 +216,55 @@ def main():
q_r = np.array([obs.get(f"{n}.q", 0) for n in jnames]) 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]) a_v = np.array([action.get(f"{n}.q", 0) for n in jnames])
cur, ts = controller.ref_cursor, controller.motion_timesteps 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) q_ref = (
log_ctx.write(f"{t_rel:.4f},{runtime.step},{cur},{ts},{int(did_blend)},{ms.mode}," + controller.motion_joint_positions[min(cur, ts - 1)] if ts > 0 else np.zeros(29)
",".join(f"{v:.6f}" for v in q_r) + "," + )
",".join(f"{v:.6f}" for v in q_ref) + "," + log_ctx.write(
",".join(f"{v:.6f}" for v in a_v) + "," + f"{t_rel:.4f},{runtime.step},{cur},{ts},{int(did_blend)},{ms.mode},"
f"{np.max(np.abs(a_v - q_r)):.6f}," + ",".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(a_v):.6f},"
f"{np.linalg.norm(controller.token):.6f}\n") f"{np.linalg.norm(controller.token):.6f}\n"
)
did_blend = False did_blend = False
now = time.time() now = time.time()
loop_ms = 1000 * (now - t0) loop_ms = 1000 * (now - t0)
if loop_ms > 50: if loop_ms > 50:
stall_src = (f"[STALL] {loop_ms:.0f}ms: " stall_src = (
f"obs={obs_t[-1]:.0f} step={step_ms:.0f} act={act_t[-1]:.0f}") 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: if loop_ms > CONTROL_DT * 1500:
slow_n += 1 slow_n += 1
if now - last_status > 2.0: if now - last_status > 2.0:
def _avg(lst): def _avg(lst):
return sum(lst) / len(lst) if lst else 0 return sum(lst) / len(lst) if lst else 0
hz = 1000 / _avg(loop_t) if _avg(loop_t) else 0 hz = 1000 / _avg(loop_t) if _avg(loop_t) else 0
print(f"\r {ms.status_line()} step={runtime.step} " print(
f"\r {ms.status_line()} step={runtime.step} "
f"ref={controller.ref_cursor}/{controller.motion_timesteps} " 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"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"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) f"slow={slow_n} blends={blend_n}",
end="",
flush=True,
)
if stall_src: if stall_src:
print(f"\n {stall_src}") print(f"\n {stall_src}")
stall_src = "" stall_src = ""
last_status = now last_status = now
loop_t = enc_t = dec_t = obs_t = act_t = [] loop_t, enc_t, dec_t, obs_t, act_t = [], [], [], [], []
slow_n = blend_n = 0 slow_n = blend_n = 0
prev_end = time.time()
gc_timer += CONTROL_DT gc_timer += CONTROL_DT
if gc_timer >= 10.0: if gc_timer >= 10.0:
gc.collect() gc.collect()
@@ -1,3 +1,19 @@
#!/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.
"""Unitree G1 locomotion controllers (Groot, Holosoma, SONIC).""" """Unitree G1 locomotion controllers (Groot, Holosoma, SONIC)."""
__all__ = [ __all__ = [
File diff suppressed because it is too large Load Diff
@@ -18,7 +18,6 @@
import logging import logging
import numpy as np
import onnxruntime as ort import onnxruntime as ort
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
@@ -32,13 +31,13 @@ from lerobot.robots.unitree_g1.controllers.sonic_pipeline import (
MovementState, MovementState,
PlannerController, PlannerController,
SonicPlanner, SonicPlanner,
_ort_providers,
_snapshot_ms,
clamp_mode_params, clamp_mode_params,
compute_kp_kd, compute_kp_kd,
lowstate_to_obs, lowstate_to_obs,
process_joystick, process_joystick,
should_replan_request, should_replan_request,
_ort_providers,
_snapshot_ms,
) )
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)