eval.py
checkpoint 加载与独立评价
原文件 · 300 行 · 9891 字节
train/eval.py
SHA-256: 3209322d0e2d0db321decfcd813a6eae46b07a1a5989759b84aea60ebcc817ce
1"""Evaluate one or more MACF actor checkpoints on a fixed scenario suite."""23import copy4import csv5import json6import re7import sys8from datetime import datetime9from pathlib import Path1011import numpy as np12import torch13import yaml1415repo_root = str(Path(__file__).resolve().parents[1])16sys.path.insert(0, repo_root)1718from config import get_args, get_config19from runner.shared.macf_runner import MACFEnvRunner20from runner.shared.gymrt2a_runner import GymRT2ARunner21from runner.qmix.macf_runner import MACFRunner as QMixRunner22from train.train import add_my_args, make_eval_env232425SUITE_CONFIG_KEYS = {26 "map_size",27 "num_uav",28 "num_ugv",29 "uav_perception_rg",30 "ugv_perception_rg",31 "target_num",32 "max_steps",33 "episode_time_limit",34 "cap_dim",35 "eval_seed",36 "eval_reset_mode",37 "eval_episodes",38 "n_eval_rollout_threads",39 "rewards_config_file",40 "gymrt2a_map_size",41 "gymrt2a_num_agents",42 "gymrt2a_objects",43 "gymrt2a_view_range",44 "gymrt2a_comm_range",45 "gymrt2a_share_intention",46}474849def parse_args(args):50 parser = get_args()51 parser.add_argument("--scenario_name", type=str, default="MyEnv")52 parser.add_argument("--num_landmarks", type=int, default=3)53 add_my_args(parser)54 parser.add_argument(55 "--models_file",56 type=str,57 required=True,58 help="YAML file containing the checkpoints and their config overrides",59 )60 parser.add_argument(61 "--output_root",62 type=str,63 default=None,64 help="override the output root declared in the model manifest",65 )66 parser.add_argument(67 "--model_name_pattern",68 default=None,69 help="optional regular expression selecting a subset of manifest models",70 )71 parser.add_argument(72 "--progress_file",73 default=None,74 help="write atomic per-episode progress for the currently evaluated model",75 )76 return parser.parse_args(args)777879def _load_manifest(path):80 manifest_path = Path(path).resolve()81 with manifest_path.open("r", encoding="utf-8") as file:82 manifest = yaml.safe_load(file) or {}83 models = manifest.get("models")84 if not isinstance(models, list) or not models:85 raise ValueError("models_file must contain a non-empty 'models' list.")86 return manifest_path, manifest878889def _resolve_checkpoint(value, manifest_dir):90 path = Path(value).expanduser()91 candidates = [path, manifest_dir / path] if not path.is_absolute() else [path]92 for candidate in candidates:93 if candidate.is_dir():94 candidate = candidate / "actor.pt"95 if candidate.is_file():96 return candidate.resolve()97 raise FileNotFoundError(f"Actor checkpoint not found: {value}")9899100def _safe_name(name):101 safe = re.sub(r"[^A-Za-z0-9_.-]+", "_", str(name)).strip("._")102 if not safe:103 raise ValueError(f"Invalid model name: {name!r}")104 return safe105106107def _select_device(config):108 if config.cuda >= 0 and torch.cuda.is_available():109 device = torch.device(f"cuda:{config.cuda}")110 if config.cuda_deterministic:111 torch.backends.cudnn.benchmark = False112 torch.backends.cudnn.deterministic = True113 else:114 device = torch.device("cpu")115 torch.set_num_threads(config.n_training_threads)116 return device117118119def _load_actor(policy, checkpoint_path, device):120 try:121 checkpoint = torch.load(122 checkpoint_path, map_location=device, weights_only=True123 )124 except TypeError:125 checkpoint = torch.load(checkpoint_path, map_location=device)126 if hasattr(policy, "load_actor_checkpoint"):127 policy.load_actor_checkpoint(checkpoint)128 else:129 policy.actor.load_state_dict(checkpoint)130131132def _close_runner(runner, envs):133 envs.close()134 runner.writter.close()135 if hasattr(runner, "writer") and runner.writer is not runner.writter:136 runner.writer.flush()137 runner.writer.close()138139140def _evaluate_model(base_config, model, checkpoint, output_dir, device):141 config = copy.deepcopy(base_config)142 overrides = model.get("overrides", {}) or {}143 if not isinstance(overrides, dict):144 raise ValueError(f"overrides for {model['name']!r} must be a mapping.")145 forbidden = SUITE_CONFIG_KEYS.intersection(overrides)146 if forbidden:147 raise ValueError(148 f"Model {model['name']!r} overrides suite settings: {sorted(forbidden)}"149 )150 for key, value in overrides.items():151 setattr(config, key, value)152153 trainer_name = str(154 getattr(config, "trainer_algorithm_name", None) or config.algorithm_name155 ).lower()156 if trainer_name in {"qmix", "vdn", "hatt_qmix"}:157 if str(config.env_name) in ("gymrt2a", "mrt2a"):158 raise NotImplementedError("QMIX-family evaluation currently targets capcoal/delivcoal.")159 config.num_agents = config.num_uav + config.num_ugv160 runner_class = QMixRunner161 elif str(config.env_name) in ("gymrt2a", "mrt2a"):162 config.num_agents = int(config.gymrt2a_num_agents)163 runner_class = GymRT2ARunner164 else:165 config.num_agents = config.num_uav + config.num_ugv166 runner_class = MACFEnvRunner167 config.use_eval = True168 config.use_render = False169 config.model_dir = None170 config.device = device171 config.experiment_name = _safe_name(model["name"])172 config.eval_output_dir = str(output_dir)173 config.eval_progress_file = getattr(base_config, "progress_file", None)174 config.run_dir = output_dir / "runner"175176 torch.manual_seed(config.seed)177 torch.cuda.manual_seed_all(config.seed)178 np.random.seed(config.seed)179180 envs = make_eval_env(config)181 config.envs = envs182 config.eval_envs = envs183 runner = None184 try:185 runner = runner_class(config)186 _load_actor(runner.policy, checkpoint, device)187 return runner.eval(total_num_steps=int(model.get("step", 0)))188 finally:189 if runner is not None:190 _close_runner(runner, envs)191 else:192 envs.close()193194195def _write_suite_results(output_dir, records, metadata):196 fields = []197 for record in records:198 for key in record:199 if key not in fields:200 fields.append(key)201 with (output_dir / "summary.csv").open(202 "w", newline="", encoding="utf-8"203 ) as file:204 writer = csv.DictWriter(file, fieldnames=fields)205 writer.writeheader()206 writer.writerows(records)207 with (output_dir / "summary.json").open("w", encoding="utf-8") as file:208 json.dump(209 {"metadata": metadata, "models": records},210 file,211 ensure_ascii=False,212 indent=2,213 )214215216def main(args=None):217 config = parse_args(sys.argv[1:] if args is None else args)218 get_config(config)219 manifest_path, manifest = _load_manifest(config.models_file)220221 suite_name = _safe_name(manifest.get("name", "macf-eval"))222 timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")223 output_root = Path(224 config.output_root or manifest.get("output_root", "eval_results")225 ).expanduser()226 output_dir = (output_root / f"{suite_name}_{timestamp}").resolve()227 output_dir.mkdir(parents=True, exist_ok=False)228229 device = _select_device(config)230 models = manifest["models"]231 if config.model_name_pattern:232 pattern = re.compile(config.model_name_pattern)233 models = [model for model in models if pattern.search(str(model.get("name", "")))]234 if not models:235 raise ValueError(236 f"model_name_pattern matched no models: {config.model_name_pattern!r}"237 )238239 records = []240 seen_names = set()241 for model in models:242 if not isinstance(model, dict) or "name" not in model or "checkpoint" not in model:243 raise ValueError("Each model requires 'name' and 'checkpoint'.")244 safe_name = _safe_name(model["name"])245 if safe_name in seen_names:246 raise ValueError(f"Duplicate model name: {model['name']!r}")247 seen_names.add(safe_name)248249 checkpoint = _resolve_checkpoint(250 model["checkpoint"], manifest_path.parent251 )252 model_output = output_dir / "models" / safe_name253 model_output.mkdir(parents=True)254 print(f"Evaluating {model['name']}: {checkpoint}")255 summary = _evaluate_model(256 config, model, checkpoint, model_output, device257 )258 records.append(259 {260 "name": str(model["name"]),261 "checkpoint": str(checkpoint),262 **summary,263 }264 )265266 metadata = {267 "name": suite_name,268 "created_at": datetime.now().isoformat(timespec="seconds"),269 "config_file": str(Path(config.config_file).resolve())270 if config.config_file271 else None,272 "models_file": str(manifest_path),273 "eval_seed": int(config.eval_seed),274 "eval_reset_mode": getattr(config, "eval_reset_mode", None),275 "eval_episodes": int(config.eval_episodes),276 "eval_deterministic": bool(277 getattr(config, "eval_deterministic", True)278 ),279 "delivery_tasks_initially_discovered": bool(280 getattr(config, "delivery_tasks_initially_discovered", False)281 ),282 "mask_explore_actions": bool(283 getattr(config, "mask_explore_actions", False)284 ),285 "max_steps": int(config.max_steps),286 "episode_time_limit": int(287 getattr(config, "episode_time_limit", None) or config.max_steps288 ),289 "device": str(device),290 "environment": {291 key: getattr(config, key) for key in sorted(SUITE_CONFIG_KEYS)292 },293 }294 _write_suite_results(output_dir, records, metadata)295 print(f"Evaluation suite written to {output_dir}")296 return output_dir297298299if __name__ == "__main__":300 main()