# Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md). # All Rights Reserved. # # SPDX-License-Identifier: BSD-3-Clause """ Visualize system identification fit by replaying real waypoints through the manager-based Sysid env (same RelCartesianOSCAction as sysid_ur5e_osc.py). Loads a CMA-ES checkpoint, applies best params to the sim, runs closed-loop replay, and plots sim vs. real joint trajectories. Usage: python scripts_v2/tools/sim2real/plot_sysid_fit.py --headless \ --checkpoint logs/sysid/YYYYMMDD_HHMMSS/checkpoint_0200.pt \ --real_data sysid_data_real.pt """ import argparse import matplotlib import numpy as np import os import torch matplotlib.use("Agg") import gymnasium as gym import matplotlib.pyplot as plt from isaaclab.app import AppLauncher parser = argparse.ArgumentParser(description="Plot sysid fit") parser.add_argument("--checkpoint", type=str, required=True, help="Path to sysid checkpoint .pt") parser.add_argument("--real_data", type=str, required=True, help="Path to real data .pt") parser.add_argument("--max_steps", type=int, default=None) AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app from isaaclab.actuators import DelayedPDActuatorCfg from isaaclab.assets import Articulation from isaaclab.utils.math import subtract_frame_transforms from uwlab_assets.robots.ur5e_robotiq_gripper.kinematics import ARM_JOINT_NAMES, EE_BODY_NAME, NUM_ARM_JOINTS import uwlab_tasks # noqa: F401 # register gym envs from uwlab_tasks.manager_based.manipulation.omnireset.config.ur5e_robotiq_2f85.sysid_cfg import SysidEnvCfg from uwlab_tasks.manager_based.manipulation.omnireset.mdp.utils import settle_robot, target_pose_to_action # ============================================================================ # Parameter application (same as sysid_ur5e_osc.py) # ============================================================================ def apply_params(robot, params, arm_joint_ids, num_joints, device): """Apply 25-element param vector (single env) to robot.""" N = 1 env_ids = torch.arange(N, device=device) p = torch.tensor(params, device=device, dtype=torch.float32).unsqueeze(0) armature_full = torch.zeros(N, num_joints, device=device) static_friction_full = torch.zeros(N, num_joints, device=device) dynamic_friction_full = torch.zeros(N, num_joints, device=device) viscous_friction_full = torch.zeros(N, num_joints, device=device) armature_full[:, arm_joint_ids] = p[:, 0:6] static_fric = p[:, 6:12] dynamic_ratio = p[:, 12:18] static_friction_full[:, arm_joint_ids] = static_fric dynamic_friction_full[:, arm_joint_ids] = dynamic_ratio * static_fric viscous_friction_full[:, arm_joint_ids] = p[:, 18:24] robot.write_joint_armature_to_sim(armature_full, env_ids=env_ids) robot.write_joint_friction_coefficient_to_sim( static_friction_full, joint_dynamic_friction_coeff=dynamic_friction_full, joint_viscous_friction_coeff=viscous_friction_full, env_ids=env_ids, ) delay_int = int(round(float(p[0, 24]))) arm_actuator = robot.actuators["arm"] delay_tensor = torch.tensor([delay_int], device=device, dtype=torch.int) arm_actuator.positions_delay_buffer.set_time_lag(delay_tensor) arm_actuator.velocities_delay_buffer.set_time_lag(delay_tensor) arm_actuator.efforts_delay_buffer.set_time_lag(delay_tensor) # ============================================================================ # Closed-loop replay # ============================================================================ def closed_loop_replay( env, wp_step_indices, wp_target_pos, wp_target_quat, initial_joint_pos, arm_joint_ids, ee_frame_idx, sim_dt, T_steps, headless=True, ): """Run closed-loop replay using env's RelCartesianOSC. Returns sim trajectory dict.""" unwrapped = env.unwrapped robot = unwrapped.scene["robot"] sim = unwrapped.sim device = unwrapped.device action_dim = unwrapped.action_manager.total_action_dim W = wp_step_indices.shape[0] default_joint_pos = robot.data.default_joint_pos.clone() default_joint_vel = robot.data.default_joint_vel.clone() default_joint_pos[:, arm_joint_ids] = initial_joint_pos.unsqueeze(0) default_joint_vel[:] = 0.0 env.reset() settle_robot(robot, sim, default_joint_pos, default_joint_vel, arm_joint_ids, sim_dt, headless=headless) sim_positions, sim_velocities, sim_ee_positions = [], [], [] wp_idx = 0 for t in range(T_steps): while wp_idx + 1 < W and t >= wp_step_indices[wp_idx + 1]: wp_idx += 1 ee_pos_w = robot.data.body_pos_w[:, ee_frame_idx] ee_quat_w = robot.data.body_quat_w[:, ee_frame_idx] ee_pos_b, ee_quat_b = subtract_frame_transforms( robot.data.root_pos_w, robot.data.root_quat_w, ee_pos_w, ee_quat_w ) target_pos = wp_target_pos[wp_idx].unsqueeze(0) target_quat = wp_target_quat[wp_idx].unsqueeze(0) action_arm = target_pose_to_action(ee_pos_b, ee_quat_b, target_pos, target_quat) action = torch.cat([action_arm, torch.zeros(1, action_dim - 6, device=device)], dim=-1) env.step(action) joint_pos = robot.data.joint_pos[:, arm_joint_ids] joint_vel = robot.data.joint_vel[:, arm_joint_ids] sim_positions.append(joint_pos[0].cpu().numpy().copy()) sim_velocities.append(joint_vel[0].cpu().numpy().copy()) sim_ee_positions.append(ee_pos_b[0].cpu().numpy().copy()) if (t + 1) % max(1, T_steps // 20) == 0: print(f" step {t+1}/{T_steps} ({100*(t+1)/T_steps:.0f}%)") return { "joint_positions": np.array(sim_positions), "joint_velocities": np.array(sim_velocities), "ee_positions": np.array(sim_ee_positions), } # ============================================================================ # Plotting # ============================================================================ JOINT_NAMES_SHORT = ["Shoulder Pan", "Shoulder Lift", "Elbow", "Wrist 1", "Wrist 2", "Wrist 3"] def plot_overlay(real_joints, sim_joints, dt, save_path="sysid_fit.png"): """Plot sim vs real joint positions.""" T = real_joints.shape[0] time_axis = np.arange(T) * dt fig, axes = plt.subplots(3, 2, figsize=(16, 10), sharex=True) axes = axes.flatten() for j in range(NUM_ARM_JOINTS): ax = axes[j] ax.plot(time_axis, np.degrees(real_joints[:, j]), "b-", linewidth=1.0, label="Real", alpha=0.8) ax.plot(time_axis, np.degrees(sim_joints[:, j]), "r-", linewidth=1.0, label="Sim", alpha=0.8) ax.set_title(JOINT_NAMES_SHORT[j], fontsize=11) ax.set_ylabel("deg") ax.legend(loc="upper right", fontsize=8) ax.grid(True, alpha=0.3) axes[-2].set_xlabel("Time (s)") axes[-1].set_xlabel("Time (s)") fig.suptitle("Sysid Fit: Sim vs Real Joint Trajectories", fontsize=13) fig.tight_layout() fig.savefig(save_path, dpi=150) print(f"Saved overlay plot: {save_path}") plt.close(fig) def plot_error(real_joints, sim_joints, dt, save_path="sysid_fit_error.png"): """Plot per-joint error over time.""" T = real_joints.shape[0] time_axis = np.arange(T) * dt error_deg = np.degrees(sim_joints - real_joints) fig, axes = plt.subplots(3, 2, figsize=(16, 10), sharex=True) axes = axes.flatten() for j in range(NUM_ARM_JOINTS): ax = axes[j] ax.plot(time_axis, error_deg[:, j], "k-", linewidth=0.8, alpha=0.7) ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5) rmse_j = np.sqrt(np.mean(error_deg[:, j] ** 2)) ax.set_title(f"{JOINT_NAMES_SHORT[j]} (RMSE={rmse_j:.2f}°)", fontsize=11) ax.set_ylabel("Error (deg)") ax.grid(True, alpha=0.3) axes[-2].set_xlabel("Time (s)") axes[-1].set_xlabel("Time (s)") fig.suptitle("Sysid Fit: Per-Joint Error", fontsize=13) fig.tight_layout() fig.savefig(save_path, dpi=150) print(f"Saved error plot: {save_path}") plt.close(fig) # ============================================================================ # Main # ============================================================================ def main(): args = args_cli device_str = args.device # Load checkpoint print(f"\nLoading checkpoint: {args.checkpoint}") ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) best_params = ckpt["best_params"] best_score = ckpt["best_score"] ckpt_args = ckpt.get("args", {}) print(f" Score (MSE): {best_score:.6f} RMSE: {np.degrees(np.sqrt(best_score)):.4f}°") print(f" Checkpoint args: sim_dt={ckpt_args.get('sim_dt', 'N/A')}") # Print best params arm = best_params[:6] sfric = best_params[6:12] dratio = best_params[12:18] vfric = best_params[18:24] delay = round(float(best_params[24])) print(f"\n {'Joint':<20s} {'Armature':>10s} {'SFric':>10s} {'DRatio':>10s} {'VFric':>10s}") for i, name in enumerate(JOINT_NAMES_SHORT): print(f" {name:<20s} {arm[i]:10.4f} {sfric[i]:10.4f} {dratio[i]:10.4f} {vfric[i]:10.4f}") print(f" Motor delay: {delay} steps") # Load real data print(f"\nLoading real data: {args.real_data}") real_data = torch.load(args.real_data, map_location="cpu", weights_only=False) real_joint_pos = real_data["joint_positions"] initial_joint_pos = real_data["initial_joint_pos"] wp_step_indices = real_data["waypoint_step_indices"] wp_target_pos = real_data["waypoint_target_pos"] wp_target_quat = real_data["waypoint_target_quat"] dt = real_data["dt"] T_steps = real_joint_pos.shape[0] if args.max_steps is not None: T_steps = min(T_steps, args.max_steps) print(f" {T_steps} steps ({T_steps*dt:.2f}s), dt={dt*1000:.1f}ms") # Move to GPU real_joint_pos_np = real_joint_pos[:T_steps].numpy() initial_joint_pos_dev = initial_joint_pos.to(device_str).float() wp_step_indices = wp_step_indices.to(device_str).long() wp_target_pos = wp_target_pos.to(device_str).float() wp_target_quat = wp_target_quat.to(device_str).float() # Create env (same as sysid_ur5e_osc.py) env_cfg = SysidEnvCfg() env_cfg.scene.num_envs = 1 env_cfg.scene.env_spacing = 2.0 delay_max = max(delay, 5) _effort_lim = { "shoulder_pan_joint": 150.0, "shoulder_lift_joint": 150.0, "elbow_joint": 150.0, "wrist_1_joint": 28.0, "wrist_2_joint": 28.0, "wrist_3_joint": 28.0, } _vel_lim = { "shoulder_pan_joint": 1.5708, "shoulder_lift_joint": 1.5708, "elbow_joint": 1.5708, "wrist_1_joint": 3.1415, "wrist_2_joint": 3.1415, "wrist_3_joint": 3.1415, } env_cfg.scene.robot.actuators["arm"] = DelayedPDActuatorCfg( joint_names_expr=["shoulder.*", "elbow.*", "wrist.*"], stiffness=0.0, damping=0.0, effort_limit=_effort_lim, velocity_limit=_vel_lim, min_delay=0, max_delay=delay_max, ) env = gym.make("OmniReset-Ur5eRobotiq2f85-Sysid-v0", cfg=env_cfg) env.reset() unwrapped = env.unwrapped robot: Articulation = unwrapped.scene["robot"] device = unwrapped.device arm_joint_ids = robot.find_joints(ARM_JOINT_NAMES)[0] ee_frame_idx = robot.find_bodies(EE_BODY_NAME)[0][0] num_joints = robot.num_joints sim_dt = env_cfg.sim.dt # Apply best params print(f"\nApplying best params (delay={delay})...") apply_params(robot, best_params, arm_joint_ids, num_joints, device) # Run closed-loop replay print(f"\nRunning closed-loop replay ({T_steps} steps)...") result = closed_loop_replay( env, wp_step_indices, wp_target_pos, wp_target_quat, initial_joint_pos_dev, arm_joint_ids, ee_frame_idx, sim_dt, T_steps, headless=args_cli.headless, ) sim_joints = result["joint_positions"] real_joints = real_joint_pos_np # Compute per-joint RMSE error_deg = np.degrees(sim_joints - real_joints) print(f"\n{'='*60}") print("Per-joint RMSE (deg)") print("=" * 60) for j in range(NUM_ARM_JOINTS): rmse_j = np.sqrt(np.mean(error_deg[:, j] ** 2)) print(f" {JOINT_NAMES_SHORT[j]:<16s}: {rmse_j:.4f}") rmse_total = np.sqrt(np.mean(error_deg**2)) mae_total = np.mean(np.abs(error_deg)) max_total = np.max(np.abs(error_deg)) print(f" TOTAL : RMSE={rmse_total:.4f} MAE={mae_total:.4f} Max={max_total:.4f}") # Sysid-equivalent score for comparison with checkpoint error_rad = sim_joints - real_joints sysid_score = np.mean(np.sum(error_rad**2, axis=1)) sysid_rmse_deg = np.degrees(np.sqrt(sysid_score)) print(f"\n Sysid-equivalent metric: score={sysid_score:.6f} RMSE={sysid_rmse_deg:.4f}°") print(f" Checkpoint metric: score={best_score:.6f} RMSE={np.degrees(np.sqrt(best_score)):.4f}°") print("=" * 60) # Plot out_dir = os.path.dirname(args.checkpoint) if os.path.dirname(args.checkpoint) else "." plot_overlay(real_joints, sim_joints, dt, save_path=os.path.join(out_dir, "sysid_fit.png")) plot_error(real_joints, sim_joints, dt, save_path=os.path.join(out_dir, "sysid_fit_error.png")) if __name__ == "__main__": main() simulation_app.close()