#!/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()