wtw，部署
This commit is contained in:
cyy_mac
897
deploy_45dim/deploy_onnx_pro_sdk.py
Normal file
897
deploy_45dim/deploy_onnx_pro_sdk.py
Normal file
@@ -0,0 +1,897 @@
|
||||
#!/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="90k_45.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()
|
||||
Reference in New Issue
Block a user