From e632c3966b095871bf72e7fa641d5d76e56bf44d Mon Sep 17 00:00:00 2001 From: lizhuoyuan Date: Sun, 9 Aug 2026 08:44:49 +0000 Subject: [PATCH] add model inference evaluation --- scripts/compare_act.py | 433 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 433 insertions(+) create mode 100644 scripts/compare_act.py diff --git a/scripts/compare_act.py b/scripts/compare_act.py new file mode 100644 index 0000000..7b234d8 --- /dev/null +++ b/scripts/compare_act.py @@ -0,0 +1,433 @@ +#!/usr/bin/env python3 +""" +Compare ACT model inference at two checkpoints against ground truth. + +Usage: + cd /home/lizhuoyuan/project/lerobot_xarm7 + python scripts/compare_act.py --num-samples 8 --episode 0 +""" + +import argparse +import csv +import os +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import torch +from matplotlib.gridspec import GridSpec + +# ── lerobot imports ────────────────────────────────────────────── +from lerobot.policies.act.modeling_act import ACTPolicy +from lerobot.processor import PolicyProcessorPipeline +from lerobot.processor.converters import ( + batch_to_transition, + transition_to_batch, + policy_action_to_transition, + transition_to_policy_action, +) +from lerobot.datasets.lerobot_dataset import LeRobotDataset + +JOINT_NAMES = ["J1", "J2", "J3", "J4", "J5", "J6", "J7", "Gripper"] + + +def parse_args(): + p = argparse.ArgumentParser(description="Compare ACT 40K vs 60K inference") + p.add_argument("--dataset-root", default="./datasets/xarm7-pick-bottle", + help="Path to dataset root") + p.add_argument("--checkpoint-40k", + default="./outputs/train/2026-08-08/09-57-03_act/checkpoints/040000/pretrained_model", + help="Path to 40K checkpoint pretrained_model dir") + p.add_argument("--checkpoint-60k", + default="./outputs/train/2026-08-08/09-57-03_act/checkpoints/060000/pretrained_model", + help="Path to 60K checkpoint pretrained_model dir") + p.add_argument("--episodes", type=str, default="0", + help="Comma-separated episode indices, e.g. '0,10,20,30,40,50'") + p.add_argument("--num-samples", type=int, default=8, help="Number of frames to sample per episode") + p.add_argument("--output-dir", default="./outputs/compare_act", + help="Output directory for results") + return p.parse_args() + + +def load_policy_and_processors(checkpoint_path: str): + """Load ACT policy, preprocessor, and postprocessor from a checkpoint.""" + print(f" Loading policy from {checkpoint_path} ...") + policy = ACTPolicy.from_pretrained(checkpoint_path) + policy.to("cpu") + policy.reset() + + preprocessor = PolicyProcessorPipeline.from_pretrained( + pretrained_model_name_or_path=checkpoint_path, + config_filename="policy_preprocessor.json", + overrides={"device_processor": {"device": "cpu"}}, + to_transition=batch_to_transition, + to_output=transition_to_batch, + ) + postprocessor = PolicyProcessorPipeline.from_pretrained( + pretrained_model_name_or_path=checkpoint_path, + config_filename="policy_postprocessor.json", + overrides={"device_processor": {"device": "cpu"}}, + to_transition=policy_action_to_transition, + to_output=transition_to_policy_action, + ) + return policy, preprocessor, postprocessor + + +def get_episode_frame_range(dataset: LeRobotDataset, episode_idx: int): + """Get (from_idx, to_idx) for a specific episode. + + dataset.meta.episodes is a HuggingFace Dataset with columns: + episode_index, dataset_from_index, dataset_to_index, length, ... + """ + eps = dataset.meta.episodes + for i in range(len(eps)): + if int(eps[i]["episode_index"]) == episode_idx: + from_idx = int(eps[i]["dataset_from_index"]) + to_idx = int(eps[i]["dataset_to_index"]) + return from_idx, to_idx + raise ValueError(f"Episode {episode_idx} not found in dataset") + + +def run_inference(policy, preprocessor, postprocessor, obs_dict: dict) -> np.ndarray: + """Run inference on a single observation dict, return predicted action as (8,) numpy.""" + batch = preprocessor(obs_dict) + with torch.inference_mode(): + action_chunk = policy.predict_action_chunk(batch) # (1, chunk_size, 8) + action = action_chunk[:, 0, :] # (1, 8) + action = postprocessor(action) + return action.cpu().numpy().squeeze(0) # (8,) + + +def run_one_episode(args, episode_idx, dataset, policy_40k, preproc_40k, postproc_40k, + policy_60k, preproc_60k, postproc_60k): + """Run inference on one episode. Returns (results, images, mae_40k, mae_60k, l2_40k, l2_60k).""" + from_idx, to_idx = get_episode_frame_range(dataset, episode_idx) + total_frames = to_idx - from_idx + print(f" Episode {episode_idx}: frames [{from_idx}, {to_idx}), total={total_frames}") + + sample_indices = np.linspace(from_idx, to_idx - 1, args.num_samples, dtype=int) + print(f" Sampling {args.num_samples} frames at indices: {list(sample_indices)}") + + results = [] + images = [] + + for i, global_idx in enumerate(sample_indices): + print(f" Frame {i+1}/{args.num_samples} (global idx={global_idx}) ...") + frame = dataset[global_idx] + + gt_action = frame["action"].numpy().squeeze() + + img = frame["observation.images.camera"].numpy() + img = np.transpose(img, (1, 2, 0)) + img_uint8 = (img * 255).clip(0, 255).astype(np.uint8) + images.append(img_uint8) + + obs_dict = { + "observation.state": frame["observation.state"], + "observation.images.camera": frame["observation.images.camera"], + } + + policy_40k.reset() + policy_60k.reset() + + pred_40k = run_inference(policy_40k, preproc_40k, postproc_40k, obs_dict) + pred_60k = run_inference(policy_60k, preproc_60k, postproc_60k, obs_dict) + + results.append({ + "frame": i, + "global_idx": int(global_idx), + "gt": gt_action, + "pred_40k": pred_40k, + "pred_60k": pred_60k, + }) + + # Compute metrics + n_joints = 8 + gt_all = np.stack([r["gt"] for r in results]) + pred_40k_all = np.stack([r["pred_40k"] for r in results]) + pred_60k_all = np.stack([r["pred_60k"] for r in results]) + + err_40k = np.abs(pred_40k_all - gt_all) + err_60k = np.abs(pred_60k_all - gt_all) + mae_40k = err_40k.mean(axis=0) + mae_60k = err_60k.mean(axis=0) + l2_40k = np.sqrt(((pred_40k_all - gt_all) ** 2).sum(axis=1)) + l2_60k = np.sqrt(((pred_60k_all - gt_all) ** 2).sum(axis=1)) + + return results, images, mae_40k, mae_60k, l2_40k.mean(), l2_60k.mean() + + +def save_episode_plots(args, episode_idx, results, images, mae_40k, mae_60k, + gt_all, pred_40k_all, pred_60k_all, out_dir): + """Generate per-episode plots and CSV.""" + n_joints = 8 + + # CSV + csv_path = os.path.join(out_dir, "comparison.csv") + err_40k = np.abs(pred_40k_all - gt_all) + err_60k = np.abs(pred_60k_all - gt_all) + l2_40k = np.sqrt(((pred_40k_all - gt_all) ** 2).sum(axis=1)) + l2_60k = np.sqrt(((pred_60k_all - gt_all) ** 2).sum(axis=1)) + + with open(csv_path, "w", newline="") as f: + writer = csv.writer(f) + header = ["frame", "global_idx"] + for jn in JOINT_NAMES: + header += [f"GT_{jn}", f"40K_{jn}", f"60K_{jn}", f"err40K_{jn}", f"err60K_{jn}"] + header += ["L2_40K", "L2_60K"] + writer.writerow(header) + for i, r in enumerate(results): + row = [r["frame"], r["global_idx"]] + for j in range(n_joints): + row += [f"{r['gt'][j]:.6f}", f"{r['pred_40k'][j]:.6f}", + f"{r['pred_60k'][j]:.6f}", + f"{err_40k[i][j]:.6f}", f"{err_60k[i][j]:.6f}"] + row += [f"{l2_40k[i]:.6f}", f"{l2_60k[i]:.6f}"] + writer.writerow(row) + + # 图1: Trajectory curves + fig1, axes = plt.subplots(2, 4, figsize=(18, 9)) + axes = axes.flatten() + colors = {"GT": "black", "40K": "#2196F3", "60K": "#FF9800"} + x = np.arange(args.num_samples) + for j in range(n_joints): + ax = axes[j] + ax.plot(x, gt_all[:, j], "o-", color=colors["GT"], label="GT", linewidth=2, markersize=5) + ax.plot(x, pred_40k_all[:, j], "s--", color=colors["40K"], label="40K", linewidth=1.5, markersize=5) + ax.plot(x, pred_60k_all[:, j], "d-.", color=colors["60K"], label="60K", linewidth=1.5, markersize=5) + ax.set_title(JOINT_NAMES[j], fontsize=12, fontweight="bold") + ax.set_xlabel("Frame index") + ax.set_ylabel("Joint value (rad)") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + fig1.suptitle(f"ACT Inference — Episode {episode_idx} ({args.num_samples} frames)", + fontsize=14, fontweight="bold") + fig1.tight_layout() + fig1.savefig(os.path.join(out_dir, "trajectory_curves.png"), dpi=150) + plt.close(fig1) + + # 图2: MAE bar chart + fig2, ax = plt.subplots(figsize=(10, 5)) + x_pos = np.arange(n_joints) + width = 0.35 + bars1 = ax.bar(x_pos - width/2, mae_40k, width, label="40K", color="#2196F3", edgecolor="white") + bars2 = ax.bar(x_pos + width/2, mae_60k, width, label="60K", color="#FF9800", edgecolor="white") + ax.set_xticks(x_pos) + ax.set_xticklabels(JOINT_NAMES) + ax.set_ylabel("MAE (rad)") + ax.set_title(f"Per-Joint MAE: 40K vs 60K — Episode {episode_idx}") + ax.legend() + ax.grid(axis="y", alpha=0.3) + for bar in bars1: + h = bar.get_height() + ax.text(bar.get_x() + bar.get_width()/2., h + 0.001, f"{h:.4f}", + ha="center", va="bottom", fontsize=7) + for bar in bars2: + h = bar.get_height() + ax.text(bar.get_x() + bar.get_width()/2., h + 0.001, f"{h:.4f}", + ha="center", va="bottom", fontsize=7) + fig2.tight_layout() + fig2.savefig(os.path.join(out_dir, "error_comparison.png"), dpi=150) + plt.close(fig2) + + # 图3: Image collage with action tables + n_cols = min(4, args.num_samples) + n_rows = (args.num_samples + n_cols - 1) // n_cols + fig3 = plt.figure(figsize=(4 * n_cols, 4.5 * n_rows)) + gs = GridSpec(n_rows * 2, n_cols, figure=fig3, height_ratios=[3, 1] * n_rows) + + for i in range(args.num_samples): + row = (i // n_cols) * 2 + col = i % n_cols + + ax_img = fig3.add_subplot(gs[row, col]) + ax_img.imshow(images[i]) + ax_img.set_title(f"Frame {i}", fontsize=10) + ax_img.axis("off") + + ax_tbl = fig3.add_subplot(gs[row + 1, col]) + ax_tbl.axis("off") + + table_data = [["Joint", "GT", "40K", "60K"]] + for j in range(n_joints): + gt_val = results[i]["gt"][j] + p40_val = results[i]["pred_40k"][j] + p60_val = results[i]["pred_60k"][j] + + def fmt(v, err, thresh=0.05): + s = f"{v:.3f}" + return f"!{s}" if err > thresh else s + + table_data.append([ + JOINT_NAMES[j], + fmt(gt_val, 0), + fmt(p40_val, abs(p40_val - gt_val)), + fmt(p60_val, abs(p60_val - gt_val)), + ]) + + tbl = ax_tbl.table(cellText=table_data, loc="center", cellLoc="center") + tbl.auto_set_font_size(False) + tbl.set_fontsize(7) + tbl.scale(1.0, 1.1) + + for j in range(n_joints): + err40 = abs(results[i]["pred_40k"][j] - results[i]["gt"][j]) + err60 = abs(results[i]["pred_60k"][j] - results[i]["gt"][j]) + if err40 > 0.05: + tbl[(j + 1, 2)].set_facecolor("#FFCDD2") + if err60 > 0.05: + tbl[(j + 1, 3)].set_facecolor("#FFCDD2") + + fig3.suptitle(f"Frame-by-Frame Comparison — Episode {episode_idx}", + fontsize=14, fontweight="bold") + fig3.tight_layout() + fig3.savefig(os.path.join(out_dir, "frame_comparison.png"), dpi=150) + plt.close(fig3) + + +def main(): + args = parse_args() + os.makedirs(args.output_dir, exist_ok=True) + + # Parse episodes + episode_list = [int(x.strip()) for x in args.episodes.split(",")] + print(f"Episodes to evaluate: {episode_list}") + + # ── 1. Load dataset (full, no episode filter) ──────────────── + print("\n" + "=" * 60) + print("Loading dataset (full) ...") + dataset = LeRobotDataset( + repo_id="xarm7-pick-bottle", + root=args.dataset_root, + ) + print(f" Total episodes: {len(dataset.meta.episodes)}") + + # ── 2. Load models (once) ──────────────────────────────────── + print("\nLoading models (shared across episodes) ...") + print(" Loading 40K model ...") + policy_40k, preproc_40k, postproc_40k = load_policy_and_processors(args.checkpoint_40k) + print(" Loading 60K model ...") + policy_60k, preproc_60k, postproc_60k = load_policy_and_processors(args.checkpoint_60k) + + # ── 3. Evaluate each episode ───────────────────────────────── + n_joints = 8 + all_mae_40k = [] + all_mae_60k = [] + all_l2_40k = [] + all_l2_60k = [] + + for ep in episode_list: + print(f"\n{'─' * 50}") + print(f"Evaluating Episode {ep}") + print(f"{'─' * 50}") + + ep_out_dir = os.path.join(args.output_dir, f"ep{ep}") + os.makedirs(ep_out_dir, exist_ok=True) + + results, images, mae_40k, mae_60k, l2_40k, l2_60k = run_one_episode( + args, ep, dataset, policy_40k, preproc_40k, postproc_40k, + policy_60k, preproc_60k, postproc_60k) + + all_mae_40k.append(mae_40k) + all_mae_60k.append(mae_60k) + all_l2_40k.append(l2_40k) + all_l2_60k.append(l2_60k) + + # Print per-episode summary + print(f"\n Episode {ep} Summary:") + print(f" {'Joint':>10} | {'40K MAE':>10} | {'60K MAE':>10} | {'Δ':>10}") + print(f" {'─' * 48}") + for j in range(n_joints): + diff = mae_40k[j] - mae_60k[j] + sign = "▼" if diff > 0 else "▲" + print(f" {JOINT_NAMES[j]:>10} | {mae_40k[j]:10.4f} | {mae_60k[j]:10.4f} | {sign}{abs(diff):9.4f}") + print(f" {'L2 mean':>10} | {l2_40k:10.4f} | {l2_60k:10.4f} |") + + # Extract arrays for plotting + gt_all = np.stack([r["gt"] for r in results]) + pred_40k_all = np.stack([r["pred_40k"] for r in results]) + pred_60k_all = np.stack([r["pred_60k"] for r in results]) + + save_episode_plots(args, ep, results, images, mae_40k, mae_60k, + gt_all, pred_40k_all, pred_60k_all, ep_out_dir) + print(f" Plots saved to {ep_out_dir}") + + # ── 4. Aggregate results across episodes ───────────────────── + mean_mae_40k = np.stack(all_mae_40k).mean(axis=0) + mean_mae_60k = np.stack(all_mae_60k).mean(axis=0) + mean_l2_40k = np.mean(all_l2_40k) + mean_l2_60k = np.mean(all_l2_60k) + + print(f"\n{'=' * 60}") + print(f"=== AGGREGATE RESULTS ({len(episode_list)} episodes) ===") + print(f"{'Joint':>10} | {'40K MAE':>10} | {'60K MAE':>10} | {'Δ':>10}") + print("-" * 48) + for j in range(n_joints): + diff = mean_mae_40k[j] - mean_mae_60k[j] + sign = "▼" if diff > 0 else "▲" + pct = (diff / mean_mae_40k[j] * 100) if mean_mae_40k[j] > 0 else 0 + print(f"{JOINT_NAMES[j]:>10} | {mean_mae_40k[j]:10.4f} | {mean_mae_60k[j]:10.4f} | {sign}{abs(diff):9.4f} ({pct:+.0f}%)") + print(f"{'L2 mean':>10} | {mean_l2_40k:10.4f} | {mean_l2_60k:10.4f} |") + l2_diff_pct = (mean_l2_40k - mean_l2_60k) / mean_l2_40k * 100 + print(f"\n Overall L2 improvement: {l2_diff_pct:.1f}%") + + # Save aggregate CSV + agg_csv = os.path.join(args.output_dir, "aggregate_summary.csv") + with open(agg_csv, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["Joint", "40K_MAE", "60K_MAE", "Diff", "Change%"]) + for j in range(n_joints): + diff = mean_mae_40k[j] - mean_mae_60k[j] + pct = (diff / mean_mae_40k[j] * 100) if mean_mae_40k[j] > 0 else 0 + writer.writerow([JOINT_NAMES[j], f"{mean_mae_40k[j]:.6f}", f"{mean_mae_60k[j]:.6f}", + f"{diff:.6f}", f"{pct:.1f}%"]) + writer.writerow(["L2_mean", f"{mean_l2_40k:.6f}", f"{mean_l2_60k:.6f}", "", f"{l2_diff_pct:.1f}%"]) + print(f"\nAggregate CSV saved to {agg_csv}") + + # Save per-episode summary CSV + eps_csv = os.path.join(args.output_dir, "per_episode_summary.csv") + with open(eps_csv, "w", newline="") as f: + writer = csv.writer(f) + writer.writerow(["Episode", "L2_40K", "L2_60K"]) + for i, ep in enumerate(episode_list): + writer.writerow([ep, f"{all_l2_40k[i]:.6f}", f"{all_l2_60k[i]:.6f}"]) + print(f"Per-episode summary saved to {eps_csv}") + + # Aggregate bar chart + fig_agg, ax = plt.subplots(figsize=(12, 6)) + x_pos = np.arange(n_joints + 1) + labels = JOINT_NAMES + ["L2"] + vals_40k = list(mean_mae_40k) + [mean_l2_40k] + vals_60k = list(mean_mae_60k) + [mean_l2_60k] + width = 0.35 + bars1 = ax.bar(x_pos - width/2, vals_40k, width, label="40K", color="#2196F3", edgecolor="white") + bars2 = ax.bar(x_pos + width/2, vals_60k, width, label="60K", color="#FF9800", edgecolor="white") + ax.set_xticks(x_pos) + ax.set_xticklabels(labels) + ax.set_ylabel("MAE / L2 (rad)") + ax.set_title(f"Aggregate Error: 40K vs 60K (avg over {len(episode_list)} episodes)") + ax.legend() + ax.grid(axis="y", alpha=0.3) + for bar in bars1: + h = bar.get_height() + ax.text(bar.get_x() + bar.get_width()/2., h + 0.0005, f"{h:.4f}", + ha="center", va="bottom", fontsize=7) + for bar in bars2: + h = bar.get_height() + ax.text(bar.get_x() + bar.get_width()/2., h + 0.0005, f"{h:.4f}", + ha="center", va="bottom", fontsize=7) + fig_agg.tight_layout() + fig_agg.savefig(os.path.join(args.output_dir, "aggregate_error.png"), dpi=150) + plt.close(fig_agg) + print("Aggregate bar chart saved: aggregate_error.png") + + print(f"\n{'=' * 60}") + print(f"All outputs saved to {args.output_dir}/") + print("Done!") + print("Done!") + + +if __name__ == "__main__": + main()