wtw,部署
This commit is contained in:
BIN
deploy_45dim/.DS_Store
vendored
Normal file
BIN
deploy_45dim/.DS_Store
vendored
Normal file
Binary file not shown.
BIN
deploy_45dim/90k_45.onnx
Normal file
BIN
deploy_45dim/90k_45.onnx
Normal file
Binary file not shown.
BIN
deploy_45dim/__pycache__/deploy_57dim_pro_sdk.cpython-310.pyc
Normal file
BIN
deploy_45dim/__pycache__/deploy_57dim_pro_sdk.cpython-310.pyc
Normal file
Binary file not shown.
Binary file not shown.
BIN
deploy_45dim/__pycache__/sim2sim_57dim_test.cpython-310.pyc
Normal file
BIN
deploy_45dim/__pycache__/sim2sim_57dim_test.cpython-310.pyc
Normal file
Binary file not shown.
BIN
deploy_45dim/__pycache__/sim2sim_test_deploy.cpython-310.pyc
Normal file
BIN
deploy_45dim/__pycache__/sim2sim_test_deploy.cpython-310.pyc
Normal file
Binary file not shown.
899
deploy_45dim/deploy_57dim_pro_sdk.py
Normal file
899
deploy_45dim/deploy_57dim_pro_sdk.py
Normal file
@@ -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()
|
||||
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()
|
||||
236
deploy_45dim/go1_sim2sim.py
Normal file
236
deploy_45dim/go1_sim2sim.py
Normal file
@@ -0,0 +1,236 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Go1 sim2sim MuJoCo viewer — Original 30k flat training.
|
||||
|
||||
Usage: python go1_sim2sim.py
|
||||
|
||||
Requires: mujoco, onnxruntime, pynput
|
||||
Install: pip install mujoco onnxruntime pynput
|
||||
|
||||
Controls:
|
||||
W/S: forward/back Q/E: strafe left/right
|
||||
A/D: rotate Space: stop R: reset Esc: quit
|
||||
"""
|
||||
import numpy as np, mujoco, onnxruntime as ort, os, time, threading, queue
|
||||
from mujoco import viewer
|
||||
from pynput import keyboard
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
ONNX = os.path.join(HERE, "90k_45.onnx")
|
||||
|
||||
# ── Parameters (original MotrixLab Go1 config) ──
|
||||
NUM_OBS = 45
|
||||
KP, KD = 80.0, 0.5 # KD=0.5 + MuJoCo joint_damping(0.5) = 1.0 = training kd
|
||||
ACTION_SCALE = 0.05
|
||||
CLIP = 23.7
|
||||
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)
|
||||
|
||||
# ── Keyboard ──
|
||||
class KB:
|
||||
def __init__(s):
|
||||
s._q = queue.Queue(); s.running = True; s.held = set()
|
||||
def _n(s, k):
|
||||
try:
|
||||
if hasattr(k, 'char') and k.char: return k.char.lower()
|
||||
except: pass
|
||||
return str(k).lower()
|
||||
def _w(s):
|
||||
while s.running:
|
||||
try:
|
||||
et, k = s._q.get(timeout=0.05)
|
||||
n = s._n(k)
|
||||
if et == 'press': s.held.add(n)
|
||||
elif et == 'release': s.held.discard(n)
|
||||
except queue.Empty: pass
|
||||
def init(s):
|
||||
s._l = keyboard.Listener(
|
||||
on_press=lambda k: s._q.put(('press', k)),
|
||||
on_release=lambda k: s._q.put(('release', k)))
|
||||
s._l.start()
|
||||
s._t = threading.Thread(target=s._w, daemon=True); s._t.start()
|
||||
def keys(s): return s.held.copy()
|
||||
def stop(s): s.running = False; s._l.stop()
|
||||
|
||||
# ── Main ──
|
||||
def main():
|
||||
# The model XML is embedded below
|
||||
xml = '''<mujoco model="go1 scene">
|
||||
<compiler angle="radian" autolimits="true"/>
|
||||
<option timestep="0.005" integrator="Euler" iterations="60">
|
||||
<flag eulerdamp="disable"/>
|
||||
</option>
|
||||
<custom>
|
||||
<numeric data="30" name="max_contact_points"/>
|
||||
<numeric data="12" name="max_geom_pairs"/>
|
||||
</custom>
|
||||
<default>
|
||||
<default class="go1">
|
||||
<geom condim="1"/>
|
||||
<joint axis="0 1 0" armature="0.005" damping="0.5"/>
|
||||
<default class="abduction">
|
||||
<joint axis="1 0 0" range="-0.863 0.863" frictionloss="0.3"/>
|
||||
</default>
|
||||
<default class="hip">
|
||||
<joint range="-0.686 4.501" frictionloss="0.3"/>
|
||||
</default>
|
||||
<default class="knee">
|
||||
<joint range="-2.818 -0.888" frictionloss="1.0"/>
|
||||
</default>
|
||||
</default>
|
||||
</default>
|
||||
<asset>
|
||||
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
|
||||
<texture name="ground" type="2d" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
|
||||
<material name="ground" texture="ground" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
|
||||
</asset>
|
||||
<worldbody>
|
||||
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
|
||||
<geom name="floor" size="0 0 0.01" type="plane" material="ground" contype="1" conaffinity="0" priority="1" friction="0.6" condim="3"/>
|
||||
<body name="trunk" pos="0 0 0.4" childclass="go1">
|
||||
<freejoint/>
|
||||
<inertial pos="0.0223 0.002 -0.0005" quat="-0.00342088 0.705204 0.000106698 0.708996" mass="5.204" diaginertia="0.0716565 0.0630105 0.0168101"/>
|
||||
<geom name="trunk_geom" contype="0" conaffinity="0" group="2" type="box" size="0.35 0.12 0.08" rgba="0.4 0.4 0.4 1"/>
|
||||
<geom name="trunk_col" contype="1" conaffinity="1" group="3" pos="0.24 0 0" size="0.05 0.05 0.05" type="box"/>
|
||||
<site name="imu" pos="-0.01592 -0.06659 -0.00617" group="5"/>
|
||||
<!-- FR leg -->
|
||||
<body name="FR_hip" pos="0.1881 -0.04675 0">
|
||||
<joint class="abduction" name="FR_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FR_thigh" pos="0 -0.08 0">
|
||||
<joint class="hip" name="FR_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FR_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="FR_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="FR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<!-- FL leg -->
|
||||
<body name="FL_hip" pos="0.1881 0.04675 0">
|
||||
<joint class="abduction" name="FL_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FL_thigh" pos="0 0.08 0">
|
||||
<joint class="hip" name="FL_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="FL_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="FL_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="FL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<!-- RR leg -->
|
||||
<body name="RR_hip" pos="-0.1881 -0.04675 0">
|
||||
<joint class="abduction" name="RR_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RR_thigh" pos="0 -0.08 0">
|
||||
<joint class="hip" name="RR_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RR_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="RR_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="RR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
<!-- RL leg -->
|
||||
<body name="RL_hip" pos="-0.1881 0.04675 0">
|
||||
<joint class="abduction" name="RL_hip_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RL_thigh" pos="0 0.08 0">
|
||||
<joint class="hip" name="RL_thigh_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<body name="RL_calf" pos="0 0 -0.213">
|
||||
<joint class="knee" name="RL_calf_joint"/>
|
||||
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
|
||||
<geom name="RL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</body>
|
||||
</worldbody>
|
||||
<actuator>
|
||||
<motor class="abduction" name="FR_hip" joint="FR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="FR_calf" joint="FR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="abduction" name="FL_hip" joint="FL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="FL_calf" joint="FL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="abduction" name="RR_hip" joint="RR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="RR_calf" joint="RR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="abduction" name="RL_hip" joint="RL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
<motor class="knee" name="RL_calf" joint="RL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
|
||||
</actuator>
|
||||
<sensor>
|
||||
<gyro site="imu" name="gyro"/>
|
||||
<velocimeter site="imu" name="local_linvel"/>
|
||||
</sensor>
|
||||
</mujoco>'''
|
||||
|
||||
model = mujoco.MjModel.from_xml_string(xml)
|
||||
data = mujoco.MjData(model)
|
||||
data.qpos[0:3] = [0, 0, 0.42]
|
||||
data.qpos[3:7] = [1, 0, 0, 0]
|
||||
data.qpos[7:19] = DEFAULT_ANGLES
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider'])
|
||||
print(f"[Go1 sim2sim] ONNX={ONNX}")
|
||||
print(f"[Go1 sim2sim] PD kp={KP} kd={KD+0.5} action_scale={ACTION_SCALE} obs={NUM_OBS}-dim")
|
||||
print(f"[Go1 sim2sim] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
|
||||
|
||||
kb = KB(); kb.init()
|
||||
view = viewer.launch_passive(model, data)
|
||||
step, vx, vy, wz = 0, 0.0, 0.0, 0.0
|
||||
last_a = np.zeros(12, dtype=np.float32)
|
||||
|
||||
while view.is_running():
|
||||
keys = kb.keys()
|
||||
if 'escape' in keys: break
|
||||
if 'r' in keys:
|
||||
data.qpos[0: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 % 2 == 0: # 100Hz control (MuJoCo dt=0.005)
|
||||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "gyro")
|
||||
adr = model.sensor_adr[sid]
|
||||
obs[0:3] = data.sensordata[adr:adr+3] * 0.25
|
||||
R = data.xmat[1].reshape(3, 3)
|
||||
obs[3:6] = (R.T @ np.array([0., 0., -1.])).astype(np.float32)
|
||||
obs[6:18] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
|
||||
obs[18:30] = data.qvel[6:18] * 0.05
|
||||
obs[30:42] = last_a
|
||||
obs[42:45] = np.array([vx, vy, wz]) * np.array([2., 2., 0.25])
|
||||
obs = np.clip(obs, -100., 100.)
|
||||
action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
|
||||
action = np.clip(action, -CLIP, CLIP)
|
||||
last_a = action.copy()
|
||||
|
||||
target = DEFAULT_ANGLES + action * ACTION_SCALE
|
||||
torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
|
||||
data.ctrl[:] = np.clip(torques, -CLIP, CLIP)
|
||||
mujoco.mj_step(model, data)
|
||||
view.sync()
|
||||
step += 1
|
||||
time.sleep(0.001)
|
||||
|
||||
kb.stop(); view.close()
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
BIN
deploy_45dim/policy.onnx
Normal file
BIN
deploy_45dim/policy.onnx
Normal file
Binary file not shown.
BIN
deploy_45dim/policy_v3.onnx
Normal file
BIN
deploy_45dim/policy_v3.onnx
Normal file
Binary file not shown.
133
deploy_45dim/sim2sim_57dim_test.py
Normal file
133
deploy_45dim/sim2sim_57dim_test.py
Normal file
@@ -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()
|
||||
421
deploy_45dim/sim2sim_test_deploy.py
Normal file
421
deploy_45dim/sim2sim_test_deploy.py
Normal file
@@ -0,0 +1,421 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
sim2sim test for deploy_onnx_pro_sdk.py using sim2sim_mujoco_example's Go1 XML.
|
||||
|
||||
Mirrors the deploy safety layer (joint limits, torque factor, position deviation protect)
|
||||
and logs all data to JSONL when --log-dir is set.
|
||||
|
||||
Usage:
|
||||
conda activate free_dog_sdk
|
||||
mjpython sim2sim_test_deploy.py --onnx policy.onnx
|
||||
mjpython sim2sim_test_deploy.py --onnx policy.onnx --log-dir logs
|
||||
|
||||
Controls: W/S=forward/back, Q/E=strafe, A/D=rotate, Space=stop, R=reset, Esc=quit
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import signal
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import mujoco
|
||||
import numpy as np
|
||||
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")
|
||||
|
||||
# ─── Policy constants (matches go1_sim2sim.py + deploy_onnx_pro_sdk.py) ───
|
||||
NUM_OBS = 45
|
||||
NUM_ACTIONS = 12
|
||||
ACTION_SCALE = 0.05
|
||||
CLIP_ACTIONS = 23.7
|
||||
CLIP_OBS = 100.0
|
||||
KP_DEFAULT = 80.0
|
||||
KD_DEFAULT = 0.5 # + joint_damping(0.5) = 1.0 total
|
||||
|
||||
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)
|
||||
|
||||
JOINT_NAMES = [
|
||||
"FR_hip", "FR_thigh", "FR_calf",
|
||||
"FL_hip", "FL_thigh", "FL_calf",
|
||||
"RR_hip", "RR_thigh", "RR_calf",
|
||||
"RL_hip", "RL_thigh", "RL_calf",
|
||||
]
|
||||
|
||||
# ─── Safety limits (mirrors deploy go1_pro_sdk safety layer) ───
|
||||
JOINT_TYPE = [("hip" if i % 3 == 0 else "thigh" if i % 3 == 1 else "knee") for i in range(12)]
|
||||
|
||||
JOINT_LIMITS = {
|
||||
"hip": (-0.78, 0.78),
|
||||
"thigh": (-0.60, 3.50),
|
||||
"knee": (-2.70, -0.95),
|
||||
}
|
||||
|
||||
TAU_MAX = {
|
||||
"hip": 23.7,
|
||||
"thigh": 23.7,
|
||||
"knee": 35.55,
|
||||
}
|
||||
|
||||
EXIT = False
|
||||
|
||||
|
||||
def _sig_handler(signum, frame):
|
||||
global EXIT
|
||||
EXIT = True
|
||||
|
||||
|
||||
signal.signal(signal.SIGINT, _sig_handler)
|
||||
signal.signal(signal.SIGTERM, _sig_handler)
|
||||
|
||||
|
||||
# ─── Safety functions (mirror deploy apply_safety) ───
|
||||
def clip_targets_to_limits(targets):
|
||||
"""PositionLimit: clamp target joint angles to JOINT_LIMITS."""
|
||||
n_clamped = 0
|
||||
safe = targets.copy()
|
||||
for i in range(12):
|
||||
lo, hi = JOINT_LIMITS[JOINT_TYPE[i]]
|
||||
if safe[i] < lo:
|
||||
safe[i] = lo; n_clamped += 1
|
||||
elif safe[i] > hi:
|
||||
safe[i] = hi; n_clamped += 1
|
||||
return safe, n_clamped
|
||||
|
||||
|
||||
def clip_torques(torques, power_factor):
|
||||
"""PowerProtect: clamp torque to TAU_MAX * power_factor / 10."""
|
||||
tau_lim = np.array([TAU_MAX[JOINT_TYPE[i]] * power_factor / 10.0 for i in range(12)],
|
||||
dtype=np.float64)
|
||||
return np.clip(torques, -tau_lim, tau_lim)
|
||||
|
||||
|
||||
def position_protect_mask(targets, current_pos, limit_rad):
|
||||
"""PositionProtect: return bool mask, True where deviation <= limit."""
|
||||
return np.abs(targets - current_pos) <= limit_rad
|
||||
|
||||
|
||||
# ─── Quaternion math (same as deploy_onnx_pro_sdk.py) ───
|
||||
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)
|
||||
|
||||
|
||||
# ─── Observations (same as deploy_onnx_pro_sdk.py) ───
|
||||
def compute_obs(model, data, commands, last_actions):
|
||||
obs = np.zeros(NUM_OBS, dtype=np.float32)
|
||||
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Gyro")
|
||||
adr = model.sensor_adr[sid]
|
||||
obs[0:3] = data.sensordata[adr:adr + 3] * 0.25
|
||||
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Quat")
|
||||
adr = model.sensor_adr[sid]
|
||||
quat = data.sensordata[adr:adr + 4]
|
||||
R = quat_to_rot_matrix(quat)
|
||||
obs[3:6] = (R.T @ np.array([0., 0., -1.], dtype=np.float64)).astype(np.float32)
|
||||
|
||||
dof_pos = np.zeros(12, dtype=np.float32)
|
||||
dof_vel = np.zeros(12, dtype=np.float32)
|
||||
for i, name in enumerate(JOINT_NAMES):
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, f"{name}_pos")
|
||||
dof_pos[i] = data.sensordata[model.sensor_adr[sid]]
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, f"{name}_vel")
|
||||
dof_vel[i] = data.sensordata[model.sensor_adr[sid]]
|
||||
|
||||
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
|
||||
|
||||
|
||||
# ─── 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"sim2sim_run_{ts}"
|
||||
self.run_dir.mkdir(parents=True, exist_ok=True)
|
||||
meta = {
|
||||
"created_at": ts,
|
||||
"args": vars(args),
|
||||
"num_obs": NUM_OBS,
|
||||
"num_actions": NUM_ACTIONS,
|
||||
"action_scale": ACTION_SCALE,
|
||||
"default_angles": DEFAULT_ANGLES.tolist(),
|
||||
"joint_names": JOINT_NAMES,
|
||||
"joint_limits": {k: list(v) for k, v in JOINT_LIMITS.items()},
|
||||
"tau_max": TAU_MAX,
|
||||
}
|
||||
(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")
|
||||
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}")
|
||||
|
||||
|
||||
# ─── Keyboard input ───
|
||||
class Keyboard:
|
||||
def __init__(self):
|
||||
self.held = set()
|
||||
|
||||
def _on_press(self, key):
|
||||
try: self.held.add(key.char.lower())
|
||||
except AttributeError: self.held.add(str(key))
|
||||
|
||||
def _on_release(self, key):
|
||||
try: self.held.discard(key.char.lower())
|
||||
except AttributeError: self.held.discard(str(key))
|
||||
|
||||
def init(self):
|
||||
from pynput import keyboard
|
||||
self._listener = keyboard.Listener(
|
||||
on_press=self._on_press, on_release=self._on_release)
|
||||
self._listener.start()
|
||||
|
||||
def keys(self):
|
||||
return self.held.copy()
|
||||
|
||||
def stop(self):
|
||||
if self._listener:
|
||||
self._listener.stop()
|
||||
|
||||
|
||||
# ─── 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_v3.onnx"))
|
||||
|
||||
# PD gains
|
||||
parser.add_argument("--kp", type=float, default=KP_DEFAULT)
|
||||
parser.add_argument("--kd", type=float, default=KD_DEFAULT)
|
||||
|
||||
# Safety (mirrors deploy args)
|
||||
parser.add_argument("--power-factor", type=int, default=7,
|
||||
help="Torque limit factor 1-10, applied as TAU_MAX * factor/10")
|
||||
parser.add_argument("--position-protect-limit", type=float, default=0.5,
|
||||
help="Max |target - actual| before zeroing torque (negative=disable)")
|
||||
parser.add_argument("--no-joint-limit", action="store_true",
|
||||
help="Disable joint limit clipping on targets")
|
||||
|
||||
# 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()
|
||||
|
||||
if not os.path.exists(SIM2SIM_XML):
|
||||
print(f"[ERROR] XML not found: {SIM2SIM_XML}"); return 1
|
||||
if not os.path.exists(args.onnx):
|
||||
print(f"[ERROR] ONNX not found: {args.onnx}"); return 1
|
||||
|
||||
print(f"[INFO] XML: {SIM2SIM_XML}")
|
||||
print(f"[INFO] ONNX: {args.onnx}")
|
||||
|
||||
model = mujoco.MjModel.from_xml_path(SIM2SIM_XML)
|
||||
data = mujoco.MjData(model)
|
||||
model.dof_damping[6:] = 0.5
|
||||
|
||||
total_kd = args.kd + model.dof_damping[6]
|
||||
print(f"[INFO] Bodies={model.nbody}, DoF={model.nq}, Actuators={model.nu}")
|
||||
print(f"[INFO] Timestep={model.opt.timestep}")
|
||||
print(f"[INFO] KP={args.kp}, KD(active)={args.kd}, KD(passive)={model.dof_damping[6]}, "
|
||||
f"total_KD={total_kd}")
|
||||
print(f"[INFO] Safety: power_factor={args.power_factor}, "
|
||||
f"position_protect={args.position_protect_limit}, "
|
||||
f"joint_limit={not args.no_joint_limit}")
|
||||
print(f"[INFO] Torque limits (factor={args.power_factor}): "
|
||||
+ ", ".join(f"{jt}={TAU_MAX[jt]*args.power_factor/10:.1f}" for jt in ["hip","thigh","knee"]))
|
||||
|
||||
# Init pose
|
||||
data.qpos[0:3] = [0.0, 0.0, 0.42]
|
||||
data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
|
||||
data.qpos[7:19] = DEFAULT_ANGLES
|
||||
data.qvel[:] = 0.0
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
# Load ONNX
|
||||
session = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
|
||||
input_name = session.get_inputs()[0].name
|
||||
print(f"[INFO] ONNX input={input_name}, shape={session.get_inputs()[0].shape}")
|
||||
print(f"[INFO] Controls: W/S=前后 Q/E=左右 A/D=旋转 Space=停 R=重置 Esc=退出")
|
||||
|
||||
logger = JsonlLogger(args.log_dir, args)
|
||||
|
||||
kb = Keyboard(); kb.init()
|
||||
|
||||
view = viewer.launch_passive(model, data)
|
||||
step = 0
|
||||
ctrl_dt = 0.01
|
||||
steps_per_inference = int(ctrl_dt / model.opt.timestep)
|
||||
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||||
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
|
||||
safety_stats = {"joint_limit_clamps": 0, "position_protect_hits": 0}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
|
||||
while view.is_running() and not EXIT:
|
||||
t_loop = time.perf_counter()
|
||||
|
||||
keys = kb.keys()
|
||||
if 'key.esc' in keys or '\x1b' in keys:
|
||||
break
|
||||
|
||||
if 'r' in keys:
|
||||
data.qpos[0:3] = [0.0, 0.0, 0.42]
|
||||
data.qpos[3:7] = [1.0, 0.0, 0.0, 0.0]
|
||||
data.qpos[7:19] = DEFAULT_ANGLES
|
||||
data.qvel[:] = 0.0
|
||||
last_actions[:] = 0.0
|
||||
action[:] = 0.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 = 0.5 if 'a' in keys else (-0.5 if 'd' in keys else 0.0)
|
||||
if ' ' in keys:
|
||||
vx = vy = wz = 0.0
|
||||
|
||||
commands = np.array([vx, vy, wz], dtype=np.float32)
|
||||
|
||||
# Inference at 100 Hz
|
||||
if step % steps_per_inference == 0:
|
||||
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 * args.action_scale
|
||||
|
||||
# ── Safety layer (mirrors deploy apply_safety) ──
|
||||
# 1. Joint limit clipping
|
||||
if not args.no_joint_limit:
|
||||
targets, n_clamped = clip_targets_to_limits(targets_raw)
|
||||
safety_stats["joint_limit_clamps"] += n_clamped
|
||||
else:
|
||||
targets = targets_raw
|
||||
|
||||
current_pos = data.qpos[7:19]
|
||||
current_vel = data.qvel[6:18]
|
||||
|
||||
# 2. Position deviation protection (zero torque where |target - actual| > limit)
|
||||
pos_ok = np.ones(12, dtype=bool)
|
||||
if args.position_protect_limit > 0:
|
||||
pos_ok = position_protect_mask(targets, current_pos, args.position_protect_limit)
|
||||
n_hit = np.sum(~pos_ok)
|
||||
safety_stats["position_protect_hits"] += n_hit
|
||||
|
||||
# 3. PD control
|
||||
torques = np.zeros(12, dtype=np.float64)
|
||||
torques[pos_ok] = (args.kp * (targets[pos_ok] - current_pos[pos_ok])
|
||||
- args.kd * current_vel[pos_ok])
|
||||
|
||||
# 4. Torque limiting (power_protect)
|
||||
torques = clip_torques(torques, args.power_factor)
|
||||
|
||||
data.ctrl[:] = torques
|
||||
|
||||
mujoco.mj_step(model, data)
|
||||
view.sync()
|
||||
|
||||
# Logging
|
||||
logger.log(
|
||||
step,
|
||||
mode="rl",
|
||||
loop_ms=(time.perf_counter() - t_loop) * 1000.0,
|
||||
commands=commands,
|
||||
obs=obs if step % steps_per_inference == 0 else np.zeros(0),
|
||||
action_raw=action,
|
||||
action_safe=action,
|
||||
joint_targets_raw=targets_raw,
|
||||
joint_targets_safe=targets,
|
||||
dof_pos=current_pos,
|
||||
dof_vel=current_vel,
|
||||
torques=torques,
|
||||
position_protect_hit_mask=(~pos_ok).astype(int),
|
||||
)
|
||||
|
||||
step += 1
|
||||
|
||||
if step % args.print_every == 0:
|
||||
z = data.qpos[2]
|
||||
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Quat")
|
||||
quat = data.sensordata[model.sensor_adr[sid]:model.sensor_adr[sid] + 4]
|
||||
pos_err = np.max(np.abs(targets - current_pos))
|
||||
n_pp = safety_stats["position_protect_hits"]
|
||||
n_jl = safety_stats["joint_limit_clamps"]
|
||||
print(f"\n[STEP {step}] z={z:.3f} cmd=[{vx:.1f},{vy:.1f},{wz:.1f}] "
|
||||
f"max_err={pos_err:.3f}")
|
||||
print(f" quat={np.round(quat, 3)} action_max={np.max(np.abs(action)):.2f}")
|
||||
print(f" target: {np.round(targets[:4], 2)}")
|
||||
print(f" actual: {np.round(current_pos[:4], 2)}")
|
||||
print(f" torque: {np.round(torques[:4], 2)}")
|
||||
if n_pp > 0 or n_jl > 0:
|
||||
print(f" safety: pos_protect_hits={n_pp} joint_limit_clamps={n_jl}")
|
||||
|
||||
# Real-time sync
|
||||
expected = (step + 1) * model.opt.timestep
|
||||
elapsed = time.perf_counter() - t0
|
||||
sleep = expected - elapsed
|
||||
if sleep > 0:
|
||||
time.sleep(sleep)
|
||||
|
||||
logger.close()
|
||||
kb.stop()
|
||||
view.close()
|
||||
print("[INFO] Done.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user