mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-22 17:32:07 +00:00
Martino Russi
706a551d26
Use SonicSimRobot for MuJoCo sim, run the planner in a background thread to avoid spawn/CUDA conflicts, and align replan/blend/GPU ORT behavior with deploy. Co-authored-by: Cursor <cursoragent@cursor.com>
1229 lines
55 KiB
Python
1229 lines
55 KiB
Python
#!/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
|
||
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)
|
||
"""
|
||
|
||
import argparse, gc, math, queue, select, sys, termios, tty
|
||
import threading, time
|
||
from dataclasses import dataclass
|
||
from enum import IntEnum
|
||
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
from huggingface_hub import hf_hub_download
|
||
|
||
from lerobot.robots.unitree_g1.config_unitree_g1 import UnitreeG1Config
|
||
from lerobot.robots.unitree_g1.unitree_g1 import UnitreeG1
|
||
from lerobot.robots.unitree_g1.g1_utils import G1_29_JointIndex
|
||
|
||
# ── Constants ────────────────────────────────────────────────────────────────
|
||
|
||
DEFAULT_ANGLES = np.array([
|
||
-0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
|
||
-0.312, 0.0, 0.0, 0.669, -0.363, 0.0,
|
||
0.0, 0.0, 0.0,
|
||
0.2, 0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
|
||
0.2, -0.2, 0.0, 0.6, 0.0, 0.0, 0.0,
|
||
], dtype=np.float32)
|
||
|
||
NATURAL_FREQ = 10.0 * 2.0 * np.pi
|
||
ARMATURE = {"5020": 0.003609725, "7520_14": 0.010177520, "7520_22": 0.025101925, "4010": 0.00425}
|
||
EFFORT = {"5020": 25.0, "7520_14": 88.0, "7520_22": 139.0, "4010": 5.0}
|
||
|
||
def _action_scale(k):
|
||
return 0.25 * EFFORT[k] / (ARMATURE[k] * NATURAL_FREQ**2)
|
||
|
||
_J = ["7520_22","7520_22","7520_14","7520_22","5020","5020"] * 2 + \
|
||
["7520_14","5020","5020"] + \
|
||
["5020","5020","5020","5020","5020","4010","4010"] * 2
|
||
ACTION_SCALE = np.array([_action_scale(k) for k in _J], dtype=np.float32)
|
||
|
||
CONTROL_DT = 0.02
|
||
DEFAULT_HEIGHT = 0.788740
|
||
TOKEN_DIM = 64
|
||
ENCODER_UPDATE_EVERY = 5
|
||
DEBUG_PRINT_EVERY = 100
|
||
MOTION_LOOK_AHEAD_STEPS = 2
|
||
INITIAL_RANDOM_SEED = 1234
|
||
MIN_TOKENS, MAX_TOKENS = 6, 16
|
||
K = MAX_TOKENS - MIN_TOKENS + 1
|
||
DEADZONE = 0.05
|
||
BLEND_FRAMES = 8
|
||
|
||
REPLAN_INTERVAL = {
|
||
"running": 0.1, "crawling": 0.2, "boxing": 1.0, "default": 1.0
|
||
}
|
||
|
||
ISAACLAB_TO_MUJOCO = np.array([
|
||
0, 3, 6, 9, 13, 17, 1, 4, 7, 10, 14, 18, 2, 5, 8,
|
||
11, 15, 19, 21, 23, 25, 27, 12, 16, 20, 22, 24, 26, 28
|
||
], dtype=np.int32)
|
||
|
||
MUJOCO_TO_ISAACLAB = np.array([
|
||
0, 6, 12, 1, 7, 13, 2, 8, 14, 3, 9, 15, 22, 4, 10,
|
||
16, 23, 5, 11, 17, 24, 18, 25, 19, 26, 20, 27, 21, 28
|
||
], dtype=np.int32)
|
||
|
||
def _to_mujoco(a): return a[MUJOCO_TO_ISAACLAB]
|
||
def _to_runtime(a): r = np.zeros(29, np.float32); r[MUJOCO_TO_ISAACLAB] = a; return r
|
||
|
||
DEFAULT_ANGLES_MUJOCO = _to_mujoco(DEFAULT_ANGLES)
|
||
ENCODER_STANDING_REF = DEFAULT_ANGLES.copy()
|
||
|
||
LOWER_BODY_IL = np.array([0,3,6,9,13,17,1,4,7,10,14,18], dtype=np.int32)
|
||
WRIST_IL = np.array([23,24,25,26,27,28], dtype=np.int32)
|
||
VR_TARGET_DEF = np.zeros(9, dtype=np.float32)
|
||
VR_ORN_DEF = np.array([1,0,0,0,1,0,0,0,1,0,0,0], dtype=np.float32)
|
||
SMPL_DEF = np.zeros(720, dtype=np.float32)
|
||
|
||
# ── PD gains ─────────────────────────────────────────────────────────────────
|
||
|
||
def _kp_kd():
|
||
s = lambda k: ARMATURE[k] * NATURAL_FREQ**2
|
||
d = lambda k: 2.0 * 2.0 * ARMATURE[k] * NATURAL_FREQ
|
||
_kp_keys = ["7520_22","7520_22","7520_14","7520_22","5020","5020"] * 2 + \
|
||
["7520_14","5020","5020"] + \
|
||
["5020","5020","5020","5020","5020","4010","4010"] * 2
|
||
_kd_keys = _kp_keys
|
||
_double = {4,5,10,11,13,14} # ankle + waist indices with factor 2
|
||
kp = np.array([2*s(k) if i in _double else s(k) for i,k in enumerate(_kp_keys)], dtype=np.float32)
|
||
kd = np.array([2*d(k) if i in _double else d(k) for i,k in enumerate(_kd_keys)], dtype=np.float32)
|
||
return kp, kd
|
||
|
||
# ── Quaternion helpers ────────────────────────────────────────────────────────
|
||
|
||
def quat_conj(q):
|
||
return np.array([q[0], -q[1], -q[2], -q[3]], dtype=np.float32)
|
||
|
||
def quat_mul(q1, q2):
|
||
w1,x1,y1,z1 = q1; w2,x2,y2,z2 = q2
|
||
return np.array([
|
||
w1*w2 - x1*x2 - y1*y2 - z1*z2,
|
||
w1*x2 + x1*w2 + y1*z2 - z1*y2,
|
||
w1*y2 - x1*z2 + y1*w2 + z1*x2,
|
||
w1*z2 + x1*y2 - y1*x2 + z1*w2,
|
||
], dtype=np.float32)
|
||
|
||
def gravity_dir(q):
|
||
q = q / (np.linalg.norm(q) + 1e-8)
|
||
qv = np.array([0, 0, 0, -1], dtype=np.float32)
|
||
return quat_mul(quat_mul(quat_conj(q), qv), q)[1:]
|
||
|
||
def quat_to_6d(q):
|
||
w,x,y,z = q
|
||
return np.array([
|
||
1-2*(y*y+z*z), 2*(x*y-z*w),
|
||
2*(x*y+z*w), 1-2*(x*x+z*z),
|
||
2*(x*z-y*w), 2*(y*z+x*w),
|
||
], dtype=np.float32)
|
||
|
||
def calc_heading(q):
|
||
w,x,y,z = q
|
||
return float(np.arctan2(2*(x*y + w*z), 1-2*(y*y+z*z)))
|
||
|
||
def heading_quat(q, sign=1.0):
|
||
a = sign * calc_heading(q) / 2.0
|
||
return np.array([np.cos(a), 0, 0, np.sin(a)], dtype=np.float64)
|
||
|
||
heading_quat_inv = lambda q: heading_quat(q, -1.0)
|
||
|
||
def quat_slerp(q0, q1, t):
|
||
q0 = q0 / (np.linalg.norm(q0)+1e-12); q1 = q1 / (np.linalg.norm(q1)+1e-12)
|
||
dot = float(np.dot(q0, q1))
|
||
if dot < 0: q1, dot = -q1, -dot
|
||
dot = min(dot, 1.0)
|
||
if dot > 0.9995:
|
||
r = q0 + t*(q1-q0); return r/(np.linalg.norm(r)+1e-12)
|
||
th = np.arccos(dot); st = np.sin(th)
|
||
return (np.sin((1-t)*th)/st)*q0 + (np.sin(t*th)/st)*q1
|
||
|
||
def quat_slerp_batch(q0, q1, t):
|
||
q0 = q0 / (np.linalg.norm(q0,axis=1,keepdims=True)+1e-12)
|
||
q1 = q1 / (np.linalg.norm(q1,axis=1,keepdims=True)+1e-12)
|
||
dot = np.sum(q0*q1, axis=1); neg = dot<0
|
||
q1=q1.copy(); q1[neg]=-q1[neg]; dot[neg]=-dot[neg]; dot=np.clip(dot,-1,1)
|
||
lin = dot>0.9995; th=np.arccos(dot); st=np.where(np.sin(th)==0,1,np.sin(th))
|
||
c0=np.sin((1-t)*th)/st; c1=np.sin(t*th)/st
|
||
c0[lin]=1-t[lin]; c1[lin]=t[lin]
|
||
r = c0[:,None]*q0 + c1[:,None]*q1
|
||
return r / (np.linalg.norm(r,axis=1,keepdims=True)+1e-12)
|
||
|
||
# ── Locomotion modes ──────────────────────────────────────────────────────────
|
||
|
||
class LocomotionMode(IntEnum):
|
||
IDLE=0; SLOW_WALK=1; WALK=2; RUN=3; SQUAT=4; KNEEL_TWO_LEGS=5; KNEEL=6
|
||
LYING_FACE_DOWN=7; CRAWLING=8; IDLE_BOXING=9; WALK_BOXING=10
|
||
LEFT_PUNCH=11; RIGHT_PUNCH=12; RANDOM_PUNCH=13; ELBOW_CRAWLING=14
|
||
LEFT_HOOK=15; RIGHT_HOOK=16; FORWARD_JUMP=17; STEALTH_WALK=18
|
||
INJURED_WALK=19; LEDGE_WALKING=20; OBJECT_CARRYING=21; STEALTH_WALK_2=22
|
||
HAPPY_DANCE_WALK=23; ZOMBIE_WALK=24; GUN_WALK=25; SCARE_WALK=26
|
||
|
||
LM = LocomotionMode
|
||
|
||
MOTION_SETS = [
|
||
("Standing", [LM.SLOW_WALK, LM.WALK, LM.RUN, LM.FORWARD_JUMP, LM.STEALTH_WALK, LM.INJURED_WALK]),
|
||
("Squat / Low", [LM.SQUAT, LM.KNEEL_TWO_LEGS, LM.KNEEL, LM.CRAWLING, LM.ELBOW_CRAWLING]),
|
||
("Boxing", [LM.IDLE_BOXING, LM.WALK_BOXING, LM.LEFT_PUNCH, LM.RIGHT_PUNCH,
|
||
LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK]),
|
||
("Styled Walks", [LM.LEDGE_WALKING, LM.OBJECT_CARRYING, LM.STEALTH_WALK_2,
|
||
LM.HAPPY_DANCE_WALK, LM.ZOMBIE_WALK, LM.GUN_WALK, LM.SCARE_WALK]),
|
||
]
|
||
|
||
STATIC_MODES = {LM.IDLE, LM.SQUAT, LM.KNEEL_TWO_LEGS, LM.KNEEL, LM.LYING_FACE_DOWN, LM.IDLE_BOXING}
|
||
STANDING_MODES = {LM.IDLE, LM.SLOW_WALK, LM.WALK, LM.RUN, LM.IDLE_BOXING, LM.WALK_BOXING,
|
||
LM.LEFT_PUNCH, LM.RIGHT_PUNCH, LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK,
|
||
LM.FORWARD_JUMP, LM.STEALTH_WALK, LM.INJURED_WALK, LM.LEDGE_WALKING,
|
||
LM.OBJECT_CARRYING, LM.STEALTH_WALK_2, LM.HAPPY_DANCE_WALK,
|
||
LM.ZOMBIE_WALK, LM.GUN_WALK, LM.SCARE_WALK}
|
||
BOXING_MODES = {LM.WALK_BOXING, LM.LEFT_PUNCH, LM.RIGHT_PUNCH,
|
||
LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK}
|
||
SPEED_RANGES = {LM.SLOW_WALK:(0.2,0.8), LM.WALK:(0.8,1.5), LM.RUN:(1.5,3.0),
|
||
LM.CRAWLING:(0.4,1.0), LM.ELBOW_CRAWLING:(0.7,1.0)}
|
||
|
||
def clamp_mode_params(ms):
|
||
m = LM(ms.mode)
|
||
ms.height = -1.0 if m in STANDING_MODES else max(0.1, min(0.8, ms.height if ms.height>=0 else 0.2))
|
||
if m in STATIC_MODES:
|
||
ms.speed = -1.0
|
||
elif m in SPEED_RANGES:
|
||
lo, hi = SPEED_RANGES[m]
|
||
ms.speed = max(lo, min(hi, ms.speed if ms.speed>=0 else lo))
|
||
elif m in BOXING_MODES:
|
||
ms.speed = max(0.7, min(1.5, ms.speed if ms.speed>=0 else 0.7))
|
||
else:
|
||
ms.speed = -1.0
|
||
|
||
def replan_interval(mode):
|
||
m = LM(mode)
|
||
if m == LM.RUN: return REPLAN_INTERVAL["running"]
|
||
if m == LM.CRAWLING: return REPLAN_INTERVAL["crawling"]
|
||
if m in {LM.LEFT_PUNCH, LM.RIGHT_PUNCH, LM.RANDOM_PUNCH, LM.LEFT_HOOK, LM.RIGHT_HOOK}:
|
||
return REPLAN_INTERVAL["boxing"]
|
||
return REPLAN_INTERVAL["default"]
|
||
|
||
|
||
def _ort_providers(force_cpu: bool = False) -> list[str]:
|
||
"""Prefer CUDA for enc/dec/planner (matches deploy when onnxruntime-gpu is installed)."""
|
||
avail = ort.get_available_providers()
|
||
if not force_cpu and "CUDAExecutionProvider" in avail:
|
||
return ["CUDAExecutionProvider", "CPUExecutionProvider"]
|
||
return ["CPUExecutionProvider"]
|
||
|
||
# ── Movement state ────────────────────────────────────────────────────────────
|
||
|
||
@dataclass
|
||
class MovementState:
|
||
mode: int = LM.SLOW_WALK # not IDLE — walking modes respond to WASD
|
||
speed: float = -1.0
|
||
height: float = -1.0
|
||
facing_angle: float = 0.0
|
||
movement_angle: float = 0.0
|
||
has_movement: bool = False
|
||
motion_set_idx: int = 0
|
||
needs_replan: bool = False
|
||
|
||
@property
|
||
def movement_direction(self):
|
||
if not self.has_movement: return (0.0, 0.0, 0.0)
|
||
return (math.cos(self.movement_angle), math.sin(self.movement_angle), 0.0)
|
||
|
||
@property
|
||
def facing_direction(self):
|
||
return (math.cos(self.facing_angle), math.sin(self.facing_angle), 0.0)
|
||
|
||
def status_line(self):
|
||
return (f"[{MOTION_SETS[self.motion_set_idx][0]}] mode={self.mode}({LM(self.mode).name}) "
|
||
f"spd={'default' if self.speed<0 else f'{self.speed:.1f}'} "
|
||
f"hgt={'default' if self.height<0 else f'{self.height:.2f}'} "
|
||
f"facing={math.degrees(self.facing_angle):.0f}° "
|
||
f"{'moving' if self.has_movement else 'still'}")
|
||
|
||
|
||
@dataclass
|
||
class MovementSnapshot:
|
||
mode: int = 0
|
||
speed: float = -1.0
|
||
height: float = -1.0
|
||
movement: tuple[float, float, float] = (0.0, 0.0, 0.0)
|
||
facing: tuple[float, float, float] = (1.0, 0.0, 0.0)
|
||
|
||
|
||
def _snapshot_ms(ms: MovementState) -> MovementSnapshot:
|
||
md, fd = ms.movement_direction, ms.facing_direction
|
||
return MovementSnapshot(ms.mode, ms.speed, ms.height, (md[0], md[1], md[2]), (fd[0], fd[1], fd[2]))
|
||
|
||
|
||
def should_replan_request(ms: MovementState, last: MovementSnapshot, replan_timer: float, step: int) -> bool:
|
||
"""Match C++ G1Deploy::Planner replan triggers (g1_deploy_onnx_ref.cpp)."""
|
||
if step <= 0:
|
||
return False
|
||
if ms.needs_replan:
|
||
return True
|
||
md, fd = ms.movement_direction, ms.facing_direction
|
||
facing_changed = fd != last.facing
|
||
height_changed = ms.height != last.height
|
||
mode_changed = ms.mode != last.mode
|
||
speed_changed = ms.speed != last.speed
|
||
dir_changed = md != last.movement
|
||
is_static = LM(ms.mode) in STATIC_MODES
|
||
if mode_changed or facing_changed or height_changed:
|
||
return True
|
||
time_to_replan = replan_timer >= replan_interval(ms.mode)
|
||
if not is_static and (speed_changed or dir_changed or (time_to_replan and ms.speed != 0)):
|
||
return True
|
||
return False
|
||
|
||
# ── Encoder / Decoder ─────────────────────────────────────────────────────────
|
||
|
||
class StandingEncoderDecoder:
|
||
def __init__(self, encoder, decoder):
|
||
self.encoder, self.decoder = encoder, decoder
|
||
self.encoder_input = encoder.get_inputs()[0].name
|
||
self.decoder_input = decoder.get_inputs()[0].name
|
||
enc_dim = int(encoder.get_inputs()[0].shape[1])
|
||
dec_dim = int(decoder.get_inputs()[0].shape[1])
|
||
if enc_dim != 1762 or dec_dim != 994:
|
||
raise RuntimeError(f"Unexpected dims encoder={enc_dim}, decoder={dec_dim}")
|
||
self.token = np.zeros(TOKEN_DIM, np.float32)
|
||
self.last_action_mj = np.zeros(29, np.float32)
|
||
self.h_q_mj = [np.zeros(29, np.float32)] * 10
|
||
self.h_dq_mj = [np.zeros(29, np.float32)] * 10
|
||
self.h_ang = [np.zeros(3, np.float32)] * 10
|
||
self.h_act_mj = [np.zeros(29, np.float32)] * 10
|
||
self.h_quat = [np.array([1,0,0,0], np.float32)] * 10
|
||
self.init_base_quat = np.array([1,0,0,0], np.float32)
|
||
self.init_ref_quat = np.array([1,0,0,0], np.float32)
|
||
self._heading_init = False
|
||
self.encode_mode = 0
|
||
self.vr_3point_local_target = VR_TARGET_DEF.copy()
|
||
self.vr_3point_local_orn_target = VR_ORN_DEF.copy()
|
||
self.smpl_joints_10frame_step1 = SMPL_DEF.copy()
|
||
self.set_zero_reference()
|
||
|
||
def update_history(self, q, dq, ang, quat):
|
||
quat = quat / (np.linalg.norm(quat)+1e-8)
|
||
q_mj = _to_mujoco(q); dq_mj = _to_mujoco(dq)
|
||
self.h_q_mj = [q_mj - DEFAULT_ANGLES_MUJOCO] + self.h_q_mj[:-1]
|
||
self.h_dq_mj = [dq_mj] + self.h_dq_mj[:-1]
|
||
self.h_ang = [ang.copy()] + self.h_ang[:-1]
|
||
self.h_act_mj = [self.last_action_mj.copy()] + self.h_act_mj[:-1]
|
||
self.h_quat = [quat.copy()] + self.h_quat[:-1]
|
||
if not self._heading_init:
|
||
self.init_base_quat = quat.copy(); self._heading_init = True
|
||
|
||
def _heading_quat(self, q):
|
||
h = calc_heading(q) / 2.0
|
||
return np.array([np.cos(h), 0, 0, np.sin(h)], np.float32)
|
||
|
||
def _heading_quat_inv(self, q):
|
||
h = calc_heading(q) / 2.0
|
||
return np.array([np.cos(-h), 0, 0, np.sin(-h)], np.float32)
|
||
|
||
def _anchor_6d(self, base_quat, ref_quat=None):
|
||
if ref_quat is None: ref_quat = self.init_ref_quat
|
||
delta = quat_mul(self._heading_quat(self.init_base_quat), self._heading_quat_inv(self.init_ref_quat))
|
||
new_ref = quat_mul(delta, ref_quat)
|
||
return quat_to_6d(quat_mul(quat_conj(base_quat), new_ref))
|
||
|
||
def set_zero_reference(self):
|
||
self.motion_joint_positions = [ENCODER_STANDING_REF.copy()]
|
||
self.motion_joint_velocities = [np.zeros(29, np.float32)]
|
||
self.motion_body_quats = [np.array([1,0,0,0], np.float32)]
|
||
self.motion_body_z = [DEFAULT_HEIGHT]
|
||
self.motion_timesteps = 1
|
||
self.freeze_ref_frame = 0
|
||
self.init_ref_quat = self.motion_body_quats[0].copy()
|
||
|
||
def build_encoder_obs(self):
|
||
obs = np.zeros(1762, np.float32)
|
||
obs[0] = float(self.encode_mode)
|
||
rf = min(self.freeze_ref_frame, self.motion_timesteps - 1)
|
||
ref_pos, ref_quat = self.motion_joint_positions[rf], self.motion_body_quats[rf]
|
||
if self.encode_mode == 0:
|
||
for f in range(10):
|
||
obs[4+29*f:4+29*(f+1)] = ref_pos
|
||
obs[601+6*f:601+6*(f+1)] = self._anchor_6d(self.h_quat[0], ref_quat)
|
||
elif self.encode_mode == 1:
|
||
ref_lower = ref_pos[LOWER_BODY_IL]
|
||
for f in range(10):
|
||
obs[661+12*f:661+12*(f+1)] = ref_lower
|
||
obs[901:910] = self.vr_3point_local_target
|
||
obs[910:922] = self.vr_3point_local_orn_target
|
||
obs[595:601] = self._anchor_6d(self.h_quat[0], ref_quat)
|
||
elif self.encode_mode == 2:
|
||
obs[922:1642] = self.smpl_joints_10frame_step1
|
||
for f in range(10):
|
||
obs[1642+6*f:1642+6*(f+1)] = self._anchor_6d(self.h_quat[0], ref_quat)
|
||
obs[1702+6*f:1702+6*(f+1)] = ref_pos[WRIST_IL]
|
||
else:
|
||
raise RuntimeError(f"Unsupported encoder mode: {self.encode_mode}")
|
||
return obs
|
||
|
||
def build_decoder_obs(self):
|
||
obs = np.zeros(994, np.float32); off = 0
|
||
obs[off:off+64] = self.token; off += 64
|
||
for h, sz in [(list(reversed(self.h_ang)),3), (list(reversed(self.h_q_mj)),29),
|
||
(list(reversed(self.h_dq_mj)),29), (list(reversed(self.h_act_mj)),29)]:
|
||
for f in range(10): obs[off:off+sz] = h[f]; off += sz
|
||
for q in reversed(self.h_quat):
|
||
obs[off:off+3] = gravity_dir(q); off += 3
|
||
assert off == 994, f"Decoder obs mismatch: {off}"
|
||
return obs
|
||
|
||
def run_encoder(self):
|
||
return self.encoder.run(None, {self.encoder_input: self.build_encoder_obs().reshape(1,-1)})[0].squeeze().astype(np.float32)
|
||
|
||
def step(self, robot_obs, update_encoder, debug=False):
|
||
jnames = [m.name for m in G1_29_JointIndex]
|
||
q = np.array([robot_obs.get(f"{n}.q", DEFAULT_ANGLES[m.value]) for m,n in zip(G1_29_JointIndex,jnames)], np.float32)
|
||
dq = np.array([robot_obs.get(f"{n}.dq", 0.0) for n in jnames], np.float32)
|
||
quat = np.array([robot_obs.get("imu.quat.w",1), robot_obs.get("imu.quat.x",0),
|
||
robot_obs.get("imu.quat.y",0), robot_obs.get("imu.quat.z",0)], np.float32)
|
||
ang = np.array([robot_obs.get(f"imu.gyro.{a}",0) for a in "xyz"], np.float32)
|
||
self.update_history(q, dq, ang, quat)
|
||
if update_encoder: self.token = self.run_encoder()
|
||
action_mj = self.decoder.run(None, {self.decoder_input: self.build_decoder_obs().reshape(1,-1)})[0].squeeze().astype(np.float32)
|
||
self.last_action_mj = action_mj.copy()
|
||
target = DEFAULT_ANGLES + action_mj[ISAACLAB_TO_MUJOCO] * ACTION_SCALE
|
||
if debug:
|
||
delta = target - q
|
||
print(f"token_norm={np.linalg.norm(self.token):.4f} action_norm={np.linalg.norm(action_mj):.4f} "
|
||
f"delta_max={np.max(np.abs(delta)):.4f} delta_rms={np.sqrt(np.mean(delta**2)):.4f}")
|
||
return {f"{m.name}.q": float(target[m.value]) for m in G1_29_JointIndex}
|
||
|
||
def print_input_diagnostics(self):
|
||
print("\n[Diag] Standing reference checks")
|
||
names = {0:"g1", 1:"teleop", 2:"smpl"}
|
||
print(f" encoder mode: {self.encode_mode} ({names.get(self.encode_mode,'unknown')})")
|
||
print(f" DEFAULT_ANGLES range: [{DEFAULT_ANGLES.min():+.4f}, {DEFAULT_ANGLES.max():+.4f}]")
|
||
print(f" anchor_6d(identity): {self._anchor_6d(np.array([1,0,0,0],np.float32), np.array([1,0,0,0],np.float32))}")
|
||
print(f" gravity(identity): {gravity_dir(np.array([1,0,0,0],np.float32))} (expect [0,0,-1])")
|
||
dec0 = self.build_decoder_obs()
|
||
print(f" decoder q-delta max: {np.max(np.abs(dec0[94:384])):.6f}")
|
||
print(f" decoder dq max: {np.max(np.abs(dec0[384:674])):.6f}")
|
||
|
||
# ── Planner motion buffer ─────────────────────────────────────────────────────
|
||
|
||
class PlannerMotion:
|
||
def __init__(self, max_frames=1500):
|
||
self.timesteps = 0
|
||
self.joint_positions = np.zeros((max_frames, 29), np.float64)
|
||
self.joint_velocities = np.zeros((max_frames, 29), np.float64)
|
||
self.body_positions = np.zeros((max_frames, 3), np.float64)
|
||
self.body_quaternions = np.zeros((max_frames, 4), np.float64)
|
||
self.body_quaternions[:, 0] = 1.0
|
||
|
||
# ── Subprocess planner ────────────────────────────────────────────────────────
|
||
|
||
def _resample_30_to_50(qpos, n30):
|
||
t50 = int(np.floor(n30 / 30.0 * 50))
|
||
f30 = np.arange(t50) / 50.0 * 30.0
|
||
f0 = np.floor(f30).astype(int)
|
||
f1 = np.minimum(f0+1, n30-1)
|
||
frac, w0 = (f30-f0).astype(np.float64), None
|
||
w0 = 1.0 - frac
|
||
jp = (w0[:,None]*qpos[f0,7:36] + frac[:,None]*qpos[f1,7:36])[:,MUJOCO_TO_ISAACLAB]
|
||
jv = np.zeros_like(jp)
|
||
if t50 >= 2: jv[:t50-1] = (jp[1:] - jp[:-1]) * 50.0; jv[-1] = jv[-2]
|
||
return {
|
||
"timesteps": t50,
|
||
"joint_positions": jp,
|
||
"joint_velocities": jv,
|
||
"body_positions": w0[:,None]*qpos[f0,:3] + frac[:,None]*qpos[f1,:3],
|
||
"body_quaternions": quat_slerp_batch(qpos[f0,3:7], qpos[f1,3:7], frac),
|
||
}
|
||
|
||
def _build_planner_inputs(ctx, ms_dict, version, seed):
|
||
inp = {
|
||
"context_mujoco_qpos": ctx.astype(np.float32).reshape(1,4,36),
|
||
"target_vel": np.array([ms_dict["speed"]], np.float32),
|
||
"mode": np.array([ms_dict["mode"]], np.int64),
|
||
"movement_direction": np.array(ms_dict["movement_direction"], np.float32).reshape(1,3),
|
||
"facing_direction": np.array(ms_dict["facing_direction"], np.float32).reshape(1,3),
|
||
"random_seed": np.array([seed], np.int64),
|
||
}
|
||
if version >= 1:
|
||
# TensorRT deploy: allow 9–11 prediction tokens only (indices 3–5 for MIN_TOKENS=6).
|
||
allowed = np.zeros((1, K), np.int64)
|
||
if K >= 6:
|
||
allowed[0, 3:6] = 1
|
||
inp.update({
|
||
"height": np.array([ms_dict["height"]], np.float32),
|
||
"has_specific_target": np.array([[0]], np.int64),
|
||
"specific_target_positions": np.zeros((1,4,3), np.float32),
|
||
"specific_target_headings": np.zeros((1,4), np.float32),
|
||
"allowed_pred_num_tokens": allowed,
|
||
})
|
||
return inp
|
||
|
||
def _planner_worker(path, req_q, res_q, stop_evt, version, seed, use_gpu):
|
||
so = ort.SessionOptions(); so.log_severity_level = 3
|
||
providers = _ort_providers(force_cpu=not use_gpu)
|
||
sess = ort.InferenceSession(path, sess_options=so, providers=providers)
|
||
while not stop_evt.is_set():
|
||
try: ctx, gf, ms_dict = req_q.get(timeout=0.05)
|
||
except Exception: continue
|
||
try:
|
||
inp = _build_planner_inputs(ctx, ms_dict, version, seed)
|
||
t0 = time.time()
|
||
qpos_out, num_pred = sess.run(None, inp)
|
||
t_inf = time.time()
|
||
n = int(num_pred.flat[0])
|
||
qpos = qpos_out[0,:n]
|
||
if np.any(np.isnan(qpos)): continue
|
||
motion = _resample_30_to_50(qpos, n)
|
||
motion["gen_frame"] = gf
|
||
print(f"[Planner] inf={1000*(t_inf-t0):.1f}ms total={1000*(time.time()-t0):.1f}ms frames={n}", flush=True)
|
||
while not res_q.empty():
|
||
try: res_q.get_nowait()
|
||
except queue.Empty: break
|
||
res_q.put(motion)
|
||
except Exception as e:
|
||
print(f"[Planner] Error: {e}", flush=True)
|
||
|
||
# ── SonicPlanner ──────────────────────────────────────────────────────────────
|
||
|
||
class SonicPlanner:
|
||
def __init__(self, session, planner_path):
|
||
self.session = session
|
||
self.planner_path = planner_path
|
||
self.gen_frame = 0
|
||
self.random_seed = INITIAL_RANDOM_SEED
|
||
self.version = 1 if len(session.get_inputs()) >= 11 else 0
|
||
self.motion_50hz = PlannerMotion()
|
||
self._snapshot = PlannerMotion()
|
||
self._req_q = self._res_q = self._stop_evt = self._planner_thread = None
|
||
self._ctrl = None
|
||
|
||
def _build_inputs(self, ctx, ms):
|
||
return _build_planner_inputs(
|
||
ctx,
|
||
{"mode": ms.mode, "speed": ms.speed, "height": ms.height,
|
||
"movement_direction": list(ms.movement_direction),
|
||
"facing_direction": list(ms.facing_direction)},
|
||
self.version, self.random_seed)
|
||
|
||
@staticmethod
|
||
def build_initial_context(joint_positions):
|
||
ctx = np.zeros((4, 36), np.float32)
|
||
jp_mj = joint_positions.astype(np.float32)[ISAACLAB_TO_MUJOCO]
|
||
for n in range(4):
|
||
ctx[n, 2] = DEFAULT_HEIGHT
|
||
ctx[n, 3] = 1.0
|
||
ctx[n, 7:36] = jp_mj
|
||
return ctx
|
||
|
||
def _context_from_controller(self, current_frame):
|
||
ctrl = self._ctrl
|
||
gen_frame = current_frame + MOTION_LOOK_AHEAD_STEPS
|
||
t_arr = gen_frame / 50.0 + np.arange(4) / 30.0
|
||
f50 = t_arr * 50.0
|
||
with ctrl.motion_lock:
|
||
ts = ctrl.motion_timesteps
|
||
if ts <= 0:
|
||
return self.build_initial_context(DEFAULT_ANGLES)
|
||
bp, bq, jp = ctrl.motion_body_pos, ctrl.motion_body_quats, ctrl.motion_joint_positions
|
||
f0 = np.minimum(np.floor(f50).astype(int), ts - 1)
|
||
f1 = np.minimum(f0 + 1, ts - 1)
|
||
frac = f50 - f0
|
||
w0 = 1.0 - frac
|
||
ctx = np.zeros((4, 36), np.float32)
|
||
ctx[:, 0:3] = w0[:, None] * bp[f0] + frac[:, None] * bp[f1]
|
||
ctx[:, 3:7] = quat_slerp_batch(bq[f0], bq[f1], frac)
|
||
ij = w0[:, None] * jp[f0] + frac[:, None] * jp[f1]
|
||
ctx[:, 7:36] = ij[:, ISAACLAB_TO_MUJOCO]
|
||
self.gen_frame = gen_frame
|
||
return ctx
|
||
|
||
def _load_motion_in_place(self, qpos, n30, target=None):
|
||
if target is None: target = self.motion_50hz
|
||
r = _resample_30_to_50(qpos, n30)
|
||
n = r["timesteps"]; target.timesteps = n
|
||
target.joint_positions[:n] = r["joint_positions"]
|
||
target.joint_velocities[:n] = r["joint_velocities"]
|
||
target.body_positions[:n] = r["body_positions"]
|
||
target.body_quaternions[:n] = r["body_quaternions"]
|
||
return target
|
||
|
||
def initialize(self, joint_positions, ms):
|
||
ctx = self.build_initial_context(joint_positions)
|
||
qpos_out, num_pred = self.session.run(None, self._build_inputs(ctx, ms))
|
||
n = int(num_pred.flat[0]); qpos = qpos_out[0,:n]
|
||
if np.any(np.isnan(qpos)): raise RuntimeError("Planner initial output contains NaN")
|
||
print(f"[Planner] Init: {n} frames @ 30 Hz")
|
||
self._load_motion_in_place(qpos, n)
|
||
print(f"[Planner] Resampled to {self.motion_50hz.timesteps} frames @ 50 Hz")
|
||
return self.motion_50hz
|
||
|
||
def request_replan(self, cursor, ms):
|
||
if self._req_q is None: return
|
||
ctx = self._context_from_controller(cursor)
|
||
ms_dict = {"mode": ms.mode, "speed": ms.speed, "height": ms.height,
|
||
"movement_direction": list(ms.movement_direction),
|
||
"facing_direction": list(ms.facing_direction)}
|
||
while not self._req_q.empty():
|
||
try: self._req_q.get_nowait()
|
||
except queue.Empty: break
|
||
self._req_q.put((ctx, self.gen_frame, ms_dict))
|
||
|
||
def try_get_new_motion(self):
|
||
if self._res_q is None: return None
|
||
result = None
|
||
while not self._res_q.empty():
|
||
try: result = self._res_q.get_nowait()
|
||
except queue.Empty: break
|
||
if result is None: return None
|
||
n, gf = result["timesteps"], result["gen_frame"]
|
||
s = self._snapshot; s.timesteps = n
|
||
s.joint_positions[:n] = result["joint_positions"]
|
||
s.joint_velocities[:n] = result["joint_velocities"]
|
||
s.body_positions[:n] = result["body_positions"]
|
||
s.body_quaternions[:n] = result["body_quaternions"]
|
||
return s, gf
|
||
|
||
def start_subprocess(self, controller, use_gpu: bool = False):
|
||
"""Run planner ONNX in a background thread (avoids mp spawn/fork + CUDA/MuJoCo issues)."""
|
||
self._ctrl = controller
|
||
self._req_q = queue.Queue()
|
||
self._res_q = queue.Queue()
|
||
self._stop_evt = threading.Event()
|
||
self._planner_thread = threading.Thread(
|
||
target=_planner_worker,
|
||
args=(self.planner_path, self._req_q, self._res_q,
|
||
self._stop_evt, self.version, self.random_seed, use_gpu),
|
||
daemon=True,
|
||
name="sonic-planner",
|
||
)
|
||
self._planner_thread.start()
|
||
print(f"[Planner] Background thread started ({'GPU' if use_gpu else 'CPU'})")
|
||
|
||
def stop_subprocess(self):
|
||
if self._stop_evt:
|
||
self._stop_evt.set()
|
||
if self._planner_thread is not None:
|
||
self._planner_thread.join(timeout=3.0)
|
||
print("[Planner] Background thread stopped")
|
||
self._planner_thread = None
|
||
self._req_q = self._res_q = self._stop_evt = None
|
||
|
||
# ── PlannerController ─────────────────────────────────────────────────────────
|
||
|
||
class PlannerController(StandingEncoderDecoder):
|
||
def __init__(self, planner, encoder, decoder):
|
||
super().__init__(encoder, decoder)
|
||
self.planner = planner
|
||
self.ref_cursor = 0
|
||
self.motion_timesteps = 0
|
||
self.motion_joint_positions = np.zeros((1500,29), np.float64)
|
||
self.motion_joint_velocities = np.zeros((1500,29), np.float64)
|
||
self.motion_body_quats = np.zeros((1500,4), np.float64); self.motion_body_quats[:,0] = 1.0
|
||
self.motion_body_pos = np.zeros((1500,3), np.float64)
|
||
self.init_ref_quat = np.array([1,0,0,0], np.float64)
|
||
self.heading_init_base_quat = np.array([1,0,0,0], np.float64)
|
||
self.delta_heading = 0.0
|
||
self.reinit_heading = False
|
||
self.playing = self.first_motion = False
|
||
self.motion_lock = threading.Lock()
|
||
|
||
def load_initial_motion(self, motion):
|
||
with self.motion_lock:
|
||
n = motion.timesteps
|
||
self.motion_timesteps = n
|
||
self.motion_joint_positions[:n] = motion.joint_positions[:n]
|
||
self.motion_joint_velocities[:n] = motion.joint_velocities[:n]
|
||
self.motion_body_quats[:n] = motion.body_quaternions[:n]
|
||
self.motion_body_pos[:n] = motion.body_positions[:n]
|
||
self.init_ref_quat = motion.body_quaternions[0].copy()
|
||
self.ref_cursor = 0; self.first_motion = True
|
||
self.playing = True; self.delta_heading = 0.0
|
||
|
||
def blend_new_motion(self, new_motion, gen_frame):
|
||
"""Blend like C++ CurrentFrameAdvancement: 8-frame cross-fade, then copy tail."""
|
||
with self.motion_lock:
|
||
cur = self.ref_cursor
|
||
new_len = gen_frame - cur + new_motion.timesteps
|
||
if new_len <= 0:
|
||
return
|
||
if self.motion_timesteps == 0:
|
||
n = new_motion.timesteps
|
||
self.motion_joint_positions[:n] = new_motion.joint_positions[:n]
|
||
self.motion_joint_velocities[:n] = new_motion.joint_velocities[:n]
|
||
self.motion_body_pos[:n] = new_motion.body_positions[:n]
|
||
self.motion_body_quats[:n] = new_motion.body_quaternions[:n]
|
||
self.motion_timesteps = n
|
||
self.ref_cursor = 0
|
||
self.init_ref_quat = self.motion_body_quats[0].copy()
|
||
self.first_motion = False
|
||
return
|
||
|
||
blend_start = max(0, gen_frame - cur)
|
||
blend_end = min(new_len, blend_start + BLEND_FRAMES)
|
||
|
||
for f in range(blend_end):
|
||
f_old = min(f + cur, self.motion_timesteps - 1)
|
||
f_new = max(0, min(f + cur - gen_frame, new_motion.timesteps - 1))
|
||
w_new = min(1.0, max(0.0, (f - blend_start) / BLEND_FRAMES))
|
||
w_old = 1.0 - w_new
|
||
self.motion_joint_positions[f] = (
|
||
w_old * self.motion_joint_positions[f_old]
|
||
+ w_new * new_motion.joint_positions[f_new]
|
||
)
|
||
self.motion_joint_velocities[f] = (
|
||
w_old * self.motion_joint_velocities[f_old]
|
||
+ w_new * new_motion.joint_velocities[f_new]
|
||
)
|
||
self.motion_body_pos[f] = (
|
||
w_old * self.motion_body_pos[f_old]
|
||
+ w_new * new_motion.body_positions[f_new]
|
||
)
|
||
self.motion_body_quats[f] = quat_slerp(
|
||
self.motion_body_quats[f_old], new_motion.body_quaternions[f_new], w_new
|
||
)
|
||
|
||
for f in range(blend_end, new_len):
|
||
f_new = max(0, min(f + cur - gen_frame, new_motion.timesteps - 1))
|
||
self.motion_joint_positions[f] = new_motion.joint_positions[f_new]
|
||
self.motion_joint_velocities[f] = new_motion.joint_velocities[f_new]
|
||
self.motion_body_pos[f] = new_motion.body_positions[f_new]
|
||
self.motion_body_quats[f] = new_motion.body_quaternions[f_new].copy()
|
||
|
||
self.motion_timesteps = new_len
|
||
self.first_motion = False
|
||
self.ref_cursor = 0
|
||
self.init_ref_quat = self.motion_body_quats[0].copy()
|
||
|
||
def _heading_apply_delta(self):
|
||
delta = quat_mul(heading_quat(self.heading_init_base_quat).astype(np.float32),
|
||
heading_quat_inv(self.init_ref_quat).astype(np.float32))
|
||
if self.delta_heading:
|
||
h = self.delta_heading / 2.0
|
||
delta = quat_mul(np.array([np.cos(h),0,0,np.sin(h)], np.float32), delta)
|
||
return delta
|
||
|
||
def _anchor_6d(self, base_quat, ref_quat=None):
|
||
if ref_quat is None: ref_quat = self.init_ref_quat
|
||
new_ref = quat_mul(self._heading_apply_delta(), ref_quat.astype(np.float32))
|
||
return quat_to_6d(quat_mul(quat_conj(base_quat.astype(np.float32)), new_ref))
|
||
|
||
def build_encoder_obs(self):
|
||
obs = np.zeros(1762, np.float32); obs[0] = float(self.encode_mode)
|
||
with self.motion_lock:
|
||
for f in range(10):
|
||
tf = min(self.ref_cursor + f*5 if self.playing else self.ref_cursor,
|
||
self.motion_timesteps - 1)
|
||
obs[4+29*f:4+29*(f+1)] = self.motion_joint_positions[tf].astype(np.float32)
|
||
if self.playing:
|
||
obs[294+29*f:294+29*(f+1)] = self.motion_joint_velocities[tf].astype(np.float32)
|
||
obs[601+6*f:601+6*(f+1)] = self._anchor_6d(
|
||
self.h_quat[0], self.motion_body_quats[tf].astype(np.float32))
|
||
return obs
|
||
|
||
def step(self, robot_obs, update_encoder, debug=False):
|
||
if robot_obs and (self.first_motion or self.reinit_heading):
|
||
q = None
|
||
if "imu.quat.w" in robot_obs:
|
||
q = np.array([
|
||
robot_obs["imu.quat.w"], robot_obs["imu.quat.x"],
|
||
robot_obs["imu.quat.y"], robot_obs["imu.quat.z"],
|
||
], np.float64)
|
||
else:
|
||
q = robot_obs.get("imu.quaternion")
|
||
if q is not None:
|
||
q = np.array(q, np.float64)
|
||
if q is not None:
|
||
self.heading_init_base_quat = np.array(q, np.float64)
|
||
with self.motion_lock:
|
||
rf = min(self.ref_cursor, self.motion_timesteps - 1)
|
||
self.init_ref_quat = self.motion_body_quats[rf].copy()
|
||
self.delta_heading = 0.0
|
||
self.first_motion = False
|
||
self.reinit_heading = False
|
||
print(f"[Heading] init quat: {self.heading_init_base_quat}")
|
||
return super().step(robot_obs, update_encoder=update_encoder, debug=debug)
|
||
|
||
def advance_cursor(self):
|
||
"""Advance one frame per 50 Hz tick (C++ current_frame_ += 1), no wall-clock catch-up."""
|
||
if not self.playing:
|
||
return
|
||
with self.motion_lock:
|
||
if self.motion_timesteps > 0:
|
||
self.ref_cursor = min(self.ref_cursor + 1, self.motion_timesteps - 1)
|
||
|
||
# ── Keyboard ──────────────────────────────────────────────────────────────────
|
||
|
||
class RawKeyboard:
|
||
def __init__(self):
|
||
self.fd = sys.stdin.fileno()
|
||
self.old = termios.tcgetattr(self.fd)
|
||
def __enter__(self): tty.setcbreak(self.fd); return self
|
||
def __exit__(self, *_): termios.tcsetattr(self.fd, termios.TCSADRAIN, self.old)
|
||
def get_key(self):
|
||
return sys.stdin.read(1) if select.select([sys.stdin],[],[],0)[0] else None
|
||
|
||
def process_keyboard(key, ms, controller=None):
|
||
if key is None: return False
|
||
if key == '\x1b': return True
|
||
if key == ' ':
|
||
ms.mode = LM.IDLE; ms.speed = ms.height = -1.0
|
||
ms.has_movement = False; ms.needs_replan = True
|
||
if controller: controller.playing = False; controller.reinit_heading = True
|
||
print("\n >> EMERGENCY STOP -> IDLE"); return False
|
||
if key in ('r','R'):
|
||
ms.needs_replan = True; print("\n >> Manual replan"); return False
|
||
if key in ('n','N','p','P'):
|
||
ms.motion_set_idx = (ms.motion_set_idx + (1 if key in ('n','N') else -1)) % len(MOTION_SETS)
|
||
name, modes = MOTION_SETS[ms.motion_set_idx]
|
||
print(f"\n >> Motion set: {name}")
|
||
[print(f" {i+1}: {m.name}") for i,m in enumerate(modes)]
|
||
return False
|
||
if key.isdigit() and key not in ('9','0'):
|
||
idx = int(key) - 1; modes = MOTION_SETS[ms.motion_set_idx][1]
|
||
if 0 <= idx < len(modes):
|
||
ms.mode = modes[idx]; ms.needs_replan = True
|
||
if controller: controller.playing = True; controller.reinit_heading = True
|
||
print(f"\n >> Mode: {LM(ms.mode).name} ({ms.mode}) [replanning...]")
|
||
return False
|
||
if key == '9':
|
||
ms.speed = max(0.0, (ms.speed if ms.speed>=0 else 1.0) - 0.1)
|
||
print(f"\n >> Speed: {ms.speed:.1f}"); return False
|
||
if key == '0':
|
||
ms.speed = min(5.0, (ms.speed if ms.speed>=0 else 1.0) + 0.1)
|
||
print(f"\n >> Speed: {ms.speed:.1f}"); return False
|
||
if key == '-':
|
||
ms.height = max(0.2, (ms.height if ms.height>=0 else DEFAULT_HEIGHT) - 0.02)
|
||
print(f"\n >> Height: {ms.height:.2f}"); return False
|
||
if key == '=':
|
||
ms.height = min(1.0, (ms.height if ms.height>=0 else DEFAULT_HEIGHT) + 0.02)
|
||
print(f"\n >> Height: {ms.height:.2f}"); return False
|
||
if key.lower() == 'w': ms.movement_angle = ms.facing_angle
|
||
elif key.lower() == 's': ms.movement_angle = ms.facing_angle + math.pi
|
||
elif key.lower() == 'a': ms.movement_angle = ms.facing_angle + math.pi/2
|
||
elif key.lower() == 'd': ms.movement_angle = ms.facing_angle - math.pi/2
|
||
if key.lower() in ('w','s','a','d'):
|
||
ms.has_movement = ms.needs_replan = True
|
||
print(f"\n >> Move {key.upper()} (replanning...)")
|
||
elif key.lower() == 'q':
|
||
ms.facing_angle += 0.1
|
||
if controller: controller.delta_heading += 0.1
|
||
print(f"\n >> Facing: {math.degrees(ms.facing_angle):.0f}°")
|
||
elif key.lower() == 'e':
|
||
ms.facing_angle -= 0.1
|
||
if controller: controller.delta_heading -= 0.1
|
||
print(f"\n >> Facing: {math.degrees(ms.facing_angle):.0f}°")
|
||
return False
|
||
|
||
_joy_prev_active = False
|
||
|
||
|
||
def _parse_wireless(wr):
|
||
"""Parse wireless_remote (bytes or int-array) into (lx, ly, rx, ry)."""
|
||
import struct as _st
|
||
if not isinstance(wr, (bytes, bytearray)):
|
||
wr = bytes(wr)
|
||
if len(wr) < 24:
|
||
return None
|
||
lx = _st.unpack("f", wr[4:8])[0]
|
||
rx = _st.unpack("f", wr[8:12])[0]
|
||
ry = _st.unpack("f", wr[12:16])[0]
|
||
ly = _st.unpack("f", wr[20:24])[0]
|
||
return lx, ly, rx, ry
|
||
|
||
|
||
def process_joystick(obs, ms, controller=None):
|
||
"""Joystick mirrors keyboard: left stick=WASD, right stick X=Q/E, right stick Y=height."""
|
||
global _joy_prev_active
|
||
wr = obs.get("wireless_remote")
|
||
if wr is None:
|
||
return
|
||
parsed = _parse_wireless(wr)
|
||
if parsed is None:
|
||
return
|
||
lx, ly, rx, ry = parsed
|
||
|
||
# Dead zone + negate both Y axes (bridge already flips them once)
|
||
lx = 0.0 if abs(lx) < DEADZONE else lx
|
||
ly = 0.0 if abs(ly) < DEADZONE else -ly
|
||
rx = 0.0 if abs(rx) < DEADZONE else rx
|
||
ry = 0.0 if abs(ry) < DEADZONE else -ry
|
||
|
||
left_active = abs(lx) > 0 or abs(ly) > 0
|
||
|
||
# Left stick → WASD (movement direction relative to facing)
|
||
if left_active:
|
||
ms.movement_angle = ms.facing_angle + math.atan2(-lx, -ly)
|
||
ms.has_movement = True
|
||
if not _joy_prev_active:
|
||
ms.needs_replan = True
|
||
_joy_prev_active = True
|
||
elif _joy_prev_active and not (abs(rx) > 0 or abs(ry) > 0):
|
||
_joy_prev_active = False
|
||
ms.has_movement = False
|
||
|
||
# Right stick X → Q/E (facing rotation, ~1 rad/s at full deflection)
|
||
if abs(rx) > 0:
|
||
delta = -0.02 * rx
|
||
ms.facing_angle += delta
|
||
if controller:
|
||
controller.delta_heading += delta
|
||
|
||
# Right stick Y → -/= (height adjustment, ~0.25/s at full deflection)
|
||
if abs(ry) > 0:
|
||
step = -0.005 * ry
|
||
ms.height = max(0.1, min(1.0, (ms.height if ms.height >= 0 else DEFAULT_HEIGHT) + step))
|
||
|
||
# ── Sonic-only sim (do not use UnitreeG1 in-process — its subscribe thread + GLFW crash) ──
|
||
|
||
def _make_hub_sim_env(headless: bool = False):
|
||
"""Load lerobot/unitree-g1-mujoco hub env for sonic.py sim only."""
|
||
import yaml as yaml_mod
|
||
|
||
from lerobot.envs.utils import _download_hub_file, _import_hub_module, _normalize_hub_result
|
||
|
||
_, _, local_file, _ = _download_hub_file(
|
||
"lerobot/unitree-g1-mujoco", trust_remote_code=True, hub_cache_dir=None
|
||
)
|
||
module = _import_hub_module(local_file, "lerobot/unitree-g1-mujoco")
|
||
kwargs = {"publish_images": False} # avoid ZMQ :5555 conflicts; use MuJoCo window instead
|
||
if headless:
|
||
orig_load = yaml_mod.safe_load
|
||
|
||
def _headless_cfg(stream):
|
||
cfg = orig_load(stream)
|
||
if isinstance(cfg, dict):
|
||
cfg = {**cfg, "ENABLE_ONSCREEN": False}
|
||
return cfg
|
||
|
||
yaml_mod.safe_load = _headless_cfg
|
||
try:
|
||
raw = module.make_env(n_envs=1, use_async_envs=False, **kwargs)
|
||
finally:
|
||
yaml_mod.safe_load = orig_load
|
||
else:
|
||
raw = module.make_env(n_envs=1, use_async_envs=False, **kwargs)
|
||
return _normalize_hub_result(raw)
|
||
|
||
|
||
def _lowstate_from_sim_obs(obs: dict):
|
||
from lerobot.robots.unitree_g1.unitree_g1 import G1_29_LowState
|
||
|
||
ls = G1_29_LowState()
|
||
for joint in G1_29_JointIndex:
|
||
i = joint.value
|
||
ls.motor_state[joint].q = float(obs["body_q"][i])
|
||
ls.motor_state[joint].dq = float(obs["body_dq"][i])
|
||
ls.motor_state[joint].tau_est = float(obs["body_tau_est"][i])
|
||
ls.imu_state.quaternion = list(obs["floating_base_pose"][3:7])
|
||
ls.imu_state.gyroscope = list(obs["floating_base_vel"][3:6])
|
||
ls.imu_state.accelerometer = list(obs["floating_base_acc"][:3])
|
||
return ls
|
||
|
||
|
||
class SonicSimRobot:
|
||
"""MuJoCo sim for sonic.py: main-thread physics + viewer (safe for GLFW)."""
|
||
|
||
def __init__(self, headless: bool = False):
|
||
self.headless = headless
|
||
self.kp = self.kd = None
|
||
self._shutdown_event = threading.Event()
|
||
self._lowstate = None
|
||
self._lowstate_lock = threading.Lock()
|
||
self._gym_env = None
|
||
self._mj = None
|
||
self._bridge = None
|
||
self._step_count = 0
|
||
self._substeps = 5
|
||
self._viewer_every = 5
|
||
self._render_every = 8
|
||
|
||
def connect(self):
|
||
wrapper = _make_hub_sim_env(headless=self.headless)
|
||
self._gym_env = wrapper["hub_env"][0].envs[0]
|
||
self._mj = self._gym_env.sim_env
|
||
self._bridge = self._mj.unitree_bridge
|
||
sim_dt = float(getattr(self._gym_env, "sim_dt", 0.004))
|
||
self._substeps = max(1, round(CONTROL_DT / sim_dt))
|
||
self._viewer_every = max(1, round(0.02 / sim_dt))
|
||
self._render_every = max(1, round(0.033333 / sim_dt))
|
||
self._advance_sim(10)
|
||
mode = "headless" if self.headless else "MuJoCo viewer (main thread)"
|
||
print(f"[SonicSim] Connected — {mode}", flush=True)
|
||
|
||
def _advance_sim(self, substeps: int | None = None):
|
||
n = self._substeps if substeps is None else substeps
|
||
for _ in range(n):
|
||
self._mj.sim_step()
|
||
self._step_count += 1
|
||
if not self.headless and self._step_count % self._viewer_every == 0:
|
||
self._mj.update_viewer()
|
||
if self._step_count % self._render_every == 0:
|
||
self._mj.update_render_caches()
|
||
with self._lowstate_lock:
|
||
self._lowstate = _lowstate_from_sim_obs(self._mj.prepare_obs())
|
||
|
||
def get_observation(self):
|
||
self._advance_sim()
|
||
with self._lowstate_lock:
|
||
ls = self._lowstate
|
||
if ls is None:
|
||
return {}
|
||
obs = {}
|
||
for motor in G1_29_JointIndex:
|
||
name, i = motor.name, motor.value
|
||
obs[f"{name}.q"] = ls.motor_state[i].q
|
||
obs[f"{name}.dq"] = ls.motor_state[i].dq
|
||
obs[f"{name}.tau"] = ls.motor_state[i].tau_est
|
||
if ls.imu_state.gyroscope:
|
||
obs["imu.gyro.x"], obs["imu.gyro.y"], obs["imu.gyro.z"] = ls.imu_state.gyroscope
|
||
if ls.imu_state.accelerometer:
|
||
obs["imu.accel.x"], obs["imu.accel.y"], obs["imu.accel.z"] = ls.imu_state.accelerometer
|
||
if ls.imu_state.quaternion:
|
||
obs["imu.quat.w"] = ls.imu_state.quaternion[0]
|
||
obs["imu.quat.x"] = ls.imu_state.quaternion[1]
|
||
obs["imu.quat.y"] = ls.imu_state.quaternion[2]
|
||
obs["imu.quat.z"] = ls.imu_state.quaternion[3]
|
||
return obs
|
||
|
||
def send_action(self, action):
|
||
with self._bridge.low_cmd_lock:
|
||
for motor in G1_29_JointIndex:
|
||
key = f"{motor.name}.q"
|
||
if key not in action:
|
||
continue
|
||
i = motor.value
|
||
self._bridge.low_cmd.motor_cmd[i].q = action[key]
|
||
self._bridge.low_cmd.motor_cmd[i].qd = 0
|
||
self._bridge.low_cmd.motor_cmd[i].kp = self.kp[i]
|
||
self._bridge.low_cmd.motor_cmd[i].kd = self.kd[i]
|
||
self._bridge.low_cmd.motor_cmd[i].tau = 0.0
|
||
self._bridge.low_cmd_received = True
|
||
self._bridge.new_low_cmd = True
|
||
return action
|
||
|
||
def reset(self, control_dt=None, default_positions=None):
|
||
if default_positions is None:
|
||
default_positions = DEFAULT_ANGLES
|
||
self._gym_env.reset()
|
||
cmd = {f"{m.name}.q": float(default_positions[m.value]) for m in G1_29_JointIndex}
|
||
self.send_action(cmd)
|
||
self._advance_sim(1)
|
||
|
||
@property
|
||
def is_connected(self):
|
||
with self._lowstate_lock:
|
||
return self._lowstate is not None
|
||
|
||
def disconnect(self):
|
||
self._shutdown_event.set()
|
||
if self._gym_env is not None:
|
||
try:
|
||
self._gym_env.close()
|
||
except Exception:
|
||
pass
|
||
self._gym_env = None
|
||
|
||
# ── Main ──────────────────────────────────────────────────────────────────────
|
||
|
||
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="Sim without MuJoCo window (default: show viewer)")
|
||
args = parser.parse_args()
|
||
|
||
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")
|
||
|
||
planner_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="planner_sonic.onnx")
|
||
encoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_encoder.onnx")
|
||
decoder_path = hf_hub_download(repo_id="nvidia/GEAR-SONIC", filename="model_decoder.onnx")
|
||
|
||
providers = _ort_providers(force_cpu=args.cpu)
|
||
use_gpu = providers[0] == "CUDAExecutionProvider"
|
||
so = ort.SessionOptions(); so.log_severity_level = 3
|
||
|
||
print(f"[ONNX] enc/dec/planner={'GPU' if use_gpu else 'CPU'}")
|
||
if not use_gpu and not args.cpu:
|
||
print("[ONNX] Tip: pip install onnxruntime-gpu for CUDA inference")
|
||
|
||
planner_sess = ort.InferenceSession(planner_path, sess_options=so, providers=providers)
|
||
encoder_sess = ort.InferenceSession(encoder_path, sess_options=so, providers=providers)
|
||
decoder_sess = ort.InferenceSession(decoder_path, sess_options=so, providers=providers)
|
||
print(f"[Planner] version={'v1+' if len(planner_sess.get_inputs())>=11 else 'v0'}")
|
||
|
||
cfg = UnitreeG1Config()
|
||
if args.ip:
|
||
cfg.is_simulation = False
|
||
cfg.robot_ip = args.ip
|
||
robot = UnitreeG1(cfg)
|
||
else:
|
||
robot = SonicSimRobot(headless=args.headless)
|
||
robot.connect()
|
||
kp, kd = _kp_kd(); robot.kp = kp.copy(); robot.kd = kd.copy()
|
||
|
||
ms = MovementState()
|
||
planner = SonicPlanner(planner_sess, planner_path)
|
||
controller = PlannerController(planner, encoder_sess, decoder_sess)
|
||
|
||
motion = planner.initialize(DEFAULT_ANGLES, ms)
|
||
controller.load_initial_motion(motion)
|
||
controller.print_input_diagnostics()
|
||
planner.start_subprocess(controller, use_gpu=use_gpu)
|
||
|
||
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)
|
||
|
||
with RawKeyboard() as kb:
|
||
try:
|
||
gc.disable(); gc_timer = 0.0
|
||
robot.reset(CONTROL_DT, DEFAULT_ANGLES); time.sleep(1.0)
|
||
|
||
step = 0; replan_timer = 0.0
|
||
last_status = time.time() - 2.1 # print status on first loop tick
|
||
loop_t = enc_t = dec_t = obs_t = act_t = []
|
||
slow_n = blend_n = 0; stall_src = ""; did_blend = False
|
||
prev_end = time.time(); t_start = time.time()
|
||
last_ms = _snapshot_ms(ms)
|
||
|
||
log_path = "/tmp/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
|
||
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 process_keyboard(kb.get_key(), ms, controller): break
|
||
|
||
obs = robot.get_observation(); t_obs = time.time()
|
||
obs_t.append(1000*(t_obs - t0))
|
||
if not obs:
|
||
step += 1; prev_end = time.time()
|
||
time.sleep(max(0.0, CONTROL_DT-(time.time()-t0))); continue
|
||
|
||
process_joystick(obs, ms, controller)
|
||
clamp_mode_params(ms)
|
||
|
||
if step > 0:
|
||
replan_timer += CONTROL_DT
|
||
do_req = should_replan_request(ms, last_ms, replan_timer, step)
|
||
if do_req:
|
||
planner.request_replan(controller.ref_cursor, ms)
|
||
replan_timer = 0.0
|
||
ms.needs_replan = False
|
||
last_ms = _snapshot_ms(ms)
|
||
|
||
do_enc = (step % ENCODER_UPDATE_EVERY == 0)
|
||
t_step = time.time()
|
||
action = controller.step(obs, update_encoder=do_enc, debug=(step % DEBUG_PRINT_EVERY == 0))
|
||
step_ms = 1000*(time.time()-t_step)
|
||
(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))
|
||
|
||
result = planner.try_get_new_motion()
|
||
t_blend = time.time()
|
||
if result:
|
||
controller.blend_new_motion(*result)
|
||
blend_ms = 1000*(time.time()-t_blend)
|
||
blend_n += 1; did_blend = True
|
||
else:
|
||
blend_ms = 0.0
|
||
|
||
if log_ctx and 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},{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)
|
||
wall_dt = now - prev_end; loop_t.append(loop_ms)
|
||
if loop_ms > 50:
|
||
stall_src = (f"[STALL] {loop_ms:.0f}ms: "
|
||
f"obs={obs_t[-1]:.0f} blend={blend_ms:.0f} step={step_ms:.0f} act={act_t[-1]:.0f}")
|
||
if loop_ms > CONTROL_DT*1500: slow_n += 1
|
||
|
||
controller.advance_cursor()
|
||
|
||
if now - last_status > 2.0:
|
||
def _avg(l): return sum(l)/len(l) if l else 0
|
||
hz = 1000/_avg(loop_t) if _avg(loop_t) else 0
|
||
print(f"\r {ms.status_line()} step={step} 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
|
||
|
||
prev_end = time.time()
|
||
gc_timer += CONTROL_DT
|
||
if gc_timer >= 10.0: gc.collect(); gc_timer = 0.0
|
||
step += 1
|
||
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}")
|
||
planner.stop_subprocess()
|
||
print("\nStopping...")
|
||
if robot.is_connected: robot.disconnect()
|
||
print("Done.")
|
||
|
||
if __name__ == "__main__":
|
||
main() |