Spaces:
Sleeping
Sleeping
File size: 3,165 Bytes
528c3d3 5c701e4 528c3d3 5c701e4 528c3d3 5c701e4 528c3d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """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()
|