# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. """ plot_results.py — generate training result charts for the submission README. Produces a 3-subplot figure saved to plots/results.png: Top: Episode return curve (x=episode, y=reward), dashed baseline, curriculum tier markers Middle: nDCG@5 before vs after training, per difficulty tier Bottom: Fakes-in-top-5 before vs after training, per difficulty tier Usage: python -m recruiter.training.plot_results \\ --baseline baseline_random.json \\ --trained trained_eval.json \\ --training training_log.jsonl \\ --output plots/results.png # Minimal (just training curve, no before/after eval): python -m recruiter.training.plot_results --training training_log.jsonl Input formats: baseline/trained JSON: output of baseline_eval.py training_log JSONL: one JSON object per update, e.g.: {"update": 0, "episode": 4, "mean_reward": -3.2, "difficulty": 1} All files are optional — subplots are skipped if data is missing. """ from __future__ import annotations import argparse import json import os from pathlib import Path from typing import Optional def _load_jsonl(path: str) -> list[dict]: records = [] with open(path) as f: for line in f: line = line.strip() if line: records.append(json.loads(line)) return records def _load_json(path: str) -> dict: with open(path) as f: return json.load(f) def _rolling_mean(values: list[float], window: int = 20) -> list[float]: out = [] for i, v in enumerate(values): start = max(0, i - window + 1) out.append(sum(values[start : i + 1]) / (i - start + 1)) return out def _tier_stats(eval_data: dict) -> dict[int, dict]: """Extract per-tier mean_ndcg and mean_fakes_in_top5 from baseline_eval output.""" per_ep = eval_data.get("per_episode", []) tier_buckets: dict[int, list[dict]] = {} for ep in per_ep: d = ep.get("difficulty", 1) tier_buckets.setdefault(d, []).append(ep) result = {} for tier, eps in tier_buckets.items(): ndcgs = [e.get("ndcg", 0.0) for e in eps] fakes = [e.get("fakes_in_top5", 0) for e in eps] result[tier] = { "mean_ndcg": sum(ndcgs) / len(ndcgs), "mean_fakes": sum(fakes) / len(fakes), "n": len(eps), } return result def plot( training_log: Optional[str] = None, baseline_path: Optional[str] = None, trained_path: Optional[str] = None, output: str = "plots/results.png", wandb_url: Optional[str] = None, ) -> None: import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import matplotlib.ticker as ticker import numpy as np Path(output).parent.mkdir(parents=True, exist_ok=True) # ── Determine number of subplots ────────────────────────────────────────── has_training = training_log is not None and Path(training_log).exists() has_before_after = ( baseline_path is not None and Path(baseline_path).exists() and trained_path is not None and Path(trained_path).exists() ) n_plots = 1 + (2 if has_before_after else 0) fig, axes = plt.subplots(n_plots, 1, figsize=(10, 4 * n_plots)) if n_plots == 1: axes = [axes] fig.suptitle("Recruiter Agent Training Results", fontsize=14, fontweight="bold", y=0.98) # ── Subplot 1: Training curve ───────────────────────────────────────────── ax = axes[0] ax.set_title("Episode Return During Training") ax.set_xlabel("Episode") ax.set_ylabel("Return (reward)") if has_training: records = _load_jsonl(training_log) episodes = [r.get("episode", i * 4) for i, r in enumerate(records)] # Support both "mean_return" (train_grpo.py key) and "mean_reward" (legacy) rewards = [r.get("mean_terminal_return", r.get("mean_return", r.get("mean_reward", 0.0))) for r in records] smooth = _rolling_mean(rewards, window=10) ax.plot(episodes, rewards, alpha=0.3, color="steelblue", linewidth=0.8, label="raw") ax.plot(episodes, smooth, color="steelblue", linewidth=2.0, label="smoothed (win=10)") # Curriculum tier markers difficulties = [r.get("difficulty", 1) for r in records] prev_d = difficulties[0] if difficulties else 1 for i, d in enumerate(difficulties): if d != prev_d: ax.axvline(x=episodes[i], color="orange", linestyle="--", alpha=0.6, linewidth=1.2) ax.text(episodes[i], ax.get_ylim()[0] if ax.get_ylim()[0] != 0 else min(rewards), f"T{d}", fontsize=8, color="darkorange", ha="left") prev_d = d else: ax.text(0.5, 0.5, "No training log provided", transform=ax.transAxes, ha="center", va="center", fontsize=12, color="grey") # Baseline reference line if baseline_path and Path(baseline_path).exists(): baseline = _load_json(baseline_path) mean_base = baseline.get("summary", {}).get("mean_reward", None) if mean_base is not None: ax.axhline(y=mean_base, color="red", linestyle="--", linewidth=1.5, label=f"random baseline ({mean_base:.2f})") ax.legend(loc="lower right", fontsize=9) ax.grid(True, alpha=0.3) if has_before_after: baseline_data = _load_json(baseline_path) trained_data = _load_json(trained_path) before_stats = _tier_stats(baseline_data) after_stats = _tier_stats(trained_data) all_tiers = sorted(set(before_stats) | set(after_stats)) tier_labels = [f"T{t}" for t in all_tiers] x = np.arange(len(all_tiers)) width = 0.35 # ── Subplot 2: nDCG@5 per tier ──────────────────────────────────────── ax2 = axes[1] ax2.set_title("nDCG@5 Before vs After Training (per difficulty tier)") ax2.set_ylabel("nDCG@5") ax2.set_ylim(0, 1.05) before_ndcg = [before_stats.get(t, {}).get("mean_ndcg", 0.0) for t in all_tiers] after_ndcg = [after_stats.get(t, {}).get("mean_ndcg", 0.0) for t in all_tiers] ax2.bar(x - width / 2, before_ndcg, width, label="Before (zero-shot)", color="salmon", alpha=0.8) ax2.bar(x + width / 2, after_ndcg, width, label="After (fine-tuned)", color="steelblue", alpha=0.8) ax2.set_xticks(x) ax2.set_xticklabels(tier_labels) ax2.legend(fontsize=9) ax2.grid(True, alpha=0.3, axis="y") # Annotate delta for i, (b, a) in enumerate(zip(before_ndcg, after_ndcg)): delta = a - b color = "green" if delta >= 0 else "red" ax2.text(x[i] + width / 2, a + 0.02, f"{delta:+.2f}", ha="center", fontsize=8, color=color, fontweight="bold") # ── Subplot 3: Fakes-in-top-5 per tier ─────────────────────────────── ax3 = axes[2] ax3.set_title("Fakes in Top-5 Before vs After Training (per difficulty tier)") ax3.set_ylabel("Mean fakes in top-5 (lower is better)") before_fakes = [before_stats.get(t, {}).get("mean_fakes", 0.0) for t in all_tiers] after_fakes = [after_stats.get(t, {}).get("mean_fakes", 0.0) for t in all_tiers] ax3.bar(x - width / 2, before_fakes, width, label="Before (zero-shot)", color="salmon", alpha=0.8) ax3.bar(x + width / 2, after_fakes, width, label="After (fine-tuned)", color="steelblue", alpha=0.8) ax3.set_xticks(x) ax3.set_xticklabels(tier_labels) ax3.legend(fontsize=9) ax3.grid(True, alpha=0.3, axis="y") # Annotate delta for i, (b, a) in enumerate(zip(before_fakes, after_fakes)): delta = a - b color = "green" if delta <= 0 else "red" ax3.text(x[i] + width / 2, a + 0.02, f"{delta:+.2f}", ha="center", fontsize=8, color=color, fontweight="bold") if wandb_url: fig.text(0.5, 0.01, f"W&B run: {wandb_url}", ha="center", fontsize=8, color="grey") plt.tight_layout(rect=[0, 0.02, 1, 0.97]) plt.savefig(output, dpi=150, bbox_inches="tight") print(f"Saved plot to {output}") def main() -> None: parser = argparse.ArgumentParser(description="Plot recruiter training results") parser.add_argument("--training", default=None, help="Path to training_log.jsonl") parser.add_argument("--baseline", default=None, help="Path to baseline_eval output JSON") parser.add_argument("--trained", default=None, help="Path to post-training eval JSON") parser.add_argument("--output", default="plots/results.png", help="Output PNG path") parser.add_argument("--wandb-url", default=None, help="W&B run URL to embed in figure") args = parser.parse_args() plot( training_log=args.training, baseline_path=args.baseline, trained_path=args.trained, output=args.output, wandb_url=args.wandb_url, ) if __name__ == "__main__": main()