eval_marl_mrt2a_delivcoal.py
MRT2A 配送留出评价
原文件 · 144 行 · 5615 字节
experiments/reproduction/eval_marl_mrt2a_delivcoal.py
SHA-256: 877c9cc30c88b9681ee9e6dd353344086316f9065ac013cdf37b43fb2561b40a
1"""用相同的一组 DelivCoal 场景评估多个 MARL-MRT2A checkpoint。"""23from __future__ import annotations45import argparse6import json7import sys8from pathlib import Path91011REPO_ROOT = Path(__file__).resolve().parents[2]12OFFICIAL_SRC = REPO_ROOT / "third_party/marl_mrt2a/marl_mrt2a/marl/src"13DEFAULT_ENV_CONFIG = REPO_ROOT / "experiments/reproduction/delivcoal_formal_base.yaml"14sys.path.insert(0, str(OFFICIAL_SRC))1516import main as epymarl_main # noqa: E402171819def scalar(value):20 return float(value.item()) if hasattr(value, "item") else float(value)212223def evaluate_one(model_root, checkpoint, env_config, eval_seed,24 eval_episodes, batch_size, episode_limit, grid_size, cuda,25 observation_layout="legacy20", share_intention="none"):26 """加载一个训练步对应的模型，并运行指定数量的评估 episode。"""27 overrides = {28 "env_args.key": repr("delivcoal-formal"),29 "env_args.config_path": repr(str(env_config)),30 "env_args.grid_size": grid_size,31 "env_args.time_limit": episode_limit,32 "env_args.observation_layout": repr(observation_layout),33 "env_args.share_intention": repr(share_intention),34 "critic_type": repr("cnn_cv_critic"),35 "critic_arch": repr("unet,8,1,2&batchNorm1d;linear,50;relu"),36 "agent": repr("cnn"),37 "agent_arch": repr("unet,8,1,2&"),38 "action_grid": True,39 "agent_distance_exp": 1.0,40 "current_target_factor": 1.0,41 "agent_reeval_rate": True,42 "filter_avail_by_objects": False,43 "gamma": 0.95,44 "standardise_returns": False,45 "standardise_rewards": True,46 "q_nstep": 5,47 "obs_agent_id": False,48 "obs_last_action": False,49 "batch_size_run": batch_size,50 "test_nepisode": eval_episodes,51 "checkpoint_path": repr(str(model_root)),52 "load_step": checkpoint,53 "evaluate": True,54 "use_cuda": cuda,55 "seed": eval_seed,56 "name": repr(f"marl_mrt2a_delivcoal_eval_c{checkpoint}_s{eval_seed}"),57 "wandb_sweep": False,58 "use_tensorboard": False,59 }60 argv = ["main.py", "--config=mappo", "--env-config=delivcoal", "with"]61 argv.extend(f"{key}={value}" for key, value in overrides.items())62 runner, _, _, _, logger = epymarl_main.main_from_arg(argv, just_build=True)63 # runner 每次同时运行 batch_size 个环境，因此向上取整调用次数。64 runs = max(1, (eval_episodes + runner.batch_size - 1) // runner.batch_size)65 for _ in range(runs):66 runner.run(test_mode=True)67 runner.close_env()68 return {69 key: scalar(values[-1][1])70 for key, values in logger.stats.items()71 if key.startswith("test_") and values72 } | {73 "loaded_step": int(runner.t_env),74 "evaluated_episodes": int(runs * runner.batch_size),75 }767778def run(args):79 """依次评估 checkpoint，并把同口径结果写入一个 JSON 文件。"""80 run_dir = Path(args.run_dir).resolve()81 model_roots = list((run_dir / "results/models").glob("*"))82 if len(model_roots) != 1:83 raise RuntimeError(f"Expected one model root, found {len(model_roots)}")84 output = Path(args.output).resolve() if args.output else run_dir / "heldout_eval"85 output.mkdir(parents=True, exist_ok=True)86 env_config = Path(args.env_config).resolve()87 results = {}88 for checkpoint in args.checkpoints:89 result = evaluate_one(90 model_roots[0], checkpoint, env_config, args.eval_seed,91 args.eval_episodes, args.batch_size_run, args.episode_limit,92 args.grid_size, args.cuda,93 args.observation_layout, args.share_intention,94 )95 results[str(checkpoint)] = result96 print(97 f"MARL-MRT2A DelivCoal {checkpoint}: "98 f"success={result.get('test_success_mean')}, "99 f"completion={result.get('test_task_completion_ratio_mean')}, "100 f"sent={result.get('test_sent_percentage_mean')}",101 flush=True,102 )103 payload = {104 "method": "MARL-MRT2A",105 "environment": "DelivCoal",106 "protocol": {107 "eval_seed": args.eval_seed,108 "eval_episodes": args.eval_episodes,109 "batch_size_run": args.batch_size_run,110 "paired_base_seed": True,111 "observation_layout": args.observation_layout,112 "share_intention": args.share_intention,113 },114 "checkpoints": results,115 }116 (output / "summary.json").write_text(117 json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"118 )119120121def parse_args():122 parser = argparse.ArgumentParser()123 parser.add_argument("--run-dir", required=True)124 parser.add_argument("--output")125 parser.add_argument("--env-config", default=str(DEFAULT_ENV_CONFIG))126 parser.add_argument("--checkpoints", nargs="+", type=int, required=True)127 parser.add_argument("--eval-seed", type=int, default=40_000)128 parser.add_argument("--eval-episodes", type=int, default=128)129 parser.add_argument("--batch-size-run", type=int, default=2)130 parser.add_argument("--episode-limit", type=int, default=750)131 parser.add_argument("--grid-size", type=int, default=21)132 parser.add_argument(133 "--observation-layout", choices=("legacy20", "full27"),134 default="legacy20",135 )136 parser.add_argument(137 "--share-intention", choices=("none", "path"), default="none"138 )139 parser.add_argument("--cuda", action="store_true")140 return parser.parse_args()141142143if __name__ == "__main__":144 run(parse_args())