Add alternating self-play training scaffolding.

Wire configurable frozen-opponent server startup and add training utilities for archive, eval harness, data collection, and freeze-alternate orchestration.

Made-with: Cursor

server/app.py

@@ -22,6 +22,9 @@ Usage:
     uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
 """
 
+import os
+import random
+
 try:
     from openenv.core.env_server.http_server import create_app
 except Exception as e:  # pragma: no cover
@@ -37,9 +40,38 @@ except (ImportError, ModuleNotFoundError):
     from server.play_environment import CyberSOCEnvironment
 
 
+class FrozenCheckpointRedPolicy:
+    """Lightweight frozen red policy shim keyed by checkpoint identity."""
+
+    def __init__(self, checkpoint: str):
+        self.checkpoint = checkpoint
+        self._rng = random.Random(hash(checkpoint))
+
+    def act(self, red_observation):
+        blue_action = red_observation.get("blue_action_type", "")
+        # Deterministic-ish heuristic controlled by checkpoint hash and episode step.
+        trigger = blue_action in {"kill_process", "isolate_segment"}
+        if trigger and self._rng.random() < 0.5:
+            return {
+                "action_type": "lateral_pivot",
+                "source_host": red_observation.get("blue_action_target", ""),
+            }
+        return {"action_type": "noop"}
+
+
+_frozen_checkpoint = os.environ.get("CYBERSOC_FROZEN_RED_CHECKPOINT", "").strip()
+_adaptive = os.environ.get("CYBERSOC_ADAPTIVE", "1").strip() not in {"0", "false", "False"}
+_red_policy = FrozenCheckpointRedPolicy(_frozen_checkpoint) if _frozen_checkpoint else None
+
+
+class ConfiguredCyberSOCEnvironment(CyberSOCEnvironment):
+    def __init__(self):
+        super().__init__(adaptive=_adaptive, neural_red_policy=_red_policy)
+
+
 # Create the app with the CyberSOCEnv environment
 app = create_app(
-    CyberSOCEnvironment,
+    ConfiguredCyberSOCEnvironment,
     SOCActionWrapper,
     SOCObservation,
     env_name="cybersocenv",

training/agent_archive.py

@@ -0,0 +1,92 @@
+"""Historical archive utilities for FSP/PFSP training."""
+
+from __future__ import annotations
+
+import json
+import random
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Dict, List, Optional
+
+
+@dataclass
+class ArchivedAgent:
+    role: str
+    version: str
+    checkpoint_path: str
+    iteration: int
+    metadata: Dict[str, float]
+
+
+class AgentArchive:
+    """Stores historical checkpoints and win-rate metadata in a JSON index."""
+
+    def __init__(self, index_path: str = "training/archive/index.json"):
+        self.index_path = Path(index_path)
+        self.index_path.parent.mkdir(parents=True, exist_ok=True)
+        self._items: List[ArchivedAgent] = self._load()
+
+    def _load(self) -> List[ArchivedAgent]:
+        if not self.index_path.exists():
+            return []
+        raw = json.loads(self.index_path.read_text(encoding="utf-8"))
+        return [ArchivedAgent(**item) for item in raw]
+
+    def save(self) -> None:
+        payload = [asdict(item) for item in self._items]
+        self.index_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+    def add(
+        self,
+        role: str,
+        version: str,
+        checkpoint_path: str,
+        iteration: int,
+        metadata: Optional[Dict[str, float]] = None,
+    ) -> None:
+        self._items.append(
+            ArchivedAgent(
+                role=role,
+                version=version,
+                checkpoint_path=checkpoint_path,
+                iteration=iteration,
+                metadata=metadata or {},
+            )
+        )
+        self.save()
+
+    def list_role(self, role: str) -> List[ArchivedAgent]:
+        return [item for item in self._items if item.role == role]
+
+    def latest(self, role: str) -> Optional[ArchivedAgent]:
+        role_items = self.list_role(role)
+        if not role_items:
+            return None
+        return max(role_items, key=lambda x: x.iteration)
+
+    def sample_fsp(self, role: str) -> Optional[ArchivedAgent]:
+        items = self.list_role(role)
+        if not items:
+            return None
+        return random.choice(items)
+
+    def sample_pfsp(self, role: str, temperature: float = 1.0) -> Optional[ArchivedAgent]:
+        """Prioritize opponents where blue has lower win-rate."""
+        items = self.list_role(role)
+        if not items:
+            return None
+
+        weights: List[float] = []
+        for item in items:
+            blue_win_rate = float(item.metadata.get("blue_win_rate", 0.5))
+            difficulty = max(0.0, min(1.0, 1.0 - blue_win_rate))
+            weights.append(max(1e-6, difficulty**temperature))
+
+        return random.choices(items, weights=weights, k=1)[0]
+
+    def must_beat_all(self, threshold: float = 0.55) -> bool:
+        """Return True when every archived red has blue_win_rate >= threshold."""
+        red_items = self.list_role("red")
+        if not red_items:
+            return True
+        return all(float(item.metadata.get("blue_win_rate", 0.0)) >= threshold for item in red_items)

training/collect_sft.py

@@ -0,0 +1,108 @@
+"""
+Collect imitation data from deterministic red-team decisions.
+
+This script runs generated scenarios and stores red decision tuples as JSONL:
+{"observation": {...}, "action": {...}}
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict, List
+
+from models import SOCActionWrapper
+from server.play_environment import CyberSOCEnvironment
+from server.tasks import get_task
+
+
+def _scripted_blue_actions(task_id: str) -> List[Dict[str, Any]]:
+    """Simple deterministic blue rollout to trigger red dynamics each step."""
+    task_def = get_task(task_id)
+    reqs = task_def.get("containment_requirements", {}) or {}
+    actions: List[Dict[str, Any]] = []
+
+    for host in reqs.get("must_forensics", []):
+        actions.append({"type": "run_forensics", "hostname": host})
+
+    for proc in reqs.get("must_kill", []):
+        actions.append(
+            {
+                "type": "kill_process",
+                "hostname": proc["hostname"],
+                "process_name": proc["process"],
+            }
+        )
+
+    for ioc in reqs.get("must_block_iocs", []):
+        ioc_type = "hash" if len(ioc) >= 32 and "." not in ioc else ("ip" if ioc.count(".") == 3 else "domain")
+        actions.append({"type": "block_ioc", "ioc_type": ioc_type, "ioc_value": ioc})
+
+    actions.append(
+        {
+            "type": "submit_containment_plan",
+            "plan": [
+                {
+                    "threat_id": t.get("threat_id", "UNKNOWN"),
+                    "actions_taken": ["run_forensics", "kill_process", "block_ioc"],
+                    "root_cause": t.get("threat_type", "unknown"),
+                    "confidence": 0.8,
+                }
+                for t in task_def.get("attack_chain", [])
+            ],
+            "executive_summary": "Automated containment sequence completed.",
+        }
+    )
+    return actions
+
+
+def collect_red_imitation_dataset(
+    output_path: Path,
+    num_tasks: int = 1000,
+    task_prefix: str = "gen_",
+) -> int:
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    records: List[Dict[str, Any]] = []
+
+    def _logger(record: Dict[str, Any]) -> None:
+        records.append(record)
+
+    env = CyberSOCEnvironment(adaptive=True, red_team_logger=_logger)
+
+    for idx in range(1, num_tasks + 1):
+        task_id = f"{task_prefix}{idx:04d}"
+        env.reset(task_id=task_id)
+        for action in _scripted_blue_actions(task_id):
+            obs = env.step(SOCActionWrapper(**action))
+            if obs.done:
+                break
+
+    with output_path.open("w", encoding="utf-8") as f:
+        for record in records:
+            f.write(json.dumps(record) + "\n")
+
+    return len(records)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Collect deterministic red-team SFT data")
+    parser.add_argument(
+        "--output",
+        default="training/data/red_imitation.jsonl",
+        help="Output JSONL path",
+    )
+    parser.add_argument("--num-tasks", type=int, default=1000, help="Number of generated scenarios")
+    parser.add_argument("--task-prefix", default="gen_", help="Task ID prefix")
+    args = parser.parse_args()
+
+    total = collect_red_imitation_dataset(
+        output_path=Path(args.output),
+        num_tasks=args.num_tasks,
+        task_prefix=args.task_prefix,
+    )
+    print(f"Saved {total} red decision examples to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

training/eval_harness.py

@@ -0,0 +1,58 @@
+"""Evaluation harness for Blue-vs-Red checkpoints."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List
+
+from models import SOCActionWrapper
+from server.play_environment import CyberSOCEnvironment
+
+
+PolicyFn = Callable[[Dict[str, Any]], Dict[str, Any]]
+
+
+@dataclass
+class EvalResult:
+    episodes: int
+    blue_win_rate: float
+    avg_blue_score: float
+    per_episode_scores: List[float]
+
+
+def run_head_to_head_eval(
+    blue_policy: PolicyFn,
+    red_policy: Any,
+    task_ids: List[str],
+    episodes_per_task: int = 50,
+) -> EvalResult:
+    """Evaluate blue against a frozen red policy across task IDs."""
+    scores: List[float] = []
+    wins = 0
+    total = 0
+
+    for task_id in task_ids:
+        for _ in range(episodes_per_task):
+            env = CyberSOCEnvironment(adaptive=True, neural_red_policy=red_policy)
+            obs = env.reset(task_id=task_id)
+
+            while not obs.done:
+                action_dict = blue_policy(obs.model_dump())
+                obs = env.step(SOCActionWrapper(**action_dict))
+
+            score = float(obs.final_score or 0.0)
+            scores.append(score)
+            wins += int(score >= 0.5)
+            total += 1
+
+    avg_score = sum(scores) / max(1, len(scores))
+    return EvalResult(
+        episodes=total,
+        blue_win_rate=wins / max(1, total),
+        avg_blue_score=avg_score,
+        per_episode_scores=scores,
+    )
+
+
+def must_beat_all_archive(win_rates: Dict[str, float], threshold: float = 0.55) -> bool:
+    return all(rate >= threshold for rate in win_rates.values())

training/freeze_alternate.py

@@ -0,0 +1,116 @@
+"""Freeze-alternate orchestration for Blue/Red GRPO training."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import subprocess
+from pathlib import Path
+from typing import Dict, Optional
+
+try:
+    from .agent_archive import AgentArchive
+    from .pfsp_scheduler import temperature_for_iteration
+except ImportError:
+    from agent_archive import AgentArchive
+    from pfsp_scheduler import temperature_for_iteration
+
+
+def _run_train_command(command: str) -> None:
+    completed = subprocess.run(command, shell=True, check=False)
+    if completed.returncode != 0:
+        raise RuntimeError(f"Training command failed: {command}")
+
+
+def _format_cmd(
+    base_cmd: str,
+    role: str,
+    output_dir: str,
+    frozen_opponent: Optional[str],
+    episodes: int,
+) -> str:
+    cmd = [
+        base_cmd,
+        f"--train-role {role}",
+        f"--output-dir {output_dir}",
+        f"--episodes {episodes}",
+    ]
+    if frozen_opponent:
+        cmd.append(f"--frozen-opponent {frozen_opponent}")
+    return " ".join(cmd)
+
+
+def run_freeze_alternate(
+    iterations: int,
+    train_blue_episodes: int,
+    train_red_episodes: int,
+    base_train_cmd: str,
+    archive_path: str,
+) -> Dict[str, str]:
+    archive = AgentArchive(archive_path)
+    latest_blue = archive.latest("blue")
+    latest_red = archive.latest("red")
+
+    for it in range(1, iterations + 1):
+        blue_version = f"blue_v{it}"
+        blue_ckpt = f"checkpoints/{blue_version}"
+        cmd_blue = _format_cmd(
+            base_cmd=base_train_cmd,
+            role="blue",
+            output_dir=blue_ckpt,
+            frozen_opponent=latest_red.checkpoint_path if latest_red else None,
+            episodes=train_blue_episodes,
+        )
+        _run_train_command(cmd_blue)
+        archive.add("blue", blue_version, blue_ckpt, iteration=it, metadata={})
+        latest_blue = archive.latest("blue")
+
+        red_version = f"red_v{it}"
+        red_ckpt = f"checkpoints/{red_version}"
+        cmd_red = _format_cmd(
+            base_cmd=base_train_cmd,
+            role="red",
+            output_dir=red_ckpt,
+            frozen_opponent=latest_blue.checkpoint_path if latest_blue else None,
+            episodes=train_red_episodes,
+        )
+        _run_train_command(cmd_red)
+        archive.add("red", red_version, red_ckpt, iteration=it, metadata={})
+        latest_red = archive.latest("red")
+
+    return {
+        "latest_blue": latest_blue.checkpoint_path if latest_blue else "",
+        "latest_red": latest_red.checkpoint_path if latest_red else "",
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Freeze-alternate Blue/Red training orchestrator")
+    parser.add_argument("--iterations", type=int, default=2)
+    parser.add_argument("--blue-episodes", type=int, default=500)
+    parser.add_argument("--red-episodes", type=int, default=300)
+    parser.add_argument("--train-cmd", default="python -m training.train_grpo")
+    parser.add_argument("--archive-path", default="training/archive/index.json")
+    parser.add_argument("--show-temp-for", type=int, default=0, help="Print PFSP temp schedule for N iterations")
+    args = parser.parse_args()
+
+    if args.show_temp_for > 0:
+        schedule = {
+            f"iter_{i + 1}": temperature_for_iteration(i, args.show_temp_for)
+            for i in range(args.show_temp_for)
+        }
+        print(json.dumps(schedule, indent=2))
+        return
+
+    outputs = run_freeze_alternate(
+        iterations=args.iterations,
+        train_blue_episodes=args.blue_episodes,
+        train_red_episodes=args.red_episodes,
+        base_train_cmd=args.train_cmd,
+        archive_path=args.archive_path,
+    )
+    print(json.dumps(outputs, indent=2))
+
+
+if __name__ == "__main__":
+    main()

training/pfsp_scheduler.py

@@ -0,0 +1,42 @@
+"""PFSP temperature scheduling and weighted opponent sampling."""
+
+from __future__ import annotations
+
+from typing import Dict, Iterable, Tuple
+
+
+def temperature_for_iteration(
+    iteration: int,
+    total_iterations: int,
+    start_temp: float = 0.5,
+    end_temp: float = 2.0,
+) -> float:
+    if total_iterations <= 1:
+        return end_temp
+    ratio = max(0.0, min(1.0, iteration / float(total_iterations - 1)))
+    return start_temp + (end_temp - start_temp) * ratio
+
+
+def pfsp_weights(
+    win_rates: Dict[str, float],
+    temperature: float,
+) -> Dict[str, float]:
+    """Compute PFSP weights = (1 - blue_win_rate) ** temperature."""
+    weights: Dict[str, float] = {}
+    for name, win_rate in win_rates.items():
+        clipped = max(0.0, min(1.0, float(win_rate)))
+        weights[name] = max(1e-6, (1.0 - clipped) ** temperature)
+    return weights
+
+
+def normalize_weights(raw_weights: Dict[str, float]) -> Dict[str, float]:
+    total = sum(raw_weights.values())
+    if total <= 0:
+        n = max(1, len(raw_weights))
+        return {k: 1.0 / n for k in raw_weights}
+    return {k: v / total for k, v in raw_weights.items()}
+
+
+def rank_hard_opponents(win_rates: Dict[str, float]) -> Iterable[Tuple[str, float]]:
+    """Yield opponents ordered from hardest (lowest win-rate) to easiest."""
+    return sorted(win_rates.items(), key=lambda kv: kv[1])

training/train_grpo.py

@@ -49,19 +49,29 @@ except ImportError:
 class EnvServer:
     """Starts and monitors the CyberSOCEnv FastAPI server as a subprocess."""
 
-    def __init__(self, cfg: TrainingConfig):
+    def __init__(self, cfg: TrainingConfig, frozen_opponent: str = ""):
         self.cfg = cfg
+        self.frozen_opponent = frozen_opponent
         self._proc: Optional[subprocess.Popen] = None
 
     def start(self) -> None:
         cmd = [
             sys.executable, "-m", "uvicorn",
-            "server:app",
+            "server.app:app",
             "--host", self.cfg.env_host,
             "--port", str(self.cfg.env_port),
             "--log-level", "warning",
         ]
-        self._proc = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
+        env = os.environ.copy()
+        if self.frozen_opponent:
+            env["CYBERSOC_FROZEN_RED_CHECKPOINT"] = self.frozen_opponent
+        env.setdefault("CYBERSOC_ADAPTIVE", "1")
+        self._proc = subprocess.Popen(
+            cmd,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.PIPE,
+            env=env,
+        )
         print(f"[server] PID {self._proc.pid} starting on :{self.cfg.env_port} ...")
         self._wait_healthy()
 
@@ -313,6 +323,20 @@ def make_all_reward_fns(env_url: str) -> List:
     return fns
 
 
+def invert_reward_fns(reward_fns: List) -> List:
+    """Invert blue-centric rewards for red training."""
+    inverted = []
+    for fn in reward_fns:
+        def _wrap(base_fn):
+            def _inv(completions: List[str], **kwargs) -> List[float]:
+                values = base_fn(completions, **kwargs)
+                return [-float(v) for v in values]
+            _inv.__name__ = f"inv_{getattr(base_fn, '__name__', 'reward')}"
+            return _inv
+        inverted.append(_wrap(fn))
+    return inverted
+
+
 # ═══════════════════════════════════════════════════════════════════════════
 # 5. Main
 # ═══════════════════════════════════════════════════════════════════════════
@@ -321,8 +345,12 @@ def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(description="CyberSOC GRPO Training")
     p.add_argument("--model",     default="", help="HuggingFace model ID (overrides config)")
     p.add_argument("--steps",     type=int, default=0, help="Max training steps (overrides config)")
+    p.add_argument("--episodes",  type=int, default=0, help="Alias for steps in freeze-alternate mode")
     p.add_argument("--tasks",     default="", help="Comma-separated task IDs, e.g. easy,medium,hard")
     p.add_argument("--hub-id",    default="", dest="hub_id", help="HF Hub model ID for checkpoint push")
+    p.add_argument("--output-dir", default="", help="Override output directory")
+    p.add_argument("--train-role", choices=["blue", "red"], default="blue", help="Policy role being optimized")
+    p.add_argument("--frozen-opponent", default="", help="Path/ID for frozen opponent checkpoint")
     p.add_argument("--no-server", action="store_true", help="Skip starting env server (already running)")
     return p.parse_args()
 
@@ -333,8 +361,10 @@ def main() -> None:
 
     if args.model:  cfg.model_name   = args.model
     if args.steps:  cfg.max_steps    = args.steps
+    if args.episodes: cfg.max_steps  = args.episodes
     if args.tasks:  cfg.task_ids     = args.tasks.split(",")
     if args.hub_id: cfg.hub_model_id = args.hub_id
+    if args.output_dir: cfg.output_dir = args.output_dir
 
     hf_token = os.environ.get("HF_TOKEN", "")
 
@@ -343,11 +373,14 @@ def main() -> None:
     print(f"  Model  : {cfg.model_name}")
     print(f"  Tasks  : {cfg.task_ids}")
     print(f"  Steps  : {cfg.max_steps}")
+    print(f"  Role   : {args.train_role}")
+    if args.frozen_opponent:
+        print(f"  Frozen : {args.frozen_opponent}")
     print(f"  Hub    : {cfg.hub_model_id or '(local only)'}")
     print("=" * 60)
 
     # 1. Start env server ────────────────────────────────────────────────────
-    server = EnvServer(cfg)
+    server = EnvServer(cfg, frozen_opponent=args.frozen_opponent)
     if not args.no_server:
         server.start()
 
@@ -382,6 +415,9 @@ def main() -> None:
         # 4. Build reward functions ───────────────────────────────────────────
         print("\n[rewards] wiring 11 reward functions (10 dims + weighted total) ...")
         reward_fns = make_all_reward_fns(cfg.env_url)
+        if args.train_role == "red":
+            reward_fns = invert_reward_fns(reward_fns)
+            print("[rewards] role=red -> inverted blue-centric rewards")
         print(f"[rewards] ✓ {len(reward_fns)} functions registered")
 
         # 5. GRPO config ──────────────────────────────────────────────────────