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())