diff --git a/.DS_Store b/.DS_Store index 5725f1f..92e9250 100644 Binary files a/.DS_Store and b/.DS_Store differ diff --git a/deploy_45dim/.DS_Store b/deploy_45dim/.DS_Store new file mode 100644 index 0000000..d56cd79 Binary files /dev/null and b/deploy_45dim/.DS_Store differ diff --git a/deploy_45dim/90k_45.onnx b/deploy_45dim/90k_45.onnx new file mode 100644 index 0000000..94002fb Binary files /dev/null and b/deploy_45dim/90k_45.onnx differ diff --git a/deploy_45dim/__pycache__/deploy_57dim_pro_sdk.cpython-310.pyc b/deploy_45dim/__pycache__/deploy_57dim_pro_sdk.cpython-310.pyc new file mode 100644 index 0000000..7d1cb28 Binary files /dev/null and b/deploy_45dim/__pycache__/deploy_57dim_pro_sdk.cpython-310.pyc differ diff --git a/deploy_45dim/__pycache__/go1_no_linevel_sim2sim_mujoco.cpython-310.pyc b/deploy_45dim/__pycache__/go1_no_linevel_sim2sim_mujoco.cpython-310.pyc new file mode 100644 index 0000000..2338c30 Binary files /dev/null and b/deploy_45dim/__pycache__/go1_no_linevel_sim2sim_mujoco.cpython-310.pyc differ diff --git a/deploy_45dim/__pycache__/sim2sim_57dim_test.cpython-310.pyc b/deploy_45dim/__pycache__/sim2sim_57dim_test.cpython-310.pyc new file mode 100644 index 0000000..4435250 Binary files /dev/null and b/deploy_45dim/__pycache__/sim2sim_57dim_test.cpython-310.pyc differ diff --git a/deploy_45dim/__pycache__/sim2sim_test_deploy.cpython-310.pyc b/deploy_45dim/__pycache__/sim2sim_test_deploy.cpython-310.pyc new file mode 100644 index 0000000..f750bec Binary files /dev/null and b/deploy_45dim/__pycache__/sim2sim_test_deploy.cpython-310.pyc differ diff --git a/deploy_45dim/deploy_57dim_pro_sdk.py b/deploy_45dim/deploy_57dim_pro_sdk.py new file mode 100644 index 0000000..45c9e13 --- /dev/null +++ b/deploy_45dim/deploy_57dim_pro_sdk.py @@ -0,0 +1,899 @@ +#!/usr/bin/env python3 +""" +deploy_onnx_pro_sdk.py + +Deploy ONNX policy on Unitree Go1 PRO using go1_pro_sdk (no official SDK needed). +Works on non-EDU Go1 PRO. + +Uses direct MCU UDP communication with Blowfish encryption. +Joint order matches ONNX directly: FR→FL→RR→RL (hip/thigh/calf per leg). +Observation: 45-dim, matches go1_sim2sim.py. + +Setup: + cd /path/to/go1_pro_sdk && pip install -e . + pip install onnxruntime numpy + +Before running, kill the Pi's sport processes: + ssh pi@192.168.123.161 + sudo pkill -9 -f keep_sport_alive + sudo pkill -9 -f Legged_sport + sudo pkill -9 -f appTransit + +ALWAYS suspend the robot for initial testing. + +Usage: + # Step 0: check 45-dim observation vector (no motors, no ONNX) + python deploy_onnx_pro_sdk.py --obs-check --kill-sport + + # Step 1: test remote controller data (no motors, no ONNX) + python deploy_onnx_pro_sdk.py --monitor --kill-sport + + # Step 2-5: state machine with RC (R2=go/stop, L2=emergency stop) + python deploy_onnx_pro_sdk.py --onnx policy.onnx --kill-sport + + # Fixed commands, no state machine + python deploy_onnx_pro_sdk.py --onnx policy.onnx --no-rc --no-sm --kill-sport +""" + +import argparse +import json +import signal +import subprocess +import sys +import time +from datetime import datetime +from enum import Enum +from pathlib import Path + +import numpy as np +import onnxruntime as ort + +from go1_pro_sdk import ( + MCUClient, LowCmd, MotorCmd, MotorMode, + apply_safety, JOINT_NAMES, +) + +# ─── Policy constants (matches go1_sim2sim.py) ─── +NUM_OBS = 57 # no-linvel: gyro(3)+gravity(3)+dof_pos(12)+dof_vel(12)+last_actions(12)+commands(3)+contacts(12) +NUM_ACTIONS = 12 +ACTION_SCALE = 0.05 +CLIP_ACTIONS = 23.7 +CLIP_OBS = 100.0 + +DEFAULT_ANGLES = np.array([ + -0.0, 0.9, -1.8, # FR + 0.0, 0.9, -1.8, # FL + -0.0, 0.9, -1.8, # RR + 0.0, 0.9, -1.8, # RL +], dtype=np.float32) + +EXIT = False + + +def _sig_handler(signum, frame): + global EXIT + EXIT = True + + +signal.signal(signal.SIGINT, _sig_handler) +signal.signal(signal.SIGTERM, _sig_handler) + + +# ─── State machine states ─── +class State(Enum): + IDLE = "IDLE" # full damping, wait for R2 + CALIBRATE = "CALIBRATE" # ramping to default pose + HOLD = "HOLD" # holding default pose, wait for R2 to start RL + RL = "RL" # running ONNX policy + + +# ─── Quaternion math ─── +def quat_to_rot_matrix(q): + w, x, y, z = q + return np.array([ + [1 - 2*y*y - 2*z*z, 2*x*y - 2*w*z, 2*x*z + 2*w*y], + [ 2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z, 2*y*z - 2*w*x], + [ 2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y], + ], dtype=np.float32) + + +def get_projected_gravity(quaternion): + R = quat_to_rot_matrix(quaternion) + return (R.T @ np.array([0., 0., -1.], dtype=np.float32)).astype(np.float32) + + +# ─── Observation ─── +def compute_obs(imu, motor_states, commands, last_actions): + """57-dim no-linvel: gyro(3)+gravity(3)+dof_pos(12)+dof_vel(12)+last_actions(12)+commands(3)+contacts(12).""" + obs = np.zeros(NUM_OBS, dtype=np.float32) + + gx, gy, gz = imu.gyroscope + obs[0:3] = np.array([gx, gy, gz], dtype=np.float32) * 0.25 + obs[3:6] = get_projected_gravity(imu.quaternion) + + dof_pos = np.array([motor_states[i].q for i in range(12)], dtype=np.float32) + dof_vel = np.array([motor_states[i].dq for i in range(12)], dtype=np.float32) + obs[6:18] = (dof_pos - DEFAULT_ANGLES) * 1.0 + obs[18:30] = dof_vel * 0.05 + + obs[30:42] = last_actions + obs[42:45] = np.array(commands, dtype=np.float32) * np.array([2.0, 2.0, 0.25], dtype=np.float32) + obs[45:57] = 0.0 # contact forces = 0 + + obs = np.clip(obs, -CLIP_OBS, CLIP_OBS) + obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0) + return obs + + +def state_ok(state): + for i in range(12): + if not np.isfinite(state.motorState[i].q): + return False, f"motorState[{i}].q non-finite" + gn = float(np.linalg.norm(get_projected_gravity(state.imu.quaternion))) + if gn < 0.5 or gn > 1.5: + return False, f"gravity norm={gn:.3f}" + return True, "ok" + + +# ─── ONNX policy ─── +class OnnxPolicy: + def __init__(self, onnx_path): + self.session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"]) + self.input_name = self.session.get_inputs()[0].name + print(f"[INFO] ONNX: {onnx_path}") + print(f"[INFO] Input : {self.input_name} {self.session.get_inputs()[0].shape}") + print(f"[INFO] Output: {[(o.name, o.shape) for o in self.session.get_outputs()]}") + + def __call__(self, obs): + out = self.session.run(None, {self.input_name: obs.reshape(1, -1).astype(np.float32)})[0][0] + return np.asarray(out, dtype=np.float32) + + +# ─── Remote controller ─── +def get_rc_commands(state, args): + r = state.remote + vx = r.ly * args.rc_vx_scale + vy = -r.lx * args.rc_vy_scale + wz = -r.rx * args.rc_wz_scale + return np.array([vx, vy, wz], dtype=np.float32) + + +class RCEdgeDetector: + """Detect rising/falling edges on RC buttons.""" + def __init__(self): + self._prev = set() + + def update(self, state): + current = set(state.remote.pressed) + rising = current - self._prev + falling = self._prev - current + self._prev = current + return rising, falling + + def rose(self, state, button): + """True on rising edge of button.""" + current = set(state.remote.pressed) + prev = self._prev + self._prev = current + return button not in prev and button in current + + +# ─── JSONL logger ─── +class JsonlLogger: + def __init__(self, log_dir, args): + self.enabled = bool(log_dir) + self.fp = None + self.run_dir = None + self.flush_every = 50 + if not self.enabled: + return + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + self.run_dir = Path(log_dir).expanduser().resolve() / f"deploy_run_{ts}" + self.run_dir.mkdir(parents=True, exist_ok=True) + meta = { + "created_at": ts, + "num_obs": NUM_OBS, + "num_actions": NUM_ACTIONS, + "action_scale": ACTION_SCALE, + "clip_actions": CLIP_ACTIONS, + "clip_obs": CLIP_OBS, + "default_angles": DEFAULT_ANGLES.tolist(), + "joint_names": list(JOINT_NAMES), + } + for k, v in vars(args).items(): + if isinstance(v, (str, int, float, bool, type(None))): + meta[k] = v + (self.run_dir / "metadata.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False)) + self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8", buffering=1) + print(f"[INFO] Log dir: {self.run_dir}") + + def log(self, step, **kw): + if not self.enabled: + return + rec = {"step": int(step), "time_wall": time.time()} + for k, v in kw.items(): + if isinstance(v, np.ndarray): + rec[k] = np.asarray(v, dtype=np.float32).reshape(-1).tolist() + elif isinstance(v, (np.float32, np.float64)): + rec[k] = float(v) + elif isinstance(v, (np.int32, np.int64)): + rec[k] = int(v) + else: + rec[k] = v + self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n") + if step % self.flush_every == 0: + self.fp.flush() + + def close(self): + if self.fp: + self.fp.flush(); self.fp.close() + print(f"[INFO] Log saved: {self.run_dir}") + + +# ─── Display helpers ─── +def fmt_rc(state): + r = state.remote + sticks = f"lx={r.lx:+.2f} ly={r.ly:+.2f} rx={r.rx:+.2f} ry={r.ry:+.2f} L2={r.L2:.2f}" + btns = ",".join(r.pressed) if r.pressed else "none" + return f"sticks:[{sticks}] btns:[{btns}]" + + +def fmt_joints(state): + parts = [] + for i in range(12): + parts.append(f"{state.motorState[i].q:+6.3f}") + return " ".join(parts) + + +# ─── Observation check mode ─── +def run_obs_check(args): + """Compute 45-dim observation and print each segment. No inference, no motors.""" + print(""" +╔══════════════════════════════════════════════════════════════╗ +║ OBS-CHECK MODE — verify 45-dim observation vector ║ +║ No inference, no motors. Compare with sim2sim reference. ║ +║ ║ +║ Ensure: ║ +║ 1. Robot SUSPENDED (sling / foot stand) ║ +║ 2. Network connected to robot (192.168.123.x) ║ +║ 3. Use --kill-sport to auto-kill Pi sport processes ║ +╚══════════════════════════════════════════════════════════════╝ +""") + input("Press Enter to start obs check...") + + if args.kill_sport: + kill_sport_processes(args.pi_host, args.pi_user) + + print("[INFO] Connecting to MCU...") + client = MCUClient() + logger = JsonlLogger(args.log_dir, args) + + try: + print("[INFO] Waking MCU...") + client.wake_mcu(n_frames=50, dt=0.01) + + state = client.recv_state(timeout=2.0) + if state is None: + print("[ERROR] No state received.") + return 1 + + print(f"[INFO] Connected. Battery={state.bms.SOC}%") + print(f"[INFO] RPY: {np.round(np.degrees(state.imu.rpy), 1)} deg") + print("[INFO] Computing 45-dim obs each step. Compare with sim2sim reference.\n") + + # Sim2sim reference (stationary, default pose, zero cmd) + print("─" * 70) + print("OBSERVATION LAYOUT (45-dim):") + print(" obs[ 0: 3] gyro * 0.25 → ~[0, 0, 0] when stationary") + print(" obs[ 3: 6] projected gravity → ~[0, 0, -1] when upright") + print(" obs[ 6:18] (dof_pos-def) * 1.0 → ~[0]*12 when at default pose") + print(" obs[18:30] dof_vel * 0.05 → ~[0]*12 when stationary") + print(" obs[30:42] last_actions → [0]*12 initially") + print(" obs[42:45] commands*[2,2,0.25] → [0,0,0] with zero cmd") + print("─" * 70) + + last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32) + step = 0 + dt = 1.0 / max(1.0, args.rate_hz) + next_t = time.perf_counter() + + while not EXIT: + new_state = client.recv_latest() + if new_state is not None: + state = new_state + if state is None: + time.sleep(0.001) + continue + + # Get commands from RC or fixed + r = state.remote + commands = np.array([r.ly * args.rc_vx_scale, + -r.lx * args.rc_vy_scale, + -r.rx * args.rc_wz_scale], dtype=np.float32) + + # Compute observation + obs = compute_obs(state.imu, state.motorState, commands, last_actions) + + dof_pos = np.array([state.motorState[i].q for i in range(12)]) + dof_vel = np.array([state.motorState[i].dq for i in range(12)]) + gx, gy, gz = state.imu.gyroscope + grav = get_projected_gravity(state.imu.quaternion) + + # Detailed print + print(f"\n{'='*70}") + print(f"[STEP {step}] bat={state.bms.SOC}% {fmt_rc(state)}") + print(f"{'='*70}") + print(f" obs[ 0: 3] gyro*0.25 = {np.round(obs[0:3], 4)}") + print(f" raw gyro (rad/s) = [{gx:+.4f}, {gy:+.4f}, {gz:+.4f}]") + print(f" obs[ 3: 6] proj gravity = {np.round(obs[3:6], 4)}") + print(f" |gravity| = {np.linalg.norm(grav):.4f} (expect ~1.0)") + print(f" obs[ 6:18] dof_pos_rel = {np.round(obs[6:18], 3)}") + diff_from_default = dof_pos - DEFAULT_ANGLES + print(f" |diff| max = {np.max(np.abs(diff_from_default)):.4f}") + print(f" actual pos = {np.round(dof_pos, 3)}") + print(f" default_angles = {np.round(DEFAULT_ANGLES, 3)}") + print(f" obs[18:30] dof_vel*0.05 = {np.round(obs[18:30], 4)}") + print(f" |dof_vel| max = {np.max(np.abs(dof_vel)):.4f}") + print(f" obs[30:42] last_actions = {np.round(obs[30:42], 4)}") + print(f" obs[42:45] commands*scale = {np.round(obs[42:45], 4)}") + print(f" raw cmd [vx,vy,wz]= {np.round(commands, 3)}") + print(f" obs min={obs.min():.3f} max={obs.max():.3f} mean={obs.mean():.3f}") + + # Warn if values look suspicious + if np.linalg.norm(grav) < 0.5 or np.linalg.norm(grav) > 1.5: + print(f" ⚠ WARN: gravity norm is off! Robot might not be upright.") + if np.max(np.abs(dof_pos)) < 1e-6: + print(f" ⚠ WARN: dof_pos is all zeros — state might be stale!") + + logger.log( + step, + mode="obs_check", + obs=obs, + dof_pos=dof_pos, + dof_vel=dof_vel, + base_ang_vel=np.array([gx, gy, gz]), + projected_gravity=grav, + imu_rpy_deg=np.degrees(state.imu.rpy), + commands=commands, + rc_buttons=r.pressed, + battery_soc=state.bms.SOC, + ) + + step += 1 + + next_t += dt + sleep = next_t - time.perf_counter() + if sleep > 0: + time.sleep(sleep) + else: + next_t = time.perf_counter() + + finally: + if logger is not None: + logger.close() + client.close() + print("[INFO] Done.") + + +# ─── Monitor mode ─── +def run_monitor(args): + """Read and display RC + IMU + joint data. No motor commands sent.""" + print(""" +╔══════════════════════════════════════════════════════════════╗ +║ MONITOR MODE — no motor commands, RC test only ║ +║ ║ +║ Ensure: ║ +║ 1. Robot SUSPENDED (sling / foot stand) ║ +║ 2. Network connected to robot (192.168.123.x) ║ +║ 3. Use --kill-sport to auto-kill Pi sport processes ║ +║ (or manually: ssh pi@192.168.123.161, pkill -9 ...) ║ +╚══════════════════════════════════════════════════════════════╝ +""") + input("Press Enter to start monitoring...") + + if args.kill_sport: + kill_sport_processes(args.pi_host, args.pi_user) + + print("[INFO] Connecting to MCU...") + client = MCUClient() + logger = JsonlLogger(args.log_dir, args) + + try: + print("[INFO] Waking MCU (damping frames)...") + client.wake_mcu(n_frames=50, dt=0.01) + + state = client.recv_state(timeout=2.0) + if state is None: + print("[ERROR] No state received. Check connection and Pi processes.") + return 1 + + print(f"[INFO] Connected. Battery={state.bms.SOC}%") + print("[INFO] Monitoring RC + IMU + joint data. Ctrl+C to exit.\n") + + edge = RCEdgeDetector() + step = 0 + dt = 1.0 / max(1.0, args.rate_hz) + next_t = time.perf_counter() + + while not EXIT: + new_state = client.recv_latest() + if new_state is not None: + state = new_state + + if state is None: + time.sleep(0.001) + continue + + rising, falling = edge.update(state) + + rpy = np.degrees(state.imu.rpy) + r = state.remote + dof_pos = np.array([state.motorState[i].q for i in range(12)]) + dof_vel = np.array([state.motorState[i].dq for i in range(12)]) + + print(f"\n─── [step {step}] bat={state.bms.SOC}% ───") + print(f" IMU roll={rpy[0]:+.1f} pitch={rpy[1]:+.1f} yaw={rpy[2]:+.1f}") + print(f" RC {fmt_rc(state)}") + if rising: + print(f" RC ↑ RISING: {sorted(rising)}") + if falling: + print(f" RC ↓ FALLING: {sorted(falling)}") + print(f" JOINT {fmt_joints(state)}") + print(f" JOINTS: {' '.join(f'{n:6s}' for n in JOINT_NAMES)}") + + logger.log( + step, + mode="monitor", + dof_pos=dof_pos, + dof_vel=dof_vel, + base_ang_vel=np.array(state.imu.gyroscope), + projected_gravity=get_projected_gravity(state.imu.quaternion), + imu_rpy_deg=np.degrees(state.imu.rpy), + rc_lx=r.lx, rc_ly=r.ly, rc_rx=r.rx, rc_ry=r.ry, + rc_buttons=r.pressed, + rc_rising=list(rising), rc_falling=list(falling), + battery_soc=state.bms.SOC, + ) + + step += 1 + + next_t += dt + sleep = next_t - time.perf_counter() + if sleep > 0: + time.sleep(sleep) + else: + next_t = time.perf_counter() + + finally: + if logger is not None: + logger.close() + client.close() + print("[INFO] Done.") + + +# ─── Ramp to default pose ─── +def ramp_to_default(client, args, state): + """Slowly interpolate from current pose to DEFAULT_ANGLES (~2s). + + Position protect is DISABLED during ramp because the movement is inherently + large (e.g. knee from -2.79 → -1.80, a 1.0 rad sweep). Only joint hard + limits and torque limiting are active. + """ + print("[INFO] Ramping to default pose (~2s)...") + current = np.array([state.motorState[i].q for i in range(12)], dtype=np.float32) + error = current - DEFAULT_ANGLES + + if np.max(np.abs(error)) < 0.05: + print("[INFO] Already near default pose.") + return state + + print(f"[INFO] Max error: {np.max(np.abs(error)):.2f} rad") + ramp_steps = 200 + step_err = error / ramp_steps + + for i in range(ramp_steps): + if EXIT: + return state + new_state = client.recv_latest() + if new_state is not None: + state = new_state + + targets = DEFAULT_ANGLES + (error - step_err * min(i + 1, ramp_steps)) + + cmd = LowCmd() + for j in range(12): + cmd.set_motor(j, MotorCmd( + mode=MotorMode.Servo, + q=float(targets[j]), dq=0.0, tau=0.0, + Kp=args.kp_cal, Kd=args.kd_cal, + )) + # Ramp: only joint limits + torque limit, NO position_protect + # (the movement is inherently >0.5 rad for many joints) + apply_safety(cmd, state, power_factor=args.power_factor, + position_limit_on=True, position_protect_limit=None) + client.send(cmd) + time.sleep(0.01) + + if i % 50 == 0: + actual = np.array([state.motorState[j].q for j in range(12)]) + print(f" ramp {i}/{ramp_steps} " + f"target_err={np.max(np.abs(targets - DEFAULT_ANGLES)):.3f} " + f"actual_err={np.max(np.abs(actual - DEFAULT_ANGLES)):.3f}") + + # Hold at default briefly (position_protect disabled — ramp continuation) + for _ in range(50): + if EXIT: + return state + new_state = client.recv_latest() + if new_state is not None: + state = new_state + + cmd = LowCmd() + for j in range(12): + cmd.set_motor(j, MotorCmd( + mode=MotorMode.Servo, + q=float(DEFAULT_ANGLES[j]), dq=0.0, tau=0.0, + Kp=args.kp, Kd=args.kd, + )) + apply_safety(cmd, state, power_factor=args.power_factor, + position_limit_on=True, position_protect_limit=None) + client.send(cmd) + time.sleep(0.01) + + print("[INFO] Default pose reached.") + return state + + +def send_hold_cmd(client, state, args): + """Send servo commands holding DEFAULT_ANGLES. + Position protect is disabled — hold is a fixed safe pose. + """ + cmd = LowCmd() + for j in range(12): + cmd.set_motor(j, MotorCmd( + mode=MotorMode.Servo, + q=float(DEFAULT_ANGLES[j]), dq=0.0, tau=0.0, + Kp=args.kp, Kd=args.kd, + )) + apply_safety(cmd, state, power_factor=args.power_factor, + position_limit_on=True, position_protect_limit=None) + client.send(cmd) + + +def send_rl_cmd(client, state, targets, args): + """Send servo commands for RL policy targets.""" + cmd = LowCmd() + for j in range(12): + cmd.set_motor(j, MotorCmd( + mode=MotorMode.Servo, + q=float(targets[j]), dq=0.0, tau=0.0, + Kp=args.kp, Kd=args.kd, + )) + pp_limit = args.position_protect_limit if args.position_protect_limit > 0 else None + apply_safety(cmd, state, power_factor=args.power_factor, + position_limit_on=True, position_protect_limit=pp_limit) + client.send(cmd) + + +# ─── Startup instructions ─── +STARTUP_BANNER = """ +╔══════════════════════════════════════════════════════════════╗ +║ Ensure: ║ +║ 1. Robot SUSPENDED (sling / foot stand) ║ +║ 2. Network connected to robot (192.168.123.x) ║ +║ 3. Use --kill-sport to auto-kill Pi sport processes ║ +║ (or manually: ssh pi@192.168.123.161, pkill -9 ...) ║ +║ ║ +║ Controls: R2=go/stop, L2=emergency-stop(→IDLE) ║ +║ Left-stick=move, Right-stick=turn ║ +╚══════════════════════════════════════════════════════════════╝ +""" + + +# ─── State machine deploy mode ─── +def run_deploy(args): + print(STARTUP_BANNER) + input("Press Enter when ready...") + + if args.kill_sport: + kill_sport_processes(args.pi_host, args.pi_user) + + print("[INFO] Connecting to MCU...") + client = MCUClient() + logger = None + + try: + print("[INFO] Waking MCU...") + client.wake_mcu(n_frames=50, dt=0.01) + + state = client.recv_state(timeout=2.0) + if state is None: + print("[ERROR] No state received. Check connection and Pi processes.") + return 1 + + print(f"[INFO] Connected. Battery={state.bms.SOC}%") + print(f"[INFO] RPY: {np.round(np.degrees(state.imu.rpy), 1)} deg") + print("[INFO] Initial joint positions (rad):") + for i in range(12): + print(f" {JOINT_NAMES[i]:6s}: {state.motorState[i].q:+7.3f}") + + # Load ONNX + policy = OnnxPolicy(args.onnx) + + # Logger + logger = JsonlLogger(args.log_dir, args) + logger.flush_every = max(1, int(args.log_flush_every)) + + # State machine + sm_state = State.IDLE + last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32) + ema_filter = type('EMA', (), {'state': np.zeros(NUM_ACTIONS, dtype=np.float32)})() + step = 0 + dt = 1.0 / args.rate_hz + next_t = time.perf_counter() + edge = RCEdgeDetector() + + # Seed edge detector with initial state + edge.update(state) + + print(f"\n[INFO] State machine: IDLE → (R2) → CALIBRATE → HOLD → (R2) → RL") + print(f"[INFO] Rate: {args.rate_hz}Hz. Ctrl+C to exit.\n") + + while not EXIT: + new_state = client.recv_latest() + if new_state is not None: + state = new_state + + if state is None: + time.sleep(0.001) + continue + + # Rising edge detection (single update per tick) + rising, falling = edge.update(state) + r2_rose = "R2" in rising + l2_rose = "L2" in rising + r = state.remote + + # ── Emergency stop: L2 → IDLE (from any state) ── + if l2_rose and sm_state != State.IDLE: + print(f"\n[L2 EMERGENCY] {sm_state.value} → IDLE") + sm_state = State.IDLE + last_actions[:] = 0.0 + # Send damping + client.send(LowCmd()) # all damping by default + + # ── State machine ── + if sm_state == State.IDLE: + # Full damping, no torques. Keep stream alive with damping frames. + if step % 10 == 0: + client.send(LowCmd()) # all damping + + if r2_rose: + print("\n[R2] IDLE → CALIBRATE") + sm_state = State.CALIBRATE + state = ramp_to_default(client, args, state) + if EXIT: + break + sm_state = State.HOLD + print(f"[STATE] → HOLD (waiting for R2 to start RL)") + + elif sm_state == State.HOLD: + send_hold_cmd(client, state, args) + + if r2_rose: + print("\n[R2] HOLD → RL") + sm_state = State.RL + last_actions[:] = 0.0 + ema_filter.state[:] = 0.0 + + elif sm_state == State.RL: + if r2_rose: + print("\n[R2] RL → HOLD") + sm_state = State.HOLD + last_actions[:] = 0.0 + # Ramp back to default + state = ramp_to_default(client, args, state) + if EXIT: + break + continue + + # Velocity commands from RC + commands = get_rc_commands(state, args) + + # Observation + inference + obs = compute_obs(state.imu, state.motorState, commands, last_actions) + + ok, reason = (True, "ok") + if not args.no_state_check: + ok, reason = state_ok(state) + + if ok: + action_raw = policy(obs) + action_clipped = np.clip(action_raw, -CLIP_ACTIONS, CLIP_ACTIONS) + # EMA smooth to suppress policy's inherent high-freq oscillation + alpha = args.action_ema_alpha + if 0 < alpha < 1 and hasattr(ema_filter, 'state'): + ema_filter.state = alpha * action_clipped + (1 - alpha) * ema_filter.state + action_smoothed = ema_filter.state.copy() + else: + action_smoothed = action_clipped + last_actions = action_smoothed.copy() + targets = DEFAULT_ANGLES + action_smoothed * ACTION_SCALE + else: + # State check failed: hold default pose + targets = DEFAULT_ANGLES.copy() + + if step >= args.warmup_steps: + send_rl_cmd(client, state, targets, args) + + # Log RL data + logger.log( + step, + mode=sm_state.value, + commands=commands, + obs_isaac=obs, + action_raw=action_raw if ok else np.zeros(NUM_ACTIONS), + action_safe=action_clipped if ok else np.zeros(NUM_ACTIONS), + joint_targets=targets, + dof_pos=np.array([state.motorState[i].q for i in range(12)]), + dof_vel=np.array([state.motorState[i].dq for i in range(12)]), + base_ang_vel=np.array(state.imu.gyroscope), + projected_gravity=get_projected_gravity(state.imu.quaternion), + imu_rpy_deg=np.degrees(state.imu.rpy), + rc_lx=state.remote.lx, rc_ly=state.remote.ly, + rc_rx=state.remote.rx, rc_ry=state.remote.ry, + rc_buttons=state.remote.pressed, + state_ok=ok, state_reason=reason, + ) + else: + # Log non-RL states (minimal) + logger.log( + step, + mode=sm_state.value, + dof_pos=np.array([state.motorState[i].q for i in range(12)]), + imu_rpy_deg=np.degrees(state.imu.rpy), + rc_buttons=state.remote.pressed, + ) + + # ── Status print ── + if step % args.print_every == 0: + rpy = np.degrees(state.imu.rpy) + dof_pos = np.array([state.motorState[i].q for i in range(12)]) + print(f"\n[STEP {step}] state={sm_state.value} bat={state.bms.SOC}% " + f"rpy=[{rpy[0]:.0f},{rpy[1]:.0f},{rpy[2]:.0f}]") + print(f" RC: {fmt_rc(state)}") + print(f" joint: {np.round(dof_pos, 2)}") + if sm_state == State.RL and 'targets' in dir(): + print(f" target:{np.round(targets, 2)}") + + step += 1 + if args.max_steps > 0 and step >= args.max_steps: + print("[INFO] max_steps reached.") + break + + # Rate limit + next_t += dt + sleep = next_t - time.perf_counter() + if sleep > 0: + time.sleep(sleep) + else: + next_t = time.perf_counter() + + finally: + if logger is not None: + logger.close() + print("[INFO] Safe stopping...") + client.safe_stop(n_frames=50, dt=0.002) + client.close() + print("[INFO] Done.") + + +# ─── Auto-kill Pi sport processes ─── +def kill_sport_processes(host, user): + """SSH into the Pi and kill sport-mode processes.""" + cmds = [ + "sudo pkill -9 -f keep_sport_alive", + "sudo pkill -9 -f Legged_sport", + "sudo pkill -9 -f appTransit", + ] + ssh_target = f"{user}@{host}" if user else host + full_cmd = " && ".join(cmds) + + print(f"[INFO] Killing sport processes on {ssh_target}...") + try: + result = subprocess.run( + ["ssh", ssh_target, full_cmd], + capture_output=True, text=True, timeout=15 + ) + if result.returncode == 0: + print("[INFO] Sport processes killed successfully.") + return True + else: + stderr = result.stderr.strip() + # pkill returns non-zero if no matching process — that's also OK + if "no process" in stderr.lower() or not stderr: + print("[INFO] No sport processes found (already killed).") + return True + print(f"[WARN] SSH returned {result.returncode}: {stderr}") + return False + except FileNotFoundError: + print("[WARN] 'ssh' command not found. Please kill processes manually.") + return False + except subprocess.TimeoutExpired: + print("[WARN] SSH timed out. Check network connection to Pi.") + return False + except Exception as e: + print(f"[WARN] Failed to kill sport processes: {e}") + return False + + +# ─── CLI ─── +def main(): + parser = argparse.ArgumentParser(description="Deploy ONNX policy on Go1 PRO (no official SDK)") + parser.add_argument("--onnx", default="better3.onnx", help="Path to ONNX model (57-dim no-linvel)") + + # Pi control + parser.add_argument("--kill-sport", action="store_true", + help="Auto-kill sport processes on Pi via SSH before starting") + parser.add_argument("--pi-host", default="192.168.123.161", help="Pi IP/hostname") + parser.add_argument("--pi-user", default="pi", help="Pi SSH user") + + # Modes + parser.add_argument("--obs-check", action="store_true", + help="Observation check mode: compute 45-dim obs, no inference, no motors") + parser.add_argument("--monitor", action="store_true", + help="Monitor mode: read RC/IMU/joint data only, no motors") + parser.add_argument("--no-rc", action="store_true", + help="Disable RC — use fixed --cmd-* velocities instead") + parser.add_argument("--no-sm", action="store_true", + help="Disable state machine — run RL continuously") + parser.add_argument("--dry-run", action="store_true", + help="Inference only, no motor commands") + + # Gains + parser.add_argument("--kp", type=float, default=30.0) + parser.add_argument("--kd", type=float, default=0.5) + parser.add_argument("--kp-cal", type=float, default=20.0) + parser.add_argument("--kd-cal", type=float, default=1.0) + parser.add_argument("--power-factor", type=int, default=7, + help="Torque limit factor 1-10. tau_lim = TAU_MAX * factor/10. hip: %.1f, thigh: %.1f, knee: %.1f Nm at factor=7" % (23.7*7/10, 23.7*7/10, 35.55*7/10)) + + # RC scale + parser.add_argument("--rc-vx-scale", type=float, default=1.0) + parser.add_argument("--rc-vy-scale", type=float, default=1.0) + parser.add_argument("--rc-wz-scale", type=float, default=1.0) + + # Fixed commands (when --no-rc) + parser.add_argument("--cmd-x", type=float, default=0.0) + parser.add_argument("--cmd-y", type=float, default=0.0) + parser.add_argument("--cmd-yaw", type=float, default=0.0) + + # Control + parser.add_argument("--rate-hz", type=float, default=100.0) + parser.add_argument("--warmup-steps", type=int, default=50) + parser.add_argument("--max-steps", type=int, default=0) + parser.add_argument("--print-every", type=int, default=50) + + # Safety + parser.add_argument("--no-state-check", action="store_true") + parser.add_argument("--position-protect-limit", type=float, default=1.0, + help="Max |target-actual| before zeroing Kp/Kd (negative=disable)") + parser.add_argument("--action-ema-alpha", type=float, default=0.0, + help="EMA smoothing factor for actions (0=no filter, 1=bypass, default 0)") + + # Logging + parser.add_argument("--log-dir", default="", help="Enable JSONL logging to this directory") + parser.add_argument("--log-flush-every", type=int, default=50) + + args = parser.parse_args() + + if args.obs_check: + return run_obs_check(args) + if args.monitor: + return run_monitor(args) + else: + return run_deploy(args) + + +if __name__ == "__main__": + main() diff --git a/deploy_onnx_pro_sdk.py b/deploy_45dim/deploy_onnx_pro_sdk.py similarity index 99% rename from deploy_onnx_pro_sdk.py rename to deploy_45dim/deploy_onnx_pro_sdk.py index c1ee77a..3b6d7ca 100644 --- a/deploy_onnx_pro_sdk.py +++ b/deploy_45dim/deploy_onnx_pro_sdk.py @@ -828,7 +828,7 @@ def kill_sport_processes(host, user): # ─── CLI ─── def main(): parser = argparse.ArgumentParser(description="Deploy ONNX policy on Go1 PRO (no official SDK)") - parser.add_argument("--onnx", default="policy.onnx", help="Path to ONNX model") + parser.add_argument("--onnx", default="90k_45.onnx", help="Path to ONNX model") # Pi control parser.add_argument("--kill-sport", action="store_true", diff --git a/go1_sim2sim.py b/deploy_45dim/go1_sim2sim.py similarity index 99% rename from go1_sim2sim.py rename to deploy_45dim/go1_sim2sim.py index d0ce6a7..6cd6ef6 100644 --- a/go1_sim2sim.py +++ b/deploy_45dim/go1_sim2sim.py @@ -15,7 +15,7 @@ from mujoco import viewer from pynput import keyboard HERE = os.path.dirname(os.path.abspath(__file__)) -ONNX = os.path.join(HERE, "policy.onnx") +ONNX = os.path.join(HERE, "90k_45.onnx") # ── Parameters (original MotrixLab Go1 config) ── NUM_OBS = 45 diff --git a/policy.onnx b/deploy_45dim/policy.onnx similarity index 100% rename from policy.onnx rename to deploy_45dim/policy.onnx diff --git a/deploy_45dim/policy_v3.onnx b/deploy_45dim/policy_v3.onnx new file mode 100644 index 0000000..94002fb Binary files /dev/null and b/deploy_45dim/policy_v3.onnx differ diff --git a/deploy_45dim/sim2sim_57dim_test.py b/deploy_45dim/sim2sim_57dim_test.py new file mode 100644 index 0000000..849b844 --- /dev/null +++ b/deploy_45dim/sim2sim_57dim_test.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 +""" +Sim2sim test for 57-dim no-linevel policy (better3.onnx). +obs: gyro(3) + gravity(3) + dof_pos(12) + dof_vel(12) + last_actions(12) + commands(3) + contacts(12) = 57 +Contacts set to zero. + +Usage: + conda activate free_dog_sdk + mjpython sim2sim_57dim_test.py + mjpython sim2sim_57dim_test.py --action-scale 0.07 --action-ema-alpha 0.3 + +Controls: W/S=前后 Q/E=左右 A/D=旋转 Space=停 R=重置 Esc=退出 +""" + +import argparse, os, signal, time +import mujoco, numpy as np, onnxruntime as ort +from mujoco import viewer + +HERE = os.path.dirname(os.path.abspath(__file__)) +XML = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml") +ONNX = os.path.join(HERE, "better3.onnx") + +NUM_OBS, NUM_ACTIONS = 57, 12 +DEFAULT_ANGLES = np.array([-0.0,0.9,-1.8, 0.0,0.9,-1.8, -0.0,0.9,-1.8, 0.0,0.9,-1.8], dtype=np.float32) + +EXIT = False +def _sig(s, f): global EXIT; EXIT = True +signal.signal(signal.SIGINT, _sig) + +class KB: + def __init__(s): s.h = set(); s._l = None + def _p(s,k): + try: s.h.add(k.char.lower()) + except: s.h.add(str(k)) + def _r(s,k): + try: s.h.discard(k.char.lower()) + except: s.h.discard(str(k)) + def init(s): + from pynput import keyboard + s._l = keyboard.Listener(on_press=s._p, on_release=s._r); s._l.start() + def keys(s): return s.h.copy() + def stop(s): + if s._l: s._l.stop() + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--action-scale", type=float, default=0.05) + p.add_argument("--action-ema-alpha", type=float, default=0.0) + p.add_argument("--kp", type=float, default=80.0) + p.add_argument("--kd", type=float, default=1.0) + args = p.parse_args() + + model = mujoco.MjModel.from_xml_path(XML) + data = mujoco.MjData(model) + + print(f"[INFO] 57-dim no-linevel ONNX: {ONNX}") + print(f"[INFO] action_scale={args.action_scale} ema_alpha={args.action_ema_alpha}") + print(f"[INFO] KP={args.kp} KD={args.kd} passive_damping={model.dof_damping[6]}") + + sess = ort.InferenceSession(ONNX, providers=["CPUExecutionProvider"]) + input_name = sess.get_inputs()[0].name + + kb = KB(); kb.init() + + data.qpos[:3] = [0,0,0.42]; data.qpos[3:7] = [1,0,0,0]; data.qpos[7:19] = DEFAULT_ANGLES + mujoco.mj_forward(model, data) + + view = viewer.launch_passive(model, data) + step = 0; si = int(0.01 / model.opt.timestep) # 100Hz + last_a = np.zeros(12, dtype=np.float32); action = np.zeros(12, dtype=np.float32) + t0 = time.perf_counter() + + while view.is_running() and not EXIT: + keys = kb.keys() + if 'key.esc' in keys: break + if 'r' in keys: + data.qpos[:3]=[0,0,0.42]; data.qpos[3:7]=[1,0,0,0] + data.qpos[7:19]=DEFAULT_ANGLES; data.qvel[:]=0; last_a[:]=0 + mujoco.mj_forward(model, data) + + vx=1.0 if 'w' in keys else (-1.0 if 's' in keys else 0.0) + vy=1.0 if 'q' in keys else (-1.0 if 'e' in keys else 0.0) + wz=1.0 if 'a' in keys else (-1.0 if 'd' in keys else 0.0) + if ' ' in keys: vx=vy=wz=0.0 + + if step % si == 0: + obs = np.zeros(NUM_OBS, dtype=np.float32) + + # gyro + sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Gyro") + obs[0:3] = data.sensordata[model.sensor_adr[sid]:model.sensor_adr[sid]+3] * 0.25 + + # gravity + R = data.xmat[1].reshape(3, 3) + obs[3:6] = (R.T @ [0., 0., -1.]).astype(np.float32) + + # dof_pos + obs[6:18] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0 + # dof_vel + obs[18:30] = data.qvel[6:18] * 0.05 + # last_actions + obs[30:42] = last_a + # commands + obs[42:45] = np.array([vx, vy, wz]) * np.array([2., 2., 0.25]) + # contact forces = 0 + obs[45:57] = 0.0 + + obs = np.clip(obs, -100., 100.) + action = sess.run(None, {input_name: obs.reshape(1, -1).astype(np.float32)})[0][0] + action = np.clip(action, -23.7, 23.7) + if 0 < args.action_ema_alpha < 1: + action = args.action_ema_alpha * action + (1 - args.action_ema_alpha) * last_a + last_a = action.copy() + + targets = DEFAULT_ANGLES + action * args.action_scale + data.ctrl[:] = np.clip(args.kp*(targets-data.qpos[7:19])-args.kd*data.qvel[6:18], -23.7, 23.7) + + mujoco.mj_step(model, data) + view.sync() + + if step % 200 == 0: + print(f"[{step}] z={data.qpos[2]:.3f} cmd=[{vx:.1f},{vy:.1f},{wz:.1f}] " + f"pos=[{data.qpos[0]:.2f},{data.qpos[1]:.2f}]") + + step += 1 + expected = (step+1)*model.opt.timestep + sleep = expected - (time.perf_counter()-t0) + if sleep > 0: time.sleep(sleep) + + kb.stop(); view.close() + +if __name__ == "__main__": + main() diff --git a/sim2sim_test_deploy.py b/deploy_45dim/sim2sim_test_deploy.py similarity index 95% rename from sim2sim_test_deploy.py rename to deploy_45dim/sim2sim_test_deploy.py index b684395..ca1e242 100644 --- a/sim2sim_test_deploy.py +++ b/deploy_45dim/sim2sim_test_deploy.py @@ -27,7 +27,7 @@ import onnxruntime as ort from mujoco import viewer HERE = os.path.dirname(os.path.abspath(__file__)) -SIM2SIM_XML = os.path.join(HERE, "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml") +SIM2SIM_XML = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml") # ─── Policy constants (matches go1_sim2sim.py + deploy_onnx_pro_sdk.py) ─── NUM_OBS = 45 @@ -227,7 +227,7 @@ class Keyboard: # ─── Main ─── def main(): parser = argparse.ArgumentParser(description="Sim2sim test for deploy_onnx_pro_sdk.py") - parser.add_argument("--onnx", default=os.path.join(HERE, "policy.onnx")) + parser.add_argument("--onnx", default=os.path.join(HERE, "policy_v3.onnx")) # PD gains parser.add_argument("--kp", type=float, default=KP_DEFAULT) @@ -244,6 +244,10 @@ def main(): # Logging parser.add_argument("--log-dir", default="", help="Enable JSONL logging to this directory") parser.add_argument("--print-every", type=int, default=200) + parser.add_argument("--action-scale", type=float, default=ACTION_SCALE, + help=f"Action scale (default {ACTION_SCALE})") + parser.add_argument("--action-ema-alpha", type=float, default=0.3, + help="EMA smoothing factor for actions (0=no filter, 1=no smooth)") args = parser.parse_args() @@ -326,10 +330,13 @@ def main(): obs = compute_obs(model, data, commands, last_actions) action = session.run(None, {input_name: obs.reshape(1, -1).astype(np.float32)})[0][0] action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS) + # EMA smooth (matches deploy --action-ema-alpha) + if 0 < args.action_ema_alpha < 1: + action = args.action_ema_alpha * action + (1 - args.action_ema_alpha) * last_actions last_actions = action.copy() # Targets (pre-safety) - targets_raw = DEFAULT_ANGLES + action * ACTION_SCALE + targets_raw = DEFAULT_ANGLES + action * args.action_scale # ── Safety layer (mirrors deploy apply_safety) ── # 1. Joint limit clipping diff --git a/deploy_wtw/__pycache__/deploy_wtw_pro_sdk.cpython-310.pyc b/deploy_wtw/__pycache__/deploy_wtw_pro_sdk.cpython-310.pyc new file mode 100644 index 0000000..e160cc1 Binary files /dev/null and b/deploy_wtw/__pycache__/deploy_wtw_pro_sdk.cpython-310.pyc differ diff --git a/deploy_wtw/__pycache__/sim2sim_wtw_test.cpython-310.pyc b/deploy_wtw/__pycache__/sim2sim_wtw_test.cpython-310.pyc new file mode 100644 index 0000000..5fa9518 Binary files /dev/null and b/deploy_wtw/__pycache__/sim2sim_wtw_test.cpython-310.pyc differ diff --git a/adaptation_module_latest.jit b/deploy_wtw/adaptation_module_latest.jit similarity index 100% rename from adaptation_module_latest.jit rename to deploy_wtw/adaptation_module_latest.jit diff --git a/body_latest.jit b/deploy_wtw/body_latest.jit similarity index 100% rename from body_latest.jit rename to deploy_wtw/body_latest.jit diff --git a/deploy_isaaclab_onnx_no_torch_wtw_r2.py b/deploy_wtw/deploy_isaaclab_onnx_no_torch_wtw_r2.py similarity index 100% rename from deploy_isaaclab_onnx_no_torch_wtw_r2.py rename to deploy_wtw/deploy_isaaclab_onnx_no_torch_wtw_r2.py diff --git a/deploy_wtw/deploy_wtw_pro_sdk.py b/deploy_wtw/deploy_wtw_pro_sdk.py new file mode 100644 index 0000000..c972dec --- /dev/null +++ b/deploy_wtw/deploy_wtw_pro_sdk.py @@ -0,0 +1,650 @@ +#!/usr/bin/env python3 +""" +deploy_wtw_pro_sdk.py + +Walk-These-Ways (RMA) deployment on Unitree Go1 PRO via go1_pro_sdk. +Direct MCU control — no LCM, no official Unitree SDK. + +Model: body_latest.jit (2102 → 12) + adaptation_module_latest.jit (2100 → 2) +Observation: 70-dim, 30-step history (2100 dims total) + +Setup: + cd /path/to/go1_pro_sdk && pip install -e . + pip install torch onnxruntime numpy + +Before running: + ssh pi@192.168.123.161 + sudo pkill -9 -f keep_sport_alive; sudo pkill -9 -f Legged_sport; sudo pkill -9 -f appTransit + +Usage: + python deploy_wtw_pro_sdk.py --kill-sport --log-dir ../logs +""" + +import argparse +import json +import signal +import subprocess +import time +from datetime import datetime +from enum import Enum +from pathlib import Path + +import numpy as np +import torch + +from go1_pro_sdk import ( + MCUClient, LowCmd, MotorCmd, MotorMode, + apply_safety, JOINT_NAMES as SDK_JOINT_NAMES, +) + +HERE = Path(__file__).parent.resolve() + +# ─── WTW policy constants ─── +NUM_OBS = 70 +NUM_ACTIONS = 12 +NUM_COMMANDS = 15 +NUM_OBS_HISTORY = 30 # 30 steps of history +OBS_HISTORY_DIM = NUM_OBS * NUM_OBS_HISTORY # 2100 +BODY_INPUT_DIM = 2102 # 2100 (history) + 2 (latent) +ADAPT_INPUT_DIM = 2100 +LATENT_DIM = 2 + +# WTW joint order: FL→FR→RL→RR (per leg: hip/thigh/calf) +# This is different from SDK order: FR→FL→RR→RL +WTW_JOINT_NAMES = [ + "FL_hip", "FL_thigh", "FL_calf", + "FR_hip", "FR_thigh", "FR_calf", + "RL_hip", "RL_thigh", "RL_calf", + "RR_hip", "RR_thigh", "RR_calf", +] + +# SDK joint order: FR→FL→RR→RL (per leg: hip/thigh/calf) +# Map: SDK index → WTW index +SDK_TO_WTW = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8], dtype=np.int64) +# Map: WTW index → SDK index +WTW_TO_SDK = np.array([3, 4, 5, 0, 1, 2, 9, 10, 11, 6, 7, 8], dtype=np.int64) + +# Default joint angles in WTW order +DEFAULT_ANGLES_WTW = np.array([ + 0.1, 0.8, -1.5, # FL + -0.1, 0.8, -1.5, # FR + 0.1, 1.0, -1.5, # RL + -0.1, 1.0, -1.5, # RR +], dtype=np.float32) + +# Default joint angles in SDK order +DEFAULT_ANGLES_SDK = DEFAULT_ANGLES_WTW[WTW_TO_SDK] + +# Observation scales (WTW standard) +OBS_SCALES = { + "lin_vel": 2.0, "ang_vel": 0.25, + "dof_pos": 1.0, "dof_vel": 0.05, + "body_height_cmd": 2.0, "footswing_height_cmd": 0.15, + "body_pitch_cmd": 0.3, "body_roll_cmd": 0.3, + "stance_width_cmd": 1.0, "stance_length_cmd": 1.0, + "aux_reward_cmd": 1.0, +} + +COMMANDS_SCALE = np.array([ + OBS_SCALES["lin_vel"], OBS_SCALES["lin_vel"], OBS_SCALES["ang_vel"], + OBS_SCALES["body_height_cmd"], 1.0, 1.0, 1.0, 1.0, 1.0, + OBS_SCALES["footswing_height_cmd"], + OBS_SCALES["body_pitch_cmd"], OBS_SCALES["body_roll_cmd"], + OBS_SCALES["stance_width_cmd"], OBS_SCALES["stance_length_cmd"], + OBS_SCALES["aux_reward_cmd"], +], dtype=np.float32)[:NUM_COMMANDS] + +ACTION_SCALE = 0.25 +HIP_SCALE_REDUCTION = 0.5 # hip joints get half action +CLIP_ACTIONS = 10.0 +CLIP_OBS = 100.0 + +EXIT = False + + +def _sig_handler(signum, frame): + global EXIT + EXIT = True + + +signal.signal(signal.SIGINT, _sig_handler) +signal.signal(signal.SIGTERM, _sig_handler) + + +# ─── State machine ─── +class State(Enum): + IDLE = "IDLE" + CALIBRATE = "CALIBRATE" + HOLD = "HOLD" + RL = "RL" + + +# ─── Quaternion math ─── +def quat_to_rot_matrix(q): + w, x, y, z = q + return np.array([ + [1 - 2*y*y - 2*z*z, 2*x*y - 2*w*z, 2*x*z + 2*w*y], + [ 2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z, 2*y*z - 2*w*x], + [ 2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y], + ], dtype=np.float32) + + +def get_projected_gravity(quaternion): + R = quat_to_rot_matrix(quaternion) + return (R.T @ np.array([0., 0., -1.], dtype=np.float32)).astype(np.float32) + + +# ─── Observation ─── +def build_commands_default(): + """Build default command vector for trotting gait.""" + cmd = np.zeros(NUM_COMMANDS, dtype=np.float32) + cmd[0:3] = [0.0, 0.0, 0.0] # vx, vy, wz + cmd[3] = 0.0 # height command (zero = nominal) + cmd[4] = 3.0 # frequency (Hz) + cmd[5] = 0.5 # phase (trot = 0.5 offset) + cmd[6] = 0.0 # offset + cmd[7] = 0.0 # bound + cmd[8] = 0.5 # duration (stance ratio) + cmd[9] = 0.15 # swing_height (matches reference footswing_height_cmd) + cmd[10] = 0.0 # body_pitch + cmd[11] = 0.0 # body_roll + cmd[12] = 0.25 # stance_width + cmd[13] = 0.4 # stance_length + cmd[14] = 0.0 # aux_reward + return cmd + + +class ClockState: + """Track gait indices and compute clock_inputs (4-dim sin per foot).""" + def __init__(self): + self.gait_indices = 0.0 + self.dt = 0.01 # 100Hz + + def step(self, commands, dt=None): + if dt is not None: + self.dt = dt + freq = commands[4] + phase = commands[5] + offset = commands[6] + bound = commands[7] if NUM_COMMANDS > 8 else 0.0 + + self.gait_indices = (self.gait_indices + self.dt * freq) % 1.0 + + foot_indices = [ + self.gait_indices + phase + offset + bound, # FL + self.gait_indices + offset, # FR + self.gait_indices + bound, # RL + self.gait_indices + phase, # RR + ] + clock = np.array([np.sin(2 * np.pi * fi) for fi in foot_indices], dtype=np.float32) + return clock + + def reset(self): + self.gait_indices = 0.0 + + +def compute_obs_wtw(imu, motor_states, commands, actions, last_actions, clock_inputs): + """Build 70-dim observation matching WTW LCM agent layout.""" + obs = np.zeros(NUM_OBS, dtype=np.float32) + + # 1. projected_gravity (3) + obs[0:3] = get_projected_gravity(imu.quaternion) + + # 2. commands * scale (15) + offset = 3 + obs[offset:offset+NUM_COMMANDS] = commands * COMMANDS_SCALE + offset += NUM_COMMANDS + + # 3. dof_pos_rel in WTW order (12) + dof_pos_sdk = np.array([motor_states[i].q for i in range(12)], dtype=np.float32) + dof_pos_wtw = dof_pos_sdk[SDK_TO_WTW] + obs[offset:offset+12] = (dof_pos_wtw - DEFAULT_ANGLES_WTW) * OBS_SCALES["dof_pos"] + offset += 12 + + # 4. dof_vel in WTW order (12) + dof_vel_sdk = np.array([motor_states[i].dq for i in range(12)], dtype=np.float32) + dof_vel_wtw = dof_vel_sdk[SDK_TO_WTW] + obs[offset:offset+12] = dof_vel_wtw * OBS_SCALES["dof_vel"] + offset += 12 + + # 5. actions clipped (12) + obs[offset:offset+12] = np.clip(actions, -CLIP_ACTIONS, CLIP_ACTIONS) + offset += 12 + + # 6. last_actions (12) + obs[offset:offset+12] = last_actions + offset += 12 + + # 7. clock_inputs (4) + obs[offset:offset+4] = clock_inputs + + obs = np.clip(obs, -CLIP_OBS, CLIP_OBS) + obs = np.nan_to_num(obs, nan=0.0, posinf=0.0, neginf=0.0) + return obs + + +# ─── Model ─── +class WTWPolicy: + def __init__(self, body_path, adapt_path): + self.body = torch.jit.load(str(body_path), map_location='cpu') + self.adapt = torch.jit.load(str(adapt_path), map_location='cpu') + self.body.eval() + self.adapt.eval() + + self.obs_history = torch.zeros(1, OBS_HISTORY_DIM, dtype=torch.float) + self.latent = torch.zeros(1, LATENT_DIM, dtype=torch.float) + + print(f"[INFO] WTW body: {body_path}") + print(f"[INFO] WTW adapt: {adapt_path}") + print(f"[INFO] History: {NUM_OBS} obs × {NUM_OBS_HISTORY} steps = {OBS_HISTORY_DIM}") + + def reset(self): + self.obs_history.zero_() + self.latent.zero_() + + def __call__(self, obs): + obs_t = torch.from_numpy(obs.reshape(1, -1)).float() + + # Update history: shift left, append new obs + self.obs_history = torch.cat( + (self.obs_history[:, NUM_OBS:], obs_t), dim=-1) + + # Adaptation module: history → latent + with torch.no_grad(): + self.latent = self.adapt(self.obs_history) + + # Body: [history, latent] → action + body_input = torch.cat((self.obs_history, self.latent), dim=-1) + with torch.no_grad(): + action = self.body(body_input) + + return action.numpy().flatten().astype(np.float32) + + +# ─── Remote controller ─── +def get_rc_commands(state, base_cmd, args): + r = state.remote + base_cmd[0] = r.ly * args.rc_vx_scale + base_cmd[1] = -r.lx * args.rc_vy_scale + base_cmd[2] = -r.rx * args.rc_wz_scale + return base_cmd + + +class RCEdgeDetector: + def __init__(self): + self._prev = set() + + def update(self, state): + current = set(state.remote.pressed) + rising = current - self._prev + falling = self._prev - current + self._prev = current + return rising, falling + + +# ─── Safety wrappers ─── +def send_hold_cmd(client, state, args): + cmd = LowCmd() + for j in range(12): + cmd.set_motor(j, MotorCmd( + mode=MotorMode.Servo, + q=float(DEFAULT_ANGLES_SDK[j]), dq=0.0, tau=0.0, + Kp=args.kp, Kd=args.kd, + )) + apply_safety(cmd, state, power_factor=args.power_factor, + position_limit_on=True, position_protect_limit=None) + client.send(cmd) + + +def send_rl_cmd(client, state, action_wtw, args): + """Convert WTW action to SDK targets and send.""" + # Scale action → position offset in WTW order + offset_wtw = action_wtw * ACTION_SCALE + # Apply hip scale reduction + for i in [0, 3, 6, 9]: # hip indices in WTW order + offset_wtw[i] *= HIP_SCALE_REDUCTION + # Target in WTW order + targets_wtw = DEFAULT_ANGLES_WTW + offset_wtw + # Convert to SDK order + targets_sdk = targets_wtw[WTW_TO_SDK] + + cmd = LowCmd() + for j in range(12): + cmd.set_motor(j, MotorCmd( + mode=MotorMode.Servo, + q=float(targets_sdk[j]), dq=0.0, tau=0.0, + Kp=args.kp, Kd=args.kd, + )) + pp_limit = args.position_protect_limit if args.position_protect_limit > 0 else None + apply_safety(cmd, state, power_factor=args.power_factor, + position_limit_on=True, position_protect_limit=pp_limit) + client.send(cmd) + + +def kill_sport_processes(host, user): + cmds = [ + "sudo pkill -9 -f keep_sport_alive", + "sudo pkill -9 -f Legged_sport", + "sudo pkill -9 -f appTransit", + ] + ssh_target = f"{user}@{host}" + print(f"[INFO] Killing sport processes on {ssh_target}...") + try: + result = subprocess.run( + ["ssh", ssh_target, " && ".join(cmds)], + capture_output=True, text=True, timeout=15) + if result.returncode == 0 or "no process" in result.stderr.lower(): + print("[INFO] Sport processes killed.") + return True + print(f"[WARN] SSH returned {result.returncode}: {result.stderr.strip()}") + return False + except Exception as e: + print(f"[WARN] Failed: {e}") + return False + + +# ─── JSONL logger ─── +class JsonlLogger: + def __init__(self, log_dir, args): + self.enabled = bool(log_dir) + self.fp = None + self.run_dir = None + self.flush_every = 50 + if not self.enabled: + return + ts = datetime.now().strftime("%Y%m%d_%H%M%S") + self.run_dir = Path(log_dir).expanduser().resolve() / f"wtw_deploy_{ts}" + self.run_dir.mkdir(parents=True, exist_ok=True) + meta = { + "created_at": ts, "num_obs": NUM_OBS, "num_actions": NUM_ACTIONS, + "num_commands": NUM_COMMANDS, "num_obs_history": NUM_OBS_HISTORY, + "action_scale": ACTION_SCALE, + "default_angles_wtw": DEFAULT_ANGLES_WTW.tolist(), + "default_angles_sdk": DEFAULT_ANGLES_SDK.tolist(), + "joint_names_wtw": WTW_JOINT_NAMES, + "joint_names_sdk": list(SDK_JOINT_NAMES), + } + for k, v in vars(args).items(): + if isinstance(v, (str, int, float, bool, type(None))): + meta[k] = v + (self.run_dir / "metadata.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False)) + self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8", buffering=1) + print(f"[INFO] Log dir: {self.run_dir}") + + def log(self, step, **kw): + if not self.enabled: + return + rec = {"step": int(step), "time_wall": time.time()} + for k, v in kw.items(): + if isinstance(v, np.ndarray): + rec[k] = np.asarray(v, dtype=np.float32).reshape(-1).tolist() + elif isinstance(v, (np.float32, np.float64)): + rec[k] = float(v) + elif isinstance(v, (np.int32, np.int64)): + rec[k] = int(v) + else: + rec[k] = v + self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n") + if step % self.flush_every == 0: + self.fp.flush() + + def close(self): + if self.fp: + self.fp.flush(); self.fp.close() + print(f"[INFO] Log saved: {self.run_dir}") + + +# ─── Ramp to default ─── +def ramp_to_default(client, args, state): + print("[INFO] Ramping to default pose (~2s)...") + current_sdk = np.array([state.motorState[i].q for i in range(12)], dtype=np.float32) + error = current_sdk - DEFAULT_ANGLES_SDK + + if np.max(np.abs(error)) < 0.05: + print("[INFO] Already near default pose.") + return state + + ramp_steps = 200 + step_err = error / ramp_steps + + for i in range(ramp_steps): + if EXIT: return state + new_state = client.recv_latest() + if new_state is not None: + state = new_state + targets = DEFAULT_ANGLES_SDK + (error - step_err * min(i + 1, ramp_steps)) + cmd = LowCmd() + for j in range(12): + cmd.set_motor(j, MotorCmd( + mode=MotorMode.Servo, + q=float(targets[j]), dq=0.0, tau=0.0, + Kp=args.kp_cal, Kd=args.kd_cal, + )) + apply_safety(cmd, state, power_factor=args.power_factor, + position_limit_on=True, position_protect_limit=None) + client.send(cmd) + time.sleep(0.01) + if i % 50 == 0: + actual = np.array([state.motorState[j].q for j in range(12)]) + print(f" ramp {i}/{ramp_steps} target_err={np.max(np.abs(targets-DEFAULT_ANGLES_SDK)):.3f} actual_err={np.max(np.abs(actual-DEFAULT_ANGLES_SDK)):.3f}") + + for _ in range(50): + if EXIT: return state + new_state = client.recv_latest() + if new_state is not None: + state = new_state + send_hold_cmd(client, state, args) + time.sleep(0.01) + + print("[INFO] Default pose reached.") + return state + + +# ─── Main ─── +def main(): + parser = argparse.ArgumentParser(description="WTW RMA deployment on Go1 PRO via go1_pro_sdk") + parser.add_argument("--body", default=str(HERE / "body_latest.jit")) + parser.add_argument("--adapt", default=str(HERE / "adaptation_module_latest.jit")) + + parser.add_argument("--kill-sport", action="store_true") + parser.add_argument("--pi-host", default="192.168.123.161") + parser.add_argument("--pi-user", default="pi") + + parser.add_argument("--kp", type=float, default=20.0) + parser.add_argument("--kd", type=float, default=0.5) + parser.add_argument("--kp-cal", type=float, default=15.0) + parser.add_argument("--kd-cal", type=float, default=0.5) + parser.add_argument("--power-factor", type=int, default=7) + parser.add_argument("--position-protect-limit", type=float, default=1.0) + + parser.add_argument("--rc-vx-scale", type=float, default=1.0) + parser.add_argument("--rc-vy-scale", type=float, default=1.0) + parser.add_argument("--rc-wz-scale", type=float, default=1.0) + parser.add_argument("--cmd-x", type=float, default=0.0) + parser.add_argument("--cmd-y", type=float, default=0.0) + parser.add_argument("--cmd-yaw", type=float, default=0.0) + + parser.add_argument("--rate-hz", type=float, default=50.0) + parser.add_argument("--warmup-steps", type=int, default=50) + parser.add_argument("--max-steps", type=int, default=0) + parser.add_argument("--print-every", type=int, default=50) + + parser.add_argument("--log-dir", default="") + parser.add_argument("--log-flush-every", type=int, default=50) + + args = parser.parse_args() + + print(""" +╔══════════════════════════════════════════════════════════════╗ +║ WTW RMA Deployment (go1_pro_sdk direct MCU) ║ +║ 1. Robot SUSPENDED ║ +║ 2. Use --kill-sport to auto-kill Pi processes ║ +║ 3. R2=go/stop, L2=estop, Left-stick=move, Right-stick=turn ║ +╚══════════════════════════════════════════════════════════════╝ +""") + input("Press Enter when ready...") + + if args.kill_sport: + kill_sport_processes(args.pi_host, args.pi_user) + + print("[INFO] Loading WTW models...") + policy = WTWPolicy(args.body, args.adapt) + + print("[INFO] Connecting to MCU...") + client = MCUClient() + logger = JsonlLogger(args.log_dir, args) + + try: + print("[INFO] Waking MCU...") + client.wake_mcu(n_frames=50, dt=0.01) + + state = client.recv_state(timeout=2.0) + if state is None: + print("[ERROR] No state received.") + return 1 + + print(f"[INFO] Connected. Battery={state.bms.SOC}%") + print(f"[INFO] RPY: {np.round(np.degrees(state.imu.rpy), 1)} deg") + + sm_state = State.IDLE + edge = RCEdgeDetector() + edge.update(state) + clock = ClockState() + base_cmd = build_commands_default() + actions = np.zeros(NUM_ACTIONS, dtype=np.float32) + last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32) + step = 0 + dt = 1.0 / args.rate_hz + next_t = time.perf_counter() + + print(f"[INFO] Rate: {args.rate_hz}Hz. Ctrl+C to exit.") + print(f"[INFO] State machine: IDLE → (R2) → CALIBRATE → HOLD → (R2) → RL") + + while not EXIT: + new_state = client.recv_latest() + if new_state is not None: + state = new_state + if state is None: + time.sleep(0.001) + continue + + rising, falling = edge.update(state) + r2_rose = "R2" in rising + l2_rose = "L2" in rising + + # Emergency stop + if l2_rose and sm_state != State.IDLE: + print(f"\n[L2 EMERGENCY] {sm_state.value} → IDLE") + sm_state = State.IDLE + actions[:] = 0.0 + last_actions[:] = 0.0 + policy.reset() + clock.reset() + client.send(LowCmd()) + + # State machine + if sm_state == State.IDLE: + if step % 10 == 0: + client.send(LowCmd()) + + if r2_rose: + print("\n[R2] IDLE → CALIBRATE") + sm_state = State.CALIBRATE + state = ramp_to_default(client, args, state) + if EXIT: break + sm_state = State.HOLD + print("[STATE] → HOLD") + + elif sm_state == State.HOLD: + send_hold_cmd(client, state, args) + + if r2_rose: + print("\n[R2] HOLD → RL") + sm_state = State.RL + actions[:] = 0.0 + last_actions[:] = 0.0 + policy.reset() + clock.reset() + + elif sm_state == State.RL: + if r2_rose: + print("\n[R2] RL → HOLD") + sm_state = State.HOLD + actions[:] = 0.0 + last_actions[:] = 0.0 + state = ramp_to_default(client, args, state) + if EXIT: break + continue + + # Update commands from RC + base_cmd = get_rc_commands(state, base_cmd, args) + + # Clock inputs + clock_inputs = clock.step(base_cmd, dt) + + # Observation + obs = compute_obs_wtw(state.imu, state.motorState, + base_cmd, actions, last_actions, clock_inputs) + + # Inference + action_raw = policy(obs) + actions = np.clip(action_raw, -CLIP_ACTIONS, CLIP_ACTIONS).astype(np.float32) + last_actions = actions.copy() + + if step >= args.warmup_steps: + send_rl_cmd(client, state, action_raw, args) + + logger.log( + step, mode="RL", + commands=base_cmd, + obs_wtw=obs, + action_raw=action_raw, + action_safe=actions, + dof_pos_sdk=np.array([state.motorState[i].q for i in range(12)]), + dof_vel_sdk=np.array([state.motorState[i].dq for i in range(12)]), + clock_inputs=clock_inputs, + imu_rpy_deg=np.degrees(state.imu.rpy), + rc_buttons=state.remote.pressed, + ) + else: + logger.log( + step, mode=sm_state.value, + dof_pos_sdk=np.array([state.motorState[i].q for i in range(12)]), + imu_rpy_deg=np.degrees(state.imu.rpy), + rc_buttons=state.remote.pressed, + ) + + # Status print + if step % args.print_every == 0: + dof_pos = np.array([state.motorState[i].q for i in range(12)]) + print(f"\n[STEP {step}] state={sm_state.value} bat={state.bms.SOC}% " + f"rpy={np.round(np.degrees(state.imu.rpy), 1)}") + print(f" RC: lx={state.remote.lx:+.2f} ly={state.remote.ly:+.2f} " + f"btns={state.remote.pressed}") + print(f" joint: {np.round(dof_pos, 2)}") + if sm_state == State.RL: + print(f" action max: {np.max(np.abs(actions)):.2f}") + + step += 1 + if args.max_steps > 0 and step >= args.max_steps: + print("[INFO] max_steps reached.") + break + + next_t += dt + sleep = next_t - time.perf_counter() + if sleep > 0: + time.sleep(sleep) + else: + next_t = time.perf_counter() + + finally: + if logger is not None: + logger.close() + print("[INFO] Safe stopping...") + client.safe_stop(n_frames=50, dt=0.002) + client.close() + print("[INFO] Done.") + + +if __name__ == "__main__": + main() diff --git a/deploy_wtw/sim2sim_wtw_test.py b/deploy_wtw/sim2sim_wtw_test.py new file mode 100644 index 0000000..cf04af9 --- /dev/null +++ b/deploy_wtw/sim2sim_wtw_test.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +""" +sim2sim test for WTW deploy_wtw_pro_sdk.py using MuJoCo Go1 XML. +Parameters verified against go1_walk_these_ways_inference.py reference. + +Usage: + conda activate free_dog_sdk + mjpython sim2sim_wtw_test.py + +Controls: W/S=前后 Q/E=左右 A/D=旋转 Space=停 R=重置 1-4=步态 Esc=退出 +""" + +import os, signal, time, queue, threading +import mujoco, numpy as np, torch +from mujoco import viewer + +HERE = os.path.dirname(os.path.abspath(__file__)) + +# ─── Paths ─── +XML_PATH = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml") +MESH_DIR = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "meshes") +BODY_JIT = os.path.join(HERE, "body_latest.jit") +ADAPT_JIT = os.path.join(HERE, "adaptation_module_latest.jit") + +# ─── WTW constants (verified against reference) ─── +NUM_OBS = 70 +NUM_ACTIONS = 12 +NUM_COMMANDS = 15 +NUM_OBS_HISTORY = 30 +OBS_BUFFER_SIZE = 2100 + +ACTION_SCALE = 0.25 +HIP_SCALE_REDUCTION = 0.5 +CLIP_ACTIONS = 10.0 +CLIP_OBS = 100.0 + +# Joint orders: +# MuJoCo/SDK: [FR_hip,FR_thigh,FR_calf, FL_hip,FL_thigh,FL_calf, RR_hip,RR_thigh,RR_calf, RL_hip,RL_thigh,RL_calf] +# WTW/Deploy: [FL_hip,FL_thigh,FL_calf, FR_hip,FR_thigh,FR_calf, RL_hip,RL_thigh,RL_calf, RR_hip,RR_thigh,RR_calf] +DEPLOY_TO_MUJOCO = np.array([3,4,5, 0,1,2, 9,10,11, 6,7,8], dtype=np.int64) + +# Default angles in WTW order (from reference) +DEFAULT_WTW = np.array([ + 0.1, 0.8, -1.5, # FL + -0.1, 0.8, -1.5, # FR + 0.1, 1.0, -1.5, # RL + -0.1, 1.0, -1.5, # RR +], dtype=np.float32) +DEFAULT_MUJOCO = np.array([ + -0.1, 0.8, -1.5, # FR + 0.1, 0.8, -1.5, # FL + -0.1, 1.0, -1.5, # RR + 0.1, 1.0, -1.5, # RL +], dtype=np.float32) + +# Commands scale (verified) +COMMANDS_SCALE = np.array([ + 2.0, 2.0, 0.25, # vx, vy, wz + 2.0, # body_height + 1, 1, 1, 1, 1, # freq, phase, offset, bound, duration + 0.15, # footswing_height + 0.3, 0.3, # body_pitch, body_roll + 1.0, 1.0, # stance_width, stance_length + 1.0, # aux_reward +], dtype=np.float32)[:15] + +OBS_SCALES = {"dof_pos":1.0, "dof_vel":0.05} + +# PD gains (from reference) +KP = 20.0; KD = 0.1 # matches reference (XML passive damping=1.0, total≈1.1) + +# Friction (training had zero floor friction) +FLOOR_FRICTION = [0.0, 0.0, 0.0] +BODY_FRICTION = [0.6, 0.3, 0.3] + +# Gait presets +GAITS = { + '1': ('Trot', 0.5, 0.0, 0.0), + '2': ('Pace', 0.0, 0.0, 0.5), + '3': ('Bound', 0.0, 0.5, 0.0), + '4': ('Pronk', 0.0, 0.0, 0.0), +} + +EXIT = False +def _sig(s, f): global EXIT; EXIT = True +signal.signal(signal.SIGINT, _sig) + + +class Keyboard: + def __init__(self): + self.held = set(); self._l = None + def _p(self, k): + try: self.held.add(k.char.lower()) + except: self.held.add(str(k)) + def _r(self, k): + try: self.held.discard(k.char.lower()) + except: self.held.discard(str(k)) + def init(self): + from pynput import keyboard + self._l = keyboard.Listener(on_press=self._p, on_release=self._r); self._l.start() + def keys(self): return self.held.copy() + def stop(self): + if self._l: self._l.stop() + + +def main(): + if not os.path.exists(BODY_JIT): + print(f"[ERROR] body not found: {BODY_JIT}"); return + if not os.path.exists(ADAPT_JIT): + print(f"[ERROR] adapt not found: {ADAPT_JIT}"); return + + # Load MuJoCo with mesh path fix + with open(XML_PATH) as f: + xml = f.read() + xml = xml.replace('meshdir="../meshes/"', f'meshdir="{MESH_DIR}"') + model = mujoco.MjModel.from_xml_string(xml) + data = mujoco.MjData(model) + + # Set friction + for i in range(model.ngeom): + model.geom_friction[i] = BODY_FRICTION + + print(f"[INFO] MuJoCo: {model.nbody} bodies, {model.nq} DoF, KP={KP}, KD(active)={KD}") + + # Load models + body = torch.jit.load(BODY_JIT, map_location='cpu').eval() + adapt = torch.jit.load(ADAPT_JIT, map_location='cpu').eval() + print(f"[INFO] WTW body+adapt loaded") + + # Init + data.qpos[0:3] = [0, 0, 0.35] + data.qpos[3:7] = [1, 0, 0, 0] + data.qpos[7:19] = DEFAULT_MUJOCO + mujoco.mj_forward(model, data) + + obs_buffer = np.zeros(OBS_BUFFER_SIZE, dtype=np.float32) + prev_action = np.zeros(12, dtype=np.float32) + last_action = np.zeros(12, dtype=np.float32) + gait_idx = 0.0 + + step, vx, vy, wz = 0, 0.0, 0.0, 0.0 + fh = 0.15 # footswing height + bp, br = 0.0, 0.0 # body pitch/roll + gait_phase, gait_offset, gait_bound, gait_dur = 0.5, 0.0, 0.0, 0.5 + current_gait = '1' + ctrl_dt = 0.02 + + kb = Keyboard(); kb.init() + view = viewer.launch_passive(model, data) + + print("[INFO] W/S=前后 Q/E=左右 A/D=旋转 1-4=步态 R=重置 Esc=退出") + + t0 = time.perf_counter() + + while view.is_running() and not EXIT: + keys = kb.keys() + if 'key.esc' in keys: break + + if 'r' in keys: + data.qpos[0:3]=[0,0,0.35]; data.qpos[3:7]=[1,0,0,0] + data.qpos[7:19]=DEFAULT_MUJOCO; data.qvel[:]=0 + obs_buffer[:]=0; prev_action[:]=0; last_action[:]=0; gait_idx=0 + mujoco.mj_forward(model, data) + + # Gait switch + for k, (name, ph, off, bd) in GAITS.items(): + if k in keys and k != current_gait: + current_gait = k + gait_phase, gait_offset, gait_bound = ph, off, bd + print(f"[INFO] Gait: {name}") + + vx=1.0 if 'w' in keys else (-1.0 if 's' in keys else 0.0) + vy=1.0 if 'q' in keys else (-1.0 if 'e' in keys else 0.0) + wz=3.0 if 'a' in keys else (-3.0 if 'd' in keys else 0.0) + if ' ' in keys: vx=vy=wz=0.0 + + # Inference at 50Hz (every 10 sim steps at dt=0.002) + if step % 10 == 0: + # Commands + raw_cmd = np.zeros(15, dtype=np.float32) + raw_cmd[0]=vx; raw_cmd[1]=vy; raw_cmd[2]=wz + raw_cmd[3]=0.0; raw_cmd[4]=3.0 + raw_cmd[5]=gait_phase; raw_cmd[6]=gait_offset; raw_cmd[7]=gait_bound; raw_cmd[8]=gait_dur + raw_cmd[9]=0.15; raw_cmd[10]=bp; raw_cmd[11]=br + raw_cmd[12]=0.25; raw_cmd[13]=0.4 + commands = raw_cmd * COMMANDS_SCALE + + # Gait index & clock + gait_idx += 0.02 * 3.0 + if gait_idx > 1.0: gait_idx -= 1.0 + p, o, b = gait_phase, gait_offset, gait_bound + fi = [gait_idx+p+o+b, gait_idx+o, gait_idx+b, gait_idx+p] + clock = np.array([np.sin(2*np.pi*f) for f in fi], dtype=np.float32) + + # Observation + obs = np.zeros(NUM_OBS, dtype=np.float32) + base_rot = data.xmat[1].reshape(3, 3) + obs[0:3] = (base_rot.T @ np.array([0., 0., -1.], dtype=np.float64)).astype(np.float32) + obs[3:18] = commands + dof_wtw = data.qpos[7:19][DEPLOY_TO_MUJOCO] + obs[18:30] = (dof_wtw - DEFAULT_WTW) * 1.0 + obs[30:42] = data.qvel[6:18][DEPLOY_TO_MUJOCO] * 0.05 + obs[42:54] = np.clip(prev_action, -CLIP_ACTIONS, CLIP_ACTIONS) + obs[54:66] = np.clip(last_action, -CLIP_ACTIONS, CLIP_ACTIONS) + obs[66:70] = clock + obs = np.clip(obs, -CLIP_OBS, CLIP_OBS) + + # History buffer + obs_buffer = np.concatenate([obs_buffer[NUM_OBS:], obs]) + + # Inference + obs_hist = torch.from_numpy(obs_buffer).float().unsqueeze(0) + with torch.inference_mode(): + latent = adapt(obs_hist) + action = body(torch.cat([obs_hist, latent], dim=1)).numpy().flatten() + action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS) + last_action = prev_action.copy() + prev_action = action.copy() + + # PD control + action_scaled = prev_action * ACTION_SCALE + for i in [0,3,6,9]: action_scaled[i] *= HIP_SCALE_REDUCTION + targets_mujoco = action_scaled[DEPLOY_TO_MUJOCO] + DEFAULT_MUJOCO + torques = KP*(targets_mujoco - data.qpos[7:19]) - KD*data.qvel[6:18] + data.ctrl[:] = np.clip(torques, -23.7, 23.7) + + mujoco.mj_step(model, data) + view.sync() + + if step % 200 == 0: + print(f"[STEP {step}] z={data.qpos[2]:.3f} cmd=[{vx:.1f},{vy:.1f},{wz:.1f}] " + f"pos=[{data.qpos[0]:.2f},{data.qpos[1]:.2f}]") + + step += 1 + expected = (step+1)*model.opt.timestep + sleep = expected - (time.perf_counter()-t0) + if sleep > 0: time.sleep(sleep) + + kb.stop(); view.close() + print("[INFO] Done.") + +if __name__ == "__main__": + main() diff --git a/test_deploy.md b/test_deploy.md index 7190fb0..613df1b 100644 --- a/test_deploy.md +++ b/test_deploy.md @@ -7,12 +7,31 @@ conda activate free_dog_sdk cd /Users/chenyouyuan/cyy_ws/deploy_go1_pro ``` +## 目录结构 + +``` +deploy_go1_pro/ +├── deploy_45dim/ # 45-dim 策略部署 +│ ├── deploy_onnx_pro_sdk.py +│ ├── policy.onnx +│ ├── go1_sim2sim.py +│ └── sim2sim_test_deploy.py +├── deploy_wtw/ # Walk-These-Ways 部署（开发中） +│ ├── deploy_isaaclab_onnx_no_torch_wtw_r2.py (LCM 版本) +│ ├── body_latest.jit +│ └── adaptation_module_latest.jit +├── sim2sim_mujoco_example/ +├── logs/ +└── test_deploy.md +``` + --- -## 推荐部署命令 +## 推荐部署命令（45-dim） ```bash -python deploy_onnx_pro_sdk.py --onnx policy.onnx --kill-sport --log-dir logs \ +cd deploy_45dim +python deploy_onnx_pro_sdk.py --onnx policy.onnx --kill-sport --log-dir ../logs \ --kp 80 --kd 1 --kp-cal 20 --kd-cal 1 \ --power-factor 7 --position-protect-limit 1.0 \ --action-ema-alpha 0.3 --print-every 10 @@ -25,7 +44,7 @@ python deploy_onnx_pro_sdk.py --onnx policy.onnx --kill-sport --log-dir logs \ ### Step 0 — 观测数据查验 ✓ ```bash -python deploy_onnx_pro_sdk.py --obs-check --kill-sport --log-dir logs +cd deploy_45dim && python deploy_onnx_pro_sdk.py --obs-check --kill-sport --log-dir ../logs ``` | obs 段 | 含义 | 期望 | 结果 | @@ -42,7 +61,7 @@ python deploy_onnx_pro_sdk.py --obs-check --kill-sport --log-dir logs ### Step 1 — 遥控器测试 ✓ ```bash -python deploy_onnx_pro_sdk.py --monitor --kill-sport --log-dir logs +cd deploy_45dim && python deploy_onnx_pro_sdk.py --monitor --kill-sport --log-dir ../logs ``` 全部按钮（R2/L2/L1/A/B/X/Y/START）上升沿/下降沿检测正常，摇杆四轴 ±1.0。