run_marl_mrt2a_delivcoal.py
MRT2A 配送训练入口
原文件 · 124 行 · 5027 字节
experiments/reproduction/run_marl_mrt2a_delivcoal.py
SHA-256: 1f1da5e3cfb5f8dba0b0328467b9274011f775dec802ec1eb862e09aeb226cf3
1"""使用 MARL-MRT2A 原项目的 PyMARL 训练流程运行 DelivCoal。"""23from __future__ import annotations45import argparse6import json7import os8import sys9from datetime import datetime10from pathlib import Path111213REPO_ROOT = Path(__file__).resolve().parents[2]14OFFICIAL_SRC = REPO_ROOT / "third_party/marl_mrt2a/marl_mrt2a/marl/src"15DEFAULT_ENV_CONFIG = REPO_ROOT / "experiments/reproduction/delivcoal_formal_base.yaml"16sys.path.insert(0, str(OFFICIAL_SRC))1718import main as epymarl_main # noqa: E402192021def run(args):22 """把命令行参数转换成 PyMARL 配置覆盖项并启动训练。"""23 output = Path(args.output).resolve()24 output.mkdir(parents=True, exist_ok=True)25 env_config = Path(args.env_config).resolve()26 name = f"marl_mrt2a_delivcoal_s{args.seed}"27 overrides = {28 # 环境适配器把 DelivCoal 的任务位置投到二维动作网格。29 "env_args.key": repr("delivcoal-formal"),30 "env_args.config_path": repr(str(env_config)),31 "env_args.grid_size": args.grid_size,32 "env_args.time_limit": args.episode_limit,33 "env_args.observation_layout": repr(args.observation_layout),34 "env_args.share_intention": repr(args.share_intention),35 "critic_type": repr("cnn_cv_critic"),36 "critic_arch": repr("unet,8,1,2&batchNorm1d;linear,50;relu"),37 "agent": repr("cnn"),38 "agent_arch": repr("unet,8,1,2&"),39 "action_grid": True,40 "agent_distance_exp": 1.0,41 "current_target_factor": 1.0,42 "agent_reeval_rate": True,43 "filter_avail_by_objects": False,44 "gamma": 0.95,45 "lr": args.lr,46 "entropy_coef": args.entropy_coef,47 "standardise_returns": False,48 "standardise_rewards": True,49 "q_nstep": 5,50 "obs_agent_id": False,51 "obs_last_action": False,52 "buffer_size": max(args.batch_size, args.buffer_size),53 "batch_size": args.batch_size,54 "batch_size_run": args.batch_size_run,55 "target_update_interval_or_tau": 200,56 "t_max": args.t_max,57 "test_nepisode": args.test_nepisode,58 "test_interval": args.test_interval,59 "log_interval": args.log_interval,60 "runner_log_interval": args.log_interval,61 "learner_log_interval": args.log_interval,62 "save_model": True,63 "save_model_interval": args.save_model_interval,64 "run_perm_importance": False,65 "save_path": repr(str(output)),66 "use_cuda": bool(args.cuda),67 "seed": args.seed,68 "name": repr(name),69 "wandb_sweep": False,70 "use_tensorboard": False,71 }72 argv = ["main.py", "--config=mappo", "--env-config=delivcoal", "with"]73 # PyMARL 使用 Sacred 的 key=value 形式接收运行时覆盖项。74 argv.extend(f"{key}={value}" for key, value in overrides.items())75 metadata = {76 "started_at": datetime.now().isoformat(timespec="seconds"),77 "method": "official MARL-MRT2A CNN policy and target revision",78 "adapter": "DelivCoal spatial task-cell adapter",79 "official_source": str(OFFICIAL_SRC),80 "environment_config": str(env_config),81 "argv": argv,82 "cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),83 }84 config_path = output / "reproduction_config.json"85 config_path.write_text(86 json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"87 )88 epymarl_main.main_from_arg(argv)89 metadata["finished_at"] = datetime.now().isoformat(timespec="seconds")90 config_path.write_text(91 json.dumps(metadata, ensure_ascii=False, indent=2), encoding="utf-8"92 )939495def parse_args():96 parser = argparse.ArgumentParser()97 parser.add_argument("--seed", type=int, required=True)98 parser.add_argument("--output", required=True)99 parser.add_argument("--env-config", default=str(DEFAULT_ENV_CONFIG))100 parser.add_argument("--t-max", type=int, default=3_600_000)101 parser.add_argument("--episode-limit", type=int, default=750)102 parser.add_argument("--grid-size", type=int, default=21)103 parser.add_argument(104 "--observation-layout", choices=("legacy20", "full27"),105 default="legacy20",106 )107 parser.add_argument(108 "--share-intention", choices=("none", "path"), default="none"109 )110 parser.add_argument("--batch-size-run", type=int, default=4)111 parser.add_argument("--batch-size", type=int, default=8)112 parser.add_argument("--buffer-size", type=int, default=16)113 parser.add_argument("--test-nepisode", type=int, default=32)114 parser.add_argument("--test-interval", type=int, default=50_000)115 parser.add_argument("--save-model-interval", type=int, default=50_000)116 parser.add_argument("--log-interval", type=int, default=10_000)117 parser.add_argument("--lr", type=float, default=5e-4)118 parser.add_argument("--entropy-coef", type=float, default=0.01)119 parser.add_argument("--cuda", action="store_true")120 return parser.parse_args()121122123if __name__ == "__main__":124 run(parse_args())