eval_cfm_delivcoal.py
CFM 采样/greedy 留出评价
原文件 · 114 行 · 4273 字节
experiments/reproduction/eval_cfm_delivcoal.py
SHA-256: 453257c21cd046339e1d3c7a57639bf8e5539fc82e277663a58f2f5eb6e93d56
1"""用相同的一组 DelivCoal 场景评估多个 CFM checkpoint。"""23from __future__ import annotations45import argparse6import csv7import json8import sys9from pathlib import Path1011import torch1213REPO_ROOT = Path(__file__).resolve().parents[2]14if str(REPO_ROOT) not in sys.path:15 sys.path.insert(0, str(REPO_ROOT))1617from experiments.reproduction.delivcoal_adapter_common import (18 DEFAULT_CONFIG,19 aggregate_rows,20 atomic_json,21 load_env_config,22)23from experiments.reproduction.train_cfm_delivcoal import AttentionNet, rollout242526def run(args):27 """依次加载 checkpoint，运行评估并汇总为 CSV 和 JSON。"""28 run_dir = Path(args.run_dir).resolve()29 output = Path(args.output).resolve() if args.output else run_dir / "heldout_eval"30 output.mkdir(parents=True, exist_ok=True)31 config = load_env_config(args.env_config)32 device = torch.device(33 f"cuda:{args.cuda}" if args.cuda >= 0 and torch.cuda.is_available() else "cpu"34 )35 seeds = list(range(args.eval_seed, args.eval_seed + args.eval_episodes))36 # seeds 对所有 checkpoint 保持一致。37 summaries = {}38 all_rows = []39 for episode in args.episodes:40 # 每个 checkpoint 使用同一组 seeds，便于直接比较训练阶段差异。41 checkpoint_path = run_dir / "models" / f"actor_ep{episode:06d}.pt"42 checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)43 network = AttentionNet(44 int(checkpoint["agent_input_dim"]),45 int(checkpoint["task_input_dim"]),46 int(checkpoint["embedding_dim"]),47 ).to(device)48 network.load_state_dict(checkpoint["model"])49 network.eval()50 # 每个 seed 只创建一个场景，greedy 模式不再引入动作采样随机性。51 rows = []52 with torch.no_grad():53 for seed in seeds:54 row, _ = rollout(55 network,56 config,57 seed,58 device,59 not args.sample_policy,60 policy_seed=(seed + args.policy_seed_offset)61 if args.sample_policy else None,62 )63 row["checkpoint_episode"] = int(episode)64 rows.append(row)65 all_rows.append(row)66 summaries[str(episode)] = aggregate_rows(rows)67 summary = summaries[str(episode)]68 print(69 f"CFM DelivCoal ep{episode}: success={summary['success']['mean']:.3f}, "70 f"completion={summary['task_completion_ratio']['mean']:.3f}, "71 f"sent={summary['sent_percentage']['mean']:.3f}",72 flush=True,73 )74 with (output / "per_episode.csv").open("w", newline="", encoding="utf-8") as file:75 # CSV 保存每个场景的原始结果，JSON 保存各指标的汇总结果。76 writer = csv.DictWriter(file, fieldnames=list(all_rows[0]))77 writer.writeheader()78 writer.writerows(all_rows)79 atomic_json(output / "summary.json", {80 "method": "CFM",81 "environment": "DelivCoal",82 "protocol": {83 "eval_seed": args.eval_seed,84 "eval_episodes": args.eval_episodes,85 "paired_seeds": True,86 "greedy": not args.sample_policy,87 "policy_seed_offset": (88 args.policy_seed_offset if args.sample_policy else None89 ),90 },91 "checkpoints": summaries,92 })939495def parse_args():96 parser = argparse.ArgumentParser()97 parser.add_argument("--run-dir", required=True)98 parser.add_argument("--output")99 parser.add_argument("--env-config", default=str(DEFAULT_CONFIG))100 parser.add_argument("--episodes", nargs="+", type=int, required=True)101 parser.add_argument("--eval-episodes", type=int, default=128)102 parser.add_argument("--eval-seed", type=int, default=40_000)103 parser.add_argument("--cuda", type=int, default=-1)104 parser.add_argument(105 "--sample-policy",106 action="store_true",107 help="Sample actions reproducibly instead of greedy decoding.",108 )109 parser.add_argument("--policy-seed-offset", type=int, default=1_000_000)110 return parser.parse_args()111112113if __name__ == "__main__":114 run(parse_args())