evaluate_source_commitment.py
共享状态库与源条件配对评价入口
原文件 · 498 行 · 20413 字节
scripts/analysis/evaluate_source_commitment.py
SHA-256: 3202c7b908cce8fed0cdbaddbc0bd7b986a12fe524b3d346592c39e343b84068
1"""Paired source-coalition exit-cost evaluation for task-conditioned AA (HATT v3.4).23This is the ``source_commitment_pair`` confirmation experiment. It tests whether a4focal agent's decision to leave its current delivery coalition responds to the5coalition's residual capacity after the agent departs:67* ``replaceable`` — the source coalition can still finish without the focal agent;8* ``critical`` — the focal agent is essential, so its departure makes the source9 coalition infeasible.1011The only difference between the two conditions is which *other* agents belong to the12source coalition (``critical`` keeps them out; ``replaceable`` adds enough idle load to13cover the remaining demand). Focal position/load, source task, target tasks and total14fleet load are all held fixed. Each model reads the same pickled state bank, so15cross-model state-distribution confounds are removed.1617The metric of interest is ``focal_stay_probability`` (and ``focal_transfer_probability``18toward the highest-deficit target). A task-conditioned A-A gate should push ``P(stay)``19up under ``critical`` relative to ``replaceable``; the v3.3AA (no-context) control should20show a smaller split.2122Two phases:23* ``--phase generate`` builds the shared state bank from base policies;24* ``--phase evaluate`` runs the paired replaceable/critical evaluation on the bank.25"""2627from __future__ import annotations2829import argparse30import copy31import csv32import gc33import json34import os35import pickle36import subprocess37import sys38import threading39from concurrent.futures import ThreadPoolExecutor, as_completed40from datetime import datetime41from pathlib import Path4243import numpy as np44import torch45import yaml4647REPO_ROOT = Path(__file__).resolve().parents[2]48sys.path.insert(0, str(REPO_ROOT))4950from scripts.analysis.diagnose_idle_opportunity import ( # noqa: E40251 _distribution,52 _make_runner,53 _parse_scenarios,54 _write_csv,55)56from train.eval import _close_runner # noqa: E4025758SCHEMA_VERSION = 1596061def _atomic_json(path, value):62 path = Path(path)63 temporary = path.with_suffix(path.suffix + ".tmp")64 temporary.write_text(65 json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8"66 )67 os.replace(temporary, path)686970def _policy_step(actor, obs, rnn_states, masks, available_actions, active, last):71 """Argmax step reusing the actor distribution; non-active agents hold their action."""72 distribution, next_rnn = _distribution(73 actor, obs, rnn_states, masks, available_actions74 )75 proposed = distribution.probs.argmax(dim=-1, keepdim=True).cpu().numpy()76 next_rnn = next_rnn.detach().cpu().numpy()77 actions = np.asarray(active) * proposed + (1.0 - np.asarray(active)) * last78 return distribution, actions, next_rnn798081def _delivery_ids(core):82 return [int(t.id) for t in core.tasks if int(t.type) == 1]838485def _committed_load(core, task_id, exclude_aid=None):86 """Sum of living member loads assigned to ``task_id``, optionally excluding an agent."""87 total = 0.088 for ag in core.agents:89 if not ag.living or (exclude_aid is not None and int(ag.aid) == int(exclude_aid)):90 continue91 if int(ag.coalition) == int(task_id):92 total += max(0.0, float(ag.load))93 return total949596def _find_focal_and_others(core):97 """Locate a focal agent and its co-members in a delivery coalition.9899 The source coalition needs at least two loaded members so that removing the co-members100 (``critical``) and keeping them (``replaceable``) changes the coalition's residual101 load without needing any outside slack. Returns ``(focal, task, others)`` or ``None``.102 """103 for task in core.tasks:104 if int(task.type) != 1 or not task.discovered or task.completed:105 continue106 members = [107 ag for ag in core.agents108 if ag.living109 and int(ag.coalition) == int(task.id)110 and float(ag.load) > 1e-6111 ]112 if len(members) < 2:113 continue114 focal = max(members, key=lambda ag: float(ag.load))115 others = [ag for ag in members if int(ag.aid) != int(focal.aid)]116 return focal, task, others117 return None118119120def _target_task(core, source_task_id):121 """Pick the discovered incomplete delivery task (not the source) with the largest122 payload deficit, as the natural transfer target."""123 best, best_deficit = None, -1.0124 for task in core.tasks:125 if int(task.type) != 1 or int(task.id) == int(source_task_id):126 continue127 if not task.discovered or task.completed:128 continue129 committed = _committed_load(core, task.id)130 deficit = max(0.0, float(task.demand) - committed)131 if deficit > best_deficit:132 best, best_deficit = task, deficit133 return best134135136def _apply_condition(core, focal, source_task, others, condition):137 """Toggle the source coalition's residual load.138139 ``critical`` moves the co-members to idle (focal is the only loaded member, so its140 departure empties the coalition). ``replaceable`` keeps the co-members in place (the141 focal has backup). Focal position/load and source task are held fixed; only the142 co-members' affiliation changes.143 """144 for ag in others:145 ag.act(int(source_task.id) if condition == "replaceable" else int(core.idle_task_id))146 return core147148149@torch.no_grad()150def _focal_distribution(actor, core, focal_id):151 obs = core._get_observations()152 available = core._get_available_actions()153 rnn = np.zeros(154 (int(actor.num_agents), 1, int(getattr(actor, "hidden_size", 1))),155 dtype=np.float32,156 )157 masks = np.ones((int(actor.num_agents), 1), dtype=np.float32)158 distribution, _ = _distribution(actor, obs, rnn, masks, available)159 probs = distribution.probs.reshape(actor.num_agents, actor.num_tasks)160 return probs[focal_id].detach().cpu().numpy()161162163# --------------------------------------------------------------------------- #164# Phase 1: shared state bank165# --------------------------------------------------------------------------- #166def _generate_state_bank(167 base_models, manifest_dir, output_dir, scenarios, cuda, event_step168):169 bank_dir = Path(output_dir) / "state_bank"170 bank_dir.mkdir(parents=True, exist_ok=True)171 metadata = []172 for model in base_models:173 name = str(model["name"])174 config, runner, envs, checkpoint = _make_runner(175 model, manifest_dir, output_dir / "generate" / name, cuda176 )177 actor = runner.policy.actor178 env = envs.envs[0]179 core = getattr(env, "env", env)180 rows = []181 try:182 for scenario_id in scenarios:183 seed = int(config.eval_seed) + int(scenario_id)184 np.random.seed(seed)185 torch.manual_seed(seed)186 if config.device.type == "cuda":187 torch.cuda.manual_seed_all(seed)188 env.seed(seed)189 obs, active, available = env.reset(scenario_id)190 core = getattr(env, "env", env)191 for task in core.tasks:192 if int(task.type) == 1:193 task.discovered = True194 task.discovered_step = 0195 obs = core._get_observations()196 available = core._get_available_actions()197 active = np.ones((config.num_agents, 1), dtype=np.float32)198 rnn = np.zeros(199 (config.num_agents, config.recurrent_N, config.hidden_size),200 dtype=np.float32,201 )202 masks = np.ones((config.num_agents, 1), dtype=np.float32)203 last = np.full(204 (config.num_agents, 1), int(core.idle_task_id), dtype=np.int64205 )206 while int(core.step_count) < int(event_step):207 _, actions, next_rnn = _policy_step(208 actor, obs, rnn, masks, available, active, last209 )210 obs, _, dones, _, active, _, available = env.step(actions)211 rnn = next_rnn212 last = actions213 if bool(np.all(np.asarray(dones, dtype=bool))):214 break215216 found = _find_focal_and_others(core)217 if found is None:218 continue219 focal, source_task, others = found220 target = _target_task(core, source_task.id)221 snap_path = bank_dir / f"{name}_{scenario_id:03d}.pkl"222 with snap_path.open("wb") as stream:223 pickle.dump(core, stream)224 metadata.append(225 {226 "schema_version": SCHEMA_VERSION,227 "base_model": name,228 "scenario_id": int(scenario_id),229 "seed": seed,230 "event_step": int(event_step),231 "snapshot": str(snap_path),232 "focal_id": int(focal.aid),233 "source_task_id": int(source_task.id),234 "source_demand": float(source_task.demand),235 "focal_load": float(focal.load),236 "other_load_without_focal": _committed_load(237 core, source_task.id, exclude_aid=focal.aid238 ),239 "other_ids": [int(h.aid) for h in others],240 "other_total_load": float(sum(float(h.load) for h in others)),241 "target_task_id": int(target.id) if target is not None else -1,242 }243 )244 finally:245 _close_runner(runner, envs)246 del runner, envs247 gc.collect()248 if config.device.type == "cuda":249 torch.cuda.empty_cache()250 _atomic_json(bank_dir / "snapshots.json", metadata)251 return bank_dir, metadata252253254# --------------------------------------------------------------------------- #255# Phase 2: paired evaluation256# --------------------------------------------------------------------------- #257@torch.no_grad()258def _evaluate_model(model, manifest_dir, output_dir, bank_dir, metadata, cuda):259 name = str(model["name"])260 config, runner, envs, checkpoint = _make_runner(261 model, manifest_dir, output_dir / "eval" / name, cuda262 )263 actor = runner.policy.actor264 rows = []265 try:266 for item in metadata:267 with Path(item["snapshot"]).open("rb") as stream:268 core = pickle.load(stream)269 focal_id = int(item["focal_id"])270 source_task_id = int(item["source_task_id"])271 target_task_id = int(item["target_task_id"])272 other_ids = [int(h) for h in item["other_ids"]]273 others = [ag for ag in core.agents if int(ag.aid) in other_ids]274 focal = next(ag for ag in core.agents if int(ag.aid) == focal_id)275 source_task = core.tasks[source_task_id]276277 pair = {}278 for condition in ("critical", "replaceable"):279 fork = pickle.loads(pickle.dumps(core))280 f_fork = next(ag for ag in fork.agents if int(ag.aid) == focal_id)281 s_fork = fork.tasks[source_task_id]282 o_fork = [ag for ag in fork.agents if int(ag.aid) in other_ids]283 _apply_condition(fork, f_fork, s_fork, o_fork, condition)284 probs = _focal_distribution(actor, fork, focal_id)285 stay = float(probs[source_task_id])286 transfer = float(probs[target_task_id]) if target_task_id >= 0 else 0.0287 leave = float(1.0 - stay)288 argmax = int(probs.argmax())289 pair[condition] = {290 "stay_probability": stay,291 "transfer_probability": transfer,292 "leave_probability": leave,293 "argmax_task": argmax,294 "argmax_is_stay": int(argmax == source_task_id),295 }296297 rows.append(298 {299 "schema_version": SCHEMA_VERSION,300 "model": name,301 "model_role": model.get("role", "unknown"),302 "ablation": model.get("ablation", "none"),303 "train_seed": int(model.get("train_seed", -1)),304 "checkpoint": str(checkpoint),305 "base_model": item["base_model"],306 "scenario_id": item["scenario_id"],307 "focal_id": focal_id,308 "source_task_id": source_task_id,309 "target_task_id": target_task_id,310 "source_demand": item["source_demand"],311 "other_load_without_focal": item["other_load_without_focal"],312 "critical_stay_probability": pair["critical"]["stay_probability"],313 "replaceable_stay_probability": pair["replaceable"]["stay_probability"],314 "stay_delta_critical_minus_replaceable": (315 pair["critical"]["stay_probability"]316 - pair["replaceable"]["stay_probability"]317 ),318 "critical_transfer_probability": pair["critical"]["transfer_probability"],319 "replaceable_transfer_probability": pair["replaceable"]["transfer_probability"],320 "critical_argmax_is_stay": pair["critical"]["argmax_is_stay"],321 "replaceable_argmax_is_stay": pair["replaceable"]["argmax_is_stay"],322 }323 )324 finally:325 _close_runner(runner, envs)326 del runner, envs327 gc.collect()328 if config.device.type == "cuda":329 torch.cuda.empty_cache()330 _write_csv(output_dir / f"{name}.csv", rows)331 return rows332333334def _summarize(rows):335 metrics = (336 "critical_stay_probability",337 "replaceable_stay_probability",338 "stay_delta_critical_minus_replaceable",339 "critical_transfer_probability",340 "replaceable_transfer_probability",341 "critical_argmax_is_stay",342 "replaceable_argmax_is_stay",343 )344 result = []345 for model in sorted({row["model"] for row in rows}):346 model_rows = [r for r in rows if r["model"] == model]347 row = {"model": model, "n": len(model_rows)}348 for metric in metrics:349 row[metric] = float(np.mean([float(r[metric]) for r in model_rows]))350 result.append(row)351 return result352353354def _worker(args):355 manifest_path = Path(args.manifest).resolve()356 manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {}357 output = Path(args.output).resolve()358 output.mkdir(parents=True, exist_ok=True)359 if args.phase == "generate":360 base = [m for m in manifest["models"] if m.get("role") == "base"]361 _generate_state_bank(362 base,363 manifest_path.parent,364 output,365 _parse_scenarios(args.scenarios),366 args.cuda,367 args.event_step,368 )369 return370 model = next(m for m in manifest["models"] if m["name"] == args.model)371 bank_dir = output / "state_bank"372 metadata = json.loads((bank_dir / "snapshots.json").read_text(encoding="utf-8"))373 _evaluate_model(model, manifest_path.parent, output, bank_dir, metadata, args.cuda)374375376def main():377 parser = argparse.ArgumentParser()378 parser.add_argument("--manifest", required=True)379 parser.add_argument("--output", required=True)380 parser.add_argument("--phase", required=True, choices=["generate", "evaluate"])381 parser.add_argument("--scenarios", default="0:32")382 parser.add_argument("--event-step", type=int, default=150)383 parser.add_argument("--cuda", type=int, default=0)384 parser.add_argument("--max-workers", type=int, default=4)385 parser.add_argument("--model")386 args = parser.parse_args()387388 # A per-model worker subprocess (``--model`` set) evaluates exactly one389 # model and exits. It must NOT fall through to the orchestrator below:390 # that path ignores ``--model`` and re-spawns the full per-model fan-out,391 # which recurses without bound (the 2026-09-06 fork-bomb / OOM root cause).392 if args.phase == "evaluate" and args.model:393 _worker(args)394 return395396 manifest_path = Path(args.manifest).resolve()397 manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {}398 output = Path(args.output).resolve()399 output.mkdir(parents=True, exist_ok=True)400 status_path = output / "status.json"401402 if args.phase == "generate":403 base = [m for m in manifest["models"] if m.get("role") == "base"]404 status = {405 "name": manifest.get("name", "source-commitment"),406 "manifest": str(manifest_path),407 "created_at": datetime.now().isoformat(timespec="seconds"),408 "state": "running",409 "phase": "generate",410 "jobs": {411 m["name"]: {"state": "running", "role": m.get("role", "unknown")}412 for m in base413 },414 }415 _atomic_json(status_path, status)416 _generate_state_bank(417 base,418 manifest_path.parent,419 output,420 _parse_scenarios(args.scenarios),421 args.cuda,422 args.event_step,423 )424 status["jobs"] = {425 m["name"]: {"state": "completed", "role": m.get("role", "unknown")}426 for m in base427 }428 status["state"] = "completed"429 status["finished_at"] = datetime.now().isoformat(timespec="seconds")430 _atomic_json(status_path, status)431 return432433 # evaluate phase434 model_names = [m["name"] for m in manifest["models"] if m.get("role") != "base"]435 status = {436 "name": manifest.get("name", "source-commitment"),437 "manifest": str(manifest_path),438 "created_at": datetime.now().isoformat(timespec="seconds"),439 "state": "running",440 "phase": "evaluate",441 "jobs": {442 name: {"state": "pending"} for name in model_names443 },444 }445 _atomic_json(status_path, status)446 lock = threading.Lock()447448 def launch(name):449 with lock:450 status["jobs"][name]["state"] = "running"451 status["jobs"][name]["started_at"] = datetime.now().isoformat(timespec="seconds")452 _atomic_json(status_path, status)453 log = output / f"{name}.log"454 with log.open("w", encoding="utf-8") as stream:455 return subprocess.run(456 [457 sys.executable, str(Path(__file__).resolve()),458 "--manifest", str(manifest_path),459 "--output", str(output),460 "--phase", "evaluate",461 "--scenarios", args.scenarios,462 "--event-step", str(args.event_step),463 "--cuda", str(args.cuda),464 "--model", name,465 ],466 cwd=REPO_ROOT,467 stdout=stream,468 stderr=subprocess.STDOUT,469 env={**os.environ, "PYTHONUNBUFFERED": "1"},470 check=False,471 ).returncode472473 with ThreadPoolExecutor(max_workers=max(1, args.max_workers)) as pool:474 futures = {pool.submit(launch, name): name for name in model_names}475 for future in as_completed(futures):476 name = futures[future]477 code = int(future.result())478 with lock:479 status["jobs"][name]["state"] = "completed" if code == 0 else "failed"480 status["jobs"][name]["return_code"] = code481 status["jobs"][name]["finished_at"] = datetime.now().isoformat(timespec="seconds")482 _atomic_json(status_path, status)483484 all_rows = []485 for name in model_names:486 path = output / f"{name}.csv"487 if path.exists():488 with path.open(encoding="utf-8", newline="") as stream:489 all_rows.extend(csv.DictReader(stream))490 _write_csv(output / "all_results.csv", all_rows)491 _write_csv(output / "summary.csv", _summarize(all_rows))492 status["state"] = "completed"493 status["finished_at"] = datetime.now().isoformat(timespec="seconds")494 _atomic_json(status_path, status)495496497if __name__ == "__main__":498 main()