#!/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 --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 --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()