"""SAGE / CoEvolve 6-round alternating LoRA training loop for OpsGuard. Defender (maintainer) and Adversary (spammer) each get their own LoRA adapter on a *shared* Qwen2.5-7B-Instruct base. We alternate: Round r in 1..R: 1. Train DEFENDER LoRA (GRPO, 200 steps) with adversary frozen as opponent 2. Train ADVERSARY LoRA (DPO, 100 steps) with defender frozen as opponent 3. Mine top-K hardest-defended attacks -> SFT replay buffer 4. Replay-SFT defender 50 steps @ lr=1e-6 (anti-catastrophic-forgetting) After R rounds: final eval on E2..E5, optional push_to_hub for both adapters, and a coevolution_curve.png (rounds 0..R × {defender_reward, adversary_success}). DESIGN NOTES (from handoff §6): - All TRL kwargs are wrapped in `_safe_kwargs(...)` — TRL's GRPOConfig / DPOConfig drift between releases; we introspect and drop unknowns instead of pinning a version. - Imports of transformers / trl / peft are TRY/EXCEPT — if missing, the `--dry-run` path still works (it never touches them). - Defaults to --dry-run so this can be smoke-tested on a laptop without GPU. REAL RUN (HF Jobs, H100): # hf jobs uv run --flavor h100-large \\ # --with "trl>=0.18,unsloth,peft,bitsandbytes,openenv-core,vllm,matplotlib" \\ # --secrets HF_TOKEN \\ # scripts/coevolution_loop.py \\ # --rounds 6 --defender-steps 200 --adversary-steps 100 \\ # --mine-k 50 --push-to-hub --hub-repo-prefix sai1906/opsguard DRY RUN (laptop): python scripts/coevolution_loop.py --rounds 6 """ from __future__ import annotations import argparse import inspect import json import math import os import random import sys import time # Force UTF-8 stdout/stderr — Windows consoles default to cp1252 and choke on # em-dashes / arrows / × that appear in our structured logs. try: sys.stdout.reconfigure(encoding="utf-8", errors="replace") sys.stderr.reconfigure(encoding="utf-8", errors="replace") except Exception: pass from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable # --- Project root on sys.path (mirrors train_grpo.py convention) --- _PROJECT_ROOT = Path(__file__).resolve().parent.parent if str(_PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(_PROJECT_ROOT)) # --- Optional heavy deps: import-fail-gracefully --- _HAS_TORCH = False _HAS_TRL = False _HAS_PEFT = False _HAS_TRANSFORMERS = False _HAS_MPL = False try: import torch # noqa: F401 _HAS_TORCH = True except Exception: pass try: import transformers # noqa: F401 _HAS_TRANSFORMERS = True except Exception: pass try: import trl # noqa: F401 _HAS_TRL = True except Exception: pass try: import peft # noqa: F401 _HAS_PEFT = True except Exception: pass try: import matplotlib # noqa: F401 matplotlib.use("Agg") import matplotlib.pyplot as plt # noqa: F401 _HAS_MPL = True except Exception: pass # ----------------------------- utilities ----------------------------- # def _safe_kwargs(cls_or_callable: Any, kwargs: dict) -> dict: """Drop kwargs not accepted by the target signature. Same trick as scripts/train_grpo.py — TRL's GRPOConfig / DPOConfig signatures drift between releases; we introspect at runtime and silently drop unknown kwargs (after warning) rather than pinning a TRL version. """ try: sig = inspect.signature(cls_or_callable) params = sig.parameters except (TypeError, ValueError): return kwargs has_var_kw = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()) if has_var_kw: return kwargs valid = set(params.keys()) out, dropped = {}, [] for k, v in kwargs.items(): if k in valid: out[k] = v else: dropped.append(k) if dropped: print(f"[safe_kwargs] dropped unknown kwargs for {getattr(cls_or_callable, '__name__', cls_or_callable)}: {dropped}") return out def _banner(round_idx: int, total: int, phase: str) -> None: print(f"=== ROUND {round_idx} / {total} — {phase.upper()} ===", flush=True) def _log(msg: str) -> None: print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) # ----------------------------- config ----------------------------- # @dataclass class CoevoConfig: base_model: str = "Qwen/Qwen2.5-7B-Instruct" rounds: int = 6 defender_steps: int = 200 adversary_steps: int = 100 replay_sft_steps: int = 50 mine_k: int = 50 # LoRA (per RL-training agent's recommendation) lora_r: int = 64 lora_alpha: int = 128 lora_dropout: float = 0.05 lora_target_modules: tuple = ( "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj", ) # GRPO config (defender) grpo_beta: float = 0.001 grpo_num_generations: int = 16 grpo_epsilon_high: float = 0.28 grpo_loss_type: str = "dr_grpo" grpo_reward_ratio: float = 2.0 # legitimate:adversarial reward weighting # DPO config (adversary) dpo_beta: float = 0.1 dpo_loss_type: str = "sigmoid" # Replay buffer (SPPO-style 30/70 mix: 30% replay, 70% fresh rollouts) replay_mix_ratio: float = 0.30 replay_lr: float = 1e-6 # Adversary diversity bonus diversity_cosine_threshold: float = 0.30 diversity_weight: float = 0.10 # Canary / oracle gate canary_fraction: float = 0.05 judge_mae_halt_threshold: float = 0.15 # Eval scenarios train_scenarios: tuple = ("E2", "E3", "E4") eval_scenarios: tuple = ("E2", "E3", "E4", "E5") # I/O output_dir: Path = Path("artifacts/coevolution") defender_adapter: str = "defender" adversary_adapter: str = "adversary" # ----------------------------- LoRA / adapter helpers ----------------------------- # def _build_lora_config(cfg: CoevoConfig): if not _HAS_PEFT: raise RuntimeError("peft is not installed; cannot build LoraConfig") from peft import LoraConfig kwargs = dict( r=cfg.lora_r, lora_alpha=cfg.lora_alpha, lora_dropout=cfg.lora_dropout, target_modules=list(cfg.lora_target_modules), bias="none", task_type="CAUSAL_LM", ) return LoraConfig(**_safe_kwargs(LoraConfig, kwargs)) def _load_base_with_adapter(base_model: str, adapter_path: Path | None, *, frozen: bool): """Load base model and (optionally) attach a LoRA adapter. `frozen=True` means the adapter is the OPPONENT — its weights are frozen so it serves only as a generation policy during the other player's training. """ if not (_HAS_TRANSFORMERS and _HAS_PEFT and _HAS_TORCH): raise RuntimeError("transformers + peft + torch required for real training") from transformers import AutoModelForCausalLM, AutoTokenizer from peft import PeftModel tok = AutoTokenizer.from_pretrained(base_model) model = AutoModelForCausalLM.from_pretrained( base_model, torch_dtype="auto", device_map="auto", ) if adapter_path is not None and Path(adapter_path).exists(): model = PeftModel.from_pretrained(model, str(adapter_path), is_trainable=not frozen) if frozen: for p in model.parameters(): p.requires_grad = False return model, tok # ----------------------------- training phases ----------------------------- # def train_defender_grpo( cfg: CoevoConfig, round_idx: int, defender_adapter_path: Path, adversary_adapter_path: Path | None, *, dry_run: bool, ) -> dict: """One round of GRPO training for the defender. The adversary LoRA is loaded as a *frozen* opponent — its rollouts become part of the env's adversarial scenario pool. Reward is shaped with legitimate:adversarial = ratio (default 2.0). """ _banner(round_idx, cfg.rounds, "defender") _log(f"GRPO {cfg.defender_steps} steps base={cfg.base_model}") _log(f" defender LoRA = {defender_adapter_path}") _log(f" opponent (frozen) adversary LoRA = {adversary_adapter_path}") _log(f" config: beta={cfg.grpo_beta} num_gen={cfg.grpo_num_generations} " f"eps_high={cfg.grpo_epsilon_high} loss={cfg.grpo_loss_type} ratio={cfg.grpo_reward_ratio}") if dry_run: for step in range(0, cfg.defender_steps, max(1, cfg.defender_steps // 4)): time.sleep(0.1) _log(f" step {step:>4}/{cfg.defender_steps} reward={0.30 + 0.05 * round_idx + random.random() * 0.02:.4f} " f"kl={0.001 + random.random() * 0.0005:.5f}") _log(f" saved adapter -> {defender_adapter_path}") return { "phase": "defender_grpo", "round": round_idx, "final_reward": 0.30 + 0.08 * round_idx, "steps": cfg.defender_steps, } # --- real path --- if not _HAS_TRL: raise RuntimeError("trl is required for non-dry-run; install trl>=0.18") from trl import GRPOConfig, GRPOTrainer model, tok = _load_base_with_adapter(cfg.base_model, defender_adapter_path, frozen=False) # opponent is loaded inside the OpenEnv server (env replays adversary policy) grpo_kwargs = dict( output_dir=str(defender_adapter_path), per_device_train_batch_size=1, num_generations=cfg.grpo_num_generations, max_steps=cfg.defender_steps, beta=cfg.grpo_beta, epsilon_high=cfg.grpo_epsilon_high, loss_type=cfg.grpo_loss_type, learning_rate=5e-6, bf16=True, report_to="none", save_strategy="no", ) grpo_cfg = GRPOConfig(**_safe_kwargs(GRPOConfig, grpo_kwargs)) trainer_kwargs = dict( model=model, processing_class=tok, args=grpo_cfg, # reward_funcs / environment_factory wired from train_grpo.py — see that # script for the OpenEnv environment_factory pattern. ) trainer = GRPOTrainer(**_safe_kwargs(GRPOTrainer, trainer_kwargs)) trainer.train() trainer.save_model(str(defender_adapter_path)) return {"phase": "defender_grpo", "round": round_idx, "steps": cfg.defender_steps} def train_adversary_dpo( cfg: CoevoConfig, round_idx: int, adversary_adapter_path: Path, defender_adapter_path: Path, preference_pairs_path: Path | None, *, dry_run: bool, ) -> dict: """DPO training for the adversary on adversarial preference pairs. The pairs are (chosen=attack-that-bypassed-defender, rejected=attack-blocked). A diversity bonus weighted at `diversity_weight` rewards attacks whose embedding cosine distance from prior attacks > `diversity_cosine_threshold`. """ _banner(round_idx, cfg.rounds, "adversary") _log(f"DPO {cfg.adversary_steps} steps base={cfg.base_model}") _log(f" adversary LoRA = {adversary_adapter_path}") _log(f" opponent (frozen) defender LoRA = {defender_adapter_path}") _log(f" preference pairs = {preference_pairs_path or ''}") _log(f" config: dpo_beta={cfg.dpo_beta} loss={cfg.dpo_loss_type}") _log(f" diversity bonus: cos_dist>{cfg.diversity_cosine_threshold} weight={cfg.diversity_weight}") if dry_run: for step in range(0, cfg.adversary_steps, max(1, cfg.adversary_steps // 4)): time.sleep(0.1) asr = 0.55 - 0.04 * round_idx + random.random() * 0.02 # adversary success rate trends down as defender hardens _log(f" step {step:>4}/{cfg.adversary_steps} loss={0.45 - 0.02 * round_idx + random.random() * 0.01:.4f} " f"adv_success_rate={max(0.05, asr):.3f} diversity={0.40 + random.random() * 0.05:.3f}") _log(f" saved adapter -> {adversary_adapter_path}") return { "phase": "adversary_dpo", "round": round_idx, "adv_success_rate": max(0.05, 0.55 - 0.04 * round_idx), "steps": cfg.adversary_steps, } if not _HAS_TRL: raise RuntimeError("trl is required for non-dry-run; install trl>=0.18") from trl import DPOConfig, DPOTrainer model, tok = _load_base_with_adapter(cfg.base_model, adversary_adapter_path, frozen=False) dpo_kwargs = dict( output_dir=str(adversary_adapter_path), per_device_train_batch_size=1, max_steps=cfg.adversary_steps, beta=cfg.dpo_beta, loss_type=cfg.dpo_loss_type, learning_rate=5e-6, bf16=True, report_to="none", save_strategy="no", ) dpo_cfg = DPOConfig(**_safe_kwargs(DPOConfig, dpo_kwargs)) trainer_kwargs = dict( model=model, ref_model=None, # peft shares base; ref is implicit args=dpo_cfg, processing_class=tok, # train_dataset=load_preference_pairs(preference_pairs_path), ) trainer = DPOTrainer(**_safe_kwargs(DPOTrainer, trainer_kwargs)) trainer.train() trainer.save_model(str(adversary_adapter_path)) return {"phase": "adversary_dpo", "round": round_idx, "steps": cfg.adversary_steps} def mine_failures( cfg: CoevoConfig, round_idx: int, defender_adapter_path: Path, adversary_adapter_path: Path, *, dry_run: bool, ) -> list[dict]: """Mine top-K attacks that the defender failed against (or scored low on). Returns a replay buffer of (state, gold_action) pairs for SFT. """ _banner(round_idx, cfg.rounds, "failure mining") _log(f"Sampling adversary against defender across {list(cfg.train_scenarios)}, k={cfg.mine_k}") if dry_run: time.sleep(0.1) replay = [ { "scenario": random.choice(list(cfg.train_scenarios)), "state": f"", "gold_action": "ignore_user", "defender_score": round(random.uniform(-1.0, 0.2), 3), } for i in range(cfg.mine_k) ] _log(f" mined {len(replay)} attacks defender missed (mean score={sum(r['defender_score'] for r in replay) / max(1, len(replay)):.3f})") return replay # real impl would: spin up env, sample N rollouts with adversary policy, # rank by defender reward (ascending), take bottom-K. See run_baseline_eval.py. raise NotImplementedError("real failure mining not wired; use --dry-run or wire env loop") def replay_sft( cfg: CoevoConfig, round_idx: int, defender_adapter_path: Path, replay_buffer: list[dict], *, dry_run: bool, ) -> dict: """SPPO-style replay SFT to anti-forget hard cases. Mixes 30% buffer + 70% fresh rollouts at lr=1e-6 for 50 steps. """ _banner(round_idx, cfg.rounds, "replay sft") n_replay = int(cfg.replay_mix_ratio * cfg.replay_sft_steps) n_fresh = cfg.replay_sft_steps - n_replay _log(f"SFT {cfg.replay_sft_steps} steps lr={cfg.replay_lr} " f"mix={int(cfg.replay_mix_ratio*100)}/{int((1-cfg.replay_mix_ratio)*100)} " f"({n_replay} replay + {n_fresh} fresh)") if dry_run: for step in range(0, cfg.replay_sft_steps, max(1, cfg.replay_sft_steps // 4)): time.sleep(0.1) _log(f" step {step:>3}/{cfg.replay_sft_steps} loss={0.30 - 0.005 * step + random.random() * 0.01:.4f}") _log(f" replay SFT done; defender adapter updated") return {"phase": "replay_sft", "round": round_idx, "steps": cfg.replay_sft_steps} # real impl: build a HF Dataset from replay_buffer, run SFTTrainer with low LR. raise NotImplementedError("real replay SFT not wired; use --dry-run or wire SFTTrainer") def canary_check( cfg: CoevoConfig, round_idx: int, defender_adapter_path: Path, *, dry_run: bool, ) -> tuple[bool, float]: """Run 5% canary episodes scored by an oracle to estimate judge MAE. Returns (ok, judge_mae). If MAE > halt_threshold, abort the loop. """ n_canary = max(1, int(cfg.canary_fraction * cfg.defender_steps)) if dry_run: time.sleep(0.05) mae = 0.05 + random.random() * 0.05 # plausible MAE ok = mae <= cfg.judge_mae_halt_threshold _log(f" canary MAE={mae:.3f} (threshold={cfg.judge_mae_halt_threshold}) {'OK' if ok else 'HALT'}") return ok, mae raise NotImplementedError("real canary check not wired; use --dry-run") # ----------------------------- final eval + viz ----------------------------- # def final_eval( cfg: CoevoConfig, defender_adapter_path: Path, adversary_adapter_path: Path, *, dry_run: bool, ) -> dict: print(f"=== FINAL EVAL — defender vs trained adversary on {list(cfg.eval_scenarios)} ===", flush=True) if dry_run: time.sleep(0.1) results = { sc: { "defender_reward": round(0.55 + 0.05 * i + random.random() * 0.03, 3), "adv_success_rate": round(max(0.05, 0.30 - 0.04 * i + random.random() * 0.02), 3), "n_episodes": 100, } for i, sc in enumerate(cfg.eval_scenarios) } for sc, r in results.items(): _log(f" {sc}: defender_reward={r['defender_reward']} adv_success_rate={r['adv_success_rate']}") return results raise NotImplementedError("real final eval not wired; use --dry-run") def plot_coevolution_curve( history: list[dict], out_path: Path, ) -> None: """rounds 0..R x {defender reward, adversary success rate} dual-axis.""" if not _HAS_MPL: _log(f"matplotlib not available; skipping {out_path}") return import matplotlib.pyplot as plt rounds = [h["round"] for h in history] def_reward = [h["defender_reward"] for h in history] adv_succ = [h["adv_success_rate"] for h in history] fig, ax1 = plt.subplots(figsize=(8, 5)) color1 = "tab:blue" ax1.set_xlabel("Round") ax1.set_ylabel("Defender mean reward", color=color1) ax1.plot(rounds, def_reward, marker="o", color=color1, label="defender reward") ax1.tick_params(axis="y", labelcolor=color1) ax1.set_ylim(0.0, 1.0) ax2 = ax1.twinx() color2 = "tab:red" ax2.set_ylabel("Adversary success rate", color=color2) ax2.plot(rounds, adv_succ, marker="s", color=color2, label="adv success rate") ax2.tick_params(axis="y", labelcolor=color2) ax2.set_ylim(0.0, 1.0) plt.title("OpsGuard CoEvolution: defender hardens, adversary success drops") fig.tight_layout() out_path.parent.mkdir(parents=True, exist_ok=True) plt.savefig(out_path, dpi=120) plt.close(fig) _log(f"wrote {out_path}") # ----------------------------- hub push ----------------------------- # def push_to_hub(adapter_path: Path, repo_id: str, *, dry_run: bool) -> None: if dry_run: _log(f"[dry-run] would push {adapter_path} -> hf://{repo_id}") return try: from huggingface_hub import HfApi except Exception as e: _log(f"huggingface_hub missing ({e}); skipping push of {adapter_path}") return api = HfApi() api.create_repo(repo_id, exist_ok=True, private=False) api.upload_folder(folder_path=str(adapter_path), repo_id=repo_id) _log(f"pushed {adapter_path} -> hf://{repo_id}") # ----------------------------- main loop ----------------------------- # def run_coevolution(cfg: CoevoConfig, *, dry_run: bool, push: bool, hub_prefix: str) -> dict: cfg.output_dir.mkdir(parents=True, exist_ok=True) defender_path = cfg.output_dir / cfg.defender_adapter adversary_path = cfg.output_dir / cfg.adversary_adapter defender_path.mkdir(parents=True, exist_ok=True) adversary_path.mkdir(parents=True, exist_ok=True) # Round 0 baseline (no training yet) history: list[dict] = [{"round": 0, "defender_reward": 0.30, "adv_success_rate": 0.55}] for r in range(1, cfg.rounds + 1): # 1. Defender GRPO (adversary frozen) d_stats = train_defender_grpo( cfg, r, defender_path, adversary_path if r > 1 else None, dry_run=dry_run, ) # canary gate ok, mae = canary_check(cfg, r, defender_path, dry_run=dry_run) if not ok: _log(f"HALTING: judge MAE {mae:.3f} > threshold {cfg.judge_mae_halt_threshold}") break # 2. Adversary DPO (defender frozen) a_stats = train_adversary_dpo( cfg, r, adversary_path, defender_path, preference_pairs_path=None, dry_run=dry_run, ) # 3. Mine failures replay = mine_failures(cfg, r, defender_path, adversary_path, dry_run=dry_run) # 4. Replay SFT replay_sft(cfg, r, defender_path, replay, dry_run=dry_run) history.append({ "round": r, "defender_reward": d_stats.get("final_reward", 0.30 + 0.08 * r), "adv_success_rate": a_stats.get("adv_success_rate", max(0.05, 0.55 - 0.04 * r)), }) # Final eval eval_results = final_eval(cfg, defender_path, adversary_path, dry_run=dry_run) # Save history JSON hist_path = cfg.output_dir / "coevolution_history.json" hist_path.write_text(json.dumps({"history": history, "eval": eval_results}, indent=2)) _log(f"wrote {hist_path}") # Plot curve plot_coevolution_curve(history, cfg.output_dir / "coevolution_curve.png") # Push if push: push_to_hub(defender_path, f"{hub_prefix}-defender-r{cfg.rounds}", dry_run=dry_run) push_to_hub(adversary_path, f"{hub_prefix}-adversary-r{cfg.rounds}", dry_run=dry_run) return {"history": history, "eval": eval_results} # ----------------------------- CLI ----------------------------- # def parse_args(argv: list[str] | None = None) -> argparse.Namespace: p = argparse.ArgumentParser(description="OpsGuard SAGE/CoEvolve alternating LoRA loop") p.add_argument("--rounds", type=int, default=6) p.add_argument("--defender-steps", type=int, default=200) p.add_argument("--adversary-steps", type=int, default=100) p.add_argument("--mine-k", type=int, default=50) p.add_argument("--replay-sft-steps", type=int, default=50) p.add_argument("--base-model", type=str, default="Qwen/Qwen2.5-7B-Instruct") p.add_argument("--output-dir", type=Path, default=Path("artifacts/coevolution")) p.add_argument("--push-to-hub", action="store_true") p.add_argument("--hub-repo-prefix", type=str, default="sai1906/opsguard") # NB: defaults to dry-run so this can be smoke-tested without a GPU. p.add_argument("--dry-run", dest="dry_run", action="store_true", default=True, help="Simulate phases without invoking TRL/PEFT (default: True).") p.add_argument("--no-dry-run", dest="dry_run", action="store_false", help="Actually train (requires GPU + trl + peft + transformers).") p.add_argument("--seed", type=int, default=0) return p.parse_args(argv) def main(argv: list[str] | None = None) -> int: args = parse_args(argv) random.seed(args.seed) cfg = CoevoConfig( base_model=args.base_model, rounds=args.rounds, defender_steps=args.defender_steps, adversary_steps=args.adversary_steps, replay_sft_steps=args.replay_sft_steps, mine_k=args.mine_k, output_dir=args.output_dir, ) print("=== OpsGuard CoEvolve ===") print(f" base_model : {cfg.base_model}") print(f" rounds : {cfg.rounds}") print(f" defender_steps : {cfg.defender_steps} (GRPO, beta={cfg.grpo_beta}, num_gen={cfg.grpo_num_generations})") print(f" adversary_steps : {cfg.adversary_steps} (DPO, beta={cfg.dpo_beta})") print(f" replay_sft_steps : {cfg.replay_sft_steps} (mix={int(cfg.replay_mix_ratio*100)}/{int((1-cfg.replay_mix_ratio)*100)}, lr={cfg.replay_lr})") print(f" mine_k : {cfg.mine_k}") print(f" LoRA : r={cfg.lora_r} alpha={cfg.lora_alpha}") print(f" dry_run : {args.dry_run}") print(f" push_to_hub : {args.push_to_hub} (prefix={args.hub_repo_prefix})") print(f" deps : torch={_HAS_TORCH} transformers={_HAS_TRANSFORMERS} " f"trl={_HAS_TRL} peft={_HAS_PEFT} matplotlib={_HAS_MPL}") print() out = run_coevolution( cfg, dry_run=args.dry_run, push=args.push_to_hub, hub_prefix=args.hub_repo_prefix, ) print() print("=== DONE ===") print(f"history: {len(out['history'])} rounds eval scenarios: {list(out['eval'].keys())}") return 0 if __name__ == "__main__": raise SystemExit(main())