scripts/coevolution_loop.py

"""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 '<auto-mined>'}")
    _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"<dummy-state-r{round_idx}-i{i}>",
                "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())