cyy/go1_pro_deploy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

新增lab

Browse Source
This commit is contained in:
cyy_mac
2026-07-26 20:52:34 +08:00
parent 6882f1abd8
commit 29de5eb8de
3 changed files with 1868 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

738
deploy_45dim_rl_gym/deploy_go1_onnx_mujoco_lab.py Normal file
View File

@@ -0,0 +1,738 @@
#!/usr/bin/env python3
"""
Go1 RobotLab ONNX Policy - MuJoCo Simulation Deployment.
Loads the RobotLab ONNX policy exported from RoboGauge and runs it in MuJoCo
with PD position control. The ONNX model uses a 10-frame history
(450-dim stacked-by-terms input, stateless).
Usage:
conda activate free_dog_sdk
MUJOCO_GL=glfw mjpython deploy_go1_onnx_mujoco_lab.py
MUJOCO_GL=glfw mjpython deploy_go1_onnx_mujoco_lab.py --terrain terrains/stairs/stairs_6.xml
Controls (matching RoboGauge keyboard convention):
↑ / ↓ forward / backward
← / → yaw left / right
, / . strafe left / right
K stop
R reset robot
Esc quit
Architecture:
ONNX input: obs [1, 450] - 10 frames x 45 dims, stacked by TERMS
ONNX output: actions [1, 12]
Control: target_q = default_q + 0.25 * action, PD with Kp=28, Kd=0.7
"""
import argparse
import json
import os
import re
import signal
import sys
import time
from collections import deque
from datetime import datetime
from pathlib import Path
import mujoco
import numpy as np
import onnxruntime as ort
from mujoco import viewer
# ── path setup ──
SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_ONNX = str(SCRIPT_DIR / "policy_robotlab_6500.onnx")
ROBOT_XML = str(SCRIPT_DIR / "go1.xml")
TERRAINS_DIR = SCRIPT_DIR / "terrains"
# ── policy constants (RoboGauge Go1Config) ──
NUM_OBS = 45
NUM_ACTIONS = 12
HISTORY_LEN = 10 # RobotLab ONNX history frames
ONNX_INPUT_DIM = 450 # 45 x 10, stacked by terms
ACTION_SCALE = 0.25
KP = 28.0
KD = 0.7
CLIP_ACTIONS = 100.0
CLIP_OBS = 100.0
ANG_VEL_SCALE = 0.25
CMD_SCALE = np.array([1.0, 1.0, 1.0], dtype=np.float32)
MAX_LIN_VEL_X = 1.0
MAX_LIN_VEL_Y = 0.5
MAX_ANG_VEL = 1.0
DEFAULT_DOF_POS = np.array([
-0.1, 0.8, -1.5, # FR_hip, FR_thigh, FR_calf
0.1, 0.8, -1.5, # FL_hip, FL_thigh, FL_calf
-0.1, 1.0, -1.5, # RR_hip, RR_thigh, RR_calf
0.1, 1.0, -1.5, # RL_hip, RL_thigh, RL_calf
], dtype=np.float32)
# ── exit flag ──
EXIT = False
def _sig_handler(signum, frame):
global EXIT
EXIT = True
signal.signal(signal.SIGINT, _sig_handler)
signal.signal(signal.SIGTERM, _sig_handler)
# ═══════════════════════════════════════════════════════════════
# XML builder — merge terrain + robot into one MuJoCo scene
# ═══════════════════════════════════════════════════════════════
def _xml_inner(text, tag):
"""Extract inner content of the first <tag>...</tag> in text."""
m = re.search(rf"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
return m.group(1).strip() if m else ""
def build_scene_xml(robot_xml_path, terrain_xml_path=None):
"""Merge robot XML with optional terrain XML into a single scene model.
The robot XML (go1.xml) has base_link as the root body with no free joint.
We add a free joint and optionally merge terrain worldbody/asset elements.
"""
robot = Path(robot_xml_path).read_text()
# 1) Add free joint to base_link
robot = robot.replace(
'<body name="base_link" pos="0 0 0.35">',
'<body name="base_link" pos="0 0 0.35">\n <joint type="free"/>'
)
if terrain_xml_path is None:
# No terrain — add a simple flat floor
floor = (
'\n <geom name="floor" size="0 0 0.05" type="plane" '
'rgba="0.5 0.9 0.9 0.1"/>\n'
)
robot = robot.replace(
'<body name="base_link"',
floor + ' <body name="base_link"'
)
return robot
terrain_text = Path(terrain_xml_path).read_text()
# 2) Merge terrain <asset> into robot <asset>
terrain_assets = _xml_inner(terrain_text, "asset")
if terrain_assets:
# Insert before closing </asset> of robot (or before <worldbody> if no asset)
robot = robot.replace('</asset>', '\n' + terrain_assets + '\n </asset>', 1)
# 3) Merge terrain <visual> settings
terrain_visual = _xml_inner(terrain_text, "visual")
if terrain_visual:
robot = robot.replace('</visual>', '\n' + terrain_visual + '\n </visual>', 1)
# 4) Merge terrain worldbody elements (lights, geoms, etc.) before base_link
terrain_wb = _xml_inner(terrain_text, "worldbody")
if terrain_wb:
# Remove <body> elements from terrain worldbody (we only want geoms/lights/cameras)
terrain_wb_no_bodies = re.sub(r'<body\b.*?</body>', '', terrain_wb, flags=re.DOTALL)
robot = robot.replace(
'<body name="base_link"',
terrain_wb_no_bodies.strip() + '\n <body name="base_link"'
)
return robot
# ═══════════════════════════════════════════════════════════════
# Observation building
# ═══════════════════════════════════════════════════════════════
def get_projected_gravity(quat_wxyz):
qw, qx, qy, qz = quat_wxyz
g = np.zeros(3, dtype=np.float32)
g[0] = 2.0 * (-qz * qx + qw * qy)
g[1] = -2.0 * (qz * qy + qw * qx)
g[2] = 1.0 - 2.0 * (qw * qw + qz * qz)
return g
def quat_rotate_inverse(q_wxyz, v):
"""Rotate vector v from world frame into body frame using quaternion q."""
q_w = q_wxyz[0]
q_vec = np.array(q_wxyz[1:], dtype=np.float32)
v = np.array(v, dtype=np.float32)
a = v * (2.0 * q_w * q_w - 1.0)
b = np.cross(q_vec, v) * q_w * 2.0
c = q_vec * np.dot(q_vec, v) * 2.0
return a - b + c
def quat_to_rpy_deg(q_wxyz):
"""Convert MuJoCo [w, x, y, z] quaternion to roll/pitch/yaw in degrees."""
qw, qx, qy, qz = [float(x) for x in q_wxyz]
sinr_cosp = 2.0 * (qw * qx + qy * qz)
cosr_cosp = 1.0 - 2.0 * (qx * qx + qy * qy)
roll = np.arctan2(sinr_cosp, cosr_cosp)
sinp = 2.0 * (qw * qy - qz * qx)
pitch = np.arcsin(np.clip(sinp, -1.0, 1.0))
siny_cosp = 2.0 * (qw * qz + qx * qy)
cosy_cosp = 1.0 - 2.0 * (qy * qy + qz * qz)
yaw = np.arctan2(siny_cosp, cosy_cosp)
return np.degrees(np.array([roll, pitch, yaw], dtype=np.float32))
def read_sensor(model, data, name, expected_dim):
"""Read a MuJoCo sensor by name, returning None if it is not present."""
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
if sid < 0:
return None
adr = int(model.sensor_adr[sid])
dim = int(model.sensor_dim[sid])
if dim != expected_dim:
raise ValueError(f"sensor {name!r} has dim {dim}, expected {expected_dim}")
return np.asarray(data.sensordata[adr:adr + dim], dtype=np.float32).copy()
class ObsBuilder:
"""Build 45-dim single-frame obs and stack 10-frame history into 450-dim ONNX input.
ONNX expects observation terms stacked by groups (not by frames):
[ang_vel(t-9..t), gravity(t-9..t), cmd(t-9..t),
dof_pos(t-9..t), dof_vel(t-9..t), last_action(t-9..t)]
Within each group, frames go from oldest (t-9) to newest (t).
"""
def __init__(self):
self.history: deque = deque(maxlen=HISTORY_LEN)
def reset(self):
self.history.clear()
def build_single_obs(self, base_ang_vel_body, base_quat_wxyz, cmd, q, dq, last_action):
"""Build 45-dim single-frame observation (RoboGauge layout)."""
obs = np.zeros(NUM_OBS, dtype=np.float32)
obs[0:3] = np.asarray(base_ang_vel_body, dtype=np.float32) * ANG_VEL_SCALE
obs[3:6] = get_projected_gravity(np.asarray(base_quat_wxyz, dtype=np.float32))
obs[6:9] = np.asarray(cmd, dtype=np.float32) * CMD_SCALE
obs[9:21] = np.asarray(q, dtype=np.float32) - DEFAULT_DOF_POS
obs[21:33] = np.asarray(dq, dtype=np.float32) * 0.05
obs[33:45] = np.asarray(last_action, dtype=np.float32)
return np.clip(obs, -CLIP_OBS, CLIP_OBS)
def build_onnx_input(self, obs_single):
"""Stack 10-frame history into 450-dim ONNX input (by-terms format)."""
self.history.append(obs_single.copy())
frames = list(self.history)
while len(frames) < HISTORY_LEN:
frames.insert(0, np.zeros(NUM_OBS, dtype=np.float32))
# term_dims: [ang_vel(3), gravity(3), cmd(3), dof_pos(12), dof_vel(12), last_action(12)]
term_dims = [3, 3, 3, 12, 12, 12]
stacked = []
offset = 0
for dim in term_dims:
for f_idx in range(HISTORY_LEN):
stacked.append(frames[f_idx][offset:offset + dim])
offset += dim
return np.concatenate(stacked, dtype=np.float32).reshape(1, -1)
# ═══════════════════════════════════════════════════════════════
# Keyboard input (pynput)
# ═══════════════════════════════════════════════════════════════
class KbReader:
"""Non-blocking keyboard reader using a background pynput listener.
Key names match RoboGauge convention:
key.up / key.down → forward/back
key.left / key.right → yaw
, / . → strafe left/right
k → stop
"""
def __init__(self):
import queue
import threading
self._q = queue.Queue()
self.held = set()
self._running = True
self._listener = None
self._thread = None
@staticmethod
def _name(key):
try:
if hasattr(key, 'char') and key.char is not None:
return key.char.lower()
except Exception:
pass
return str(key).lower()
def _worker(self):
while self._running:
try:
evt, key = self._q.get(timeout=0.05)
except Exception:
continue
name = self._name(key)
if evt == 'press':
self.held.add(name)
else:
self.held.discard(name)
def start(self):
from pynput import keyboard
self._listener = keyboard.Listener(
on_press=lambda k: self._q.put(('press', k)),
on_release=lambda k: self._q.put(('release', k)),
)
self._listener.start()
import threading
self._thread = threading.Thread(target=self._worker, daemon=True)
self._thread.start()
print("[INFO] keyboard started (↑↓←→ , . K R Esc)")
def is_held(self, key):
return self._name(key) in self.held
def snapshot(self):
return set(self.held)
def stop(self):
self._running = False
if self._listener:
self._listener.stop()
def get_command(held):
"""Extract velocity command from held key set.
Returns (vx, vy, yaw_rate).
"""
if 'k' in held:
return 0.0, 0.0, 0.0
def axis(pos_key, neg_key, limit):
p = pos_key in held
n = neg_key in held
if p == n:
return 0.0
return limit if p else -limit
vx = axis('key.up', 'key.down', MAX_LIN_VEL_X)
vy = axis(',', '.', MAX_LIN_VEL_Y)
yaw = axis('key.left', 'key.right', MAX_ANG_VEL)
return vx, vy, yaw
class SampleLogger:
def __init__(self, log_dir, terrain_path, cmd, args):
ts = datetime.now().strftime("%Y%m%d_%H%M%S")
self.run_dir = Path(log_dir).expanduser().resolve() / f"mujoco_stairs_sample_{ts}"
self.run_dir.mkdir(parents=True, exist_ok=True)
self.fp = open(self.run_dir / "steps.jsonl", "a", encoding="utf-8", buffering=1)
meta = {
"created_at": ts,
"terrain": str(terrain_path) if terrain_path else "flat",
"sample_seconds": args.sample_seconds,
"sample_cmd": list(cmd),
"spawn": [args.spawn_x, args.spawn_y, args.spawn_z],
"action_scale": ACTION_SCALE,
"clip_actions": args.clip_actions,
"max_target_step": args.max_target_step,
"kp": KP,
"kd": KD,
"control_dt": 0.02,
"sim_dt": 0.002,
"num_obs": NUM_OBS,
"history_len": HISTORY_LEN,
"onnx_input_dim": ONNX_INPUT_DIM,
}
(self.run_dir / "metadata.json").write_text(
json.dumps(meta, indent=2, ensure_ascii=False))
print(f"[INFO] sample log dir: {self.run_dir}")
def write(self, rec):
self.fp.write(json.dumps(rec, ensure_ascii=False) + "\n")
def close(self):
self.fp.flush()
self.fp.close()
def summarize_samples(samples):
if not samples:
return {}
def max_abs_field(field):
return float(max(max(abs(x) for x in s[field]) for s in samples))
def p99_abs_field(field):
vals = sorted(abs(x) for s in samples for x in s[field])
if not vals:
return 0.0
idx = min(len(vals) - 1, int(0.99 * (len(vals) - 1)))
return float(vals[idx])
rpy = np.array([s["rpy_deg"] for s in samples], dtype=np.float32)
xyz = np.array([s["base_pos"] for s in samples], dtype=np.float32)
return {
"num_samples": len(samples),
"duration": float(samples[-1]["time_sim"] - samples[0]["time_sim"]),
"base_x_start": float(xyz[0, 0]),
"base_x_end": float(xyz[-1, 0]),
"base_x_progress": float(xyz[-1, 0] - xyz[0, 0]),
"base_z_min": float(np.min(xyz[:, 2])),
"base_z_max": float(np.max(xyz[:, 2])),
"rpy_min_deg": np.min(rpy, axis=0).astype(float).tolist(),
"rpy_max_deg": np.max(rpy, axis=0).astype(float).tolist(),
"action_raw_maxabs": max_abs_field("action_raw"),
"action_raw_p99abs": p99_abs_field("action_raw"),
"action_applied_maxabs": max_abs_field("action"),
"action_applied_p99abs": p99_abs_field("action"),
"target_offset_maxabs": max_abs_field("target_offset"),
"target_offset_p99abs": p99_abs_field("target_offset"),
"dof_vel_maxabs": max_abs_field("dof_vel"),
"torque_maxabs": max_abs_field("torques"),
"fallen": bool(any(s["fallen"] for s in samples)),
}
# ═══════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════
def main():
parser = argparse.ArgumentParser(description="Go1 RobotLab ONNX MuJoCo Deployment")
parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX, help="ONNX model path")
parser.add_argument("--terrain", type=str, default=None,
help="Terrain XML path (e.g. terrains/stairs/stairs_6.xml)")
parser.add_argument("--sample", action="store_true",
help="Run a non-interactive fixed-command sampling rollout")
parser.add_argument("--sample-stairs-level", type=int, default=None,
help="Shortcut for --sample --terrain terrains/stairs/stairs_LEVEL.xml")
parser.add_argument("--sample-seconds", type=float, default=20.0)
parser.add_argument("--sample-cmd-x", type=float, default=MAX_LIN_VEL_X)
parser.add_argument("--sample-cmd-y", type=float, default=0.0)
parser.add_argument("--sample-cmd-yaw", type=float, default=0.0)
parser.add_argument("--sample-log-dir", type=str, default=str(SCRIPT_DIR / "mujoco_logs"))
parser.add_argument("--spawn-x", type=float, default=None)
parser.add_argument("--spawn-y", type=float, default=None)
parser.add_argument("--spawn-z", type=float, default=None)
parser.add_argument("--clip-actions", type=float, default=CLIP_ACTIONS,
help="Clip ONNX actions before applying action_scale")
parser.add_argument("--max-target-step", type=float, default=0.0,
help="Optional per-control-step target q slew limit in rad")
args = parser.parse_args()
sample_mode = args.sample or args.sample_stairs_level is not None
if args.sample_stairs_level is not None:
args.sample = True
args.terrain = f"terrains/stairs/stairs_{args.sample_stairs_level}.xml"
if args.spawn_x is None:
args.spawn_x = -0.6 if args.sample_stairs_level is not None else 0.0
if args.spawn_y is None:
args.spawn_y = 0.0
if args.spawn_z is None:
args.spawn_z = 0.34
args.sample_log_dir = str(Path(args.sample_log_dir).expanduser().resolve())
# Resolve terrain path
terrain_path = None
if args.terrain:
terrain_path = Path(args.terrain)
if not terrain_path.is_absolute():
terrain_path = SCRIPT_DIR / terrain_path
if not terrain_path.exists():
print(f"[ERROR] terrain not found: {terrain_path}")
return 1
# ── Build scene XML ──
print(f"[INFO] robot: {ROBOT_XML}")
print(f"[INFO] terrain: {terrain_path or 'flat floor'}")
scene_xml = build_scene_xml(ROBOT_XML, str(terrain_path) if terrain_path else None)
# ── Load MuJoCo model ──
os.chdir(str(SCRIPT_DIR)) # meshdir="assets" is relative to go1.xml
model = mujoco.MjModel.from_xml_string(scene_xml)
data = mujoco.MjData(model)
model.opt.timestep = 0.002
print(f"[INFO] model: {model.nbody} bodies, qpos={model.nq}, actuators={model.nu}")
print(f"[INFO] timestep: {model.opt.timestep}s")
has_imu_gyro = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Gyro") >= 0
has_imu_quat = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "Body_Quat") >= 0
print(f"[INFO] IMU sensors: Body_Gyro={has_imu_gyro}, Body_Quat={has_imu_quat}")
# ── Load ONNX ──
session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
inp = session.get_inputs()[0]
print(f"[INFO] ONNX input : {inp.name} {inp.shape}")
for o in session.get_outputs():
print(f"[INFO] ONNX output: {o.name} {o.shape}")
# ── Init state ──
data.qpos[0:3] = np.array([args.spawn_x, args.spawn_y, args.spawn_z], dtype=np.float64)
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
data.qpos[7:19] = DEFAULT_DOF_POS.astype(np.float64)
data.qvel[:] = 0.0
data.ctrl[:] = 0.0
mujoco.mj_forward(model, data)
# ── Control state ──
ctrl_dt = 0.02 # 50 Hz
sim_dt = model.opt.timestep
steps_per_inference = max(1, int(ctrl_dt / sim_dt))
print(f"[INFO] control: {ctrl_dt}s ({1/ctrl_dt:.0f}Hz), "
f"sim steps per inference: {steps_per_inference}")
last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
action_raw = np.zeros(NUM_ACTIONS, dtype=np.float32)
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
target_pos = DEFAULT_DOF_POS.copy()
prev_target_pos = DEFAULT_DOF_POS.copy()
obs_builder = ObsBuilder()
step_count = 0
inference_step = 0
if sample_mode:
cmd = np.array([
args.sample_cmd_x,
args.sample_cmd_y,
args.sample_cmd_yaw,
], dtype=np.float32)
cmd = np.clip(cmd, [-MAX_LIN_VEL_X, -MAX_LIN_VEL_Y, -MAX_ANG_VEL],
[MAX_LIN_VEL_X, MAX_LIN_VEL_Y, MAX_ANG_VEL])
logger = SampleLogger(args.sample_log_dir, terrain_path, cmd, args)
samples = []
max_sim_steps = int(args.sample_seconds / sim_dt)
print(
f"[INFO] sampling: seconds={args.sample_seconds:.1f} "
f"cmd=({cmd[0]:.2f},{cmd[1]:.2f},{cmd[2]:.2f}) "
f"spawn=({args.spawn_x:.2f},{args.spawn_y:.2f},{args.spawn_z:.2f})"
)
while step_count < max_sim_steps and not EXIT:
if inference_step == 0:
imu_quat = read_sensor(model, data, "Body_Quat", 4)
imu_ang_vel = read_sensor(model, data, "Body_Gyro", 3)
quat_wxyz = imu_quat if imu_quat is not None else data.qpos[3:7].copy()
if imu_ang_vel is not None:
base_ang_vel_body = imu_ang_vel
else:
world_ang_vel = data.qvel[3:6].copy()
base_ang_vel_body = quat_rotate_inverse(quat_wxyz, world_ang_vel)
q = data.qpos[7:19].copy()
dq = data.qvel[6:18].copy()
obs_single = obs_builder.build_single_obs(
base_ang_vel_body, quat_wxyz, cmd, q, dq, last_action)
onnx_input = obs_builder.build_onnx_input(obs_single)
outputs = session.run(None, {'obs': onnx_input})
action_raw = outputs[0][0].astype(np.float32)
action = np.clip(action_raw, -args.clip_actions, args.clip_actions)
last_action = action.copy()
target_pos_raw = DEFAULT_DOF_POS + action * ACTION_SCALE
if args.max_target_step > 0.0:
delta = np.clip(
target_pos_raw - prev_target_pos,
-args.max_target_step,
args.max_target_step,
)
target_pos = prev_target_pos + delta
else:
target_pos = target_pos_raw
prev_target_pos = target_pos.copy()
current_pos = data.qpos[7:19]
current_vel = data.qvel[6:18]
torques = KP * (target_pos - current_pos) - KD * current_vel
torques = np.clip(torques, -33.5, 33.5)
data.ctrl[:] = torques.astype(np.float64)
mujoco.mj_step(model, data)
if inference_step == 0:
quat_wxyz = data.qpos[3:7].copy()
rpy_deg = quat_to_rpy_deg(quat_wxyz)
fallen = bool(
data.qpos[2] < 0.16
or abs(rpy_deg[0]) > 60.0
or abs(rpy_deg[1]) > 60.0
or not np.all(np.isfinite(data.qpos))
)
rec = {
"step": int(step_count),
"control_step": int(step_count // steps_per_inference),
"time_sim": float(data.time),
"cmd": cmd.astype(float).tolist(),
"base_pos": data.qpos[0:3].astype(float).tolist(),
"base_quat": quat_wxyz.astype(float).tolist(),
"rpy_deg": rpy_deg.astype(float).tolist(),
"base_lin_vel": data.qvel[0:3].astype(float).tolist(),
"base_ang_vel": data.qvel[3:6].astype(float).tolist(),
"dof_pos": current_pos.astype(float).tolist(),
"dof_vel": current_vel.astype(float).tolist(),
"action_raw": action_raw.astype(float).tolist(),
"action": action.astype(float).tolist(),
"target_pos": target_pos.astype(float).tolist(),
"target_offset": (target_pos - DEFAULT_DOF_POS).astype(float).tolist(),
"torques": torques.astype(float).tolist(),
"fallen": fallen,
}
logger.write(rec)
samples.append(rec)
if rec["control_step"] % 50 == 0:
print(
f"[sample {rec['control_step']:04d}] "
f"t={rec['time_sim']:.2f} x={rec['base_pos'][0]:.2f} "
f"z={rec['base_pos'][2]:.2f} rpy={np.round(rpy_deg, 1)} "
f"act_max={np.max(np.abs(action)):.2f} "
f"target_off={np.max(np.abs(target_pos - DEFAULT_DOF_POS)):.2f}"
)
if fallen:
print(f"[WARN] sample stopped: fallen at t={data.time:.2f}s")
break
step_count += 1
inference_step = (inference_step + 1) % steps_per_inference
summary = summarize_samples(samples)
(logger.run_dir / "summary.json").write_text(
json.dumps(summary, indent=2, ensure_ascii=False))
logger.close()
print("[INFO] sample summary:")
print(json.dumps(summary, indent=2, ensure_ascii=False))
print(f"[INFO] sample saved: {logger.run_dir}")
return 0
# ── Keyboard ──
kb = KbReader()
kb.start()
# ── Viewer ──
view = viewer.launch_passive(model, data)
# Track the robot body
body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
if body_id >= 0:
view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
view.cam.trackbodyid = body_id
view.cam.distance = 2.5
view.cam.elevation = -20
view.cam.azimuth = 60
print("[INFO] viewer launched")
loop_start = time.time()
while view.is_running() and not EXIT:
held = kb.snapshot()
# ── Quit ──
if 'key.esc' in held:
break
# ── Command from keyboard ──
vx, vy, yaw = get_command(held)
cmd = np.array([vx, vy, yaw], dtype=np.float32)
# ── Reset ──
if 'r' in held:
data.qpos[0:3] = np.array([0.0, 0.0, 0.34], dtype=np.float64)
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float64)
data.qpos[7:19] = DEFAULT_DOF_POS.astype(np.float64)
data.qvel[:] = 0.0
data.ctrl[:] = 0.0
last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
action_raw = np.zeros(NUM_ACTIONS, dtype=np.float32)
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
target_pos = DEFAULT_DOF_POS.copy()
prev_target_pos = DEFAULT_DOF_POS.copy()
obs_builder.reset()
mujoco.mj_forward(model, data)
print("[RESET]")
# ── Inference ──
if inference_step == 0:
# Match RoboGauge: policy obs uses XML IMU gyro and framequat sensors.
# Fall back to qpos/qvel only for XMLs without those sensors.
imu_quat = read_sensor(model, data, "Body_Quat", 4)
imu_ang_vel = read_sensor(model, data, "Body_Gyro", 3)
quat_wxyz = imu_quat if imu_quat is not None else data.qpos[3:7].copy()
if imu_ang_vel is not None:
base_ang_vel_body = imu_ang_vel
else:
world_ang_vel = data.qvel[3:6].copy()
base_ang_vel_body = quat_rotate_inverse(quat_wxyz, world_ang_vel)
# Joint state (qpos[7:19] is FR,FL,RR,RL — matches policy order)
q = data.qpos[7:19].copy()
dq = data.qvel[6:18].copy()
# Build obs
obs_single = obs_builder.build_single_obs(
base_ang_vel_body, quat_wxyz, cmd, q, dq, last_action)
onnx_input = obs_builder.build_onnx_input(obs_single)
# Run ONNX
outputs = session.run(None, {'obs': onnx_input})
action_raw = outputs[0][0].astype(np.float32) # [12]
action = np.clip(action_raw, -args.clip_actions, args.clip_actions)
last_action = action.copy()
target_pos_raw = DEFAULT_DOF_POS + action * ACTION_SCALE
if args.max_target_step > 0.0:
delta = np.clip(
target_pos_raw - prev_target_pos,
-args.max_target_step,
args.max_target_step,
)
target_pos = prev_target_pos + delta
else:
target_pos = target_pos_raw
prev_target_pos = target_pos.copy()
# ── PD control ──
current_pos = data.qpos[7:19]
current_vel = data.qvel[6:18]
torques = KP * (target_pos - current_pos) - KD * current_vel
torques = np.clip(torques, -33.5, 33.5)
data.ctrl[:] = torques.astype(np.float64)
mujoco.mj_step(model, data)
view.sync()
# Real-time sync — use sim_dt because step_count increments every sim step
expected_time = step_count * sim_dt
elapsed = time.time() - loop_start
if 0 < expected_time - elapsed < ctrl_dt:
time.sleep(expected_time - elapsed)
step_count += 1
inference_step = (inference_step + 1) % steps_per_inference
# Periodic status
if step_count % 200 == 0:
z = data.qpos[2]
lin_vel_abs = np.linalg.norm(data.qvel[0:3])
print(f"[{step_count}] cmd=({vx:.1f},{vy:.1f},{yaw:.1f}) "
f"z={z:.3f} |v|={lin_vel_abs:.2f} "
f"act[0:4]={np.round(action[:4], 3)}")
kb.stop()
view.close()
print("[INFO] done.")
return 0
if __name__ == "__main__":
sys.exit(main())

1130
deploy_45dim_rl_gym/deploy_go1_rlgym_pro_sdk_lab.py Normal file
View File

File diff suppressed because it is too large Load Diff

BIN
deploy_45dim_rl_gym/policy_robotlab_6500.onnx Normal file
View File

Binary file not shown.