Spaces:

Rushhaabhhh
/

HONEST-RL-Calibrator

Running

HONEST-RL-Calibrator / bin /plot_training_curves.py

HONEST-RL-Calibrator-v0

3040767 verified 6 months ago

11.9 kB

#!/usr/bin/env python
"""Render committed training-curve evidence (loss, reward, KL).
Two operating modes
-------------------
1. ``--trainer-state PATH`` — read TRL's canonical
   ``trainer_state.json`` (saved automatically inside the trainer's
   ``output_dir``) and emit *real* curves. This is the path
   ``docs/RUNBOOK.md`` instructs operators to take after a full GPU run::
       python bin/plot_training_curves.py \
              --trainer-state ./honest-qwen3b-grpo/trainer_state.json \
              --label "qwen3b · 350 steps · L4"
2. ``--demo`` — synthesise a labelled, deterministic, *representative*
   trajectory grounded in HONEST's actual reward formula:
   * Brier dominates: ``-1.5 * (c-y)^2``.
   * Initial overconfidence ⇒ reward ≈ -0.40.
   * Calibrated state ⇒ reward asymptotes near format bonus (+0.15).
   * KL is held below the early-stop threshold (0.5) by the adaptive-beta
     callback (see ``training/train_grpo.py``).
   Demo plots are committed to the repo so judges / readers see the
   *shape* of the training curve even before they reproduce the run.
   They are clearly tagged "DEMO TRACE" in the title and watermark.
Usage examples
--------------
    # Regenerate demo plots committed to docs/training/
    python bin/plot_training_curves.py --demo --out docs/training
    # After a real run, overwrite with real data
    python bin/plot_training_curves.py \
        --trainer-state ./honest-qwen3b-grpo/trainer_state.json \
        --out docs/training \
        --label "qwen3b · 350 steps · L4"
"""

from __future__ import annotations

import argparse
import json
import math
import random
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import matplotlib

matplotlib.use("Agg")  # non-interactive; never tries to open an X display
import matplotlib.pyplot as plt
import numpy as np


# ---------------------------------------------------------------------------
# trainer_state.json reader
# ---------------------------------------------------------------------------


def _read_trainer_state(path: Path) -> List[Dict[str, Any]]:
    """Return TRL's ``log_history`` array (one dict per logged step)."""
    if not path.exists():
        raise FileNotFoundError(
            f"trainer_state.json not found at {path}. "
            "Run a real training session first or pass --demo."
        )
    with path.open(encoding="utf-8") as fh:
        state = json.load(fh)
    history = state.get("log_history")
    if not isinstance(history, list) or not history:
        raise ValueError(f"{path} contains no log_history entries.")
    return history


def _series(history: List[Dict[str, Any]], key: str) -> Tuple[List[int], List[float]]:
    """Extract ``(steps, values)`` for a given metric key, skipping eval rows."""
    steps: List[int] = []
    values: List[float] = []
    for row in history:
        v = row.get(key)
        s = row.get("step") or row.get("global_step")
        if v is None or s is None:
            continue
        try:
            values.append(float(v))
            steps.append(int(s))
        except (TypeError, ValueError):
            continue
    return steps, values


# ---------------------------------------------------------------------------
# Demonstration trajectory generator (deterministic, seeded)
# ---------------------------------------------------------------------------


def _demo_history(num_steps: int = 350, seed: int = 42) -> List[Dict[str, Any]]:
    """Synthesise a representative GRPO log_history.
    The shape follows the project's documented reward dynamics:
    * Reward
        - starts at the empirical Brier of an over-confident base model
          (``c≈0.95, y≈0.5`` ⇒ ``-1.5*(0.45)^2 ≈ -0.30``),
        - climbs as confidence aligns with empirical correctness,
        - asymptotes near the format bonus (``+0.15``) once the model is
          calibrated.
    * Loss (policy loss surrogate)
        - decays from ~0.50 to ~0.05 with noise, modulated by the cosine
          LR schedule's warmup-then-decay envelope.
    * KL(π||π_ref)
        - stays well below the ``KLEarlyStopCallback`` threshold (0.5)
          thanks to ``AdaptiveBetaCallback``; oscillates around 0.05–0.15.
    The trace is **strictly synthetic** and the plots produced from it
    are watermarked accordingly.
    """
    rng = random.Random(seed)
    np_rng = np.random.default_rng(seed)

    history: List[Dict[str, Any]] = []
    log_every = max(1, num_steps // 70)  # ~70 logged points

    for step in range(0, num_steps + 1, log_every):
        progress = step / max(1, num_steps)

        # Reward: -0.40 -> +0.12 with diminishing returns.
        # Sigmoid envelope matches a cosine-LR + Brier curriculum.
        env = 1.0 / (1.0 + math.exp(-6.0 * (progress - 0.35)))
        reward_mean = -0.40 + (0.55 * env)
        reward_noise = float(np_rng.normal(0.0, 0.06 * (1.0 - 0.4 * progress)))
        reward = reward_mean + reward_noise

        # Reward std drops as the policy concentrates.
        reward_std = max(0.05, 0.45 - 0.30 * progress + 0.05 * abs(reward_noise))

        # Policy loss: starts ~0.55, exponentially decays to ~0.05.
        loss_mean = 0.55 * math.exp(-3.0 * progress) + 0.06
        loss = max(0.0, loss_mean + float(np_rng.normal(0.0, 0.04 * (1.0 - 0.5 * progress))))

        # KL: small, with a mid-training bump that adaptive-beta tames.
        kl_mean = 0.04 + 0.10 * math.exp(-12.0 * (progress - 0.55) ** 2)
        kl = max(0.005, kl_mean + float(np_rng.normal(0.0, 0.015)))

        # Cosine LR schedule with 5% warmup, peak 2e-6.
        peak_lr = 2.0e-6
        warmup_frac = 0.05
        if progress < warmup_frac:
            lr = peak_lr * (progress / warmup_frac)
        else:
            after = (progress - warmup_frac) / (1.0 - warmup_frac)
            lr = peak_lr * 0.5 * (1.0 + math.cos(math.pi * after))

        history.append(
            {
                "step": step,
                "epoch": progress,
                "loss": round(loss, 4),
                "reward": round(reward, 4),
                "reward_std": round(reward_std, 4),
                "kl": round(kl, 4),
                "learning_rate": float(f"{lr:.3e}"),
            }
        )

    rng.shuffle  # silence unused-import linter if any; keep rng for future seeding
    return history


# ---------------------------------------------------------------------------
# Plotting helpers
# ---------------------------------------------------------------------------


def _annotate_demo(ax: plt.Axes) -> None:
    """Stamp 'DEMO TRACE' so anyone looking at the PNG knows it is synthetic."""
    ax.text(
        0.99,
        0.02,
        "DEMO TRACE — replace via\nbin/plot_training_curves.py --trainer-state ...",
        transform=ax.transAxes,
        ha="right",
        va="bottom",
        fontsize=8,
        color="#999999",
        alpha=0.85,
        fontstyle="italic",
    )


def _smooth(values: List[float], window: int = 7) -> np.ndarray:
    """Centred moving-average for a calmer curve overlay."""
    arr = np.asarray(values, dtype=float)
    if arr.size < 2 or window <= 1:
        return arr
    pad = window // 2
    padded = np.pad(arr, pad, mode="edge")
    kernel = np.ones(window) / window
    return np.convolve(padded, kernel, mode="valid")


def _plot_curve(
    steps: List[int],
    values: List[float],
    *,
    out_path: Path,
    title: str,
    ylabel: str,
    color: str,
    is_demo: bool,
    label: Optional[str],
    band: Optional[Tuple[List[float], List[float]]] = None,
    ylim: Optional[Tuple[float, float]] = None,
) -> None:
    """Render a single training-curve PNG and write it to disk."""
    fig, ax = plt.subplots(figsize=(8.0, 4.5), dpi=150)
    ax.plot(steps, values, color=color, alpha=0.35, linewidth=1.2, label="raw")
    smooth = _smooth(values, window=max(3, len(values) // 30 or 3))
    ax.plot(steps, smooth, color=color, linewidth=2.2, label="smoothed")

    if band is not None:
        lo, hi = band
        ax.fill_between(steps, lo, hi, color=color, alpha=0.10, label="±1 std")

    suffix = "" if not label else f" — {label}"
    ax.set_title(f"{title}{suffix}", fontsize=12, weight="bold")
    ax.set_xlabel("training step")
    ax.set_ylabel(ylabel)
    ax.grid(True, alpha=0.25)
    if ylim is not None:
        ax.set_ylim(*ylim)
    ax.legend(loc="best", fontsize=8, framealpha=0.85)

    if is_demo:
        _annotate_demo(ax)

    fig.tight_layout()
    out_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out_path, bbox_inches="tight")
    plt.close(fig)
    print(f"[plot_training_curves] wrote {out_path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    src = parser.add_mutually_exclusive_group(required=True)
    src.add_argument(
        "--trainer-state",
        type=Path,
        help="Path to TRL trainer_state.json (real run).",
    )
    src.add_argument(
        "--demo",
        action="store_true",
        help="Synthesise a labelled demonstration trajectory.",
    )
    parser.add_argument(
        "--out",
        type=Path,
        default=Path("docs/training"),
        help="Output directory for PNGs (default: docs/training).",
    )
    parser.add_argument(
        "--label",
        type=str,
        default=None,
        help="Suffix appended to plot titles (e.g. 'qwen3b · 350 steps · L4').",
    )
    parser.add_argument(
        "--demo-steps",
        type=int,
        default=350,
        help="Number of demo training steps (default: 350, matches qwen3b preset).",
    )
    parser.add_argument(
        "--demo-seed",
        type=int,
        default=42,
        help="RNG seed for the demo trace (deterministic across runs).",
    )
    args = parser.parse_args()

    if args.demo:
        history = _demo_history(num_steps=args.demo_steps, seed=args.demo_seed)
        is_demo = True
        label = args.label or "demo trace"
    else:
        history = _read_trainer_state(args.trainer_state)
        is_demo = False
        label = args.label

    out_dir = args.out

    # Reward
    steps, rewards = _series(history, "reward")
    if steps:
        _, std = _series(history, "reward_std")
        band = None
        if len(std) == len(rewards) and len(std) > 0:
            band = (
                [r - s for r, s in zip(rewards, std)],
                [r + s for r, s in zip(rewards, std)],
            )
        _plot_curve(
            steps,
            rewards,
            out_path=out_dir / "reward_curve.png",
            title="GRPO mean reward (Brier-shaped)",
            ylabel="reward = -1.5·(c-y)² + format/abstain",
            color="#1f77b4",
            is_demo=is_demo,
            label=label,
            band=band,
        )

    # Loss
    steps, losses = _series(history, "loss")
    if steps:
        _plot_curve(
            steps,
            losses,
            out_path=out_dir / "loss_curve.png",
            title="GRPO policy loss",
            ylabel="loss",
            color="#d62728",
            is_demo=is_demo,
            label=label,
        )

    # KL (optional but useful evidence of stability under adaptive beta)
    steps, kls = _series(history, "kl")
    if steps:
        _plot_curve(
            steps,
            kls,
            out_path=out_dir / "kl_curve.png",
            title="KL(π‖π_ref) — bounded by AdaptiveBetaCallback",
            ylabel="KL divergence",
            color="#2ca02c",
            is_demo=is_demo,
            label=label,
            ylim=(0.0, max(0.6, max(kls) * 1.2)),
        )


if __name__ == "__main__":
    main()

	#!/usr/bin/env python
	"""Render committed training-curve evidence (loss, reward, KL).

	Two operating modes
	-------------------

	1. ``--trainer-state PATH`` — read TRL's canonical
	``trainer_state.json`` (saved automatically inside the trainer's
	``output_dir``) and emit real curves. This is the path
	``docs/RUNBOOK.md`` instructs operators to take after a full GPU run::

	python bin/plot_training_curves.py \
	--trainer-state ./honest-qwen3b-grpo/trainer_state.json \
	--label "qwen3b · 350 steps · L4"

	2. ``--demo`` — synthesise a labelled, deterministic, representative
	trajectory grounded in HONEST's actual reward formula:

	* Brier dominates: ``-1.5 * (c-y)^2``.
	* Initial overconfidence ⇒ reward ≈ -0.40.
	* Calibrated state ⇒ reward asymptotes near format bonus (+0.15).
	* KL is held below the early-stop threshold (0.5) by the adaptive-beta
	callback (see ``training/train_grpo.py``).

	Demo plots are committed to the repo so judges / readers see the
	shape of the training curve even before they reproduce the run.
	They are clearly tagged "DEMO TRACE" in the title and watermark.

	Usage examples
	--------------

	# Regenerate demo plots committed to docs/training/
	python bin/plot_training_curves.py --demo --out docs/training

	# After a real run, overwrite with real data
	python bin/plot_training_curves.py \
	--trainer-state ./honest-qwen3b-grpo/trainer_state.json \
	--out docs/training \
	--label "qwen3b · 350 steps · L4"
	"""

	from __future__ import annotations

	import argparse
	import json
	import math
	import random
	from pathlib import Path
	from typing import Any, Dict, List, Optional, Tuple

	import matplotlib

	matplotlib.use("Agg") # non-interactive; never tries to open an X display
	import matplotlib.pyplot as plt
	import numpy as np


	# ---------------------------------------------------------------------------
	# trainer_state.json reader
	# ---------------------------------------------------------------------------


	def _read_trainer_state(path: Path) -> List[Dict[str, Any]]:
	"""Return TRL's ``log_history`` array (one dict per logged step)."""
	if not path.exists():
	raise FileNotFoundError(
	f"trainer_state.json not found at {path}. "
	"Run a real training session first or pass --demo."
	)
	with path.open(encoding="utf-8") as fh:
	state = json.load(fh)
	history = state.get("log_history")
	if not isinstance(history, list) or not history:
	raise ValueError(f"{path} contains no log_history entries.")
	return history


	def _series(history: List[Dict[str, Any]], key: str) -> Tuple[List[int], List[float]]:
	"""Extract ``(steps, values)`` for a given metric key, skipping eval rows."""
	steps: List[int] = []
	values: List[float] = []
	for row in history:
	v = row.get(key)
	s = row.get("step") or row.get("global_step")
	if v is None or s is None:
	continue
	try:
	values.append(float(v))
	steps.append(int(s))
	except (TypeError, ValueError):
	continue
	return steps, values


	# ---------------------------------------------------------------------------
	# Demonstration trajectory generator (deterministic, seeded)
	# ---------------------------------------------------------------------------


	def _demo_history(num_steps: int = 350, seed: int = 42) -> List[Dict[str, Any]]:
	"""Synthesise a representative GRPO log_history.

	The shape follows the project's documented reward dynamics:

	* Reward
	- starts at the empirical Brier of an over-confident base model
	(``c≈0.95, y≈0.5`` ⇒ ``-1.5*(0.45)^2 ≈ -0.30``),
	- climbs as confidence aligns with empirical correctness,
	- asymptotes near the format bonus (``+0.15``) once the model is
	calibrated.
	* Loss (policy loss surrogate)
	- decays from ~0.50 to ~0.05 with noise, modulated by the cosine
	LR schedule's warmup-then-decay envelope.
	* KL(π\|\|π_ref)
	- stays well below the ``KLEarlyStopCallback`` threshold (0.5)
	thanks to ``AdaptiveBetaCallback``; oscillates around 0.05–0.15.

	The trace is strictly synthetic and the plots produced from it
	are watermarked accordingly.
	"""
	rng = random.Random(seed)
	np_rng = np.random.default_rng(seed)

	history: List[Dict[str, Any]] = []
	log_every = max(1, num_steps // 70) # ~70 logged points

	for step in range(0, num_steps + 1, log_every):
	progress = step / max(1, num_steps)

	# Reward: -0.40 -> +0.12 with diminishing returns.
	# Sigmoid envelope matches a cosine-LR + Brier curriculum.
	env = 1.0 / (1.0 + math.exp(-6.0 * (progress - 0.35)))
	reward_mean = -0.40 + (0.55 * env)
	reward_noise = float(np_rng.normal(0.0, 0.06 * (1.0 - 0.4 * progress)))
	reward = reward_mean + reward_noise

	# Reward std drops as the policy concentrates.
	reward_std = max(0.05, 0.45 - 0.30 * progress + 0.05 * abs(reward_noise))

	# Policy loss: starts ~0.55, exponentially decays to ~0.05.
	loss_mean = 0.55 * math.exp(-3.0 * progress) + 0.06
	loss = max(0.0, loss_mean + float(np_rng.normal(0.0, 0.04 * (1.0 - 0.5 * progress))))

	# KL: small, with a mid-training bump that adaptive-beta tames.
	kl_mean = 0.04 + 0.10 * math.exp(-12.0 * (progress - 0.55) ** 2)
	kl = max(0.005, kl_mean + float(np_rng.normal(0.0, 0.015)))

	# Cosine LR schedule with 5% warmup, peak 2e-6.
	peak_lr = 2.0e-6
	warmup_frac = 0.05
	if progress < warmup_frac:
	lr = peak_lr * (progress / warmup_frac)
	else:
	after = (progress - warmup_frac) / (1.0 - warmup_frac)
	lr = peak_lr * 0.5 * (1.0 + math.cos(math.pi * after))

	history.append(
	{
	"step": step,
	"epoch": progress,
	"loss": round(loss, 4),
	"reward": round(reward, 4),
	"reward_std": round(reward_std, 4),
	"kl": round(kl, 4),
	"learning_rate": float(f"{lr:.3e}"),
	}
	)

	rng.shuffle # silence unused-import linter if any; keep rng for future seeding
	return history


	# ---------------------------------------------------------------------------
	# Plotting helpers
	# ---------------------------------------------------------------------------


	def _annotate_demo(ax: plt.Axes) -> None:
	"""Stamp 'DEMO TRACE' so anyone looking at the PNG knows it is synthetic."""
	ax.text(
	0.99,
	0.02,
	"DEMO TRACE — replace via\nbin/plot_training_curves.py --trainer-state ...",
	transform=ax.transAxes,
	ha="right",
	va="bottom",
	fontsize=8,
	color="#999999",
	alpha=0.85,
	fontstyle="italic",
	)


	def _smooth(values: List[float], window: int = 7) -> np.ndarray:
	"""Centred moving-average for a calmer curve overlay."""
	arr = np.asarray(values, dtype=float)
	if arr.size < 2 or window <= 1:
	return arr
	pad = window // 2
	padded = np.pad(arr, pad, mode="edge")
	kernel = np.ones(window) / window
	return np.convolve(padded, kernel, mode="valid")


	def _plot_curve(
	steps: List[int],
	values: List[float],
	*,
	out_path: Path,
	title: str,
	ylabel: str,
	color: str,
	is_demo: bool,
	label: Optional[str],
	band: Optional[Tuple[List[float], List[float]]] = None,
	ylim: Optional[Tuple[float, float]] = None,
	) -> None:
	"""Render a single training-curve PNG and write it to disk."""
	fig, ax = plt.subplots(figsize=(8.0, 4.5), dpi=150)
	ax.plot(steps, values, color=color, alpha=0.35, linewidth=1.2, label="raw")
	smooth = _smooth(values, window=max(3, len(values) // 30 or 3))
	ax.plot(steps, smooth, color=color, linewidth=2.2, label="smoothed")

	if band is not None:
	lo, hi = band
	ax.fill_between(steps, lo, hi, color=color, alpha=0.10, label="±1 std")

	suffix = "" if not label else f" — {label}"
	ax.set_title(f"{title}{suffix}", fontsize=12, weight="bold")
	ax.set_xlabel("training step")
	ax.set_ylabel(ylabel)
	ax.grid(True, alpha=0.25)
	if ylim is not None:
	ax.set_ylim(*ylim)
	ax.legend(loc="best", fontsize=8, framealpha=0.85)

	if is_demo:
	_annotate_demo(ax)

	fig.tight_layout()
	out_path.parent.mkdir(parents=True, exist_ok=True)
	fig.savefig(out_path, bbox_inches="tight")
	plt.close(fig)
	print(f"[plot_training_curves] wrote {out_path}")


	# ---------------------------------------------------------------------------
	# Main
	# ---------------------------------------------------------------------------


	def main() -> None:
	parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
	src = parser.add_mutually_exclusive_group(required=True)
	src.add_argument(
	"--trainer-state",
	type=Path,
	help="Path to TRL trainer_state.json (real run).",
	)
	src.add_argument(
	"--demo",
	action="store_true",
	help="Synthesise a labelled demonstration trajectory.",
	)
	parser.add_argument(
	"--out",
	type=Path,
	default=Path("docs/training"),
	help="Output directory for PNGs (default: docs/training).",
	)
	parser.add_argument(
	"--label",
	type=str,
	default=None,
	help="Suffix appended to plot titles (e.g. 'qwen3b · 350 steps · L4').",
	)
	parser.add_argument(
	"--demo-steps",
	type=int,
	default=350,
	help="Number of demo training steps (default: 350, matches qwen3b preset).",
	)
	parser.add_argument(
	"--demo-seed",
	type=int,
	default=42,
	help="RNG seed for the demo trace (deterministic across runs).",
	)
	args = parser.parse_args()

	if args.demo:
	history = _demo_history(num_steps=args.demo_steps, seed=args.demo_seed)
	is_demo = True
	label = args.label or "demo trace"
	else:
	history = _read_trainer_state(args.trainer_state)
	is_demo = False
	label = args.label

	out_dir = args.out

	# Reward
	steps, rewards = _series(history, "reward")
	if steps:
	_, std = _series(history, "reward_std")
	band = None
	if len(std) == len(rewards) and len(std) > 0:
	band = (
	[r - s for r, s in zip(rewards, std)],
	[r + s for r, s in zip(rewards, std)],
	)
	_plot_curve(
	steps,
	rewards,
	out_path=out_dir / "reward_curve.png",
	title="GRPO mean reward (Brier-shaped)",
	ylabel="reward = -1.5·(c-y)² + format/abstain",
	color="#1f77b4",
	is_demo=is_demo,
	label=label,
	band=band,
	)

	# Loss
	steps, losses = _series(history, "loss")
	if steps:
	_plot_curve(
	steps,
	losses,
	out_path=out_dir / "loss_curve.png",
	title="GRPO policy loss",
	ylabel="loss",
	color="#d62728",
	is_demo=is_demo,
	label=label,
	)

	# KL (optional but useful evidence of stability under adaptive beta)
	steps, kls = _series(history, "kl")
	if steps:
	_plot_curve(
	steps,
	kls,
	out_path=out_dir / "kl_curve.png",
	title="KL(π‖π_ref) — bounded by AdaptiveBetaCallback",
	ylabel="KL divergence",
	color="#2ca02c",
	is_demo=is_demo,
	label=label,
	ylim=(0.0, max(0.6, max(kls) * 1.2)),
	)


	if __name__ == "__main__":
	main()