898 lines
34 KiB
Python
898 lines
34 KiB
Python
#!/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 = 45
|
|
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):
|
|
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 = 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="policy.onnx", help="Path to ONNX model")
|
|
|
|
# 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.3,
|
|
help="EMA smoothing factor for actions (0=no smooth, 1=no filter, default 0.3)")
|
|
|
|
# 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()
|