feat: 5 training plots (reward±std, KL/loss, completion stats) + Drive backup in Cell 8

notebooks/colab_runner.ipynb

@@ -315,7 +315,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# EVALUATE + PLOTS \u2014 Phase 6\nimport os, glob, shutil\n\nOUTPUT_DIR = '/content/drive/MyDrive/injectarena/run_v1'\nCHECKPOINT  = f'{OUTPUT_DIR}/final'\n\n%cd /content/injectarena\n\n# Pull latest code (gets make_plots.py etc.)\n!git pull origin main\n\n# 1. Evaluate trained checkpoint against eval split.\n!python train/eval.py \\\n    --checkpoint {CHECKPOINT} \\\n    --output-json docs/eval_results.json\n\n# 2. Copy training logs from Drive \u2192 repo.\nos.makedirs('logs', exist_ok=True)\nfor f in glob.glob(f'{OUTPUT_DIR}/logs/*.jsonl'):\n    shutil.copy(f, 'logs/')\nprint(f\"Copied {len(glob.glob(OUTPUT_DIR+'/logs/*.jsonl'))} log files\")\n\n# 3. Generate plots.\n!pip install matplotlib --quiet\n!python scripts/make_plots.py --logs logs/ --out docs/plots/ --eval docs/eval_results.json --trainer-state {OUTPUT_DIR}/trainer_state.json\n\n# 4. Commit and push using GH_TOKEN if available.\nfrom google.colab import userdata\nimport subprocess\n\ngh_token = ''\ntry:\n    gh_token = userdata.get('GH_TOKEN')\nexcept Exception:\n    pass\n\n!git config user.email \"colab@bot\"\n!git config user.name \"colab\"\n!git add docs/plots docs/eval_results.json\n\n# Stage logs only if any were copied\nimport os as _os\nif _os.path.exists('logs') and any(_os.listdir('logs')):\n    !git add logs/\n\n!git status\n\nif gh_token:\n    remote_url = !git remote get-url origin\n    remote_url = remote_url[0].replace('https://', f'https://{gh_token}@')\n    !git commit -m \"Phase 6: training results and plots\" || echo \"Nothing to commit\"\n    !git push {remote_url} main\nelse:\n    !git commit -m \"Phase 6: training results and plots\" || echo \"Nothing to commit\"\n    print(\"\u26a0 GH_TOKEN not set \u2014 commit created locally but not pushed.\")\n    print(\"  Add GH_TOKEN to Colab secrets and re-run this cell, or run:\")\n    print(\"  git push from your Mac after pulling.\")\n\nprint(\"\\n\u2713 Cell 8 done.\")"
+    "# EVALUATE + PLOTS \u2014 Phase 6\nimport os, glob, shutil\nfrom pathlib import Path\n\nOUTPUT_DIR = '/content/drive/MyDrive/injectarena/run_v1'\nCHECKPOINT  = f'{OUTPUT_DIR}/final'\nPLOTS_DRIVE = f'{OUTPUT_DIR}/plots'   # also save plots to Drive\n\n%cd /content/injectarena\n\n# Pull latest code (gets updated make_plots.py etc.)\n!git pull origin main\n\n# 1. Evaluate trained checkpoint against eval split.\n!python train/eval.py \\\n    --checkpoint {CHECKPOINT} \\\n    --output-json docs/eval_results.json\n\n# 2. Generate all 5 plots from trainer_state.json + eval results.\n#    trainer_state.json is written by TRL directly into the output dir.\n!pip install matplotlib --quiet\nos.makedirs('docs/plots', exist_ok=True)\n!python scripts/make_plots.py \\\n    --trainer-state {OUTPUT_DIR}/trainer_state.json \\\n    --logs logs/ \\\n    --eval docs/eval_results.json \\\n    --out docs/plots/\n\n# 3. Copy all plots to Drive so they survive session resets.\nos.makedirs(PLOTS_DRIVE, exist_ok=True)\ncopied = 0\nfor src in Path('docs/plots').glob('*.png'):\n    shutil.copy(src, PLOTS_DRIVE)\n    copied += 1\nprint(f\"Copied {copied} plots to Drive: {PLOTS_DRIVE}\")\n\n# 4. Commit everything and push using GH_TOKEN if available.\nfrom google.colab import userdata\nimport subprocess\n\ngh_token = ''\ntry:\n    gh_token = userdata.get('GH_TOKEN')\nexcept Exception:\n    pass\n\n!git config user.email \"colab@bot\"\n!git config user.name \"colab\"\n!git add docs/plots docs/eval_results.json\n!git status\n\nif gh_token:\n    remote_url = !git remote get-url origin\n    remote_url = remote_url[0].replace('https://', f'https://{gh_token}@')\n    !git commit -m \"Phase 6: training results and plots\" || echo \"Nothing to commit\"\n    !git push {remote_url} main\nelse:\n    !git commit -m \"Phase 6: training results and plots\" || echo \"Nothing to commit\"\n    print(\"\u26a0 GH_TOKEN not set \u2014 commit created locally but not pushed.\")\n    print(\"  Add GH_TOKEN to Colab secrets and re-run this cell, or push from Mac.\")\n\nprint(\"\\n\u2713 Cell 8 done. Plots at docs/plots/ and backed up to Drive.\")\nprint(\"Plots generated:\")\nfor p in sorted(Path('docs/plots').glob('*.png')):\n    print(f\"  {p}\")\n"
    ]
   }
  ],

scripts/make_plots.py

@@ -1,29 +1,48 @@
-"""Generate training plots from JSONL logs and eval results.
-
-Usage:
-    python scripts/make_plots.py --logs logs/ --out docs/plots/
-    python scripts/make_plots.py --logs logs/ --out docs/plots/ --eval docs/eval_results.json
+"""Generate training plots from TRL trainer_state.json or JSONL logs + eval results.
+
+Plots produced
+--------------
+1. reward_curve.png   — mean reward ± std band + smoothed trend
+2. kl_loss_curve.png  — KL divergence & policy loss on twin axes
+3. completion_stats.png — mean completion length + clipped-ratio line
+4. bypass_bars.png    — RL-trained vs handcrafted baseline (eval results)
+5. per_category.png   — per-scenario-category breakdown (eval results)
+
+Usage
+-----
+    python scripts/make_plots.py \\
+        --trainer-state /path/to/trainer_state.json \\
+        --eval docs/eval_results.json \\
+        --out docs/plots/
 """
 
 from __future__ import annotations
 
 import argparse
 import json
-import os
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional, Tuple
+
 
+# ---------------------------------------------------------------------------
+# Argument parsing
+# ---------------------------------------------------------------------------
 
 def _parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser()
-    p.add_argument("--logs", type=str, default="logs/")
+    p.add_argument("--logs", type=str, default="logs/",
+                   help="Directory of JSONL log files (legacy fallback)")
+    p.add_argument("--trainer-state", type=str, default=None,
+                   help="Path to TRL trainer_state.json (preferred)")
     p.add_argument("--out", type=str, default="docs/plots/")
     p.add_argument("--eval", type=str, default="docs/eval_results.json")
-    p.add_argument("--trainer-state", type=str, default=None,
-                   help="Path to TRL trainer_state.json (fallback when no JSONL logs)")
     return p.parse_args()
 
 
+# ---------------------------------------------------------------------------
+# Data loading
+# ---------------------------------------------------------------------------
+
 def _load_jsonl(path: Path) -> List[Dict[str, Any]]:
     rows = []
     with open(path) as f:
@@ -38,7 +57,7 @@ def _load_jsonl(path: Path) -> List[Dict[str, Any]]:
 
 
 def _load_all_logs(logs_dir: Path) -> List[Dict[str, Any]]:
-    rows = []
+    rows: List[Dict[str, Any]] = []
     for p in sorted(logs_dir.glob("*.jsonl")):
         rows.extend(_load_jsonl(p))
     rows.sort(key=lambda r: r.get("step", r.get("global_step", 0)))
@@ -46,149 +65,283 @@ def _load_all_logs(logs_dir: Path) -> List[Dict[str, Any]]:
 
 
 def _load_trainer_state(state_path: Path) -> List[Dict[str, Any]]:
-    """Read TRL's trainer_state.json log_history into the same row format."""
+    """Parse TRL trainer_state.json log_history into a normalised row list."""
     if not state_path.exists():
+        print(f"trainer_state.json not found at {state_path}")
         return []
     with open(state_path) as f:
         data = json.load(f)
-    rows = []
+
+    # TRL GRPO key → normalised key mapping
+    KEY_MAP = {
+        "reward":                     "reward/mean",
+        "rewards/mean":               "reward/mean",
+        "reward/mean":                "reward/mean",
+        "reward_std":                 "reward/std",
+        "rewards/std":                "reward/std",
+        "reward/std":                 "reward/std",
+        "kl":                         "kl",
+        "loss":                       "loss",
+        "train/loss":                 "loss",
+        "learning_rate":              "lr",
+        "completions/mean_length":    "completion/mean_length",
+        "completions/clipped_ratio":  "completion/clipped_ratio",
+    }
+
+    rows: List[Dict[str, Any]] = []
     for entry in data.get("log_history", []):
         step = entry.get("step")
         if step is None:
             continue
         row: Dict[str, Any] = {"step": step}
-        # TRL key variants
-        for src, dst in [
-            ("reward", "reward/mean"),
-            ("rewards/mean", "reward/mean"),
-            ("reward/mean", "reward/mean"),
-            ("loss", "loss"),
-            ("train/loss", "loss"),
-        ]:
+        for src, dst in KEY_MAP.items():
             if src in entry:
                 row[dst] = entry[src]
-        rows.append(row)
+        if len(row) > 1:   # has at least one metric besides step
+            rows.append(row)
     rows.sort(key=lambda r: r["step"])
+    print(f"Loaded {len(rows)} log entries from {state_path}")
     return rows
 
 
+def _extract(rows: List[Dict[str, Any]], key: str) -> Tuple[List[int], List[float]]:
+    steps, vals = [], []
+    for r in rows:
+        v = r.get(key)
+        if v is not None:
+            steps.append(r["step"])
+            vals.append(float(v))
+    return steps, vals
+
+
+# ---------------------------------------------------------------------------
+# Plot helpers
+# ---------------------------------------------------------------------------
+
+def _smooth(vals: List[float], window: int) -> List[float]:
+    import numpy as np
+    if window <= 1 or len(vals) < window:
+        return vals
+    return list(np.convolve(vals, np.ones(window) / window, mode="valid"))
+
+
+BLUE   = "#3b82f6"
+DBLUE  = "#1d4ed8"
+RED    = "#ef4444"
+GREEN  = "#22c55e"
+ORANGE = "#f97316"
+PURPLE = "#a855f7"
+GRAY   = "#94a3b8"
+
+
+# ---------------------------------------------------------------------------
+# Plot 1: Reward curve with ±std band
+# ---------------------------------------------------------------------------
+
 def _plot_reward_curve(rows: List[Dict[str, Any]], out_dir: Path) -> None:
     import matplotlib.pyplot as plt
     import numpy as np
 
-    steps, rewards, losses = [], [], []
-    for r in rows:
-        step = r.get("step", r.get("global_step"))
-        if step is None:
-            continue
-        reward = r.get("reward/mean", r.get("mean_reward"))
-        loss = r.get("loss", r.get("train/loss"))
-        if reward is not None:
-            steps.append(step)
-            rewards.append(reward)
-        if loss is not None and step not in [s for s in steps[:-1]]:
-            losses.append((step, loss))
-
-    if not steps:
-        print("No reward data found in logs — skipping reward curve.")
+    steps_r, rewards = _extract(rows, "reward/mean")
+    steps_s, stds    = _extract(rows, "reward/std")
+
+    if not steps_r:
+        print("No reward data — skipping reward_curve.png")
         return
 
-    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8), sharex=True)
-
-    # Smooth reward with rolling window
-    window = max(1, len(rewards) // 20)
-    smoothed = np.convolve(rewards, np.ones(window) / window, mode="valid")
-    smooth_steps = steps[window - 1:]
-
-    ax1.plot(steps, rewards, alpha=0.3, color="#3b82f6", linewidth=0.8, label="raw")
-    ax1.plot(smooth_steps, smoothed, color="#1d4ed8", linewidth=2, label=f"smoothed (w={window})")
-    ax1.axhline(0, color="gray", linestyle="--", linewidth=0.5)
-    ax1.set_ylabel("Mean Reward")
-    ax1.set_title("InjectArena — GRPO Training: Reward Curve")
-    ax1.legend(loc="lower right")
-    ax1.grid(alpha=0.3)
-
-    if losses:
-        l_steps = [x[0] for x in losses]
-        l_vals = [x[1] for x in losses]
-        ax2.plot(l_steps, l_vals, color="#ef4444", linewidth=1.5)
-        ax2.set_ylabel("Training Loss")
-        ax2.set_xlabel("Step")
-        ax2.grid(alpha=0.3)
-    else:
-        ax2.set_visible(False)
+    window = max(1, len(rewards) // 15)
+    smoothed = _smooth(rewards, window)
+    smooth_steps = steps_r[window - 1:] if window > 1 else steps_r
+
+    fig, ax = plt.subplots(figsize=(10, 5))
+
+    # std band
+    if steps_s and len(steps_s) == len(steps_r):
+        r_arr = np.array(rewards)
+        s_arr = np.array(stds)
+        ax.fill_between(steps_r, r_arr - s_arr, r_arr + s_arr,
+                        alpha=0.15, color=BLUE, label="±1 std")
+
+    ax.plot(steps_r, rewards, alpha=0.35, color=BLUE, linewidth=0.9, label="raw reward")
+    ax.plot(smooth_steps, smoothed, color=DBLUE, linewidth=2.2,
+            label=f"smoothed (w={window})")
+    ax.axhline(0, color="gray", linestyle="--", linewidth=0.6)
+
+    ax.set_xlabel("Training Step")
+    ax.set_ylabel("Mean Reward")
+    ax.set_title("InjectArena — GRPO Reward Curve (300 steps)")
+    ax.legend(loc="lower right")
+    ax.set_ylim(bottom=0)
+    ax.grid(alpha=0.25)
 
-    plt.tight_layout()
     out_path = out_dir / "reward_curve.png"
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {out_path}")
+
+
+# ---------------------------------------------------------------------------
+# Plot 2: KL divergence + policy loss on twin axes
+# ---------------------------------------------------------------------------
+
+def _plot_kl_loss(rows: List[Dict[str, Any]], out_dir: Path) -> None:
+    import matplotlib.pyplot as plt
+
+    steps_kl, kls   = _extract(rows, "kl")
+    steps_l,  losses = _extract(rows, "loss")
+
+    if not steps_kl and not steps_l:
+        print("No KL/loss data — skipping kl_loss_curve.png")
+        return
+
+    fig, ax1 = plt.subplots(figsize=(10, 4))
+
+    if steps_kl:
+        ax1.plot(steps_kl, kls, color=PURPLE, linewidth=1.8, label="KL divergence")
+        ax1.set_ylabel("KL Divergence", color=PURPLE)
+        ax1.tick_params(axis="y", labelcolor=PURPLE)
+
+    if steps_l:
+        ax2 = ax1.twinx()
+        ax2.plot(steps_l, losses, color=RED, linewidth=1.8, linestyle="--", label="Policy loss")
+        ax2.set_ylabel("Policy Loss", color=RED)
+        ax2.tick_params(axis="y", labelcolor=RED)
+
+    ax1.set_xlabel("Training Step")
+    ax1.set_title("InjectArena — KL Divergence & Policy Loss")
+    ax1.grid(alpha=0.2)
+
+    # Combined legend
+    lines = []
+    if steps_kl:
+        lines += ax1.get_lines()
+    if steps_l:
+        lines += ax2.get_lines()
+    if lines:
+        ax1.legend(lines, [l.get_label() for l in lines], loc="upper right")
+
+    out_path = out_dir / "kl_loss_curve.png"
+    plt.tight_layout()
     plt.savefig(out_path, dpi=150, bbox_inches="tight")
     plt.close()
     print(f"Saved {out_path}")
 
 
+# ---------------------------------------------------------------------------
+# Plot 3: Completion length + clipped ratio
+# ---------------------------------------------------------------------------
+
+def _plot_completion_stats(rows: List[Dict[str, Any]], out_dir: Path) -> None:
+    import matplotlib.pyplot as plt
+
+    steps_l, lengths = _extract(rows, "completion/mean_length")
+    steps_c, clipped = _extract(rows, "completion/clipped_ratio")
+
+    if not steps_l and not steps_c:
+        print("No completion stats — skipping completion_stats.png")
+        return
+
+    fig, ax1 = plt.subplots(figsize=(10, 4))
+
+    if steps_l:
+        ax1.plot(steps_l, lengths, color=ORANGE, linewidth=1.8, label="Mean completion length (tokens)")
+        ax1.set_ylabel("Mean Length (tokens)", color=ORANGE)
+        ax1.tick_params(axis="y", labelcolor=ORANGE)
+
+    if steps_c:
+        ax2 = ax1.twinx()
+        ax2.plot(steps_c, clipped, color=RED, linewidth=1.8, linestyle="--",
+                 label="Clipped ratio (hit max_len)")
+        ax2.set_ylabel("Clipped Ratio", color=RED)
+        ax2.set_ylim(0, 1.05)
+        ax2.tick_params(axis="y", labelcolor=RED)
+        ax2.axhline(1.0, color=RED, linestyle=":", linewidth=0.7, alpha=0.5)
+
+    ax1.set_xlabel("Training Step")
+    ax1.set_title("InjectArena — Completion Length & Clipping")
+    ax1.grid(alpha=0.2)
+
+    lines = []
+    if steps_l:
+        lines += ax1.get_lines()
+    if steps_c:
+        lines += ax2.get_lines()
+    if lines:
+        ax1.legend(lines, [l.get_label() for l in lines], loc="upper right")
+
+    out_path = out_dir / "completion_stats.png"
+    plt.tight_layout()
+    plt.savefig(out_path, dpi=150, bbox_inches="tight")
+    plt.close()
+    print(f"Saved {out_path}")
+
+
+# ---------------------------------------------------------------------------
+# Plot 4: Bypass bars (eval results vs baseline)
+# ---------------------------------------------------------------------------
+
 def _plot_bypass_bars(eval_path: Path, out_dir: Path) -> None:
     import matplotlib.pyplot as plt
     import numpy as np
 
     if not eval_path.exists():
-        print(f"Eval results not found at {eval_path} — skipping bypass bar chart.")
+        print(f"Eval results not found at {eval_path} — skipping bypass_bars.png")
         return
 
     with open(eval_path) as f:
         data = json.load(f)
 
     metrics = {
-        "PG2 Bypass": data.get("pg2_bypass_rate", 0),
-        "FW Bypass": data.get("fw_bypass_rate", 0),
-        "Task Success": data.get("task_success_rate", 0),
+        "PG2 Bypass":      data.get("pg2_bypass_rate", 0),
+        "FW Bypass":       data.get("fw_bypass_rate", 0),
+        "Task Success":    data.get("task_success_rate", 0),
         "Composed Bypass": data.get("composed_bypass_rate", 0),
     }
-
-    # Hardcoded approximate baselines (handcrafted corpus on same eval split)
-    # These are filled in once the zero_shot eval is run.
     baselines = {
-        "PG2 Bypass": 0.15,
-        "FW Bypass": 0.20,
-        "Task Success": 0.05,
-        "Composed Bypass": 0.02,
+        "PG2 Bypass": 0.15, "FW Bypass": 0.20,
+        "Task Success": 0.05, "Composed Bypass": 0.02,
     }
 
     x = np.arange(len(metrics))
     width = 0.35
-
     fig, ax = plt.subplots(figsize=(9, 5))
-    bars1 = ax.bar(x - width / 2,
-                   [baselines[k] for k in metrics],
-                   width, label="Handcrafted Baseline", color="#94a3b8", edgecolor="white")
-    bars2 = ax.bar(x + width / 2,
-                   [metrics[k] for k in metrics],
-                   width, label="InjectArena (RL-trained)", color="#3b82f6", edgecolor="white")
+
+    bars1 = ax.bar(x - width / 2, [baselines[k] for k in metrics],
+                   width, label="Handcrafted Baseline", color=GRAY, edgecolor="white")
+    bars2 = ax.bar(x + width / 2, [metrics[k] for k in metrics],
+                   width, label="InjectArena (RL-trained)", color=BLUE, edgecolor="white")
 
     ax.set_ylabel("Rate")
     ax.set_title("InjectArena — Attacker Performance vs Baseline")
     ax.set_xticks(x)
     ax.set_xticklabels(list(metrics.keys()))
-    ax.set_ylim(0, 1.0)
+    ax.set_ylim(0, 1.05)
     ax.legend()
     ax.grid(axis="y", alpha=0.3)
 
     for bar in bars1:
         h = bar.get_height()
-        if h > 0.02:
+        if h > 0.01:
             ax.text(bar.get_x() + bar.get_width() / 2, h + 0.01,
-                    f"{h:.0%}", ha="center", va="bottom", fontsize=9, color="#64748b")
+                    f"{h:.0%}", ha="center", va="bottom", fontsize=9, color="#475569")
     for bar in bars2:
         h = bar.get_height()
-        if h > 0.02:
+        if h > 0.01:
             ax.text(bar.get_x() + bar.get_width() / 2, h + 0.01,
-                    f"{h:.0%}", ha="center", va="bottom", fontsize=9, color="#1d4ed8")
+                    f"{h:.0%}", ha="center", va="bottom", fontsize=9, color=DBLUE)
 
-    plt.tight_layout()
     out_path = out_dir / "bypass_bars.png"
+    plt.tight_layout()
     plt.savefig(out_path, dpi=150, bbox_inches="tight")
     plt.close()
     print(f"Saved {out_path}")
 
 
+# ---------------------------------------------------------------------------
+# Plot 5: Per-category breakdown
+# ---------------------------------------------------------------------------
+
 def _plot_per_category(eval_path: Path, out_dir: Path) -> None:
     import matplotlib.pyplot as plt
     import numpy as np
@@ -201,40 +354,42 @@ def _plot_per_category(eval_path: Path, out_dir: Path) -> None:
 
     per_cat = data.get("per_category", {})
     if not per_cat:
+        print("No per_category data in eval results — skipping per_category.png")
         return
 
     cats = list(per_cat.keys())
-    task_rates = [per_cat[c]["task_success"] for c in cats]
+    task_rates   = [per_cat[c]["task_success"] for c in cats]
     bypass_rates = [per_cat[c]["composed_bypass"] for c in cats]
 
     x = np.arange(len(cats))
     width = 0.35
-
     fig, ax = plt.subplots(figsize=(8, 5))
-    ax.bar(x - width / 2, task_rates, width, label="Task Success", color="#22c55e", edgecolor="white")
-    ax.bar(x + width / 2, bypass_rates, width, label="Composed Bypass", color="#3b82f6", edgecolor="white")
+    ax.bar(x - width / 2, task_rates,   width, label="Task Success",    color=GREEN, edgecolor="white")
+    ax.bar(x + width / 2, bypass_rates, width, label="Composed Bypass", color=BLUE,  edgecolor="white")
 
     ax.set_ylabel("Rate")
     ax.set_title("InjectArena — Per-Category Breakdown")
     ax.set_xticks(x)
-    ax.set_xticklabels(cats, rotation=15)
-    ax.set_ylim(0, 1.0)
+    ax.set_xticklabels(cats, rotation=15, ha="right")
+    ax.set_ylim(0, 1.05)
     ax.legend()
     ax.grid(axis="y", alpha=0.3)
 
-    plt.tight_layout()
     out_path = out_dir / "per_category.png"
+    plt.tight_layout()
     plt.savefig(out_path, dpi=150, bbox_inches="tight")
     plt.close()
     print(f"Saved {out_path}")
 
 
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
 def main() -> None:
     args = _parse_args()
-    logs_dir = Path(args.logs)
-    out_dir = Path(args.out)
+    out_dir   = Path(args.out)
     eval_path = Path(args.eval)
-
     out_dir.mkdir(parents=True, exist_ok=True)
 
     try:
@@ -244,24 +399,32 @@ def main() -> None:
         print("matplotlib not installed — pip install matplotlib")
         return
 
+    # Load training log rows (trainer_state preferred, JSONL fallback)
     rows: List[Dict[str, Any]] = []
-    if logs_dir.exists():
-        rows = _load_all_logs(logs_dir)
-        print(f"Loaded {len(rows)} log rows from {logs_dir}")
-
-    if not rows and args.trainer_state:
-        state_path = Path(args.trainer_state)
-        rows = _load_trainer_state(state_path)
-        print(f"Loaded {len(rows)} log rows from trainer_state {state_path}")
-
+    if args.trainer_state:
+        rows = _load_trainer_state(Path(args.trainer_state))
     if not rows:
-        print("No log data found — skipping reward curve.")
-    else:
+        logs_dir = Path(args.logs)
+        if logs_dir.exists():
+            rows = _load_all_logs(logs_dir)
+            if rows:
+                print(f"Loaded {len(rows)} log rows from {logs_dir}")
+
+    # Training plots (require rows)
+    if rows:
         _plot_reward_curve(rows, out_dir)
+        _plot_kl_loss(rows, out_dir)
+        _plot_completion_stats(rows, out_dir)
+    else:
+        print("No training log data found — skipping reward/KL/completion plots.")
 
+    # Eval plots (require eval results JSON)
     _plot_bypass_bars(eval_path, out_dir)
     _plot_per_category(eval_path, out_dir)
-    print("Done.")
+
+    print("\nAll plots done.")
+    for p in sorted(out_dir.glob("*.png")):
+        print(f"  {p}")
 
 
 if __name__ == "__main__":