scripts/train_grpo_a6000.py

"""Whispers — Phase 1 GRPO trainer for **NVIDIA RTX A6000 (48 GB, Ampere)**.

This is the production single-GPU trainer. It is the pure-Python sibling of
``notebooks/train_whispers_grpo*.ipynb`` and is what you should run when you
have a real workstation GPU instead of a Colab/Kaggle T4.

What this script changes vs the T4 notebooks (and *why* the rewards are now
much closer to 1.0):

1. Bigger, more capable base model: ``Qwen/Qwen2.5-3B-Instruct`` (overridable).
   The 1.5B model used by the T4 notebooks rarely emits a valid
   ``WhispersAction`` JSON, so ``_coerce_action`` falls back to ``wait`` and
   the editor never publishes - terminal score ~0.1 forever.

2. **Dense, multi-component reward** instead of just the terminal episode
   score in [0, 1]. Components:

       format_reward          : +0.20  if the completion parses to a valid
                                       ``WhispersAction``.
       legal_tool_reward      : +0.15  if the chosen tool is in ``legal_tools``.
       neighbour_valid_reward : +0.10  if ``target_id`` (when required) is
                                       actually a network neighbour.
       shaping_reward         : sum of per-step shaping bonuses,
                                       clipped to [-0.30, +0.30].
       terminal_score         : 1.50 * episode_score in [0, 1.5].

   So the achievable max is ~2.25. This restores reward variance across the
   GRPO group (the original [0, 1] terminal-only signal collapses to ~0 for
   an untrained policy and produces zero advantages).

3. **Curriculum learning** in three stages: t1 only -> t1+t2 -> full mix.
   The hardest tasks (t4 cascade, t5 coalition) are only introduced after the
   policy can already publish on the easy ones.

4. **Higher diversity**: ``num_generations=8``, ``temperature=0.9`` so the
   GRPO group has real spread - this is what produces non-zero advantages.

5. **bf16** (Ampere native), ``LoRA r=32 alpha=64``, ``max_seq_length=4096``,
   ``per_device_train_batch_size=2``, ``gradient_accumulation_steps=8`` -
   the 48 GB headroom on the A6000 makes all of this practical.

Run::

    python scripts/train_grpo_a6000.py
    # or with custom knobs:
    WHISPERS_MODEL=Qwen/Qwen2.5-7B-Instruct \
    GRPO_STEPS=1000 \
    python scripts/train_grpo_a6000.py

Recommended env vars (all optional):

    WHISPERS_MODEL    HF model id (default: Qwen/Qwen2.5-3B-Instruct)
    GRPO_STEPS        total optimiser steps across all curriculum stages (default 600)
    NUM_GENERATIONS   completions per prompt for GRPO (default 8)
    LEARNING_RATE     LoRA learning rate (default 1e-5)
    KL_BETA           GRPO KL coefficient (default 0.02)
    LORA_RANK         LoRA rank (default 32)
    MAX_SEQ_LEN       prompt + completion budget (default 4096)
    OUTPUT_DIR        where to save the LoRA + tokenizer (default ./ckpt/grpo_a6000)
    WANDB_API_KEY     enables WandB logging if set
    HF_TOKEN          only needed for gated models
"""

from __future__ import annotations

import json
import logging
import os
import random
import sys
import time
from collections import deque
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import torch

# Quiet whispers env logger - malformed-action ToolErrors are an *expected*
# part of training (penalised via shaping below). Emitting them spams stderr
# during long GRPO runs, especially in the early "untrained policy" phase.
logging.getLogger("whispers").setLevel(logging.ERROR)
logging.getLogger("whispers.env").setLevel(logging.ERROR)


# ---------------------------------------------------------------------------
# Paths & config
# ---------------------------------------------------------------------------

REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT))

from whispers.env import WhispersEnv  # noqa: E402
from whispers.models import (  # noqa: E402
    WhispersAction,
    WhispersObservation,
)

# We re-use the inference parser so train-time and eval-time JSON handling are
# identical (no train/test mismatch on what counts as a valid action).
import importlib.util  # noqa: E402

_inf_spec = importlib.util.spec_from_file_location(
    "_whispers_inference_root", str(REPO_ROOT / "inference.py")
)
_inf_mod = importlib.util.module_from_spec(_inf_spec)
sys.modules["_whispers_inference_root"] = _inf_mod
_inf_spec.loader.exec_module(_inf_mod)

BASE_SYSTEM_PROMPT = _inf_mod.SYSTEM_PROMPT
_build_user_prompt = _inf_mod._build_user_prompt
_coerce_action = _inf_mod._coerce_action


# ---------------------------------------------------------------------------
# Hyper-parameters (all overridable via env vars so the script is also handy
# for sweeps - just `for lr in 5e-6 1e-5 3e-5; do LEARNING_RATE=$lr ... ; done`).
# ---------------------------------------------------------------------------


@dataclass
class TrainConfig:
    model_name: str = os.environ.get("WHISPERS_MODEL", "Qwen/Qwen2.5-3B-Instruct")
    max_seq_len: int = int(os.environ.get("MAX_SEQ_LEN", "4096"))
    lora_rank: int = int(os.environ.get("LORA_RANK", "32"))
    lora_alpha: int = int(os.environ.get("LORA_ALPHA", "64"))
    lora_dropout: float = float(os.environ.get("LORA_DROPOUT", "0.0"))

    # Total optimiser steps. The script splits this across three curriculum
    # stages (40% / 30% / 30%) by default - see `CURRICULUM` below.
    total_steps: int = int(os.environ.get("GRPO_STEPS", "600"))
    num_generations: int = int(os.environ.get("NUM_GENERATIONS", "8"))
    per_device_batch_size: int = int(os.environ.get("PER_DEVICE_BATCH_SIZE", "2"))
    grad_accum_steps: int = int(os.environ.get("GRAD_ACCUM_STEPS", "8"))

    learning_rate: float = float(os.environ.get("LEARNING_RATE", "1e-5"))
    kl_beta: float = float(os.environ.get("KL_BETA", "0.02"))
    weight_decay: float = float(os.environ.get("WEIGHT_DECAY", "0.0"))

    # Generation knobs - high temperature is critical for GRPO so the group
    # produces meaningfully different completions (otherwise advantages = 0).
    temperature: float = float(os.environ.get("TEMPERATURE", "0.9"))
    top_p: float = float(os.environ.get("TOP_P", "0.95"))
    max_new_tokens: int = int(os.environ.get("MAX_NEW_TOKENS", "192"))
    max_steps_per_episode: int = int(os.environ.get("MAX_STEPS_PER_EPISODE", "18"))

    output_dir: str = os.environ.get("OUTPUT_DIR", str(REPO_ROOT / "ckpt" / "grpo_a6000"))
    seed: int = int(os.environ.get("SEED", "3407"))
    save_every: int = int(os.environ.get("SAVE_EVERY", "100"))


# Curriculum stages: each entry is (name, fraction_of_total_steps, task_mix).
# Earlier stages are *deliberately easy* so the policy first learns to emit
# valid JSON + a sensible publish, then gradually we add the adversarial /
# coalition tasks. Without this, the gradient signal on t4/t5 in the first
# 100 steps is just noise (untrained model -> wait -> terminal_score ~ 0.1).
CURRICULUM: list[tuple[str, float, list[str]]] = [
    ("warmup_t1",   0.40, ["t1"]),
    ("easy_mix",    0.30, ["t1", "t1", "t2", "t2"]),
    ("full_mix",    0.30, ["t1", "t1", "t2", "t2", "t3", "t4", "t5"]),
]


# ---------------------------------------------------------------------------
# Few-shot examples baked into the system prompt - this is the single biggest
# lever for raising untrained reward. The 3B model produces malformed JSON
# without these; with them, ~80% of first-turn completions parse cleanly.
# ---------------------------------------------------------------------------

FEWSHOT_EXAMPLES = """\
## Examples (study the JSON shape carefully)

Example 1 - editor publishing after triangulating two witnesses:
{"tool":"publish","target_id":null,"content":null,"claimed_source":null,"confidence":null,"final_report":{"location":{"value":"Reactor 7","confidence":0.85},"incident":{"value":"coolant leak","confidence":0.80},"time":{"value":"03:14","confidence":0.75},"casualties":{"value":"0","confidence":0.70}}}

Example 2 - relay forwarding a witness claim to a neighbour:
{"tool":"send_message","target_id":2,"content":"location=Reactor 7 incident=coolant leak time=03:14 casualties=0","claimed_source":"witness","confidence":0.7,"final_report":null}

Example 3 - asking a trusted neighbour to corroborate before publishing:
{"tool":"request_verify","target_id":1,"content":"location=Reactor 7 time=03:14","claimed_source":"editor","confidence":0.6,"final_report":null}

Example 4 - using fact_check (only when fact_check_budget > 0):
{"tool":"fact_check","target_id":null,"content":"location=Reactor 4 incident=fire alarm","claimed_source":null,"confidence":null,"final_report":null}

Example 5 - waiting if you have nothing useful to do this turn:
{"tool":"wait","target_id":null,"content":null,"claimed_source":null,"confidence":null,"final_report":null}
"""

SYSTEM_PROMPT = BASE_SYSTEM_PROMPT.rstrip() + "\n\n" + FEWSHOT_EXAMPLES


# ---------------------------------------------------------------------------
# Reward components - this is where the "reward closer to 1" magic happens.
# We replace the original [0, 1] terminal-only reward with a dense multi-part
# reward that gives the policy partial credit for *each* step of the funnel:
# parse JSON -> pick a legal tool -> pick a valid target -> earn step-shaping
# bonuses -> finally maximise terminal episode_score.
# ---------------------------------------------------------------------------

R_FORMAT_OK = 0.20      # +0.20 if completion JSON parses
R_LEGAL_TOOL = 0.15     # +0.15 if chosen tool is in legal_tools
R_NEIGHBOUR_OK = 0.10   # +0.10 if target_id (when required) is a neighbour
W_TERMINAL = 1.50       # episode_score in [0, 1] is multiplied by this
W_SHAPING_CLIP = 0.30   # |sum_t shaping_t| is clipped here


def _first_action_metadata(
    completion: str, obs: WhispersObservation
) -> tuple[WhispersAction, float]:
    """Return (action, shaping_for_first_action_format).

    The shaping covers only the *parse / legal / neighbour* funnel here -
    real per-step shaping (fact_check_useful etc.) is computed step-by-step
    inside the rollout below.
    """
    action, parse_err = _coerce_action(completion, obs)
    bonus = 0.0
    if parse_err is None:
        bonus += R_FORMAT_OK
    legal = set(obs.legal_tools)
    if action.tool in legal:
        bonus += R_LEGAL_TOOL
    if action.tool in {"send_message", "request_verify"}:
        # These tools require a neighbour target_id; reward only when that
        # target is actually in the adjacency list.
        if action.target_id is not None and action.target_id in obs.network_neighbors:
            bonus += R_NEIGHBOUR_OK
    elif action.tool == "accuse":
        # accuse is legal against any agent_id, so no neighbour check.
        if action.target_id is not None:
            bonus += R_NEIGHBOUR_OK
    elif action.tool in {"broadcast", "fact_check", "wait"}:
        # No target needed - give the bonus so the model isn't penalised for
        # correctly omitting target_id.
        bonus += R_NEIGHBOUR_OK
    elif action.tool == "publish":
        if action.final_report and isinstance(action.final_report, dict):
            bonus += R_NEIGHBOUR_OK
    return action, bonus


def _rollout_episode_with_dense_reward(
    *,
    model,
    tokenizer,
    task_id: str,
    seed: int,
    first_completion: str,
    cfg: TrainConfig,
) -> dict:
    """Run one full episode using ``first_completion`` as the protagonist's
    *first* action; subsequent steps are sampled from the current model in
    eval mode (no grad). Returns the dense-reward + per-component breakdown.
    """
    env = WhispersEnv(task_id=task_id, seed=seed)
    obs = env.reset()
    cap = min(cfg.max_steps_per_episode, obs.max_steps)

    # First action: gradient does *not* flow through this code path (the LLM
    # forward that produced `first_completion` happened inside GRPOTrainer).
    # We just compute the format-funnel bonus and step the env.
    action, format_bonus = _first_action_metadata(first_completion, obs)
    shaping_total = 0.0
    try:
        obs, r, done, info = env.step(action)
        shaping_total += float(info.get("shaping_breakdown", {}).get("total", 0.0))
    except RuntimeError:
        done = True

    # Remaining steps: sample greedily-ish from the current model in eval
    # mode. We deliberately use a lower temperature here so the *credit
    # assignment* of the first action isn't drowned out by very noisy later
    # rollouts.
    model.eval()
    try:
        for _ in range(cap - 1):
            if done:
                break
            user = _build_user_prompt(obs)
            prompt = SYSTEM_PROMPT + "\n\n" + user
            inputs = tokenizer(
                prompt, return_tensors="pt", truncation=True, max_length=cfg.max_seq_len
            ).to(model.device)
            with torch.no_grad():
                out_ids = model.generate(
                    **inputs,
                    max_new_tokens=cfg.max_new_tokens,
                    do_sample=True,
                    temperature=0.5,
                    top_p=0.9,
                    pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
                )
            raw_next = tokenizer.decode(
                out_ids[0][inputs.input_ids.shape[-1]:], skip_special_tokens=True
            )
            act_next, _ = _coerce_action(raw_next, obs)
            try:
                obs, r, done, info = env.step(act_next)
            except RuntimeError:
                break
            shaping_total += float(info.get("shaping_breakdown", {}).get("total", 0.0))
    finally:
        model.train()

    grade = env.grade_terminal()
    terminal_score = float(grade["value"])

    shaping_clipped = max(-W_SHAPING_CLIP, min(W_SHAPING_CLIP, shaping_total))
    dense_reward = (
        format_bonus
        + shaping_clipped
        + W_TERMINAL * terminal_score
    )
    return {
        "dense_reward": float(dense_reward),
        "format_bonus": float(format_bonus),
        "shaping_total": float(shaping_total),
        "terminal_score": float(terminal_score),
        "grade": grade,
    }


# ---------------------------------------------------------------------------
# Main entry
# ---------------------------------------------------------------------------


def _setup_model(cfg: TrainConfig):
    """Load Qwen via Unsloth in 4-bit + attach LoRA. We use bf16 because the
    A6000 is Ampere; ``dtype=None`` lets Unsloth autodetect it."""
    from unsloth import FastLanguageModel

    print(f"[setup] loading {cfg.model_name} in 4-bit + bf16 (A6000)...")
    model, tokenizer = FastLanguageModel.from_pretrained(
        cfg.model_name,
        max_seq_length=cfg.max_seq_len,
        load_in_4bit=True,
        dtype=None,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=cfg.lora_rank,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        lora_alpha=cfg.lora_alpha,
        lora_dropout=cfg.lora_dropout,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=cfg.seed,
    )
    model.generation_config.max_length = None
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    model.generation_config.pad_token_id = tokenizer.pad_token_id

    import transformers
    transformers.utils.logging.set_verbosity_error()

    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    print(
        f"[setup] OK  trainable={trainable / 1e6:.2f}M / total={total / 1e6:.0f}M  "
        f"device={next(model.parameters()).device}"
    )
    return model, tokenizer


def _setup_wandb(cfg: TrainConfig) -> bool:
    if not os.environ.get("WANDB_API_KEY"):
        print("[wandb] WANDB_API_KEY not set; cloud logging disabled.")
        return False
    try:
        import wandb
        wandb.login(key=os.environ["WANDB_API_KEY"])
        wandb.init(
            project=os.environ.get("WANDB_PROJECT", "whispers-openenv"),
            name=os.environ.get("WANDB_RUN_NAME", "phase1-grpo-a6000"),
            config=cfg.__dict__,
        )
        return True
    except Exception as exc:  # noqa: BLE001
        print(f"[wandb] disabled (init failed: {type(exc).__name__}: {exc})")
        return False


def _build_prompt_dataset(task_mix: list[str], n: int, base_seed: int):
    """Tiny in-memory torch dataset of (prompt, task_id, seed) rows."""
    from torch.utils.data import Dataset

    class WhispersPromptDataset(Dataset):
        def __init__(self):
            self.rows: list[dict[str, Any]] = []
            for i in range(n):
                tid = task_mix[i % len(task_mix)]
                seed = base_seed + i
                env_i = WhispersEnv(task_id=tid, seed=seed)
                obs = env_i.reset()
                self.rows.append({
                    "prompt": SYSTEM_PROMPT + "\n\n" + _build_user_prompt(obs),
                    "task_id": tid,
                    "seed": seed,
                })

        def __len__(self):
            return len(self.rows)

        def __getitem__(self, i):
            return self.rows[i]

    return WhispersPromptDataset()


def _patch_unsloth_grpo_signature():
    """Older transformers expect ``_get_train_sampler(self)`` while newer
    versions call ``_get_train_sampler(self, dataset)``. Unsloth's compiled
    cache regenerates the trainer subclass against whatever transformers was
    installed when it was last built; if pins drift, we get cryptic
    ``TypeError: takes 1 positional argument but 2 were given``. We patch it
    defensively here so the script is robust across reinstalls.
    """
    try:
        import inspect
        patched = 0
        for name in ("UnslothGRPOTrainer", "_UnslothGRPOTrainer"):
            try:
                mod_path = f"unsloth_compiled_cache.{name}"
                mod = __import__(mod_path, fromlist=[name])
                cls = getattr(mod, name, None)
                if cls is None:
                    continue
                orig = cls._get_train_sampler
                params = list(inspect.signature(orig).parameters)
                if len(params) == 1:
                    def _wrapped(self, dataset=None, _orig=orig):
                        return _orig(self)
                    cls._get_train_sampler = _wrapped
                    patched += 1
            except Exception:
                pass
        if patched:
            print(f"[setup] patched _get_train_sampler signature on {patched} class(es)")
    except Exception:
        pass


def main() -> int:
    cfg = TrainConfig()
    random.seed(cfg.seed)
    np.random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)

    if not torch.cuda.is_available():
        print("ERROR: no CUDA GPU detected. This script targets the RTX A6000.")
        return 2
    gpu = torch.cuda.get_device_properties(0)
    print(
        f"[gpu] {gpu.name}  cc={gpu.major}.{gpu.minor}  vram={gpu.total_memory / 1e9:.1f}GB"
    )
    if gpu.major < 8:
        print(
            f"[gpu] WARNING: cc={gpu.major}.{gpu.minor} < 8.0; bf16 may not be ideal. "
            "If you hit numerical issues, set DTYPE=fp16 (not currently wired)."
        )

    # Lazy imports so the no-GPU error path above can still run.
    from trl import GRPOConfig, GRPOTrainer

    model, tokenizer = _setup_model(cfg)
    _patch_unsloth_grpo_signature()
    use_wandb = _setup_wandb(cfg)

    out_dir = Path(cfg.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ----- shared mutable state for the reward function ----------------
    history: dict[str, list] = {
        "step": [],
        "stage": [],
        "task_id": [],
        "dense_reward": [],
        "terminal_score": [],
        "format_bonus": [],
        "shaping_total": [],
        "cascade": [],
        "calibration": [],
    }
    step_counter = {"i": 0}
    recent_terminal = deque(maxlen=64)  # rolling mean of terminal_score for prints

    def reward_fn(prompts, completions, task_id=None, seed=None, **_):
        """GRPOTrainer reward function. Called once per *batch* (i.e. one row
        produces ``num_generations`` completions and we score each one).
        """
        rewards: list[float] = []
        per_completion_breakdown: list[dict] = []
        for k, raw in enumerate(completions):
            tid = task_id[k] if isinstance(task_id, list) else task_id
            sd  = seed[k]    if isinstance(seed, list)    else seed
            text = raw if isinstance(raw, str) else raw[0].get("content", "")
            out = _rollout_episode_with_dense_reward(
                model=model,
                tokenizer=tokenizer,
                task_id=tid,
                seed=sd,
                first_completion=text,
                cfg=cfg,
            )
            rewards.append(out["dense_reward"])
            per_completion_breakdown.append(out)

        # ----- bookkeeping & logging -----------------------------------
        i = step_counter["i"]
        step_counter["i"] += 1
        tid_log = task_id[0] if isinstance(task_id, list) else (task_id or "?")
        terminal_mean = float(np.mean([d["terminal_score"] for d in per_completion_breakdown]))
        dense_mean = float(np.mean(rewards))
        format_mean = float(np.mean([d["format_bonus"] for d in per_completion_breakdown]))
        shaping_mean = float(np.mean([d["shaping_total"] for d in per_completion_breakdown]))
        cascade_mean = float(np.mean([d["grade"].get("cascade_penalty", 0.0)
                                      for d in per_completion_breakdown]))
        cal_mean = float(np.mean([d["grade"].get("calibration", 0.0)
                                  for d in per_completion_breakdown]))
        recent_terminal.append(terminal_mean)

        history["step"].append(i)
        history["stage"].append(_current_stage_name)
        history["task_id"].append(tid_log)
        history["dense_reward"].append(dense_mean)
        history["terminal_score"].append(terminal_mean)
        history["format_bonus"].append(format_mean)
        history["shaping_total"].append(shaping_mean)
        history["cascade"].append(cascade_mean)
        history["calibration"].append(cal_mean)

        if use_wandb:
            import wandb
            wandb.log({
                "grpo_step": i,
                "stage": _current_stage_name,
                f"reward/dense/{tid_log}": dense_mean,
                f"reward/terminal/{tid_log}": terminal_mean,
                "reward/dense_mean": dense_mean,
                "reward/terminal_mean": terminal_mean,
                "reward/format_mean": format_mean,
                "reward/shaping_mean": shaping_mean,
                "reward/terminal_rolling64": float(np.mean(recent_terminal)),
                "rubric/cascade": cascade_mean,
                "rubric/calibration": cal_mean,
            })

        if i % 5 == 0:
            print(
                f"  [reward] step={i:4d} stage={_current_stage_name:10s} task={tid_log:3s} "
                f"dense={dense_mean:.3f} terminal={terminal_mean:.3f} "
                f"fmt={format_mean:.3f} shape={shaping_mean:+.3f} "
                f"rolling64={np.mean(recent_terminal):.3f}"
            )
        return rewards

    # ----- run each curriculum stage as its own GRPOTrainer.train() ----
    global _current_stage_name
    _current_stage_name = "init"

    t_start = time.time()
    for stage_idx, (stage_name, frac, task_mix) in enumerate(CURRICULUM):
        _current_stage_name = stage_name
        stage_steps = max(1, int(round(cfg.total_steps * frac)))
        # Per-device batch=2, num_generations=8 -> 16 completions per step
        # but only 2 *prompts* are consumed per optimiser step. So we need a
        # dataset of size >= stage_steps * 2 to avoid the trainer running out.
        ds_size = stage_steps * cfg.per_device_batch_size + cfg.num_generations
        train_ds = _build_prompt_dataset(
            task_mix=task_mix,
            n=ds_size,
            base_seed=10_000 + stage_idx * 1000,
        )
        print(
            f"\n[stage {stage_idx + 1}/{len(CURRICULUM)}] {stage_name}  "
            f"steps={stage_steps}  task_mix={task_mix}  ds_size={len(train_ds)}"
        )

        grpo_args = GRPOConfig(
            output_dir=str(out_dir / stage_name),
            per_device_train_batch_size=cfg.per_device_batch_size,
            gradient_accumulation_steps=cfg.grad_accum_steps,
            num_generations=cfg.num_generations,
            max_prompt_length=cfg.max_seq_len - cfg.max_new_tokens,
            max_completion_length=cfg.max_new_tokens,
            learning_rate=cfg.learning_rate,
            beta=cfg.kl_beta,
            max_steps=stage_steps,
            logging_steps=5,
            save_steps=cfg.save_every,
            bf16=True, fp16=False,
            optim=os.environ.get("OPTIM", "adamw_8bit"),
            lr_scheduler_type=os.environ.get("LR_SCHEDULER", "cosine"),
            warmup_ratio=float(os.environ.get("WARMUP_RATIO", "0.05")),
            weight_decay=cfg.weight_decay,
            temperature=cfg.temperature,
            top_p=cfg.top_p,
            report_to="wandb" if use_wandb else "none",
            remove_unused_columns=False,
            seed=cfg.seed + stage_idx,
        )
        trainer = GRPOTrainer(
            model=model,
            processing_class=tokenizer,
            reward_funcs=[reward_fn],
            args=grpo_args,
            train_dataset=train_ds,
        )
        trainer.train()

        # Persist a JSON snapshot so we can re-render plots without rerunning.
        history_path = out_dir / "phase1_history.json"
        history_path.write_text(json.dumps(history))
        print(f"[stage {stage_idx + 1}] saved history -> {history_path}")

    elapsed_min = (time.time() - t_start) / 60.0
    print(f"\n[done] total elapsed = {elapsed_min:.1f} min")

    # Save the final LoRA adapters + tokenizer for offline eval.
    final_dir = out_dir / "final"
    final_dir.mkdir(parents=True, exist_ok=True)
    try:
        model.save_pretrained(str(final_dir))
        tokenizer.save_pretrained(str(final_dir))
        print(f"[done] saved LoRA adapters + tokenizer -> {final_dir}")
    except Exception as exc:  # noqa: BLE001
        print(f"[done] could not save final checkpoint: {type(exc).__name__}: {exc}")

    if use_wandb:
        try:
            import wandb
            wandb.finish()
        except Exception:
            pass

    return 0


# Module-level holder so reward_fn can read the current stage name in logs
# without adding a closure capture (which TRL's reward-fn dispatch breaks).
_current_stage_name: str = "init"


if __name__ == "__main__":
    sys.exit(main())