"""Self-play training loop for RecallTrace. Runs episodes where the Adversary picks intervention placements and the Investigator tries to find them. Both agents update after each episode. """ from __future__ import annotations import random import time from typing import Any, Dict, List from env.env import RecallTraceEnv from selfplay.adversary import AdversaryAgent, GRAPH_REGIONS from selfplay.investigator import InvestigatorAgent from selfplay.scenario_gen import apply_intervention, compute_f1, generate_graph class SelfPlayTrainer: """Orchestrates adversarial self-play between Investigator and Adversary.""" def __init__(self, num_nodes: int = 12): self.num_nodes = num_nodes self.adversary = AdversaryAgent(temperature=2.0, min_temperature=0.3) self.investigator = InvestigatorAgent() self.all_stats: List[Dict[str, Any]] = [] def run_episode(self, episode_num: int, seed: int | None = None) -> Dict[str, Any]: """Run a single self-play episode. Returns episode stats dict.""" rng = random.Random(seed) # 1) Generate a fresh supply-chain graph graph_scenario = generate_graph(num_nodes=self.num_nodes, seed=seed) # 2) Adversary picks intervention intervention_type, target_node, num_hops = self.adversary.choose_intervention( graph_scenario, rng=rng, ) # Determine graph region of target node graph_region = graph_scenario.get("_node_regions", {}).get(target_node, "downstream") # 3) Apply intervention to scenario scenario = apply_intervention( graph_scenario, intervention_type, target_node, num_hops, rng=rng, ) # 4) Create environment and reset env = RecallTraceEnv(scenario_data=scenario) observation = env.reset() # 5) Investigator runs the episode self.investigator.reset_episode() total_reward = 0.0 steps = 0 done = False while not done and steps < scenario["max_steps"]: action = self.investigator.act(observation, rng=rng) observation, reward, done, info = env.step(action) total_reward += reward steps += 1 # Force finalize if not done if not done: action = self.investigator.act(observation, rng=rng) if action.type.value != "finalize": from env.models import RecallAction action = RecallAction(type="finalize", rationale="Budget exhausted.") observation, reward, done, info = env.step(action) total_reward += reward steps += 1 # 6) Compute F1 from quarantine results quarantined_nodes = list(set(self.investigator.nodes_quarantined)) # Also check env state for quarantined inventory env_state = env.state() for node_id, node_data in env_state.state_data.get("nodes", {}).items(): q_inv = node_data.get("quarantined_inventory", {}) if q_inv and node_id not in quarantined_nodes: quarantined_nodes.append(node_id) f1, f1_details = compute_f1(scenario, quarantined_nodes) # 7) Compute investigator reward with the specified reward structure inv_reward = 0.0 tp = f1_details["tp"] fp = f1_details["fp"] inv_reward += tp * 2.0 # +2.0 per correctly quarantined unsafe node inv_reward += fp * (-1.5) # -1.5 per safe node wrongly blocked inv_reward += steps * (-0.05) # -0.05 per step # Belief calibration bonus if f1 > 0.6: inv_reward += 0.3 # 8) Update both agents adversary_reward = self.adversary.update(intervention_type, graph_region, f1) self.investigator.update(inv_reward, f1, steps) # 9) Build stats dict inv_summary = self.investigator.get_episode_summary() correctly_identified = ( inv_summary["intervention_guess"] == intervention_type if inv_summary["intervention_guess"] is not None else False ) stats = { "episode": episode_num, "investigator_f1": round(f1, 4), "adversary_reward": round(adversary_reward, 4), "investigator_reward": round(inv_reward, 4), "num_quarantined": len(quarantined_nodes), "intervention_type": intervention_type, "graph_region": graph_region, "target_node": target_node, "num_hops": num_hops, "steps_taken": steps, "nodes_visited": inv_summary["nodes_visited"], "nodes_quarantined_list": sorted(set(quarantined_nodes)), "belief_confidence": inv_summary["belief_confidence"], "quarantine_threshold": inv_summary["quarantine_threshold"], "exploration_rate": inv_summary["exploration_rate"], "intervention_guess": inv_summary["intervention_guess"], "intervention_correctly_identified": correctly_identified, "f1_details": f1_details, } return stats def train(self, num_episodes: int = 200) -> List[Dict[str, Any]]: """Run the full self-play training loop.""" print(f"\n{'='*70}") print(f" RecallTrace — Adversarial Self-Play Training") print(f" Episodes: {num_episodes} | Nodes per graph: {self.num_nodes}") print(f"{'='*70}\n") self.all_stats = [] start_time = time.time() for ep in range(1, num_episodes + 1): stats = self.run_episode(episode_num=ep, seed=ep * 42) self.all_stats.append(stats) # Progress logging every 20 episodes if ep % 20 == 0 or ep == 1: recent = self.all_stats[-20:] if len(self.all_stats) >= 20 else self.all_stats avg_f1 = sum(s["investigator_f1"] for s in recent) / len(recent) avg_adv = sum(s["adversary_reward"] for s in recent) / len(recent) avg_q = sum(s["num_quarantined"] for s in recent) / len(recent) avg_steps = sum(s["steps_taken"] for s in recent) / len(recent) elapsed = time.time() - start_time print( f" Episode {ep:>4d} | " f"F1: {avg_f1:.3f} | " f"Adv Reward: {avg_adv:+.3f} | " f"Quarantined: {avg_q:.1f} | " f"Steps: {avg_steps:.1f} | " f"Time: {elapsed:.1f}s" ) elapsed = time.time() - start_time print(f"\n Training complete in {elapsed:.1f}s") print(f" Adversary strategy: {self.adversary.get_strategy_summary()}") # Print summary early = self.all_stats[:20] late = self.all_stats[-20:] print(f"\n Early avg F1: {sum(s['investigator_f1'] for s in early)/len(early):.3f}") print(f" Late avg F1: {sum(s['investigator_f1'] for s in late)/len(late):.3f}") print(f" Early avg quarantined: {sum(s['num_quarantined'] for s in early)/len(early):.1f}") print(f" Late avg quarantined: {sum(s['num_quarantined'] for s in late)/len(late):.1f}") print() return self.all_stats @staticmethod def get_training_curves(stats: List[Dict[str, Any]]) -> Dict[str, List[float]]: """Extract plottable series from training stats.""" return { "episodes": [s["episode"] for s in stats], "investigator_f1": [s["investigator_f1"] for s in stats], "adversary_reward": [s["adversary_reward"] for s in stats], "num_quarantined": [s["num_quarantined"] for s in stats], "steps_taken": [s["steps_taken"] for s in stats], "quarantine_threshold": [s["quarantine_threshold"] for s in stats], "exploration_rate": [s["exploration_rate"] for s in stats], "belief_confidence": [s["belief_confidence"] for s in stats], }