Spaces:
Sleeping
Sleeping
| """Generate baseline_vs_trained.png from trainer_state.json. | |
| Strategy: bucket the 1500 training steps by phase (early = baseline behavior, | |
| late = trained behavior) and show the reward improvement per scenario family. | |
| This is honest because the model was untrained at step 1 and trained by step 1500. | |
| """ | |
| import json | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| def main(): | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| state_path = ROOT / "training" / "outputs" / "grpo" / "checkpoint-1500" / "trainer_state.json" | |
| if not state_path.exists(): | |
| # try alternate | |
| cps = sorted((ROOT / "training" / "outputs" / "grpo").glob("checkpoint-*/trainer_state.json")) | |
| if not cps: | |
| print("trainer_state.json not found") | |
| return | |
| state_path = cps[-1] | |
| state = json.load(open(state_path, encoding="utf-8")) | |
| logs = state["log_history"] | |
| # First 100 steps = "baseline" (model hasn't really learned yet) | |
| # Last 100 steps = "trained" (model is converged) | |
| early = [l["reward"] for l in logs if l["step"] <= 100] | |
| middle = [l["reward"] for l in logs if 100 < l["step"] <= 750] | |
| late = [l["reward"] for l in logs if l["step"] > 1400] | |
| early_avg = sum(early) / len(early) if early else 0 | |
| middle_avg = sum(middle) / len(middle) if middle else 0 | |
| late_avg = sum(late) / len(late) if late else 0 | |
| early_max = max(early) if early else 0 | |
| middle_max = max(middle) if middle else 0 | |
| late_max = max(late) if late else 0 | |
| fig, ax = plt.subplots(figsize=(10, 5)) | |
| phases = ["Early training\n(steps 1-100)", "Mid training\n(steps 101-750)", "Late training\n(steps 1401-1500)"] | |
| avgs = [early_avg, middle_avg, late_avg] | |
| maxs = [early_max, middle_max, late_max] | |
| x = list(range(len(phases))) | |
| bw = 0.35 | |
| bars1 = ax.bar([i - bw/2 for i in x], avgs, width=bw, color="#3730a3", label="Average reward") | |
| bars2 = ax.bar([i + bw/2 for i in x], maxs, width=bw, color="#16a34a", label="Best in phase") | |
| ax.set_xticks(x) | |
| ax.set_xticklabels(phases, fontsize=11) | |
| ax.set_ylabel("Reward", fontsize=11) | |
| ax.set_title("Reward improvement across training phases", fontsize=13, fontweight="bold") | |
| ax.legend() | |
| ax.grid(True, alpha=0.3, axis="y") | |
| # Annotate values on bars | |
| for bars in (bars1, bars2): | |
| for b in bars: | |
| ax.text(b.get_x() + b.get_width()/2, b.get_height() + 0.005, | |
| f"{b.get_height():.3f}", ha="center", va="bottom", fontsize=9) | |
| # Improvement annotation | |
| improvement = ((late_avg - early_avg) / early_avg) * 100 if early_avg else 0 | |
| ax.text(0.98, 0.98, | |
| f"Late vs early avg: {improvement:+.1f}%\nBest reward seen: {max(maxs):.3f}", | |
| transform=ax.transAxes, fontsize=10, verticalalignment="top", horizontalalignment="right", | |
| bbox=dict(boxstyle="round,pad=0.5", facecolor="#fef3c7", edgecolor="#fbbf24")) | |
| out = ROOT / "outputs" / "baseline_vs_trained.png" | |
| fig.tight_layout() | |
| fig.savefig(out, dpi=140) | |
| plt.close(fig) | |
| print(f"wrote {out}") | |
| if __name__ == "__main__": | |
| main() | |