#!/usr/bin/env python3 """Build the aggregate showcase plot + summary table comparing every Membrane GRPO run we have metrics for. Outputs: docs/plots/grpo_warmstart_ablation.png docs/plots/grpo_warmstart_ablation.svg docs/plots/grpo_warmstart_summary.csv docs/plots/grpo_warmstart_summary.md """ from __future__ import annotations import csv import json from pathlib import Path import matplotlib.pyplot as plt import numpy as np from huggingface_hub import hf_hub_download ROOT = Path(__file__).resolve().parents[2] PLOTS = ROOT / "docs" / "plots" PLOTS.mkdir(parents=True, exist_ok=True) def read_csv_curve(path: Path) -> tuple[list[int], list[float]]: steps, rewards = [], [] with path.open() as fp: for row in csv.DictReader(fp): steps.append(int(row["step"])) rewards.append(float(row["reward_mean"])) return steps, rewards def load_hero() -> tuple[list[int], list[float]]: p = hf_hub_download( "Tejasghatule/membrane-qwen25-1p5b-grpo-lora", "membrane_grpo_existing_checkpoint_1000/final_adapter/trainer_state.json", ) log = json.load(open(p))["log_history"] steps = [r["step"] for r in log if "reward" in r] rewards = [r["reward"] for r in log if "reward" in r] return steps, rewards COLD_DIR = ROOT / "docs" / "hf_runs" WARM_DIR = COLD_DIR / "continue_warm_start" cold_runs = [ ("cold-deep / seed 3408 (collapsed)", COLD_DIR / "repeat_seed_3408" / "training_metrics.csv"), ("cold-deep / seed 3409 (collapsed)", COLD_DIR / "repeat_seed_3409" / "training_metrics.csv"), ("cold-multi / seed 3410 (collapsed)", COLD_DIR / "multi_seed_3410" / "training_metrics.csv"), ] warm_runs = [ # (label_for_summary, csv_path, plot_label, color, lw) ("warm-deep / lr=2e-6 (conservative)", WARM_DIR / "continue_deep_seed_5821_lr2e-6" / "training_metrics.csv", "Warm start, single task, slow LR → best run", "#1f77b4", 2.4), ("warm-deep / lr=5e-6 (aggressive, saturates)", WARM_DIR / "continue_deep_seed_5823_lr5e-6" / "training_metrics.csv", "Warm start, single task, fast LR → saturates", "#9ec5e8", 1.9), ("warm-multi / lr=3e-6 (conservative, 7 tasks)", WARM_DIR / "continue_multifull_seed_5822_lr3e-6" / "training_metrics.csv", "Warm start, all 7 tasks, slow LR", "#d62728", 2.0), ("warm-multi / lr=5e-6 (aggressive, 7 tasks)", WARM_DIR / "continue_multifull_seed_5825_lr5e-6" / "training_metrics.csv", "Warm start, all 7 tasks, fast LR", "#f4a6a4", 1.7), ] def smooth(xs: list[float], w: int = 3) -> list[float]: if len(xs) < w: return xs pad = w // 2 out = [] for i in range(len(xs)): lo = max(0, i - pad) hi = min(len(xs), i + pad + 1) out.append(sum(xs[lo:hi]) / (hi - lo)) return out def build_plot() -> None: fig, ax = plt.subplots(1, 1, figsize=(12, 7)) # ---- Cold-start: collapse all 3 into a single shaded band, single legend entry. cold_curves = [] for _, csv_path in cold_runs: if csv_path.exists(): cold_curves.append(read_csv_curve(csv_path)) if cold_curves: all_steps = sorted({s for steps, _ in cold_curves for s in steps}) if all_steps: arr = np.full((len(cold_curves), len(all_steps)), np.nan) for i, (steps, rewards) in enumerate(cold_curves): for s, r in zip(steps, rewards): arr[i, all_steps.index(s)] = r band_lo = np.nanmin(arr, axis=0) band_hi = np.nanmax(arr, axis=0) band_mid = np.nanmean(arr, axis=0) ax.fill_between(all_steps, band_lo, band_hi, color="#cccccc", alpha=0.55, lw=0, label="3 cold-start runs → stuck below 0.02") ax.plot(all_steps, band_mid, color="#888888", lw=1.0, ls="--", alpha=0.85) # ---- Hero h_steps, h_rewards = load_hero() ax.plot(h_steps, smooth(h_rewards), color="black", lw=2.4, label="Colab hero → cold start that converged (lr=5e-6, 1000 steps)", zorder=10) ax.scatter([h_steps[-1]], [h_rewards[-1]], color="black", zorder=11) # ---- Warm runs best_warm_step = None best_warm_rew = -1.0 sat_steps = sat_rew = None for _, csv_path, plot_label, color, lw in warm_runs: if not csv_path.exists(): continue s, r = read_csv_curve(csv_path) rs = smooth(r) ax.plot(s, rs, color=color, lw=lw, label=plot_label) # Track best warm-deep / lr=2e-6 peak (the headline) if plot_label.startswith("Warm start, single task, slow LR"): best_idx = int(np.argmax(rs)) best_warm_step, best_warm_rew = s[best_idx], rs[best_idx] # Track saturating warm-deep aggressive descent point if plot_label.startswith("Warm start, single task, fast LR"): # Find the descent - peak then a later lower point peak_idx = int(np.argmax(rs)) if peak_idx < len(rs) - 5: sat_steps, sat_rew = s[-1], rs[-1] sat_peak_step, sat_peak_rew = s[peak_idx], rs[peak_idx] # ---- Hero vs warm reference lines ax.axhline(1.0, color="#bbbbbb", lw=0.7, ls=":") ax.text(2010, 1.0, " perfect = 1.0", fontsize=8.5, color="#888", va="center") ax.axhline(0.974, color="black", lw=0.6, ls=":", alpha=0.5) ax.text(2010, 0.974, " hero peak 0.974", fontsize=8.5, color="black", va="center", alpha=0.7) # ---- Annotations on the plot itself if best_warm_step is not None: ax.annotate( f"Warm start beats hero\n(peak {best_warm_rew:.3f})", xy=(best_warm_step, best_warm_rew), xytext=(700, 0.42), fontsize=10.5, fontweight="bold", color="#1f3a8a", ha="center", arrowprops=dict(arrowstyle="->", color="#1f3a8a", lw=1.2, connectionstyle="arc3,rad=-0.15"), ) if sat_steps is not None: ax.annotate( "Fast LR → policy saturates,\nadvantage signal disappears,\nreward drifts down", xy=(sat_steps, sat_rew), xytext=(1450, 0.30), fontsize=10, color="#3a3a3a", ha="center", arrowprops=dict(arrowstyle="->", color="#3a3a3a", lw=1.0, connectionstyle="arc3,rad=0.15"), ) # Annotate the cold-start band so the takeaway is on the chart, not just the legend ax.annotate( "Cold-start runs never escape zero -\n" "Membrane's reward is sparse, so a\n" "fresh policy almost never gets a\n" "non-zero signal to learn from.", xy=(900, 0.012), xytext=(60, 0.16), fontsize=9.5, color="#555", ha="left", arrowprops=dict(arrowstyle="->", color="#888", lw=0.9, connectionstyle="arc3,rad=-0.25"), ) # ---- Labels and legend ax.set_xlabel("Training step", fontsize=11) ax.set_ylabel("Membrane reward (0 = invalid / leak, 1 = perfect refusal + commit)", fontsize=11) ax.set_ylim(-0.05, 1.08) ax.set_xlim(0, 2200) ax.grid(True, alpha=0.25) leg = ax.legend(loc="lower right", fontsize=9.5, framealpha=0.97, title="Run", title_fontsize=10) leg.get_title().set_fontweight("bold") # Title + subtitle as separate axes-level objects so they stack cleanly. ax.set_title("How a 1.5 B Qwen learns Membrane", fontsize=15, fontweight="bold", pad=34) ax.text( 0.5, 1.015, "Cold start collapses near zero. The Colab hero converges. Warm starts loaded from the hero adapter\n" "climb above it on the training task, and lift onto the 7-task curriculum.", transform=ax.transAxes, ha="center", va="bottom", fontsize=10.5, color="#444", ) fig.tight_layout() fig.savefig(PLOTS / "grpo_warmstart_ablation.png", dpi=170) fig.savefig(PLOTS / "grpo_warmstart_ablation.svg") plt.close(fig) print("wrote", PLOTS / "grpo_warmstart_ablation.png") def write_summary() -> None: rows = [] h_steps, h_rewards = load_hero() rows.append({ "run": "Colab hero (cold start, lr=5e-6, 1000 steps)", "category": "hero", "first_reward": round(h_rewards[0], 4), "final_reward": round(h_rewards[-1], 4), "best_reward": round(max(h_rewards), 4), "best_step": h_steps[h_rewards.index(max(h_rewards))], "steps_logged": len(h_rewards), }) def _summarize(name: str, csv_path: Path, category: str) -> dict | None: if not csv_path.exists(): return None s, r = read_csv_curve(csv_path) return { "run": name, "category": category, "first_reward": round(r[0], 4), "final_reward": round(r[-1], 4), "best_reward": round(max(r), 4), "best_step": s[r.index(max(r))], "steps_logged": len(r), } for name, p in cold_runs: d = _summarize(name, p, "cold-start") if d: rows.append(d) for entry in warm_runs: name, p = entry[0], entry[1] d = _summarize(name, p, "warm-start") if d: rows.append(d) csv_path = PLOTS / "grpo_warmstart_summary.csv" with csv_path.open("w", newline="") as fp: w = csv.DictWriter(fp, fieldnames=list(rows[0].keys())) w.writeheader() w.writerows(rows) print("wrote", csv_path) md_lines = [ "# Membrane GRPO - warm-start ablation", "", "Aggregate of every Membrane GRPO run that has produced metrics:", "", "- 1 Colab run (the hero) that converged from a cold start.", "- 3 cold-start Hugging Face Job runs that did not converge.", "- 4 warm-start Hugging Face Job runs that loaded the hero adapter as initial", " weights and continued training. The four form a 2 × 2 ablation across", " learning rate (conservative vs aggressive) and task mix (single task vs", " full 7-scenario curriculum).", "", "| run | category | first reward | final reward | best reward | best step | steps logged |", "|---|---|---|---|---|---|---|", ] for r in rows: md_lines.append( f"| {r['run']} | {r['category']} | {r['first_reward']:.3f} | " f"{r['final_reward']:.3f} | **{r['best_reward']:.3f}** | " f"{r['best_step']} | {r['steps_logged']} |" ) md_lines += [ "", "## Findings", "", "1. **Cold-start GRPO does not learn Membrane in the budgets we tested.** Three", " independent HF Job runs (single-task seeds 3408 / 3409 at 900 steps,", " multi-task seed 3410 at 1400 steps) all stay below 0.02 mean reward.", " Membrane's reward is sparse on purpose - any malformed JSONL action zeroes", " the episode - so a freshly-initialised policy almost never produces a", " non-zero advantage signal long enough for GRPO to bootstrap.", "", "2. **The Colab T4 run converged with the same recipe that collapsed on the", " A10G HF Job runs.** Same script, same hyperparameters, same seed family.", " The only difference is the RNG stream from a different GPU and the", " Unsloth/TRL versions that the version pin now reproduces. This is why", " warm-starting was worth doing: a known-good policy whose weights could be", " redeployed as initial conditions on the HF compute backend.", "", "3. **Warm-starting from the hero adapter beats the hero on the same task.**", " `warm-deep / lr=2e-6` lifts mean reward from 0.880 (the hero adapter's", " starting score) to **0.971 final / 0.988 peak**, surpassing both the", " hero's 0.935 final and 0.974 peak. The conservative learning rate is the", " key - see finding 4.", "", "4. **Aggressive learning rate saturates GRPO on a single task.**", " `warm-deep / lr=5e-6` (the *hero's* learning rate) climbs from 0.849 to a", " peak of **0.988 by step 240**, then drifts back down to 0.959 by", " step 1500. This is not a model failure: `frac_reward_zero_std` rises", " from 0.2 to ≥ 0.7, meaning ≥ 70 % of GRPO prompt groups produce identical", " rewards across all 4 completions. With zero per-group advantage there is", " no gradient signal, and `grad_norm` falls to exactly 0.0. Conservative", " lr=2e-6 keeps per-group variance alive longer and continues improving.", "", "5. **A single-task warm-start transfers to the full 7-scenario curriculum.**", " `warm-multi / lr=3e-6` (must-refuse, must-comply, long, triad,", " round-robin, and two more held-out scenarios) starts at 0.495 - the model", " has never seen 6 of those 7 tasks during the original Colab training -", " and climbs to **0.793** by step 2000 without collapsing. The aggressive", " variant peaks higher (0.854 at step 800) but slides to 0.785 as it", " over-fits the task it was warm-started on.", "", "## Source data", "", "- Application source: ", "- Per-run metrics: `docs/hf_runs//training_metrics.csv` and", " `docs/hf_runs//run_summary.json`.", "- Aggregate plot: `grpo_warmstart_ablation.png` / `.svg`.", "- Headline CSV: `grpo_warmstart_summary.csv`.", "- All adapters:", " .", "", ] md_path = PLOTS / "grpo_warmstart_summary.md" md_path.write_text("\n".join(md_lines)) print("wrote", md_path) def main() -> None: build_plot() write_summary() if __name__ == "__main__": main()