47372443ff
Upstream: https://github.com/pollen-robotics/microduck_rl Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63 Upstream-Branch: develop
612 lines
24 KiB
Python
612 lines
24 KiB
Python
#!/usr/bin/env python3
|
||
"""Sim2real validation for the XL330 test bench.
|
||
|
||
Runs the same ONNX policy on a fixed sequence of target angles, either in
|
||
MuJoCo (with the BAM M6 actuator model) or on the real XL330 (via rustypot),
|
||
logs the joint trajectory, and plots sim vs real for comparison.
|
||
|
||
Example workflow
|
||
----------------
|
||
# 1) Record in sim:
|
||
uv run python scripts/testbench_sim2real.py --mode sim --onnx policy.onnx --out sim.npz
|
||
# 2) Plug the real bench via USB, then record on hardware:
|
||
uv run python scripts/testbench_sim2real.py --mode real --onnx policy.onnx --out real.npz \
|
||
--port /dev/ttyUSB0 --motor-id 1
|
||
# 3) Compare the two traces:
|
||
uv run python scripts/testbench_sim2real.py --compare sim.npz real.npz --out-plot comparison.png
|
||
|
||
Observation layout (must match the training env): [joint_pos, joint_vel, last_action, command]
|
||
Action: 1-D position offset in radians, scaled by 1.0, added to default pose (0.0).
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import math
|
||
import time
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import onnxruntime as ort
|
||
from mjlab_microduck.robot.testbench_constants import (
|
||
TESTBENCH_ARM_MASS,
|
||
TESTBENCH_XML,
|
||
_set_arm_mass,
|
||
)
|
||
|
||
|
||
# --- Match training env ---
|
||
CONTROL_DT = 0.02 # decimation=4 × timestep=0.005 (policy rate = 50 Hz)
|
||
SIM_DT = 0.005 # (logging rate = 200 Hz — one sample per inner sim step)
|
||
LOG_DT = SIM_DT
|
||
DEFAULT_POS = 0.0
|
||
MAX_ANGLE = math.radians(80.0)
|
||
|
||
# XL330 present_velocity is returned by rustypot as raw ticks (i32, NOT converted).
|
||
# Each tick = 0.229 RPM (per Dynamixel XL330 spec). rad/s = ticks * 0.229 * 2π/60.
|
||
DXL_VEL_TICK_TO_RAD_S = 0.229 * 2.0 * math.pi / 60.0 # ≈ 0.02398 rad/s per tick
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Shared: target schedule + policy wrapper
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def make_target_schedule(
|
||
total_time: float,
|
||
hold_time: float = 4.0,
|
||
seed: int = 0,
|
||
) -> np.ndarray:
|
||
"""Return one target angle per control step."""
|
||
rng = np.random.default_rng(seed)
|
||
n_steps = int(round(total_time / CONTROL_DT))
|
||
steps_per_hold = int(round(hold_time / CONTROL_DT))
|
||
targets = np.zeros(n_steps, dtype=np.float32)
|
||
i = 0
|
||
while i < n_steps:
|
||
angle = float(rng.uniform(-MAX_ANGLE, MAX_ANGLE))
|
||
end = min(i + steps_per_hold, n_steps)
|
||
targets[i:end] = angle
|
||
i = end
|
||
return targets
|
||
|
||
|
||
class PolicyRunner:
|
||
def __init__(self, onnx_path: str, action_scale: float = 1.0):
|
||
print(f"Loading policy: {onnx_path} (action_scale={action_scale})")
|
||
self.session = ort.InferenceSession(onnx_path)
|
||
self.in_name = self.session.get_inputs()[0].name
|
||
in_shape = self.session.get_inputs()[0].shape
|
||
print(f" input {self.in_name} shape={in_shape}")
|
||
self.action_scale = action_scale
|
||
self.last_action = np.zeros(1, dtype=np.float32)
|
||
|
||
def reset(self):
|
||
self.last_action[:] = 0.0
|
||
|
||
def step(self, q: float, qd: float, target: float) -> float:
|
||
# Matches the testbench env's policy obs layout:
|
||
# [joint_pos_rel, joint_vel_rel, last_action, command] (4-d).
|
||
obs = np.array(
|
||
[q - DEFAULT_POS, qd, self.last_action[0], target],
|
||
dtype=np.float32,
|
||
)[None, :]
|
||
action = self.session.run(None, {self.in_name: obs})[0].reshape(-1)
|
||
self.last_action = action.astype(np.float32)
|
||
return DEFAULT_POS + float(action[0]) * self.action_scale
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Sim rollout (mujoco, same XL330 testbench XML as training)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def rollout_sim_bam(onnx_path: str, total_time: float, seed: int, action_scale: float) -> dict:
|
||
"""Sim rollout using bam's MujocoController on a vanilla MuJoCo step loop.
|
||
|
||
Pros: 200 Hz inner-step logging, no torch/mjwarp. Cons: not the exact
|
||
actuator that was trained against (uses bam upstream, not mjlab's M6).
|
||
"""
|
||
import mujoco # local import so --mode real works without mujoco
|
||
|
||
from bam.actuators import actuators as bam_actuators
|
||
from bam.model import models as bam_models
|
||
from bam.mujoco import MujocoController
|
||
|
||
# Load the fitted XL330 m6 params from the canonical bam bundle (identical to
|
||
# the values that used to live in mjlab_microduck.actuator.bam_params).
|
||
import json as _json
|
||
from bam.model import _resolve_json_path
|
||
with open(_resolve_json_path(None, "xl330", "m6")) as _f:
|
||
DEFAULT_XL330_M6 = _json.load(_f)
|
||
|
||
VIN = 7.4
|
||
KP_FW = 200.0
|
||
ACTUATOR_NAME = "1"
|
||
|
||
# Build BAM's M6 model + XL330 voltage-controlled actuator. The
|
||
# MujocoController below drives the joint via this model on every step,
|
||
# writing torque to data.ctrl and updating dof_frictionloss/dof_damping
|
||
# so MuJoCo's solver applies BAM's Stribeck+load+quadratic friction.
|
||
bam_model = bam_models["m6"]()
|
||
bam_model.set_actuator(bam_actuators["xl330"]())
|
||
bam_model.actuator.kp = KP_FW
|
||
bam_model.actuator.vin = VIN
|
||
bam_model.load_parameters_from_dict(DEFAULT_XL330_M6)
|
||
|
||
kt = bam_model.kt.value
|
||
R = bam_model.R.value
|
||
|
||
spec = mujoco.MjSpec.from_file(str(TESTBENCH_XML))
|
||
_set_arm_mass(spec, TESTBENCH_ARM_MASS)
|
||
|
||
# MujocoController needs a torque-controlled motor; the XL330 entry in the
|
||
# XML is a position actuator, so convert it and set the voltage-bounded
|
||
# force range. Armature is set on the dof by MujocoController.__init__.
|
||
for act in spec.actuators:
|
||
act.set_to_motor()
|
||
act.forcelimited = False
|
||
fl = VIN * kt / R
|
||
act.forcerange = (-fl, fl)
|
||
act.gear = [1.0, 0, 0, 0, 0, 0]
|
||
for joint in spec.joints:
|
||
if joint.type == mujoco.mjtJoint.mjJNT_HINGE:
|
||
joint.damping = 0.0
|
||
joint.frictionloss = 0.0
|
||
|
||
model = spec.compile()
|
||
data = mujoco.MjData(model)
|
||
model.opt.timestep = SIM_DT
|
||
|
||
joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "1")
|
||
dof_id = int(model.jnt_dofadr[joint_id])
|
||
qpos_id = int(model.jnt_qposadr[joint_id])
|
||
|
||
data.qpos[qpos_id] = 0.0
|
||
data.qvel[dof_id] = 0.0
|
||
mujoco.mj_forward(model, data)
|
||
|
||
bam_ctrl = MujocoController(bam_model, ACTUATOR_NAME, model, data)
|
||
bam_ctrl.reset(data.qpos)
|
||
|
||
runner = PolicyRunner(onnx_path, action_scale=action_scale)
|
||
policy_targets = make_target_schedule(total_time, seed=seed)
|
||
decim = int(round(CONTROL_DT / SIM_DT))
|
||
|
||
# Logging at SIM_DT (200 Hz): decim samples per policy step.
|
||
N_log = len(policy_targets) * decim
|
||
rec = {k: np.zeros(N_log, dtype=np.float32)
|
||
for k in ("t", "target", "q", "qd", "action", "ctrl")}
|
||
|
||
t = 0.0
|
||
log_i = 0
|
||
for policy_i, target in enumerate(policy_targets):
|
||
q = float(data.qpos[qpos_id])
|
||
qd = float(data.qvel[dof_id])
|
||
goal = runner.step(q, qd, float(target))
|
||
action_raw = float(runner.last_action[0])
|
||
|
||
for _ in range(decim):
|
||
q = float(data.qpos[qpos_id])
|
||
dq = float(data.qvel[dof_id])
|
||
|
||
# ---- log at 200 Hz ----
|
||
rec["t"][log_i] = t
|
||
rec["target"][log_i] = target
|
||
rec["q"][log_i] = q
|
||
rec["qd"][log_i] = dq
|
||
rec["action"][log_i] = action_raw
|
||
rec["ctrl"][log_i] = goal
|
||
log_i += 1
|
||
|
||
# BAM owns control/torque/friction: set the target, then update()
|
||
# writes torque to data.ctrl and pushes friction/damping onto the
|
||
# dof so MuJoCo's solver applies them on the next step.
|
||
bam_ctrl.set_q_target(ACTUATOR_NAME, goal)
|
||
bam_ctrl.update()
|
||
mujoco.mj_step(model, data)
|
||
t += SIM_DT
|
||
|
||
return rec
|
||
|
||
|
||
def rollout_sim_mjlab(onnx_path: str, total_time: float, seed: int, action_scale: float) -> dict:
|
||
"""Sim rollout via the actual mjlab testbench env (same BAM M6 the policy was trained against).
|
||
|
||
Boots make_testbench_env_cfg() with num_envs=1, overrides the target_angle
|
||
command with our deterministic schedule each policy tick, and steps the env
|
||
with the policy action. We replicate ManagerBasedRlEnv.step's inner
|
||
decimation loop manually so we can log q/qd at SIM_DT (200 Hz) between
|
||
sub-steps, matching the bam backend's logging rate.
|
||
"""
|
||
import torch
|
||
|
||
from mjlab.envs import ManagerBasedRlEnv
|
||
|
||
from mjlab_microduck.tasks.testbench_env_cfg import make_testbench_env_cfg
|
||
|
||
env_cfg = make_testbench_env_cfg(play=True)
|
||
env_cfg.scene.num_envs = 1
|
||
# Disable auto-resampling and auto-reset so our deterministic schedule and
|
||
# initial pose hold for the entire rollout.
|
||
env_cfg.commands["target_angle"].resampling_time_range = (1e6, 1e6)
|
||
env_cfg.episode_length_s = max(total_time + 10.0, env_cfg.episode_length_s)
|
||
# Drop observation noise so the mjlab path is a fair sim2real reference
|
||
# (matches the bam path which doesn't inject noise either).
|
||
env_cfg.observations["policy"].enable_corruption = False
|
||
|
||
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
||
env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
|
||
env.reset(seed=seed)
|
||
|
||
cmd_term = env.command_manager.get_term("target_angle")
|
||
robot = env.scene["robot"]
|
||
|
||
runner = PolicyRunner(onnx_path, action_scale=action_scale)
|
||
policy_targets = make_target_schedule(total_time, seed=seed)
|
||
decim = env.cfg.decimation
|
||
physics_dt = env.physics_dt
|
||
N_log = len(policy_targets) * decim
|
||
rec = {k: np.zeros(N_log, dtype=np.float32)
|
||
for k in ("t", "target", "q", "qd", "action", "ctrl")}
|
||
|
||
t = 0.0
|
||
log_i = 0
|
||
for target in policy_targets:
|
||
# Inject deterministic target and recompute obs so the policy sees it
|
||
# this tick (the env's TargetAngleCommand otherwise samples randomly).
|
||
cmd_term._target[0, 0] = float(target)
|
||
# update_history=True is critical: the testbench env's joint_vel obs
|
||
# has a 1-tick delay, so the history buffer must advance each policy
|
||
# tick or the policy sees stale velocity.
|
||
obs_buf = env.observation_manager.compute(update_history=True)
|
||
policy_obs = obs_buf["policy"][0].detach().cpu().numpy().astype(np.float32)
|
||
ort_out = runner.session.run(None, {runner.in_name: policy_obs[None, :]})[0].reshape(-1)
|
||
runner.last_action = ort_out.astype(np.float32)
|
||
action_raw = float(ort_out[0])
|
||
goal = DEFAULT_POS + action_raw * action_scale
|
||
|
||
# Manually run the decimation loop ManagerBasedRlEnv.step uses, so we
|
||
# can sample joint state at the physics rate (200 Hz).
|
||
action = torch.as_tensor(ort_out, device=device).reshape(1, -1)
|
||
env.action_manager.process_action(action)
|
||
for _ in range(decim):
|
||
# Log the pre-step state to mirror the bam backend (which records
|
||
# q/qd right before each mj_step).
|
||
rec["t"][log_i] = t
|
||
rec["target"][log_i] = float(target)
|
||
rec["q"][log_i] = float(robot.data.joint_pos[0, 0].item())
|
||
rec["qd"][log_i] = float(robot.data.joint_vel[0, 0].item())
|
||
rec["action"][log_i] = action_raw
|
||
rec["ctrl"][log_i] = goal
|
||
log_i += 1
|
||
|
||
env.action_manager.apply_action()
|
||
env.scene.write_data_to_sim()
|
||
env.sim.step()
|
||
env.scene.update(dt=physics_dt)
|
||
t += physics_dt
|
||
|
||
env.close()
|
||
return rec
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Real rollout (rustypot XL330)
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def rollout_real(
|
||
onnx_path: str,
|
||
total_time: float,
|
||
seed: int,
|
||
port: str,
|
||
motor_id: int,
|
||
baudrate: int,
|
||
kp: int,
|
||
action_scale: float,
|
||
) -> dict:
|
||
from rustypot import Xl330PyController
|
||
|
||
ctrl = Xl330PyController(port, baudrate, 0.05)
|
||
assert ctrl.ping(motor_id), f"motor id={motor_id} not responding on {port}"
|
||
|
||
# Match the firmware gain used in sim (BAM kp_fw=200).
|
||
ctrl.write_torque_enable(motor_id, False)
|
||
ctrl.write_operating_mode(motor_id, 3) # position control
|
||
ctrl.write_position_p_gain(motor_id, kp)
|
||
ctrl.write_position_i_gain(motor_id, 0)
|
||
ctrl.write_position_d_gain(motor_id, 0)
|
||
# Read back to confirm the gain actually landed (firmware silently clamps
|
||
# out-of-range values, so verifying catches mismatches early).
|
||
readback = ctrl.read_position_p_gain(motor_id)
|
||
if isinstance(readback, (list, tuple)):
|
||
readback = readback[0]
|
||
print(f" XL330 position P-gain: requested={kp}, readback={readback}")
|
||
ctrl.write_goal_position(motor_id, 0.0)
|
||
ctrl.write_torque_enable(motor_id, True)
|
||
time.sleep(1.0) # let it settle at zero
|
||
|
||
runner = PolicyRunner(onnx_path, action_scale=action_scale)
|
||
policy_targets = make_target_schedule(total_time, seed=seed)
|
||
|
||
decim = int(round(CONTROL_DT / LOG_DT)) # samples per policy tick (4 at 200 Hz / 50 Hz)
|
||
N_log = len(policy_targets) * decim
|
||
rec = {k: np.zeros(N_log, dtype=np.float32)
|
||
for k in ("t", "target", "q", "qd", "action", "ctrl")}
|
||
|
||
def _scalar(x) -> float:
|
||
if isinstance(x, (list, tuple)):
|
||
return float(x[0])
|
||
return float(x)
|
||
|
||
t_start = time.perf_counter()
|
||
prev_q = 0.0
|
||
log_i = 0
|
||
goal = 0.0
|
||
action_raw = 0.0
|
||
|
||
for policy_i, target in enumerate(policy_targets):
|
||
tick_start = time.perf_counter()
|
||
target_f = float(target)
|
||
|
||
# Read once, run policy, write goal — all at the start of the 20 ms window.
|
||
q = _scalar(ctrl.read_present_position(motor_id))
|
||
try:
|
||
qd = _scalar(ctrl.read_present_velocity(motor_id)) * DXL_VEL_TICK_TO_RAD_S
|
||
except Exception:
|
||
qd = (q - prev_q) / CONTROL_DT
|
||
|
||
goal = runner.step(q, qd, target_f)
|
||
action_raw = float(runner.last_action[0])
|
||
# ctrl.write_goal_position(motor_id, float(np.clip(goal, -MAX_ANGLE, MAX_ANGLE)))
|
||
ctrl.write_goal_position(motor_id, float(goal))
|
||
|
||
# First 200 Hz sample uses the values we just read (no extra USB round-trip).
|
||
rec["t"][log_i] = time.perf_counter() - t_start
|
||
rec["target"][log_i] = target_f
|
||
rec["q"][log_i] = q
|
||
rec["qd"][log_i] = qd
|
||
rec["action"][log_i] = action_raw
|
||
rec["ctrl"][log_i] = goal
|
||
prev_q = q
|
||
log_i += 1
|
||
|
||
# Remaining (decim-1) samples inside the policy window: read only.
|
||
for k in range(1, decim):
|
||
sample_deadline = tick_start + (k + 1) * LOG_DT
|
||
while time.perf_counter() < sample_deadline - 0.001:
|
||
time.sleep(0.0005)
|
||
q = _scalar(ctrl.read_present_position(motor_id))
|
||
try:
|
||
qd = _scalar(ctrl.read_present_velocity(motor_id)) * DXL_VEL_TICK_TO_RAD_S
|
||
except Exception:
|
||
qd = (q - prev_q) / LOG_DT
|
||
prev_q = q
|
||
|
||
rec["t"][log_i] = time.perf_counter() - t_start
|
||
rec["target"][log_i] = target_f
|
||
rec["q"][log_i] = q
|
||
rec["qd"][log_i] = qd
|
||
rec["action"][log_i] = action_raw
|
||
rec["ctrl"][log_i] = goal
|
||
log_i += 1
|
||
|
||
# Live status on every new segment plus a heartbeat.
|
||
new_segment = policy_i == 0 or policy_targets[policy_i] != policy_targets[policy_i - 1]
|
||
if new_segment or policy_i % 25 == 0:
|
||
print(
|
||
f"\r t={rec['t'][log_i-1]:6.2f}s target={math.degrees(target_f):+6.1f}° "
|
||
f"q={math.degrees(q):+6.1f}° err={math.degrees(q - target_f):+6.1f}° "
|
||
f"goal={math.degrees(goal):+6.1f}°",
|
||
end="" if not new_segment else "\n",
|
||
flush=True,
|
||
)
|
||
|
||
# Hold the remaining time of the policy window if we got here early.
|
||
dt_left = CONTROL_DT - (time.perf_counter() - tick_start)
|
||
if dt_left > 0:
|
||
time.sleep(dt_left)
|
||
print()
|
||
|
||
ctrl.write_torque_enable(motor_id, False)
|
||
return rec
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Plotting / analytics
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _mae(a: np.ndarray, b: np.ndarray) -> float:
|
||
n = min(len(a), len(b))
|
||
return float(np.mean(np.abs(a[:n] - b[:n])))
|
||
|
||
|
||
def npz_to_bam_log(npz_path: str, json_path: str, *, mass: float, length: float,
|
||
kp: int, vin: float) -> None:
|
||
"""Convert a rollout .npz (written by rollout_sim/rollout_real) to a BAM log json.
|
||
|
||
BAM log format (see ~/Rhoban/bam/bam/logs.py):
|
||
top-level: mass, length, kp, vin, motor, trajectory, dt
|
||
entries: position, speed, load, input_volts, temp, goal_position, torque_enable, timestamp
|
||
Can be fed to `python -m bam.plot --logdir <dir> --actuator xl330`.
|
||
"""
|
||
import json
|
||
|
||
d = dict(np.load(npz_path))
|
||
t = d["t"]
|
||
# Prefer the actual recorded timestamps for dt to handle small jitter;
|
||
# fall back to the fixed control period if there are fewer than 2 samples.
|
||
dt = float(np.mean(np.diff(t))) if len(t) > 1 else CONTROL_DT
|
||
|
||
entries = []
|
||
for i in range(len(t)):
|
||
entries.append({
|
||
"position": float(d["q"][i]),
|
||
"speed": float(d["qd"][i]),
|
||
"load": 0.0,
|
||
"input_volts": vin,
|
||
"temp": 25.0,
|
||
"goal_position": float(d["ctrl"][i]),
|
||
"torque_enable": True,
|
||
"timestamp": float(t[i]),
|
||
})
|
||
|
||
log = {
|
||
"mass": mass,
|
||
"length": length,
|
||
"kp": kp,
|
||
"vin": vin,
|
||
"motor": "xl330",
|
||
"trajectory": "rl_policy",
|
||
"dt": dt,
|
||
"entries": entries,
|
||
}
|
||
|
||
out = Path(json_path)
|
||
out.parent.mkdir(parents=True, exist_ok=True)
|
||
with open(out, "w") as f:
|
||
json.dump(log, f)
|
||
print(f"Wrote BAM log: {out} ({len(entries)} entries, dt={dt:.4f}s, mass={mass}kg, kp={kp})")
|
||
print(f" Replay with: (cd ~/Rhoban/bam && python -m bam.plot --logdir {out.parent} --actuator xl330)")
|
||
|
||
|
||
def compare_and_plot(sim_file: str, real_file: str, out_path: str) -> None:
|
||
import matplotlib.pyplot as plt
|
||
|
||
sim = dict(np.load(sim_file))
|
||
real = dict(np.load(real_file))
|
||
n = min(len(sim["t"]), len(real["t"]))
|
||
t = sim["t"][:n]
|
||
|
||
err_sim = sim["q"][:n] - sim["target"][:n]
|
||
err_real = real["q"][:n] - real["target"][:n]
|
||
|
||
print("\n=== Analytics ===")
|
||
print(f" steps compared : {n}")
|
||
print(f" MAE q (sim vs real) : {_mae(sim['q'], real['q']):.4f} rad "
|
||
f"({math.degrees(_mae(sim['q'], real['q'])):.2f}°)")
|
||
print(f" sim tracking MAE : {float(np.mean(np.abs(err_sim))):.4f} rad "
|
||
f"({math.degrees(float(np.mean(np.abs(err_sim)))):.2f}°)")
|
||
print(f" real tracking MAE : {float(np.mean(np.abs(err_real))):.4f} rad "
|
||
f"({math.degrees(float(np.mean(np.abs(err_real)))):.2f}°)")
|
||
print(f" sim qd RMS : {float(np.sqrt(np.mean(sim['qd'][:n]**2))):.3f} rad/s")
|
||
print(f" real qd RMS : {float(np.sqrt(np.mean(real['qd'][:n]**2))):.3f} rad/s")
|
||
print(f" action MAE : {_mae(sim['action'], real['action']):.4f} rad")
|
||
|
||
fig, axes = plt.subplots(4, 1, figsize=(11, 10), sharex=True)
|
||
|
||
axes[0].plot(t, sim["target"][:n], "k-", lw=1, label="target", alpha=0.4)
|
||
axes[0].plot(t, sim["q"][:n], "b-", lw=1.2, label="sim q")
|
||
axes[0].plot(t, real["q"][:n], "r-", lw=1.2, label="real q")
|
||
axes[0].plot(t, sim["ctrl"][:n], "b:", lw=0.8, alpha=0.6, label="sim goal (policy)")
|
||
axes[0].plot(t, real["ctrl"][:n], "r:", lw=0.8, alpha=0.6, label="real goal (policy)")
|
||
axes[0].set_ylabel("position [rad]")
|
||
axes[0].legend(loc="upper right", fontsize=8)
|
||
axes[0].grid(alpha=0.3)
|
||
axes[0].set_title(f"Testbench sim2real — MAE(sim, real) = {_mae(sim['q'], real['q']):.4f} rad")
|
||
|
||
axes[1].plot(t, np.degrees(err_sim), "b-", lw=1, label="sim")
|
||
axes[1].plot(t, np.degrees(err_real), "r-", lw=1, label="real")
|
||
axes[1].axhline(0, color="k", lw=0.5)
|
||
axes[1].set_ylabel("tracking error [deg]")
|
||
axes[1].legend(fontsize=8)
|
||
axes[1].grid(alpha=0.3)
|
||
|
||
axes[2].plot(t, sim["qd"][:n], "b-", lw=1, label="sim")
|
||
axes[2].plot(t, real["qd"][:n], "r-", lw=1, label="real")
|
||
axes[2].set_ylabel("velocity [rad/s]")
|
||
axes[2].legend(fontsize=8)
|
||
axes[2].grid(alpha=0.3)
|
||
|
||
axes[3].plot(t, sim["action"][:n], "b-", lw=1, label="sim action")
|
||
axes[3].plot(t, real["action"][:n], "r-", lw=1, label="real action")
|
||
axes[3].set_ylabel("policy action [rad]")
|
||
axes[3].set_xlabel("time [s]")
|
||
axes[3].legend(fontsize=8)
|
||
axes[3].grid(alpha=0.3)
|
||
|
||
plt.tight_layout()
|
||
plt.savefig(out_path, dpi=140)
|
||
print(f"\nSaved plot: {out_path}")
|
||
plt.show()
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def main():
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--mode", choices=["sim", "real"], help="Rollout mode")
|
||
ap.add_argument("--sim-backend", choices=["bam", "mjlab"], default="bam",
|
||
help="Sim backend: 'bam' uses bam.MujocoController on a vanilla "
|
||
"mujoco loop (200 Hz log, lightweight); 'mjlab' boots the actual "
|
||
"make_testbench_env_cfg() mjlab env with its BamM6Actuator (50 Hz log).")
|
||
ap.add_argument("--onnx", type=str, help="Path to trained ONNX policy")
|
||
ap.add_argument("--out", type=str, help="Output .npz log file")
|
||
ap.add_argument("--duration", type=float, default=30.0, help="Total time [s]")
|
||
ap.add_argument("--seed", type=int, default=0, help="Target schedule seed")
|
||
# real-only
|
||
ap.add_argument("--port", type=str, default="/dev/ttyUSB0")
|
||
ap.add_argument("--motor-id", type=int, default=1)
|
||
ap.add_argument("--baudrate", type=int, default=1_000_000)
|
||
ap.add_argument("--kp", type=int, default=200, help="XL330 position P gain")
|
||
ap.add_argument("--action-scale", type=float, default=1.0,
|
||
help="Multiplier applied to the policy action before offsetting "
|
||
"by the default pose (must match training env action scale)")
|
||
# compare mode
|
||
ap.add_argument("--compare", nargs=2, metavar=("SIM_NPZ", "REAL_NPZ"),
|
||
help="Plot two logged runs side by side")
|
||
ap.add_argument("--out-plot", type=str, default="testbench_sim2real.png")
|
||
# BAM log export
|
||
ap.add_argument("--to-bam", nargs=2, metavar=("NPZ", "JSON"),
|
||
help="Convert a rollout NPZ to BAM log format "
|
||
"(run with: python -m bam.plot --logdir <dir> --actuator xl330)")
|
||
ap.add_argument("--bam-mass", type=float, default=TESTBENCH_ARM_MASS, help="Payload mass [kg]")
|
||
ap.add_argument("--bam-length", type=float, default=0.1, help="Arm length [m]")
|
||
ap.add_argument("--bam-vin", type=float, default=7.4, help="Supply voltage [V]")
|
||
|
||
args = ap.parse_args()
|
||
|
||
if args.compare:
|
||
compare_and_plot(args.compare[0], args.compare[1], args.out_plot)
|
||
return
|
||
|
||
if args.to_bam:
|
||
npz_to_bam_log(
|
||
args.to_bam[0],
|
||
args.to_bam[1],
|
||
mass=args.bam_mass,
|
||
length=args.bam_length,
|
||
kp=args.kp,
|
||
vin=args.bam_vin,
|
||
)
|
||
return
|
||
|
||
if not (args.mode and args.onnx and args.out):
|
||
ap.error("--mode, --onnx and --out are required for a rollout")
|
||
|
||
if args.mode == "sim":
|
||
sim_fn = rollout_sim_mjlab if args.sim_backend == "mjlab" else rollout_sim_bam
|
||
rec = sim_fn(args.onnx, args.duration, args.seed, args.action_scale)
|
||
else:
|
||
rec = rollout_real(
|
||
args.onnx, args.duration, args.seed,
|
||
args.port, args.motor_id, args.baudrate, args.kp,
|
||
args.action_scale,
|
||
)
|
||
|
||
out = Path(args.out)
|
||
out.parent.mkdir(parents=True, exist_ok=True)
|
||
np.savez(out, **rec)
|
||
err = rec["q"] - rec["target"]
|
||
print(f"\nSaved {len(rec['t'])} samples to {out}")
|
||
print(f" tracking MAE: {float(np.mean(np.abs(err))):.4f} rad ({math.degrees(float(np.mean(np.abs(err)))):.2f}°)")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|