Spaces:
Sleeping
Sleeping
| """Generate the README plots from results.tsv and any wandb export. | |
| Outputs (in outputs/): | |
| - baseline_vs_trained.png (bar chart per scenario) | |
| - reward_curve.png (line chart over training steps) | |
| - memory_growth.png (lessons.md size vs episode) | |
| - curriculum_progression.png (sampling weights over episodes) | |
| - results_tsv_summary.png (per-episode score over time) | |
| Designed to run AFTER training completes. Produces placeholder plots | |
| when underlying data is missing so the README never has broken images. | |
| """ | |
| import csv | |
| import json | |
| import os | |
| from pathlib import Path | |
| from typing import Dict, List | |
| ROOT = Path(__file__).resolve().parents[1] | |
| def _safe_import(): | |
| try: | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| return plt | |
| except ImportError: | |
| return None | |
| def _load_results_tsv() -> List[Dict]: | |
| path = ROOT / "results.tsv" | |
| if not path.exists(): | |
| return [] | |
| rows: List[Dict] = [] | |
| with open(path, "r", encoding="utf-8", newline="") as f: | |
| reader = csv.DictReader(f, delimiter="\t") | |
| for r in reader: | |
| rows.append(r) | |
| return rows | |
| def plot_results_tsv_summary(plt): | |
| rows = _load_results_tsv() | |
| fig, ax = plt.subplots(figsize=(8, 4)) | |
| if not rows: | |
| ax.text(0.5, 0.5, "No episodes logged yet", ha="center", va="center") | |
| else: | |
| scores = [float(r.get("score", 0)) for r in rows] | |
| ax.plot(range(1, len(scores) + 1), scores, marker="o", linewidth=1) | |
| ax.set_xlabel("Episode") | |
| ax.set_ylabel("Final Score") | |
| ax.set_ylim(0, 1) | |
| ax.set_title("APIShift — Per-Episode Score Over Training") | |
| ax.grid(True, alpha=0.3) | |
| out = ROOT / "outputs" / "results_tsv_summary.png" | |
| out.parent.mkdir(exist_ok=True) | |
| fig.tight_layout() | |
| fig.savefig(out, dpi=120) | |
| plt.close(fig) | |
| print(f"wrote {out}") | |
| def plot_baseline_vs_trained(plt): | |
| base_path = ROOT / "outputs" / "baseline_scores.json" | |
| trained_path = ROOT / "outputs" / "trained_scores.json" | |
| fig, ax = plt.subplots(figsize=(9, 4.5)) | |
| if not base_path.exists() or not trained_path.exists(): | |
| ax.text(0.5, 0.5, "Run training/run_baseline.py and trained eval first.", | |
| ha="center", va="center") | |
| else: | |
| base = json.load(open(base_path))["scores"] | |
| trained = json.load(open(trained_path))["scores"] | |
| keys = list(base.keys()) | |
| x = range(len(keys)) | |
| bw = 0.35 | |
| ax.bar([i - bw/2 for i in x], [base[k] for k in keys], width=bw, label="untrained") | |
| ax.bar([i + bw/2 for i in x], [trained.get(k, 0) for k in keys], width=bw, label="trained") | |
| ax.set_xticks(list(x)) | |
| ax.set_xticklabels(keys, rotation=20, ha="right") | |
| ax.set_ylim(0, 1) | |
| ax.set_ylabel("Final Score") | |
| ax.set_title("APIShift — Untrained vs Trained Manager (Held-Out)") | |
| ax.legend() | |
| out = ROOT / "outputs" / "baseline_vs_trained.png" | |
| fig.tight_layout() | |
| fig.savefig(out, dpi=120) | |
| plt.close(fig) | |
| print(f"wrote {out}") | |
| def plot_memory_growth(plt): | |
| """Memory growth over episodes — derived from results.tsv + | |
| snapshots of lessons.md size if available.""" | |
| rows = _load_results_tsv() | |
| fig, ax = plt.subplots(figsize=(8, 4)) | |
| if not rows: | |
| ax.text(0.5, 0.5, "No episodes yet", ha="center", va="center") | |
| else: | |
| # Approximate lessons growth: count distinct lessons_used over time | |
| seen = set() | |
| cum = [] | |
| for r in rows: | |
| for lid in (r.get("lessons_used") or "").split(","): | |
| lid = lid.strip() | |
| if lid and lid != "(none)": | |
| seen.add(lid) | |
| cum.append(len(seen)) | |
| ax.plot(range(1, len(cum) + 1), cum, marker="o", linewidth=1) | |
| ax.set_xlabel("Episode") | |
| ax.set_ylabel("Distinct lessons referenced") | |
| ax.set_title("APIShift — Memory Growth Over Training") | |
| ax.grid(True, alpha=0.3) | |
| out = ROOT / "outputs" / "memory_growth.png" | |
| fig.tight_layout() | |
| fig.savefig(out, dpi=120) | |
| plt.close(fig) | |
| print(f"wrote {out}") | |
| def _load_trainer_state(): | |
| """Find and load the trainer_state.json from training outputs.""" | |
| candidates = [ | |
| ROOT / "training" / "outputs" / "grpo" / "checkpoint-1500" / "trainer_state.json", | |
| ROOT / "training" / "outputs" / "grpo" / "trainer_state.json", | |
| ] | |
| for p in candidates: | |
| if p.exists(): | |
| return json.load(open(p, encoding="utf-8")) | |
| # Fallback: latest checkpoint in outputs/grpo/ | |
| grpo_dir = ROOT / "training" / "outputs" / "grpo" | |
| if grpo_dir.exists(): | |
| cps = sorted(grpo_dir.glob("checkpoint-*/trainer_state.json")) | |
| if cps: | |
| return json.load(open(cps[-1], encoding="utf-8")) | |
| return None | |
| def plot_reward_curve(plt): | |
| """The reward curve: per-step reward + rolling average + trend line.""" | |
| state = _load_trainer_state() | |
| fig, ax = plt.subplots(figsize=(10, 5)) | |
| if state is None: | |
| ax.text(0.5, 0.5, "trainer_state.json not found, run training first", | |
| ha="center", va="center") | |
| ax.axis("off") | |
| else: | |
| logs = state["log_history"] | |
| steps = [l["step"] for l in logs] | |
| rewards = [l.get("reward", 0) for l in logs] | |
| # Rolling mean window of 50 | |
| window = 50 | |
| rolling = [] | |
| for i in range(len(rewards)): | |
| lo = max(0, i - window + 1) | |
| rolling.append(sum(rewards[lo:i+1]) / (i - lo + 1)) | |
| # Linear trend | |
| n = len(steps) | |
| if n > 1: | |
| mean_x = sum(steps) / n | |
| mean_y = sum(rewards) / n | |
| num = sum((steps[i] - mean_x) * (rewards[i] - mean_y) for i in range(n)) | |
| den = sum((steps[i] - mean_x) ** 2 for i in range(n)) | |
| slope = num / den if den else 0 | |
| intercept = mean_y - slope * mean_x | |
| trend = [slope * x + intercept for x in steps] | |
| ax.plot(steps, rewards, color="#a5b4fc", linewidth=0.8, alpha=0.5, label="Per step") | |
| ax.plot(steps, rolling, color="#3730a3", linewidth=2.2, label=f"Rolling avg ({window})") | |
| if n > 1: | |
| ax.plot(steps, trend, color="#dc2626", linewidth=1.5, linestyle="--", | |
| label=f"Trend ({slope*100:+.3f} per step)") | |
| ax.axhline(rewards[0], color="#9ca3af", linewidth=0.8, linestyle=":", alpha=0.6) | |
| ax.set_xlabel("Training step", fontsize=11) | |
| ax.set_ylabel("Reward", fontsize=11) | |
| ax.set_title("APIShift Manager, GRPO Training Reward Curve", fontsize=13, fontweight="bold") | |
| ax.grid(True, alpha=0.3) | |
| ax.legend(loc="lower right") | |
| # Footnote | |
| avg = sum(rewards) / len(rewards) | |
| best = max(rewards) | |
| ax.text(0.02, 0.98, | |
| f"Steps: {len(steps)} | Avg reward: {avg:.3f} | Best: {best:.3f}", | |
| transform=ax.transAxes, fontsize=9, verticalalignment="top", | |
| bbox=dict(boxstyle="round,pad=0.4", facecolor="#fef3c7", edgecolor="#fbbf24")) | |
| out = ROOT / "outputs" / "reward_curve.png" | |
| out.parent.mkdir(exist_ok=True) | |
| fig.tight_layout() | |
| fig.savefig(out, dpi=140) | |
| plt.close(fig) | |
| print(f"wrote {out}") | |
| def _kde(values, x_grid, bandwidth=0.012): | |
| """Tiny gaussian KDE so we don't depend on scipy.""" | |
| import math | |
| n = len(values) | |
| if n == 0: | |
| return [0.0] * len(x_grid) | |
| out = [] | |
| coef = 1.0 / (n * bandwidth * math.sqrt(2 * math.pi)) | |
| for x in x_grid: | |
| s = 0.0 | |
| for v in values: | |
| d = (x - v) / bandwidth | |
| s += math.exp(-0.5 * d * d) | |
| out.append(coef * s) | |
| return out | |
| def plot_reward_distribution(plt): | |
| """KDE density of training rewards: early training behavior vs late training | |
| behavior of the same model. Honest framing, no claim of separate untrained baseline.""" | |
| state = _load_trainer_state() | |
| fig, ax = plt.subplots(figsize=(10, 5)) | |
| if state is None: | |
| ax.text(0.5, 0.5, "trainer_state.json not found", ha="center", va="center") | |
| ax.axis("off") | |
| else: | |
| logs = state["log_history"] | |
| rewards = [l.get("reward", 0) for l in logs] | |
| early = rewards[:200] | |
| late = rewards[-200:] | |
| early_avg = sum(early) / len(early) | |
| late_avg = sum(late) / len(late) | |
| x_grid = [0.05 + 0.005 * i for i in range(101)] # 0.05 -> 0.55 | |
| early_d = _kde(early, x_grid, bandwidth=0.018) | |
| late_d = _kde(late, x_grid, bandwidth=0.018) | |
| ax.fill_between(x_grid, early_d, alpha=0.35, color="#dc2626", label=f"Early training, steps 1-200 (avg {early_avg:.3f})") | |
| ax.plot(x_grid, early_d, color="#991b1b", linewidth=2) | |
| ax.fill_between(x_grid, late_d, alpha=0.35, color="#16a34a", label=f"Late training, steps 1301-1500 (avg {late_avg:.3f})") | |
| ax.plot(x_grid, late_d, color="#15803d", linewidth=2) | |
| # Mean markers | |
| ax.axvline(early_avg, color="#991b1b", linestyle=":", linewidth=1.5, alpha=0.8) | |
| ax.axvline(late_avg, color="#15803d", linestyle=":", linewidth=1.5, alpha=0.8) | |
| # Improvement annotation | |
| delta_pct = ((late_avg - early_avg) / early_avg) * 100 if early_avg else 0 | |
| ax.text(0.98, 0.96, | |
| f"Mean reward shifted right by {delta_pct:+.1f}%\nLate training spends more weight on high-quality episodes", | |
| transform=ax.transAxes, fontsize=10, verticalalignment="top", horizontalalignment="right", | |
| bbox=dict(boxstyle="round,pad=0.5", facecolor="#fef3c7", edgecolor="#fbbf24")) | |
| ax.set_xlabel("Reward", fontsize=11) | |
| ax.set_ylabel("Density", fontsize=11) | |
| ax.set_title("Reward density shift: early training vs late training", fontsize=13, fontweight="bold") | |
| ax.legend(loc="upper left", fontsize=9) | |
| ax.grid(True, alpha=0.3) | |
| out = ROOT / "outputs" / "reward_distribution.png" | |
| fig.tight_layout() | |
| fig.savefig(out, dpi=140) | |
| plt.close(fig) | |
| print(f"wrote {out}") | |
| def plot_reward_components(plt): | |
| """Stacked area showing how each of the 5 reward components contributed. | |
| For training, TRL only logs the aggregate. To get per-component data we | |
| run a quick eval pass at multiple checkpoints. For now, generate a | |
| representative chart from the final state.component_rewards dict via | |
| a rollout pass.""" | |
| fig, ax = plt.subplots(figsize=(9, 4.5)) | |
| # Run an in-process rollout of 10 episodes against current env to get | |
| # per-component breakdown. This is fast (no LLM, deterministic optimal plan). | |
| try: | |
| sys = __import__("sys") | |
| sys.path.insert(0, str(ROOT)) | |
| from server.environment import APIShiftEnvironment | |
| from models import APIShiftAction | |
| from scenarios import SCENARIO_IDS | |
| plan = [ | |
| {"command": "dispatch_diff", "rationale": ""}, | |
| {"command": "dispatch_patch", "payload": {"change_id": "ch_001"}, "rationale": ""}, | |
| {"command": "dispatch_test", "rationale": ""}, | |
| {"command": "dispatch_rollback", "rationale": ""}, | |
| {"command": "submit", "rationale": ""}, | |
| ] | |
| comps_by_episode = [] | |
| labels = [] | |
| for sid in SCENARIO_IDS[:10]: | |
| env = APIShiftEnvironment() | |
| env.reset(task_id=sid) | |
| for a in plan: | |
| env.step(APIShiftAction.model_validate(a)) | |
| comps_by_episode.append(dict(env.state.component_rewards)) | |
| labels.append(sid[:18]) | |
| keys = ["breaking_change_detection", "migration_patch_correctness", | |
| "backward_compat_preservation", "rollback_plan_completeness", "simplicity_bonus"] | |
| weights = [0.33, 0.28, 0.24, 0.10, 0.05] | |
| colors = ["#3730a3", "#0891b2", "#16a34a", "#f59e0b", "#dc2626"] | |
| nice_labels = ["Breaking change detection (33%)", "Migration patch correctness (28%)", | |
| "Backward compat preservation (24%)", "Rollback plan completeness (10%)", | |
| "Simplicity bonus (5%)"] | |
| bottom = [0] * len(comps_by_episode) | |
| x = list(range(len(comps_by_episode))) | |
| for k, w, c, nl in zip(keys, weights, colors, nice_labels): | |
| vals = [comps_by_episode[i].get(k, 0) * w for i in range(len(comps_by_episode))] | |
| ax.bar(x, vals, bottom=bottom, color=c, label=nl, edgecolor="white", linewidth=0.5) | |
| bottom = [bottom[i] + vals[i] for i in range(len(bottom))] | |
| ax.set_xticks(x) | |
| ax.set_xticklabels(labels, rotation=30, ha="right", fontsize=8) | |
| ax.set_ylabel("Weighted reward contribution", fontsize=11) | |
| ax.set_title("Reward Components Per Scenario, Optimal Plan", fontsize=13, fontweight="bold") | |
| ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.25), ncol=2, fontsize=8) | |
| except Exception as e: | |
| ax.text(0.5, 0.5, f"Could not generate component plot: {e}", | |
| ha="center", va="center") | |
| ax.axis("off") | |
| out = ROOT / "outputs" / "reward_components.png" | |
| fig.tight_layout() | |
| fig.savefig(out, dpi=140) | |
| plt.close(fig) | |
| print(f"wrote {out}") | |
| def main(): | |
| plt = _safe_import() | |
| if plt is None: | |
| print("matplotlib not installed; install with `pip install matplotlib`.") | |
| return | |
| plot_reward_curve(plt) | |
| plot_reward_distribution(plt) | |
| plot_reward_components(plt) | |
| plot_baseline_vs_trained(plt) | |
| plot_results_tsv_summary(plt) | |
| plot_memory_growth(plt) | |
| if __name__ == "__main__": | |
| main() | |