"""4-arm memory ablation across E2/E3/E4 to prove memory CAUSALLY changes outcomes. Arms: A — keyword_security_triager (no memory queries — just heuristic) B — keyword_security_triager w/ memory store disabled (env returns no hits; OPSGUARD_MEMORY_DISABLED=1) C — memory_aware policy (current: periodic query_history) D — memory_aware + ContributorProfile features actively consulted in policy (already wired into observation; this arm uses them) Outputs (under --out, default `eval_outputs/memory_ablation/`): rollouts.jsonl per-rollout records summary.md arm × scenario reward (mean ± std) + attack_recall delta.png bar plot of arm reward delta vs Arm B (no-memory baseline) pmi.json PMI between memory hit in last 5 steps and a correct attack-detection action (counterfactual proxy) Usage: python scripts/memory_ablation.py python scripts/memory_ablation.py --scenarios E2_social_eng_buildup --seeds 0 1 """ from __future__ import annotations import argparse import json import math import os import sys import time from pathlib import Path # Make repo importable when run as a script. sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) from eval.harness import rollout from eval.policies import build_policy, make_keyword_security_triager from models import ActionType, OpsguardAction, OpsguardObservation, SecurityVerdict from server.opsguard_environment import OpsguardEnvironment DEFAULT_SCENARIOS = [ "E2_social_eng_buildup", "E3_compromised_maintainer", "E4_multi_vector", ] DEFAULT_SEEDS = [0, 1, 2] CATCH_ACTIONS = {"reject_pr", "close_spam", "flag_security"} PMI_LOOKBACK = 5 # ------------------------------------------------------------------------------ # Arm D: memory-aware policy that ACTUALLY consults ContributorProfile features. # ------------------------------------------------------------------------------ def make_memory_aware_profile_triager(): """Memory-aware triager + per-author behavioural change-point detection. Periodically issues `query_history` like the base memory_aware policy, and additionally inspects `current_issue.author_profile` to flag PRs whose behavioural delta or anomaly score sharply diverges from the contributor's established baseline (Jia-Tan-style change-point). """ base = make_keyword_security_triager() state = {"last_query_step": -10} def act(obs: OpsguardObservation) -> OpsguardAction: ci = obs.current_issue if ci is None: return OpsguardAction(action_type=ActionType.WAIT) # Periodic memory query, same cadence as memory_aware. if obs.step > 0 and obs.step - state["last_query_step"] >= 8: state["last_query_step"] = obs.step return OpsguardAction( action_type=ActionType.QUERY_HISTORY, query=ci.author_login, reasoning="check prior author activity", ) # Profile-driven change-point detection BEFORE falling through to the # keyword triager. A long-trusted author (high trust_score, several prior # PRs) suddenly producing a PR with a large behavioural delta or anomaly # spike is the supply-chain signal we want to catch. prof = ci.author_profile or {} n_seen = prof.get("n_prs_seen", 0) trust = prof.get("trust_score", 0.5) anomaly = prof.get("anomaly_score_ewma", 0.0) delta = prof.get("behavioral_delta", 0.0) if ci.is_pr and n_seen >= 3 and trust >= 0.7 and (anomaly >= 1.5 or delta >= 200.0): return OpsguardAction( action_type=ActionType.FLAG_SECURITY, target_issue_id=ci.issue_id, comment_body=( f"behavioural change-point on trusted author " f"(anomaly_ewma={anomaly:.2f}, delta={delta:.1f})" ), reasoning="contributor profile change-point", security_verdict=SecurityVerdict.SUSPICIOUS, ) return base(obs) return act # ------------------------------------------------------------------------------ # Arm specs. # ------------------------------------------------------------------------------ ARMS = { "A_no_memory_query": { "policy_factory": lambda: build_policy("keyword_security_triager"), "env_overrides": {}, "description": "keyword triager; never issues query_history", }, "B_memory_disabled": { "policy_factory": lambda: build_policy("keyword_security_triager"), "env_overrides": {"OPSGUARD_MEMORY_DISABLED": "1"}, "description": "keyword triager; env honors OPSGUARD_MEMORY_DISABLED → memory_hits always []", }, "C_memory_aware": { "policy_factory": lambda: build_policy("memory_aware"), "env_overrides": {}, "description": "memory_aware: periodic query_history", }, "D_memory_aware_profile": { "policy_factory": make_memory_aware_profile_triager, "env_overrides": {}, "description": "memory_aware + ContributorProfile change-point features", }, } # ------------------------------------------------------------------------------ # Custom rollout that mirrors eval.harness.rollout but ALSO records the per-step # trace we need for PMI: action, was_attack, was_caught, memory_hits_in_obs. # ------------------------------------------------------------------------------ def rollout_with_trace( env: OpsguardEnvironment, policy_name: str, policy_fn, scenario_id: str, seed: int, ): t0 = time.time() obs = env.reset(scenario_id=scenario_id, seed=seed) cum = 0.0 last_meta: dict = {} n_steps = 0 trace: list[dict] = [] budget = env._episode.scenario.step_budget + 5 # type: ignore[union-attr] while not obs.done and n_steps < budget: # Snapshot memory-hit count BEFORE acting (the obs the policy sees). mem_hits_now = len(obs.memory_hits or []) action = policy_fn(obs) target_id = action.target_issue_id if target_id is None and obs.current_issue is not None: target_id = obs.current_issue.issue_id obs = env.step(action) if obs.reward is not None: cum += obs.reward n_steps += 1 if obs.metadata: last_meta = obs.metadata # The most recent action lands at the end of recent_actions; pull # ground-truth `is_attack` from there (env populated it). recent = (obs.recent_actions or []) last_rec = recent[-1] if recent else {} action_str = last_rec.get("action", action.action_type.value if hasattr(action.action_type, "value") else str(action.action_type)) was_attack = bool(last_rec.get("is_attack", False)) was_catch_action = action_str in CATCH_ACTIONS # "Caught" = a catch action applied to a true attack. caught = was_catch_action and was_attack trace.append({ "step": n_steps, "action": action_str, "issue_id": last_rec.get("issue_id", target_id), "is_attack": was_attack, "catch_action": was_catch_action, "caught": caught, "memory_hits_in_obs": mem_hits_now, }) record = { "policy": policy_name, "scenario_id": scenario_id, "seed": seed, "cumulative_reward": round(cum, 4), "n_steps": n_steps, "n_resolved": last_meta.get("legit_resolved", 0), "n_total": last_meta.get("legit_total", 0), "n_spam_caught": last_meta.get("attacks_caught", 0), "n_spam_total": last_meta.get("attacks_total", 0), "elapsed_sec": round(time.time() - t0, 2), "final_breakdown": last_meta, "trace": trace, } return record # ------------------------------------------------------------------------------ # Aggregation, PMI, summary, plot. # ------------------------------------------------------------------------------ def _mean_std(xs: list[float]) -> tuple[float, float]: if not xs: return 0.0, 0.0 m = sum(xs) / len(xs) if len(xs) < 2: return m, 0.0 v = sum((x - m) ** 2 for x in xs) / (len(xs) - 1) return m, math.sqrt(v) def aggregate_records(records: list[dict]) -> dict: cells = {} for r in records: cells.setdefault((r["policy"], r["scenario_id"]), []).append(r) out = [] for (arm, scen), rs in sorted(cells.items()): rewards = [r["cumulative_reward"] for r in rs] recall = [ r["n_spam_caught"] / r["n_spam_total"] for r in rs if r["n_spam_total"] ] rm, rs_ = _mean_std(rewards) recm, _recs = _mean_std(recall) out.append({ "arm": arm, "scenario": scen, "n": len(rs), "reward_mean": round(rm, 3), "reward_std": round(rs_, 3), "attack_recall_mean": round(recm, 3) if recall else None, "spam_caught_mean": round(sum(r["n_spam_caught"] for r in rs) / len(rs), 2), "spam_total_mean": round(sum(r["n_spam_total"] for r in rs) / len(rs), 2), }) return {"cells": out} def compute_pmi(records: list[dict]) -> dict: """Per-arm PMI between "memory hit retrieved in prior 5 steps" and "this step is a correct attack-detection (catch action on a real attack)". For each catch-action step (action ∈ {reject_pr, close_spam, flag_security}), we ask: - Was there ANY memory hit observed in the prior PMI_LOOKBACK steps? - Did this catch a true attack? PMI_proxy = P(catch | retrieved_recent) / P(catch) Reported per arm. Values > 1 ⇒ memory retrieval covaries with successful attack detection above chance. """ by_arm: dict[str, list[dict]] = {} for r in records: by_arm.setdefault(r["policy"], []).append(r) out: dict[str, dict] = {} for arm, recs in by_arm.items(): # Pool catch-action events across all rollouts for this arm. catch_events: list[dict] = [] # one entry per catch action for r in recs: trace = r.get("trace", []) for i, step in enumerate(trace): if not step["catch_action"]: continue window = trace[max(0, i - PMI_LOOKBACK):i] retrieved_recent = any(s["memory_hits_in_obs"] > 0 for s in window) catch_events.append({ "caught": bool(step["caught"]), "retrieved_recent": retrieved_recent, }) n = len(catch_events) n_caught = sum(1 for e in catch_events if e["caught"]) n_retrieved = sum(1 for e in catch_events if e["retrieved_recent"]) n_caught_and_retrieved = sum( 1 for e in catch_events if e["caught"] and e["retrieved_recent"] ) p_catch = n_caught / n if n else 0.0 p_catch_given_retrieved = ( n_caught_and_retrieved / n_retrieved if n_retrieved else 0.0 ) # PMI proxy as a ratio (paper-style "lift"). if p_catch > 0 and n_retrieved > 0: pmi_ratio = p_catch_given_retrieved / p_catch pmi_log2 = math.log2(pmi_ratio) if pmi_ratio > 0 else float("-inf") else: pmi_ratio = None pmi_log2 = None out[arm] = { "n_catch_events": n, "n_caught": n_caught, "n_with_recent_memory_hit": n_retrieved, "n_caught_and_retrieved": n_caught_and_retrieved, "p_catch": round(p_catch, 4), "p_catch_given_retrieved": round(p_catch_given_retrieved, 4), "pmi_ratio_p_catch_given_retrieved_over_p_catch": ( round(pmi_ratio, 4) if pmi_ratio is not None else None ), "pmi_log2": round(pmi_log2, 4) if pmi_log2 is not None else None, "lookback_steps": PMI_LOOKBACK, } return out def write_summary_md(agg: dict, pmi: dict, out_path: Path) -> None: cells = agg["cells"] arms = sorted({c["arm"] for c in cells}) scenarios = sorted({c["scenario"] for c in cells}) lines = ["# OpsGuard memory ablation\n"] lines.append("4-arm ablation showing the causal contribution of memory.\n") lines.append("## Arms\n") for k, v in ARMS.items(): lines.append(f"- **{k}** — {v['description']}") lines.append("") for s in scenarios: lines.append(f"\n## {s}\n") lines.append("| Arm | reward (mean ± std) | attack_recall | spam_caught / total |") lines.append("|---|---:|---:|---:|") for a in arms: cell = next((c for c in cells if c["arm"] == a and c["scenario"] == s), None) if not cell: continue ar = ( f"{cell['attack_recall_mean']:.2f}" if cell["attack_recall_mean"] is not None else "n/a" ) lines.append( f"| {a} | {cell['reward_mean']:+.2f} ± {cell['reward_std']:.2f} | " f"{ar} | {cell['spam_caught_mean']:.1f} / {cell['spam_total_mean']:.1f} |" ) # Δ-vs-B table — the headline number. lines.append("\n## Reward delta vs Arm B (memory disabled), per scenario\n") lines.append("| Scenario | A − B | C − B | D − B |") lines.append("|---|---:|---:|---:|") for s in scenarios: b_cell = next((c for c in cells if c["arm"] == "B_memory_disabled" and c["scenario"] == s), None) b_r = b_cell["reward_mean"] if b_cell else 0.0 deltas = [] for a in ("A_no_memory_query", "C_memory_aware", "D_memory_aware_profile"): cell = next((c for c in cells if c["arm"] == a and c["scenario"] == s), None) if cell is None: deltas.append("n/a") else: d = cell["reward_mean"] - b_r deltas.append(f"{d:+.2f}") lines.append(f"| {s} | {deltas[0]} | {deltas[1]} | {deltas[2]} |") lines.append("\n## PMI: P(catch | recent memory hit) / P(catch), per arm\n") lines.append(f"Lookback = {PMI_LOOKBACK} steps. Values > 1.0 ⇒ memory retrieval lifts catch probability.\n") lines.append("| Arm | n_catch | n_caught | n_with_recent_hit | P(catch) | P(catch|hit) | PMI (ratio) | PMI (log2) |") lines.append("|---|---:|---:|---:|---:|---:|---:|---:|") for a in arms: p = pmi.get(a, {}) ratio = p.get("pmi_ratio_p_catch_given_retrieved_over_p_catch") log2 = p.get("pmi_log2") ratio_s = f"{ratio:.3f}" if ratio is not None else "n/a" log2_s = f"{log2:+.3f}" if log2 is not None else "n/a" lines.append( f"| {a} | {p.get('n_catch_events', 0)} | {p.get('n_caught', 0)} | " f"{p.get('n_with_recent_memory_hit', 0)} | " f"{p.get('p_catch', 0):.3f} | {p.get('p_catch_given_retrieved', 0):.3f} | " f"{ratio_s} | {log2_s} |" ) out_path.write_text("\n".join(lines) + "\n", encoding="utf-8") def plot_delta_vs_b(agg: dict, out_path: Path) -> None: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt cells = agg["cells"] scenarios = sorted({c["scenario"] for c in cells}) arms_to_show = ["A_no_memory_query", "C_memory_aware", "D_memory_aware_profile"] deltas: dict[str, list[float]] = {a: [] for a in arms_to_show} for s in scenarios: b_cell = next((c for c in cells if c["arm"] == "B_memory_disabled" and c["scenario"] == s), None) b_r = b_cell["reward_mean"] if b_cell else 0.0 for a in arms_to_show: cell = next((c for c in cells if c["arm"] == a and c["scenario"] == s), None) deltas[a].append((cell["reward_mean"] - b_r) if cell else 0.0) fig, ax = plt.subplots(figsize=(10, 5)) width = 0.8 / len(arms_to_show) for i, a in enumerate(arms_to_show): x = [j + i * width - 0.4 + width / 2 for j in range(len(scenarios))] ax.bar(x, deltas[a], width=width * 0.95, label=a) ax.set_xticks(range(len(scenarios))) ax.set_xticklabels(scenarios, rotation=15, ha="right") ax.set_ylabel("Reward delta vs Arm B (memory disabled)") ax.set_title("OpsGuard memory ablation: reward lift over no-memory baseline") ax.axhline(0, color="gray", linestyle="--", linewidth=0.8) ax.legend(loc="best", fontsize=9) ax.grid(axis="y", alpha=0.3) plt.tight_layout() plt.savefig(out_path, dpi=130) plt.close(fig) # ------------------------------------------------------------------------------ # Main. # ------------------------------------------------------------------------------ def main(): ap = argparse.ArgumentParser() ap.add_argument("--scenarios", nargs="*", default=DEFAULT_SCENARIOS) ap.add_argument("--seeds", nargs="*", type=int, default=DEFAULT_SEEDS) ap.add_argument("--out", default="eval_outputs/memory_ablation") args = ap.parse_args() out_dir = Path(args.out) out_dir.mkdir(parents=True, exist_ok=True) rollouts_path = out_dir / "rollouts.jsonl" records: list[dict] = [] saved_env: dict[str, str | None] = {} print( f"[memory_ablation] arms={list(ARMS)} scenarios={args.scenarios} " f"seeds={args.seeds} -> {out_dir}" ) t_total = time.time() with open(rollouts_path, "w", encoding="utf-8") as fout: for arm_name, spec in ARMS.items(): # Apply env overrides for this arm; restore on exit. for k, v in spec["env_overrides"].items(): saved_env[k] = os.environ.get(k) os.environ[k] = v try: for sid in args.scenarios: for seed in args.seeds: env = OpsguardEnvironment() policy_fn = spec["policy_factory"]() rec = rollout_with_trace(env, arm_name, policy_fn, sid, seed) records.append(rec) fout.write(json.dumps(rec) + "\n") fout.flush() recall = ( rec["n_spam_caught"] / rec["n_spam_total"] if rec["n_spam_total"] else float("nan") ) print( f" {arm_name:>26} | {sid:<26} seed={seed} | " f"reward={rec['cumulative_reward']:>+8.2f} " f"recall={recall:.2f} " f"steps={rec['n_steps']:>3} " f"({rec['elapsed_sec']:.1f}s)", flush=True, ) finally: # Restore env vars. for k in spec["env_overrides"]: prev = saved_env.get(k) if prev is None: os.environ.pop(k, None) else: os.environ[k] = prev agg = aggregate_records(records) pmi = compute_pmi(records) (out_dir / "summary.json").write_text(json.dumps(agg, indent=2), encoding="utf-8") (out_dir / "pmi.json").write_text(json.dumps(pmi, indent=2), encoding="utf-8") write_summary_md(agg, pmi, out_dir / "summary.md") plot_delta_vs_b(agg, out_dir / "delta.png") print( f"\n[done] {len(records)} rollouts in {time.time() - t_total:.1f}s -> {out_dir}" ) print(f" rollouts: {rollouts_path}") print(f" summary: {out_dir / 'summary.md'}") print(f" pmi: {out_dir / 'pmi.json'}") print(f" plot: {out_dir / 'delta.png'}") if __name__ == "__main__": main()