eval_task_slot_permutation.py
任务槽重排、动作映射与配对统计
原文件 · 469 行 · 16528 字节
scripts/analysis/eval_task_slot_permutation.py
SHA-256: f6626ecb6e0c455edd7510bb8283363e26532b25d5ca37b2ae64276916963b9f
1"""Paired evaluation under a fixed per-episode permutation of task slots.23The physical environment remains in canonical task order. Only the actor4observation/action interface is permuted, including coalition IDs and action5masks. A task-slot permutation-equivariant policy should therefore preserve6its physical decisions and performance. ``permutation_scope=deliver`` keeps7the historical homogeneous-Deliver test; ``all`` also moves the heterogeneous8Explore/Idle/Resupply slots.9"""1011import argparse12import copy13import csv14import json15import subprocess16import sys17import traceback18from pathlib import Path1920import numpy as np21import torch22import yaml2324REPO_ROOT = Path(__file__).resolve().parents[2]25sys.path.insert(0, str(REPO_ROOT))2627from config import get_args, get_config28from runner.shared.macf_runner import MACFEnvRunner, _t2n29from train.eval import _close_runner, _load_actor, _safe_name30from train.train import add_my_args, make_eval_env313233def _build_config(config_file):34 parser = get_args()35 parser.add_argument("--scenario_name", type=str, default="MyEnv")36 parser.add_argument("--num_landmarks", type=int, default=3)37 add_my_args(parser)38 config = parser.parse_args(["--config_file", str(config_file)])39 get_config(config)40 return config414243def _resolve_checkpoint(value, manifest_dir):44 path = Path(value).expanduser()45 if not path.is_absolute():46 path = manifest_dir / path47 if path.is_dir():48 path = path / "actor.pt"49 if not path.is_file():50 raise FileNotFoundError(path)51 return path.resolve()525354def _task_permutation(env, scenario_id, permutation_seed, permutation_scope):55 task_num = len(env.env.tasks)56 presented_to_canonical = np.arange(task_num, dtype=np.int64)57 if permutation_scope in ("deliver", "explore"):58 selected_type = 1 if permutation_scope == "deliver" else 059 permuted_slots = np.asarray(60 [61 index62 for index, task in enumerate(env.env.tasks)63 if int(task.type) == selected_type64 ],65 dtype=np.int64,66 )67 elif permutation_scope == "all":68 permuted_slots = np.arange(task_num, dtype=np.int64)69 else:70 raise ValueError(f"Unsupported permutation_scope={permutation_scope!r}")71 shuffled = permuted_slots.copy()72 np.random.default_rng(permutation_seed + scenario_id).shuffle(shuffled)73 presented_to_canonical[permuted_slots] = shuffled74 canonical_to_presented = np.empty_like(presented_to_canonical)75 canonical_to_presented[presented_to_canonical] = np.arange(task_num)76 return presented_to_canonical, canonical_to_presented777879def _map_normalized_coalitions(values, canonical_to_presented, valid=None):80 task_num = len(canonical_to_presented)81 canonical = np.rint(values * max(1, task_num - 1)).astype(np.int64)82 canonical = np.clip(canonical, 0, task_num - 1)83 mapped = canonical_to_presented[canonical] / float(max(1, task_num - 1))84 if valid is not None:85 mapped = np.where(valid, mapped, values)86 return mapped.astype(values.dtype, copy=False)878889def _present_actor_inputs(90 obs,91 available_actions,92 presented_to_canonical,93 canonical_to_presented,94 cap_dim,95):96 presented_obs = {key: np.array(value, copy=True) for key, value in obs.items()}97 presented_obs["tasks"] = np.take(98 presented_obs["tasks"], presented_to_canonical, axis=-299 )100 ego_coalition_index = 2 + cap_dim101 presented_obs["ego"][..., ego_coalition_index] = _map_normalized_coalitions(102 presented_obs["ego"][..., ego_coalition_index],103 canonical_to_presented,104 )105 other_coalition_index = cap_dim + 4106 valid_others = np.abs(presented_obs["others"]).sum(axis=-1) > 0107 presented_obs["others"][..., other_coalition_index] = (108 _map_normalized_coalitions(109 presented_obs["others"][..., other_coalition_index],110 canonical_to_presented,111 valid=valid_others,112 )113 )114 presented_available = np.take(115 np.asarray(available_actions), presented_to_canonical, axis=-1116 )117 return presented_obs, presented_available118119120@torch.no_grad()121def _run_episode(122 config,123 runner,124 env,125 scenario_id,126 permuted,127 permutation_seed,128 permutation_scope,129):130 scenario_seed = int(config.eval_seed) + int(scenario_id)131 numpy_state = np.random.get_state()132 np.random.seed(scenario_seed)133 try:134 env.seed(scenario_seed)135 obs, active_masks, available_actions = env.reset(scenario_id)136 task_num = len(env.env.tasks)137 if permuted:138 presented_to_canonical, canonical_to_presented = _task_permutation(139 env, scenario_id, permutation_seed, permutation_scope140 )141 else:142 presented_to_canonical = np.arange(task_num, dtype=np.int64)143 canonical_to_presented = presented_to_canonical.copy()144145 rnn_states = np.zeros(146 (config.num_agents, config.recurrent_N, config.hidden_size),147 dtype=np.float32,148 )149 masks = np.ones((config.num_agents, 1), dtype=np.float32)150 last_actions = np.zeros((config.num_agents, 1), dtype=np.float32)151 episode_reward = 0.0152 final_info = None153154 runner.trainer.prep_rollout()155 for step in range(config.max_steps):156 actor_obs, actor_available = _present_actor_inputs(157 obs,158 available_actions,159 presented_to_canonical,160 canonical_to_presented,161 config.cap_dim,162 )163 actions, next_rnn_states = runner.trainer.policy.act(164 actor_obs,165 rnn_states,166 masks,167 available_actions=actor_available,168 deterministic=True,169 )170 presented_actions = _t2n(actions).astype(np.int64)171 next_rnn_states = _t2n(next_rnn_states)172 canonical_actions = presented_to_canonical[173 presented_actions.reshape(-1)174 ].reshape(presented_actions.shape)175 canonical_actions = (176 active_masks * canonical_actions177 + (1.0 - active_masks) * last_actions178 )179180 (181 obs,182 rewards,183 dones,184 info,185 active_masks,186 _,187 available_actions,188 ) = env.step(canonical_actions)189 final_info = info190 episode_reward += float(np.mean(rewards))191192 done_mask = np.asarray(dones, dtype=bool)193 next_rnn_states[done_mask] = 0.0194 masks = np.ones((config.num_agents, 1), dtype=np.float32)195 masks[done_mask] = 0.0196 rnn_states = next_rnn_states197 last_actions = canonical_actions198 if bool(np.all(done_mask)) or info.get("end", 0):199 break200 finally:201 np.random.set_state(numpy_state)202203 if final_info is None:204 raise RuntimeError("Episode ended without environment info.")205 delivery_tasks = [206 task for task in final_info["tasks"] if int(task["type"]) == 1207 ]208 completed = int(sum(bool(task["completed"]) for task in delivery_tasks))209 return {210 "scenario_id": int(scenario_id),211 "scenario_seed": scenario_seed,212 "mode": "permuted" if permuted else "original",213 "success": int(completed == len(delivery_tasks)),214 "tasks_completed": completed,215 "task_completion_ratio": completed / max(1, len(delivery_tasks)),216 "episode_length": int(step + 1),217 "episode_reward": float(episode_reward),218 "permutation": presented_to_canonical.tolist(),219 "permutation_scope": permutation_scope,220 }221222223def _evaluate_one(task):224 (225 model,226 base_config_file,227 manifest_dir,228 output_root,229 episodes,230 permutation_seed,231 permutation_scope,232 cuda,233 ) = task234 base_config = _build_config(base_config_file)235 config = copy.deepcopy(base_config)236 for key, value in (model.get("overrides") or {}).items():237 setattr(config, key, value)238 config.cuda = int(cuda)239 config.num_agents = config.num_uav + config.num_ugv240 config.use_eval = True241 config.use_render = False242 config.model_dir = None243 config.device = (244 torch.device(f"cuda:{cuda}")245 if int(cuda) >= 0246 else torch.device("cpu")247 )248 config.experiment_name = _safe_name(model["name"] + "_slot_permutation")249 model_output = Path(output_root) / _safe_name(model["name"])250 model_output.mkdir(parents=True, exist_ok=True)251 config.eval_output_dir = str(model_output)252 config.run_dir = model_output / "runner"253254 torch.manual_seed(config.seed)255 torch.cuda.manual_seed_all(config.seed)256 np.random.seed(config.seed)257 torch.set_num_threads(config.n_training_threads)258259 envs = make_eval_env(config)260 config.envs = envs261 config.eval_envs = envs262 runner = MACFEnvRunner(config)263 checkpoint = _resolve_checkpoint(model["checkpoint"], Path(manifest_dir))264 _load_actor(runner.policy, checkpoint, config.device)265 records = []266 try:267 env = envs.envs[0]268 for scenario_id in range(episodes):269 records.append(270 _run_episode(271 config,272 runner,273 env,274 scenario_id,275 False,276 permutation_seed,277 permutation_scope,278 )279 )280 records.append(281 _run_episode(282 config,283 runner,284 env,285 scenario_id,286 True,287 permutation_seed,288 permutation_scope,289 )290 )291 finally:292 _close_runner(runner, envs)293294 for record in records:295 record["model"] = model["name"]296 with (model_output / "episodes.json").open("w", encoding="utf-8") as file:297 json.dump(records, file, ensure_ascii=False, indent=2)298 return {"model": model["name"], "checkpoint": str(checkpoint), "records": records}299300301def _write_csv(path, records):302 if not records:303 return304 fields = []305 for record in records:306 for key in record:307 if key not in fields:308 fields.append(key)309 with path.open("w", newline="", encoding="utf-8") as file:310 writer = csv.DictWriter(file, fieldnames=fields)311 writer.writeheader()312 writer.writerows(records)313314315def _summarize(result):316 by_mode = {317 mode: {318 record["scenario_id"]: record319 for record in result["records"]320 if record["mode"] == mode321 }322 for mode in ("original", "permuted")323 }324 scenario_ids = sorted(set(by_mode["original"]) & set(by_mode["permuted"]))325 original = [by_mode["original"][index] for index in scenario_ids]326 permuted = [by_mode["permuted"][index] for index in scenario_ids]327 success_to_failure = sum(328 a["success"] == 1 and b["success"] == 0329 for a, b in zip(original, permuted)330 )331 failure_to_success = sum(332 a["success"] == 0 and b["success"] == 1333 for a, b in zip(original, permuted)334 )335 return {336 "model": result["model"],337 "checkpoint": result["checkpoint"],338 "num_episodes": len(scenario_ids),339 "original_success": float(np.mean([row["success"] for row in original])),340 "permuted_success": float(np.mean([row["success"] for row in permuted])),341 "success_delta": float(342 np.mean([row["success"] for row in permuted])343 - np.mean([row["success"] for row in original])344 ),345 "success_to_failure": int(success_to_failure),346 "failure_to_success": int(failure_to_success),347 "original_completion": float(348 np.mean([row["task_completion_ratio"] for row in original])349 ),350 "permuted_completion": float(351 np.mean([row["task_completion_ratio"] for row in permuted])352 ),353 "original_length": float(354 np.mean([row["episode_length"] for row in original])355 ),356 "permuted_length": float(357 np.mean([row["episode_length"] for row in permuted])358 ),359 }360361362def main():363 parser = argparse.ArgumentParser()364 parser.add_argument("--config_file", required=True)365 parser.add_argument("--models_file", required=True)366 parser.add_argument("--output_dir", required=True)367 parser.add_argument("--episodes", type=int, default=128)368 parser.add_argument("--permutation_seed", type=int, default=760000)369 parser.add_argument(370 "--permutation_scope",371 choices=("deliver", "explore", "all"),372 default="deliver",373 )374 parser.add_argument("--cuda", type=int, default=1)375 parser.add_argument("--max_workers", type=int, default=2)376 parser.add_argument("--model_name_pattern")377 parser.add_argument("--detach", action="store_true")378 args = parser.parse_args()379380 config_file = Path(args.config_file).resolve()381 models_file = Path(args.models_file).resolve()382 output_dir = Path(args.output_dir).resolve()383 if args.detach:384 output_dir.parent.mkdir(parents=True, exist_ok=True)385 command = [386 sys.executable,387 str(Path(__file__).resolve()),388 *[argument for argument in sys.argv[1:] if argument != "--detach"],389 ]390 log_path = output_dir.parent / f"{output_dir.name}.log"391 with log_path.open("w", encoding="utf-8") as log_file:392 process = subprocess.Popen(393 command,394 cwd=REPO_ROOT,395 stdin=subprocess.DEVNULL,396 stdout=log_file,397 stderr=subprocess.STDOUT,398 start_new_session=True,399 )400 pid_path = output_dir.parent / f"{output_dir.name}.pid"401 pid_path.write_text(f"{process.pid}\n", encoding="utf-8")402 print(403 f"Detached permutation eval pid={process.pid} "404 f"output_dir={output_dir}"405 )406 return407 output_dir.mkdir(parents=True, exist_ok=False)408 with models_file.open("r", encoding="utf-8") as file:409 manifest = yaml.safe_load(file) or {}410 models = manifest["models"]411 if args.model_name_pattern:412 models = [413 model414 for model in models415 if args.model_name_pattern in model["name"]416 ]417 if not models:418 raise ValueError(419 f"No model matched {args.model_name_pattern!r}."420 )421422 # Initialize the shared CUDA context before worker threads construct423 # independent runners. Concurrent lazy initialization can otherwise make424 # one thread observe an incomplete CUDA driver state.425 if args.cuda >= 0:426 if not torch.cuda.is_available():427 raise RuntimeError(f"CUDA device {args.cuda} requested but unavailable.")428 torch.cuda.set_device(args.cuda)429 torch.empty(1, device=f"cuda:{args.cuda}")430431 tasks = [432 (433 model,434 str(config_file),435 str(models_file.parent),436 str(output_dir),437 args.episodes,438 args.permutation_seed,439 args.permutation_scope,440 args.cuda,441 )442 for model in models443 ]444 # Evaluate models sequentially inside each process. Each runner mutates445 # global NumPy/Torch RNG state and owns a CUDA model, so threads would make446 # paired scenarios nondeterministic and race CUDA initialization. Campaign447 # supervisors can still run independent permutation scopes in parallel.448 results = []449 for task in tasks:450 try:451 result = _evaluate_one(task)452 except Exception:453 (output_dir / "failure.log").write_text(454 traceback.format_exc(), encoding="utf-8"455 )456 raise457 results.append(result)458 print(f"Completed {task[0]['name']}", flush=True)459460 summaries = [_summarize(result) for result in results]461 summaries.sort(key=lambda row: row["model"])462 _write_csv(output_dir / "summary.csv", summaries)463 with (output_dir / "summary.json").open("w", encoding="utf-8") as file:464 json.dump(summaries, file, ensure_ascii=False, indent=2)465 print(f"Wrote {output_dir}", flush=True)466467468if __name__ == "__main__":469 main()