"""Validate the BAM M6 actuator kernel against real testbench data. Loads real testbench recordings, replays them in MuJoCo with the BAM M6 actuator, and compares simulated vs real position traces. Also runs BAM's own Python simulator as a reference. Usage: uv run python3 scripts/validate_bam_testbench.py [--plot] [--max-files N] """ import argparse import json import os import sys from copy import copy from pathlib import Path import mujoco import numpy as np # ── Paths ── BAM_DIR = Path(os.path.expanduser("~/Rhoban/bam")) DATA_DIR = BAM_DIR / "bam" / "data" / "processed" PARAMS_FILE = BAM_DIR / "params" / "xl330" / "m6_new.json" TESTBENCH_XML = ( Path(__file__).resolve().parent.parent / "src" / "mjlab_microduck" / "robot" / "xl330_test_bench" / "scene.xml" ) # ── Load M6 params ── with open(PARAMS_FILE) as f: M6 = json.load(f) # XL330 firmware constants ERROR_GAIN = (4096 / (2 * np.pi)) / (256 * 885) VIN = 7.4 MAX_PWM = 1.0 def bam_python_rollout(log: dict) -> list[float]: """Reference: BAM's own Python simulator.""" sys.path.insert(0, str(BAM_DIR)) from bam.model import load_model from bam.simulate import Simulator # BAM expects arm_mass in the log (mass of the arm itself, not the payload) if "arm_mass" not in log: log = dict(log) log["arm_mass"] = 0.0 model = load_model(str(PARAMS_FILE)) sim = Simulator(model) result = sim.rollout_log(log, simulate_control=True) return result[0] # positions def compute_m6_friction(motor_torque, external_torque, dq): """M6 friction computation matching our kernel (and BAM's model.py).""" p = M6 stribeck_coeff = np.exp(-(np.abs(dq / p["dtheta_stribeck"]) ** p["alpha"])) gearbox_torque = np.abs( external_torque * p["load_friction_external"] - motor_torque * p["load_friction_motor"] ) gearbox_torque_stribeck = np.abs( external_torque * p["load_friction_external_stribeck"] - motor_torque * p["load_friction_motor_stribeck"] ) frictionloss = p["friction_base"] frictionloss += gearbox_torque frictionloss += stribeck_coeff * p["friction_stribeck"] frictionloss += gearbox_torque_stribeck * stribeck_coeff # quadratic (tiny, skip for clarity) damping = p["friction_viscous"] friction_budget = frictionloss + damping * np.abs(dq) return friction_budget def mujoco_rollout(log: dict) -> list[float]: """Run the testbench in MuJoCo with our BAM M6 actuator logic.""" mass = log["mass"] kp_fw = log["kp"] dt = log["dt"] entries = log["entries"] # Load and modify the testbench model spec = mujoco.MjSpec.from_file(str(TESTBENCH_XML)) # Convert actuator to motor (same as our kernel's edit_spec) for act in spec.actuators: act.set_to_motor() act.forcelimited = True force_limit = VIN * M6["kt"] / M6["R"] act.forcerange = (-force_limit, force_limit) act.gear = [1.0, 0, 0, 0, 0, 0] # Zero out MuJoCo joint friction (we handle it) for joint in spec.joints: if joint.type == mujoco.mjtJoint.mjJNT_HINGE: joint.damping = 0.0 joint.frictionloss = 0.0 joint.armature = M6["armature"] # Set the arm mass to match the BAM recording for body in spec.bodies: if body.name == "arm": # Scale mass and inertia proportionally original_mass = body.mass scale = mass / original_mass if original_mass > 0 else 1.0 body.mass = mass # Scale inertia proportionally to mass body.fullinertia = [x * scale for x in body.fullinertia] break model = spec.compile() data = mujoco.MjData(model) model.opt.timestep = dt # Find joint and actuator IDs joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "1") dof_id = model.jnt_dofadr[joint_id] # Initialize state data.qpos[dof_id] = entries[0]["position"] data.qvel[dof_id] = entries[0].get("speed", 0.0) mujoco.mj_forward(model, data) positions = [] for entry in entries: positions.append(float(data.qpos[dof_id])) if not entry["torque_enable"]: data.ctrl[0] = 0.0 mujoco.mj_step(model, data) continue goal = entry["goal_position"] q = data.qpos[dof_id] dq = data.qvel[dof_id] # ── BAM M6 actuator logic (same as our kernel) ── # 1. Firmware control law duty = (goal - q) * kp_fw * ERROR_GAIN duty = np.clip(duty, -MAX_PWM, MAX_PWM) voltage = VIN * duty # 2. DC motor torque motor_torque = M6["kt"] * voltage / M6["R"] - M6["kt"] ** 2 * dq / M6["R"] # 3. External torque (from MuJoCo bias forces) # BAM convention: bias_torque = m*g*l*sin(q) with g=-9.81 (gravity negative) # MuJoCo convention: qfrc_bias has opposite sign external_torque = -data.qfrc_bias[dof_id] # 4. M6 friction friction_budget = compute_m6_friction(motor_torque, external_torque, dq) # 5. Static friction clipping eff_inertia = 1.0 / model.dof_invweight0[dof_id] if model.dof_invweight0[dof_id] > 0 else 1e6 net_no_friction = motor_torque + external_torque tau_stop = (eff_inertia / dt) * dq + net_no_friction friction_mag = min(abs(tau_stop), friction_budget) friction_torque = -np.sign(tau_stop) * friction_mag # 6. Set ctrl = motor + friction (MuJoCo adds qfrc_bias) data.ctrl[0] = motor_torque + friction_torque mujoco.mj_step(model, data) return positions def main(): parser = argparse.ArgumentParser() parser.add_argument("--plot", action="store_true", help="Show plots") parser.add_argument("--max-files", type=int, default=5) args = parser.parse_args() data_files = sorted(DATA_DIR.glob("*.json")) if args.max_files: data_files = data_files[: args.max_files] print(f"Validating BAM M6 kernel against {len(data_files)} testbench recordings") print(f"M6 params: kt={M6['kt']:.4f} R={M6['R']:.4f}") print(f"Testbench XML: {TESTBENCH_XML}") print() results = [] for fpath in data_files: log = json.load(open(fpath)) name = f"{log['trajectory']}_m{log['mass']}_kp{log['kp']}" print(f" {name}...", end=" ", flush=True) real_pos = [e["position"] for e in log["entries"]] # BAM Python reference bam_pos = bam_python_rollout(log) # Our MuJoCo M6 kernel mj_pos = mujoco_rollout(log) # Compute MAE real_np = np.array(real_pos) bam_np = np.array(bam_pos) mj_np = np.array(mj_pos[: len(real_np)]) mae_bam = np.mean(np.abs(bam_np - real_np)) mae_mj = np.mean(np.abs(mj_np - real_np)) mae_bam_vs_mj = np.mean(np.abs(bam_np - mj_np)) print( f"MAE bam_vs_real={mae_bam:.5f} mj_vs_real={mae_mj:.5f} bam_vs_mj={mae_bam_vs_mj:.5f}" ) results.append( { "name": name, "real": real_np, "bam": bam_np, "mj": mj_np, "mae_bam": mae_bam, "mae_mj": mae_mj, "mae_bam_vs_mj": mae_bam_vs_mj, } ) print() avg_bam = np.mean([r["mae_bam"] for r in results]) avg_mj = np.mean([r["mae_mj"] for r in results]) avg_diff = np.mean([r["mae_bam_vs_mj"] for r in results]) print(f"Average MAE bam_vs_real={avg_bam:.5f} mj_vs_real={avg_mj:.5f} bam_vs_mj={avg_diff:.5f}") if avg_diff > 0.01: print("\n⚠ BAM and MuJoCo diverge significantly — likely a kernel bug!") elif avg_mj > avg_bam * 1.5: print("\n⚠ MuJoCo worse than BAM — MuJoCo dynamics differ from BAM's simple integrator") else: print("\n✓ BAM and MuJoCo agree — kernel is correct") if args.plot: try: import matplotlib.pyplot as plt n = len(results) fig, axes = plt.subplots(n, 1, figsize=(12, 3 * n), sharex=False) if n == 1: axes = [axes] for ax, r in zip(axes, results): t = np.arange(len(r["real"])) * 0.005 ax.plot(t, r["real"], "k-", lw=1.5, label="Real") ax.plot(t, r["bam"], "b--", lw=1.2, label=f'BAM (MAE={r["mae_bam"]:.4f})') ax.plot(t, r["mj"], "r:", lw=1.2, label=f'MuJoCo M6 (MAE={r["mae_mj"]:.4f})') ax.set_title(r["name"]) ax.set_ylabel("Position (rad)") ax.legend(fontsize=8) ax.grid(alpha=0.3) axes[-1].set_xlabel("Time (s)") plt.tight_layout() plt.savefig("bam_validation.png", dpi=150) print("Saved bam_validation.png") plt.show() except ImportError: print("matplotlib not available, skipping plots") if __name__ == "__main__": main()