"""Paired source-coalition exit-cost evaluation for task-conditioned AA (HATT v3.4). This is the ``source_commitment_pair`` confirmation experiment. It tests whether a focal agent's decision to leave its current delivery coalition responds to the coalition's residual capacity after the agent departs: * ``replaceable`` — the source coalition can still finish without the focal agent; * ``critical`` — the focal agent is essential, so its departure makes the source coalition infeasible. The only difference between the two conditions is which *other* agents belong to the source coalition (``critical`` keeps them out; ``replaceable`` adds enough idle load to cover the remaining demand). Focal position/load, source task, target tasks and total fleet load are all held fixed. Each model reads the same pickled state bank, so cross-model state-distribution confounds are removed. The metric of interest is ``focal_stay_probability`` (and ``focal_transfer_probability`` toward the highest-deficit target). A task-conditioned A-A gate should push ``P(stay)`` up under ``critical`` relative to ``replaceable``; the v3.3AA (no-context) control should show a smaller split. Two phases: * ``--phase generate`` builds the shared state bank from base policies; * ``--phase evaluate`` runs the paired replaceable/critical evaluation on the bank. """ from __future__ import annotations import argparse import copy import csv import gc import json import os import pickle import subprocess import sys import threading from concurrent.futures import ThreadPoolExecutor, as_completed from datetime import datetime from pathlib import Path import numpy as np import torch import yaml REPO_ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(REPO_ROOT)) from scripts.analysis.diagnose_idle_opportunity import ( # noqa: E402 _distribution, _make_runner, _parse_scenarios, _write_csv, ) from train.eval import _close_runner # noqa: E402 SCHEMA_VERSION = 1 def _atomic_json(path, value): path = Path(path) temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text( json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8" ) os.replace(temporary, path) def _policy_step(actor, obs, rnn_states, masks, available_actions, active, last): """Argmax step reusing the actor distribution; non-active agents hold their action.""" distribution, next_rnn = _distribution( actor, obs, rnn_states, masks, available_actions ) proposed = distribution.probs.argmax(dim=-1, keepdim=True).cpu().numpy() next_rnn = next_rnn.detach().cpu().numpy() actions = np.asarray(active) * proposed + (1.0 - np.asarray(active)) * last return distribution, actions, next_rnn def _delivery_ids(core): return [int(t.id) for t in core.tasks if int(t.type) == 1] def _committed_load(core, task_id, exclude_aid=None): """Sum of living member loads assigned to ``task_id``, optionally excluding an agent.""" total = 0.0 for ag in core.agents: if not ag.living or (exclude_aid is not None and int(ag.aid) == int(exclude_aid)): continue if int(ag.coalition) == int(task_id): total += max(0.0, float(ag.load)) return total def _find_focal_and_others(core): """Locate a focal agent and its co-members in a delivery coalition. The source coalition needs at least two loaded members so that removing the co-members (``critical``) and keeping them (``replaceable``) changes the coalition's residual load without needing any outside slack. Returns ``(focal, task, others)`` or ``None``. """ for task in core.tasks: if int(task.type) != 1 or not task.discovered or task.completed: continue members = [ ag for ag in core.agents if ag.living and int(ag.coalition) == int(task.id) and float(ag.load) > 1e-6 ] if len(members) < 2: continue focal = max(members, key=lambda ag: float(ag.load)) others = [ag for ag in members if int(ag.aid) != int(focal.aid)] return focal, task, others return None def _target_task(core, source_task_id): """Pick the discovered incomplete delivery task (not the source) with the largest payload deficit, as the natural transfer target.""" best, best_deficit = None, -1.0 for task in core.tasks: if int(task.type) != 1 or int(task.id) == int(source_task_id): continue if not task.discovered or task.completed: continue committed = _committed_load(core, task.id) deficit = max(0.0, float(task.demand) - committed) if deficit > best_deficit: best, best_deficit = task, deficit return best def _apply_condition(core, focal, source_task, others, condition): """Toggle the source coalition's residual load. ``critical`` moves the co-members to idle (focal is the only loaded member, so its departure empties the coalition). ``replaceable`` keeps the co-members in place (the focal has backup). Focal position/load and source task are held fixed; only the co-members' affiliation changes. """ for ag in others: ag.act(int(source_task.id) if condition == "replaceable" else int(core.idle_task_id)) return core @torch.no_grad() def _focal_distribution(actor, core, focal_id): obs = core._get_observations() available = core._get_available_actions() rnn = np.zeros( (int(actor.num_agents), 1, int(getattr(actor, "hidden_size", 1))), dtype=np.float32, ) masks = np.ones((int(actor.num_agents), 1), dtype=np.float32) distribution, _ = _distribution(actor, obs, rnn, masks, available) probs = distribution.probs.reshape(actor.num_agents, actor.num_tasks) return probs[focal_id].detach().cpu().numpy() # --------------------------------------------------------------------------- # # Phase 1: shared state bank # --------------------------------------------------------------------------- # def _generate_state_bank( base_models, manifest_dir, output_dir, scenarios, cuda, event_step ): bank_dir = Path(output_dir) / "state_bank" bank_dir.mkdir(parents=True, exist_ok=True) metadata = [] for model in base_models: name = str(model["name"]) config, runner, envs, checkpoint = _make_runner( model, manifest_dir, output_dir / "generate" / name, cuda ) actor = runner.policy.actor env = envs.envs[0] core = getattr(env, "env", env) rows = [] try: for scenario_id in scenarios: seed = int(config.eval_seed) + int(scenario_id) np.random.seed(seed) torch.manual_seed(seed) if config.device.type == "cuda": torch.cuda.manual_seed_all(seed) env.seed(seed) obs, active, available = env.reset(scenario_id) core = getattr(env, "env", env) for task in core.tasks: if int(task.type) == 1: task.discovered = True task.discovered_step = 0 obs = core._get_observations() available = core._get_available_actions() active = np.ones((config.num_agents, 1), dtype=np.float32) rnn = np.zeros( (config.num_agents, config.recurrent_N, config.hidden_size), dtype=np.float32, ) masks = np.ones((config.num_agents, 1), dtype=np.float32) last = np.full( (config.num_agents, 1), int(core.idle_task_id), dtype=np.int64 ) while int(core.step_count) < int(event_step): _, actions, next_rnn = _policy_step( actor, obs, rnn, masks, available, active, last ) obs, _, dones, _, active, _, available = env.step(actions) rnn = next_rnn last = actions if bool(np.all(np.asarray(dones, dtype=bool))): break found = _find_focal_and_others(core) if found is None: continue focal, source_task, others = found target = _target_task(core, source_task.id) snap_path = bank_dir / f"{name}_{scenario_id:03d}.pkl" with snap_path.open("wb") as stream: pickle.dump(core, stream) metadata.append( { "schema_version": SCHEMA_VERSION, "base_model": name, "scenario_id": int(scenario_id), "seed": seed, "event_step": int(event_step), "snapshot": str(snap_path), "focal_id": int(focal.aid), "source_task_id": int(source_task.id), "source_demand": float(source_task.demand), "focal_load": float(focal.load), "other_load_without_focal": _committed_load( core, source_task.id, exclude_aid=focal.aid ), "other_ids": [int(h.aid) for h in others], "other_total_load": float(sum(float(h.load) for h in others)), "target_task_id": int(target.id) if target is not None else -1, } ) finally: _close_runner(runner, envs) del runner, envs gc.collect() if config.device.type == "cuda": torch.cuda.empty_cache() _atomic_json(bank_dir / "snapshots.json", metadata) return bank_dir, metadata # --------------------------------------------------------------------------- # # Phase 2: paired evaluation # --------------------------------------------------------------------------- # @torch.no_grad() def _evaluate_model(model, manifest_dir, output_dir, bank_dir, metadata, cuda): name = str(model["name"]) config, runner, envs, checkpoint = _make_runner( model, manifest_dir, output_dir / "eval" / name, cuda ) actor = runner.policy.actor rows = [] try: for item in metadata: with Path(item["snapshot"]).open("rb") as stream: core = pickle.load(stream) focal_id = int(item["focal_id"]) source_task_id = int(item["source_task_id"]) target_task_id = int(item["target_task_id"]) other_ids = [int(h) for h in item["other_ids"]] others = [ag for ag in core.agents if int(ag.aid) in other_ids] focal = next(ag for ag in core.agents if int(ag.aid) == focal_id) source_task = core.tasks[source_task_id] pair = {} for condition in ("critical", "replaceable"): fork = pickle.loads(pickle.dumps(core)) f_fork = next(ag for ag in fork.agents if int(ag.aid) == focal_id) s_fork = fork.tasks[source_task_id] o_fork = [ag for ag in fork.agents if int(ag.aid) in other_ids] _apply_condition(fork, f_fork, s_fork, o_fork, condition) probs = _focal_distribution(actor, fork, focal_id) stay = float(probs[source_task_id]) transfer = float(probs[target_task_id]) if target_task_id >= 0 else 0.0 leave = float(1.0 - stay) argmax = int(probs.argmax()) pair[condition] = { "stay_probability": stay, "transfer_probability": transfer, "leave_probability": leave, "argmax_task": argmax, "argmax_is_stay": int(argmax == source_task_id), } rows.append( { "schema_version": SCHEMA_VERSION, "model": name, "model_role": model.get("role", "unknown"), "ablation": model.get("ablation", "none"), "train_seed": int(model.get("train_seed", -1)), "checkpoint": str(checkpoint), "base_model": item["base_model"], "scenario_id": item["scenario_id"], "focal_id": focal_id, "source_task_id": source_task_id, "target_task_id": target_task_id, "source_demand": item["source_demand"], "other_load_without_focal": item["other_load_without_focal"], "critical_stay_probability": pair["critical"]["stay_probability"], "replaceable_stay_probability": pair["replaceable"]["stay_probability"], "stay_delta_critical_minus_replaceable": ( pair["critical"]["stay_probability"] - pair["replaceable"]["stay_probability"] ), "critical_transfer_probability": pair["critical"]["transfer_probability"], "replaceable_transfer_probability": pair["replaceable"]["transfer_probability"], "critical_argmax_is_stay": pair["critical"]["argmax_is_stay"], "replaceable_argmax_is_stay": pair["replaceable"]["argmax_is_stay"], } ) finally: _close_runner(runner, envs) del runner, envs gc.collect() if config.device.type == "cuda": torch.cuda.empty_cache() _write_csv(output_dir / f"{name}.csv", rows) return rows def _summarize(rows): metrics = ( "critical_stay_probability", "replaceable_stay_probability", "stay_delta_critical_minus_replaceable", "critical_transfer_probability", "replaceable_transfer_probability", "critical_argmax_is_stay", "replaceable_argmax_is_stay", ) result = [] for model in sorted({row["model"] for row in rows}): model_rows = [r for r in rows if r["model"] == model] row = {"model": model, "n": len(model_rows)} for metric in metrics: row[metric] = float(np.mean([float(r[metric]) for r in model_rows])) result.append(row) return result def _worker(args): manifest_path = Path(args.manifest).resolve() manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} output = Path(args.output).resolve() output.mkdir(parents=True, exist_ok=True) if args.phase == "generate": base = [m for m in manifest["models"] if m.get("role") == "base"] _generate_state_bank( base, manifest_path.parent, output, _parse_scenarios(args.scenarios), args.cuda, args.event_step, ) return model = next(m for m in manifest["models"] if m["name"] == args.model) bank_dir = output / "state_bank" metadata = json.loads((bank_dir / "snapshots.json").read_text(encoding="utf-8")) _evaluate_model(model, manifest_path.parent, output, bank_dir, metadata, args.cuda) def main(): parser = argparse.ArgumentParser() parser.add_argument("--manifest", required=True) parser.add_argument("--output", required=True) parser.add_argument("--phase", required=True, choices=["generate", "evaluate"]) parser.add_argument("--scenarios", default="0:32") parser.add_argument("--event-step", type=int, default=150) parser.add_argument("--cuda", type=int, default=0) parser.add_argument("--max-workers", type=int, default=4) parser.add_argument("--model") args = parser.parse_args() # A per-model worker subprocess (``--model`` set) evaluates exactly one # model and exits. It must NOT fall through to the orchestrator below: # that path ignores ``--model`` and re-spawns the full per-model fan-out, # which recurses without bound (the 2026-09-06 fork-bomb / OOM root cause). if args.phase == "evaluate" and args.model: _worker(args) return manifest_path = Path(args.manifest).resolve() manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} output = Path(args.output).resolve() output.mkdir(parents=True, exist_ok=True) status_path = output / "status.json" if args.phase == "generate": base = [m for m in manifest["models"] if m.get("role") == "base"] status = { "name": manifest.get("name", "source-commitment"), "manifest": str(manifest_path), "created_at": datetime.now().isoformat(timespec="seconds"), "state": "running", "phase": "generate", "jobs": { m["name"]: {"state": "running", "role": m.get("role", "unknown")} for m in base }, } _atomic_json(status_path, status) _generate_state_bank( base, manifest_path.parent, output, _parse_scenarios(args.scenarios), args.cuda, args.event_step, ) status["jobs"] = { m["name"]: {"state": "completed", "role": m.get("role", "unknown")} for m in base } status["state"] = "completed" status["finished_at"] = datetime.now().isoformat(timespec="seconds") _atomic_json(status_path, status) return # evaluate phase model_names = [m["name"] for m in manifest["models"] if m.get("role") != "base"] status = { "name": manifest.get("name", "source-commitment"), "manifest": str(manifest_path), "created_at": datetime.now().isoformat(timespec="seconds"), "state": "running", "phase": "evaluate", "jobs": { name: {"state": "pending"} for name in model_names }, } _atomic_json(status_path, status) lock = threading.Lock() def launch(name): with lock: status["jobs"][name]["state"] = "running" status["jobs"][name]["started_at"] = datetime.now().isoformat(timespec="seconds") _atomic_json(status_path, status) log = output / f"{name}.log" with log.open("w", encoding="utf-8") as stream: return subprocess.run( [ sys.executable, str(Path(__file__).resolve()), "--manifest", str(manifest_path), "--output", str(output), "--phase", "evaluate", "--scenarios", args.scenarios, "--event-step", str(args.event_step), "--cuda", str(args.cuda), "--model", name, ], cwd=REPO_ROOT, stdout=stream, stderr=subprocess.STDOUT, env={**os.environ, "PYTHONUNBUFFERED": "1"}, check=False, ).returncode with ThreadPoolExecutor(max_workers=max(1, args.max_workers)) as pool: futures = {pool.submit(launch, name): name for name in model_names} for future in as_completed(futures): name = futures[future] code = int(future.result()) with lock: status["jobs"][name]["state"] = "completed" if code == 0 else "failed" status["jobs"][name]["return_code"] = code status["jobs"][name]["finished_at"] = datetime.now().isoformat(timespec="seconds") _atomic_json(status_path, status) all_rows = [] for name in model_names: path = output / f"{name}.csv" if path.exists(): with path.open(encoding="utf-8", newline="") as stream: all_rows.extend(csv.DictReader(stream)) _write_csv(output / "all_results.csv", all_rows) _write_csv(output / "summary.csv", _summarize(all_rows)) status["state"] = "completed" status["finished_at"] = datetime.now().isoformat(timespec="seconds") _atomic_json(status_path, status) if __name__ == "__main__": main()