"""Evaluate frozen policies on a controlled task-priority reveal event. Each model first creates its own pre-event trajectory on a scenario with nine visible delivery tasks and one reserved hidden task. At ``event_step`` the same serialized environment state is forked into three conditions: * ``high``: reveal the reserved task with priority 1.0; * ``low``: reveal it with priority 0.5; * ``high_keep_current``: reveal the high-priority task while retaining the previous joint action at the event-triggered decision. The high/low pair measures selective policy response. ``high_keep_current`` is an event-response counterfactual baseline, not an oracle or a separate policy. The environment next invokes the policy at its usual redecision event, so this branch measures the value of reacting at the reveal itself. All continuations are deterministic so the three branches share the same pre-event state and differ only through the registered intervention. """ from __future__ import annotations import argparse import copy import csv import gc import hashlib import json import os 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 _physical_state_hash(core): payload = { "step": int(core.step_count), "agents": [ { "id": int(agent.aid), "position": np.asarray(agent.position, dtype=np.float64).tolist(), "living": bool(agent.living), "load": float(agent.load), "battery": float(agent.battery), "coalition": int(agent.coalition), } for agent in core.agents ], "tasks": [ { "id": int(task.id), "position": np.asarray(task.position, dtype=np.float64).tolist(), "demand": float(task.demand), "priority": float(task.req_vec[0]), "urgency": float(task.urgency), "discovered": bool(task.discovered), "completed": bool(task.completed), } for task in core.tasks ], } rendered = json.dumps(payload, sort_keys=True, separators=(",", ":")) return hashlib.sha256(rendered.encode("utf-8")).hexdigest() def _policy_step(actor, obs, rnn_states, masks, available_actions, active, last): 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 _set_reserved_visibility(core, task_id, discovered): task = core.tasks[int(task_id)] task.discovered = bool(discovered) task.completed = False if discovered: task.discovered_step = int(core.step_count) core.event_new_task_found = 1 else: task.discovered_step = None task.first_assigned_step = None task.response_time = None def _event_observation(env, task_id, priority): core = getattr(env, "env", env) task = core.tasks[int(task_id)] task.req_vec[0] = float(priority) _set_reserved_visibility(core, task_id, True) obs = core._get_observations() available = core._get_available_actions() active = np.asarray( [1.0 if agent.living else 0.0 for agent in core.agents], dtype=np.float32, )[:, None] core._current_decision_mask = active[:, 0].copy() return obs, active, available def _source_departure_diagnostics(core, actions): rows = [] for agent in core.agents: old = int(agent.coalition) new = int(actions[int(agent.aid), 0]) if ( new == old or old < 0 or old >= len(core.tasks) or int(core.tasks[old].type) != 1 ): continue task = core.tasks[old] remaining_other_load = sum( max(0.0, float(other.load)) for other in core.agents if other.living and int(other.aid) != int(agent.aid) and int(other.coalition) == old ) rows.append( { "agent": int(agent.aid), "source_task": old, "source_demand": float(task.demand), "source_other_committed_load": remaining_other_load, "source_fragile": int( remaining_other_load + 1e-6 < float(task.demand) ), } ) return rows def _metrics(core, task_id, start_step, reward_sum): deliver = [task for task in core.tasks if int(task.type) == 1] completed = sum(bool(task.completed) for task in deliver) weighted_total = sum(max(0.0, float(task.req_vec[0])) for task in deliver) weighted_done = sum( max(0.0, float(task.req_vec[0])) for task in deliver if task.completed ) target = core.tasks[int(task_id)] return { "success": int(completed == len(deliver)), "tasks_completed": int(completed), "task_completion_ratio": float(completed / max(1, len(deliver))), "priority_weighted_completion": float( weighted_done / max(1e-8, weighted_total) ), "remaining_demand": float( sum(max(0.0, float(task.demand)) for task in deliver) ), "event_task_completed": int(target.completed), "event_task_response_delay": ( int(target.first_assigned_step - start_step) if target.first_assigned_step is not None else -1 ), "episode_length": int(core.step_count), "team_reward_after_event": float(reward_sum), } @torch.no_grad() def _continue( config, actor, env, task_id, obs, active, available, rnn, masks, last, first ): core = getattr(env, "env", env) actions = np.asarray(first).copy() next_rnn = np.asarray(rnn).copy() reward_sum = 0.0 start_step = int(core.step_count) while int(core.step_count) < int(config.max_steps): obs, rewards, dones, _, active, _, available = env.step(actions) reward_sum += float(np.asarray(rewards, dtype=np.float64).sum()) done = np.asarray(dones, dtype=bool) next_rnn[done] = 0.0 masks = np.ones((config.num_agents, 1), dtype=np.float32) masks[done] = 0.0 rnn = next_rnn last = actions if bool(np.all(done)): break _, actions, next_rnn = _policy_step( actor, obs, rnn, masks, available, active, last ) return _metrics(core, task_id, start_step, reward_sum) @torch.no_grad() def _run_model(model, manifest_dir, output_dir, scenarios, cuda, event_step): name = str(model["name"]) config, runner, envs, checkpoint = _make_runner( model, manifest_dir, output_dir, cuda ) actor = runner.policy.actor env = envs.envs[0] 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) deliver_ids = [int(t.id) for t in core.tasks if int(t.type) == 1] task_id = int(deliver_ids[0]) for task in core.tasks: if int(task.type) == 1: task.discovered = True task.discovered_step = 0 _set_reserved_visibility(core, task_id, False) 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 ) valid = True 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 core.tasks[task_id].discovered: _set_reserved_visibility(core, task_id, False) core.event_new_task_found = 0 obs = core._get_observations() available = core._get_available_actions() all_still = all( (not agent.living) or bool(agent.idle) for agent in core.agents ) active = core._get_reconfiguration_active_mask( bool(core.event_any_task_completed), all_still )[:, None] core._current_decision_mask = active[:, 0].copy() if bool(np.all(dones)): valid = False break if not valid: rows.append( { "model": name, "model_role": model.get("role", "unknown"), "ablation": model.get("ablation", "none"), "scenario_id": int(scenario_id), "scenario_seed": seed, "condition": "skipped_early_terminal", } ) continue base = copy.deepcopy(env) pre_hash = _physical_state_hash(core) for condition, priority, defer in ( ("high", 1.0, False), ("low", 0.5, False), ("high_keep_current_at_event", 1.0, True), ): branch = copy.deepcopy(base) branch_core = getattr(branch, "env", branch) branch_obs, branch_active, branch_available = _event_observation( branch, task_id, priority ) distribution, policy_actions, next_rnn = _policy_step( actor, branch_obs, rnn, masks, branch_available, branch_active, last, ) probs = distribution.probs.detach().cpu().numpy() executed = np.asarray(last if defer else policy_actions).copy() source_rows = _source_departure_diagnostics( branch_core, policy_actions ) before = [int(agent.coalition) for agent in branch_core.agents] immediate_switches = sum( int(policy_actions[i, 0]) != before[i] for i in range(config.num_agents) if branch_core.agents[i].living ) immediate_target = sum( int(policy_actions[i, 0]) == task_id for i in range(config.num_agents) if branch_core.agents[i].living ) outcome = _continue( config, actor, branch, task_id, branch_obs, branch_active, branch_available, rnn, masks, last, executed, ) rows.append( { "schema_version": SCHEMA_VERSION, "model": name, "model_role": model.get("role", "unknown"), "training_condition": model.get( "training_condition", "unknown" ), "ablation": model.get("ablation", "none"), "train_seed": int(model.get("train_seed", -1)), "checkpoint_role": model.get( "checkpoint_kind", "validation-best" ), "checkpoint": str(checkpoint), "scenario_id": int(scenario_id), "scenario_seed": seed, "event_step": int(event_step), "condition": condition, "priority": priority, "pre_event_state_hash": pre_hash, "mean_event_task_probability": float( probs[:, task_id].mean() ), "max_event_task_probability": float( probs[:, task_id].max() ), "immediate_target_assignments": immediate_target, "immediate_policy_switches": immediate_switches, "departures_from_fragile_sources": int( sum(row["source_fragile"] for row in source_rows) ), "departures_from_delivery_sources": len(source_rows), **outcome, } ) print(f"{name} scenario={scenario_id} complete", flush=True) 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): grouped = {} for row in rows: if not str(row.get("condition", "")).startswith(("high", "low")): continue key = (row["model"], row["condition"]) grouped.setdefault(key, []).append(row) result = [] metrics = ( "mean_event_task_probability", "immediate_target_assignments", "immediate_policy_switches", "departures_from_fragile_sources", "success", "task_completion_ratio", "priority_weighted_completion", "event_task_completed", "event_task_response_delay", "remaining_demand", "episode_length", ) for (model, condition), items in sorted(grouped.items()): row = {"model": model, "condition": condition, "n": len(items)} for metric in metrics: values = [float(item[metric]) for item in items if metric in item] row[metric] = float(np.mean(values)) if values else None result.append(row) return result def _paired_effects(rows): indexed = { (row.get("model"), row.get("scenario_id"), row.get("condition")): row for row in rows } metrics = ( "mean_event_task_probability", "immediate_target_assignments", "immediate_policy_switches", "departures_from_fragile_sources", "success", "task_completion_ratio", "priority_weighted_completion", "event_task_completed", "remaining_demand", "episode_length", "team_reward_after_event", ) result = [] pairs = ( ("high_minus_low", "high", "low"), ( "react_minus_keep_current", "high", "high_keep_current_at_event", ), ) for model, scenario_id, condition in sorted(indexed): if condition != "high": continue for contrast, treatment, control in pairs: left = indexed.get((model, scenario_id, treatment)) right = indexed.get((model, scenario_id, control)) if left is None or right is None: continue item = { "model": model, "scenario_id": scenario_id, "contrast": contrast, "pre_event_state_hash": left.get("pre_event_state_hash"), } for metric in metrics: item[metric] = float(left[metric]) - float(right[metric]) result.append(item) return result def _worker(args): manifest_path = Path(args.manifest).resolve() manifest = yaml.safe_load(manifest_path.read_text(encoding="utf-8")) or {} model = next(item for item in manifest["models"] if item["name"] == args.model) output = Path(args.output).resolve() output.mkdir(parents=True, exist_ok=True) _run_model( model, manifest_path.parent, output, _parse_scenarios(args.scenarios), args.cuda, args.event_step, ) def main(): parser = argparse.ArgumentParser() parser.add_argument("--manifest", required=True) parser.add_argument("--output", required=True) parser.add_argument("--scenarios", default="0:32") parser.add_argument("--event-step", type=int, default=150) parser.add_argument("--cuda", type=int, default=1) parser.add_argument("--max-workers", type=int, default=4) parser.add_argument("--model") args = parser.parse_args() if 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" status = { "name": manifest.get("name", "priority-reconfiguration"), "manifest": str(manifest_path), "created_at": datetime.now().isoformat(timespec="seconds"), "state": "running", "jobs": { item["name"]: { "state": "pending", "role": item.get("role", "unknown"), "training_condition": item.get("training_condition", "unknown"), "ablation": item.get("ablation", "none"), } for item in manifest["models"] }, } _atomic_json(status_path, status) status_lock = threading.Lock() command_base = [ sys.executable, str(Path(__file__).resolve()), "--manifest", str(manifest_path), "--output", str(output), "--scenarios", args.scenarios, "--event-step", str(args.event_step), "--cuda", str(args.cuda), ] def launch(name): with status_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( command_base + ["--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 = {} for item in manifest["models"]: name = item["name"] futures[pool.submit(launch, name)] = name for future in as_completed(futures): name = futures[future] code = int(future.result()) with status_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 item in manifest["models"]: path = output / f"{item['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) summary = _summarize(all_rows) _write_csv(output / "summary.csv", summary) _write_csv(output / "paired_effects.csv", _paired_effects(all_rows)) status["state"] = ( "completed" if all(job["state"] == "completed" for job in status["jobs"].values()) else "partial_failure" ) status["finished_at"] = datetime.now().isoformat(timespec="seconds") _atomic_json(status_path, status) if __name__ == "__main__": main()