training/train_grpo.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch>=2.2",
#   "transformers>=4.46",
#   "peft>=0.13",
#   "accelerate>=1.0",
#   "bitsandbytes>=0.43; platform_system != 'Darwin'",
#   "huggingface_hub>=0.26",
#   "trackio>=0.1.4",
#   "openenv-core[core]>=0.2.1",
#   "pydantic>=2.6",
#   "pydantic-settings>=2.0",
#   "fastapi>=0.115",
#   "uvicorn>=0.30",
#   "python-dotenv",
#   "openai>=1.40",
#   "matplotlib>=3.8",
# ]
# ///
"""
InvoiceGuard Round 2 - Trajectory-level GRPO trainer (HF Jobs UV script).

Trains a small instruction-tuned LM on the InvoiceGuard OpenEnv with a
hand-written multi-step GRPO loop:

  for each iteration over train tasks:
      sample G trajectories per task (stochastic policy)
      reward per trajectory  = env cumulative reward + alpha * grader_score
      advantage per trajectory = (reward - group_mean) / (group_std + eps)
      apply PPO-clipped policy gradient on every (obs, action) pair in
          each trajectory, weighted by that trajectory's advantage,
          regularised by KL against a frozen reference policy.

The trainer is deliberately small (no TRL GRPOTrainer dep) because TRL's
GRPO assumes single-turn rewards; our env is multi-turn agentic.

Launch on HF Jobs:
    See `invoice_guard/training/launch_hf_job.py` for the recommended
    submission flow (uploads `invoice_guard/` to a code repo on the Hub
    and points this script at it via INVOICEGUARD_CODE_REPO).

Run a tiny local smoke test (CPU/GPU, no Hub push):
    cd invoice_guard
    python -m training.train_grpo \
        --model-name Qwen/Qwen2.5-0.5B-Instruct \
        --num-iterations 1 --group-size 2 --max-train-tasks 2 \
        --no-push

Required env vars on HF Jobs:
    HF_TOKEN                -- write-scoped token (passed via `secrets=`)
    HF_USERNAME             -- pushes adapter to {HF_USERNAME}/{HUB_MODEL_ID}

Optional env vars:
    INVOICEGUARD_CODE_REPO  -- model/dataset repo containing the env code;
                               cloned into /tmp at startup if set
    HUB_MODEL_ID            -- name of the LoRA adapter repo to create
    BASE_MODEL              -- HF model id of the base policy
    TRACKIO_PROJECT         -- defaults to "invoiceguard-round2"
"""

from __future__ import annotations

import argparse
import json
import os
import random
import subprocess
import sys
import time
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import List, Optional

# -----------------------------------------------------------------------------
# 0. Bootstrap: make `invoice_guard/` importable on HF Jobs.
# -----------------------------------------------------------------------------

def _hf_token() -> Optional[str]:
    return os.environ.get("HF_TOKEN") or os.environ.get("API_TOKEN_HF")


def _bootstrap_invoice_guard_path() -> Path:
    """Ensure `inference`, `models`, `tasks`, `server` modules can be imported.

    Priority order:
      1. INVOICEGUARD_CODE_DIR       -> already on disk
      2. INVOICEGUARD_CODE_REPO      -> hf_hub_download / snapshot_download
      3. INVOICEGUARD_GIT_URL        -> git clone --depth=1
      4. fall back to the parent dir of this file (local dev)
    """
    code_dir = os.environ.get("INVOICEGUARD_CODE_DIR")
    if code_dir and Path(code_dir).is_dir():
        sys.path.insert(0, code_dir)
        return Path(code_dir)

    repo = os.environ.get("INVOICEGUARD_CODE_REPO")
    if repo:
        from huggingface_hub import snapshot_download
        local = snapshot_download(
            repo_id=repo,
            repo_type="model",
            token=_hf_token(),
        )
        sys.path.insert(0, local)
        return Path(local)

    git_url = os.environ.get("INVOICEGUARD_GIT_URL")
    if git_url:
        target = Path("/tmp/invoiceguard_src")
        if not target.is_dir():
            subprocess.check_call(
                ["git", "clone", "--depth=1", git_url, str(target)],
            )
        sub = target / "invoice_guard"
        sys.path.insert(0, str(sub if sub.is_dir() else target))
        return sub if sub.is_dir() else target

    here = Path(__file__).resolve().parent.parent  # invoice_guard/
    sys.path.insert(0, str(here))
    return here


_CODE_ROOT = _bootstrap_invoice_guard_path()


# -----------------------------------------------------------------------------
# 1. Heavy imports (after sys.path is set).
# -----------------------------------------------------------------------------

import torch
import torch.nn.functional as F
from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

from models import TaskID  # type: ignore
from server.invoice_guard_environment import InvoiceGuardEnvironment  # type: ignore
from tasks import HARD_TASK_LIST, TASK_LIST  # type: ignore
from inference import SYSTEM_PROMPT, build_action, build_observation_prompt  # type: ignore
from training.rollout import Trajectory, TrajectoryStep, rollout_episode  # type: ignore


# -----------------------------------------------------------------------------
# 2. Config.
# -----------------------------------------------------------------------------

@dataclass
class TrainConfig:
    base_model: str = os.environ.get("BASE_MODEL", "Qwen/Qwen3-4B-Instruct-2507")
    hub_username: Optional[str] = os.environ.get("HF_USERNAME")
    hub_model_id: str = os.environ.get("HUB_MODEL_ID", "invoiceguard-qwen3-4b-grpo")
    trackio_project: str = os.environ.get("TRACKIO_PROJECT", "invoiceguard-round2")
    trackio_run_name: str = os.environ.get("TRACKIO_RUN_NAME", "qwen3-4b-grpo")
    artifact_dir: str = os.environ.get("ARTIFACT_DIR", "/tmp/invoiceguard-training-artifacts")

    seed: int = 42
    num_iterations: int = 3            # full passes over train tasks
    group_size: int = 4                # G trajectories per task per iteration
    max_train_tasks: Optional[int] = None  # truncate train set (smoke runs)
    eval_holdout_canonical: int = 3
    eval_holdout_hard: int = 3

    # Optimisation
    lr: float = 1e-5
    grad_clip: float = 1.0
    ppo_clip: float = 0.2
    kl_coef: float = 0.05
    grader_bonus: float = 1.0          # weight on terminal grader_score
    micro_batch_size: int = 1          # (obs, action) pairs per fwd/bwd
    bf16: bool = torch.cuda.is_available()
    use_4bit: bool = True
    gradient_checkpointing: bool = True

    # Sampling
    sample_temperature: float = 1.0
    sample_top_p: float = 0.95
    max_new_tokens: int = 384
    max_prompt_tokens: int = 2048

    # Tiny behavior warm-start. Smoke showed the raw model sometimes echoes the
    # observation instead of emitting JSON; this teaches format before RL.
    format_warmup: bool = True
    format_warmup_tasks: int = 8
    format_warmup_lr: float = 5e-5
    save_format_warmup_checkpoint: bool = True
    format_warmup_model_id: Optional[str] = os.environ.get("FORMAT_WARMUP_MODEL_ID")
    resume_adapter: Optional[str] = os.environ.get("RESUME_ADAPTER")

    # LoRA
    lora_r: int = 16
    lora_alpha: int = 32
    lora_dropout: float = 0.05
    lora_target_modules: tuple = (
        "q_proj", "k_proj", "v_proj", "o_proj",
    )

    push_to_hub: bool = True


# -----------------------------------------------------------------------------
# 3. Train / eval task split.
# -----------------------------------------------------------------------------

def split_tasks(cfg: TrainConfig) -> tuple[list[TaskID], list[TaskID]]:
    """Deterministic seeded split. Held-out tasks are NEVER trained on."""
    rng = random.Random(cfg.seed)

    canonical = list(TASK_LIST)
    hard = list(HARD_TASK_LIST)

    rng.shuffle(canonical)
    rng.shuffle(hard)

    eval_c = canonical[: cfg.eval_holdout_canonical]
    eval_h = hard[: cfg.eval_holdout_hard]
    train = canonical[cfg.eval_holdout_canonical:] + hard[cfg.eval_holdout_hard:]

    if cfg.max_train_tasks is not None:
        train = train[: cfg.max_train_tasks]

    eval_set = eval_c + eval_h
    return train, eval_set


# -----------------------------------------------------------------------------
# 4. Log-prob computation.
# -----------------------------------------------------------------------------

def _completion_logprobs(
    model,
    prompt_ids: torch.Tensor,
    completion_ids: torch.Tensor,
    device: torch.device,
) -> torch.Tensor:
    """Sum log p(completion | prompt) under `model`. Returns scalar tensor."""
    input_ids = torch.cat([prompt_ids, completion_ids], dim=0).unsqueeze(0).to(device)
    attention_mask = torch.ones_like(input_ids)

    out = model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
    # Shift: predict token t from logits at t-1.
    logits = out.logits[0, :-1, :]              # (L-1, V)
    targets = input_ids[0, 1:]                  # (L-1,)
    logprobs = F.log_softmax(logits.float(), dim=-1)
    token_lp = logprobs.gather(-1, targets.unsqueeze(-1)).squeeze(-1)  # (L-1,)

    # Only sum log-probs over the completion tokens.
    comp_len = completion_ids.shape[0]
    return token_lp[-comp_len:].sum()


# -----------------------------------------------------------------------------
# 5. Trajectory advantage computation.
# -----------------------------------------------------------------------------

def trajectory_reward(traj: Trajectory, grader_bonus: float) -> float:
    """Single scalar that GRPO will rank within a group."""
    return traj.cumulative_reward + grader_bonus * traj.grader_score


def compute_group_advantages(
    trajectories: List[Trajectory], grader_bonus: float
) -> List[float]:
    rewards = [trajectory_reward(t, grader_bonus) for t in trajectories]
    if len(rewards) < 2:
        return [max(min(r, 2.0), -2.0) for r in rewards]
    mean = sum(rewards) / len(rewards)
    var = sum((r - mean) ** 2 for r in rewards) / len(rewards)
    if var < 1e-10:
        return [max(min(r, 2.0), -2.0) for r in rewards]
    std = var ** 0.5
    return [(r - mean) / std for r in rewards]


_ALL_INVESTIGATION_ACTIONS = [
    {"action_type": "inspect_purchase_order"},
    {"action_type": "inspect_goods_receipt_note"},
    {"action_type": "inspect_invoice_line_items"},
    {"action_type": "inspect_vendor_profile"},
    {"action_type": "compare_quantity"},
    {"action_type": "compare_price"},
    {"action_type": "compare_totals"},
    {"action_type": "check_for_duplicate_invoice"},
    {"action_type": "inspect_policy_rules"},
]


def _format_warmup_actions(
    env: InvoiceGuardEnvironment,
    task_id: TaskID,
    max_investigation_steps: int = 9,
) -> list[dict]:
    case = getattr(env, "_case", None)
    if case is None:
        env.reset(task_id=task_id.value)
        case = getattr(env, "_case", None)
    assert case is not None
    gt = case.ground_truth
    investigation = _ALL_INVESTIGATION_ACTIONS[:max_investigation_steps]
    used_names = [a["action_type"] for a in investigation]
    evidence = list(dict.fromkeys([*used_names, *gt.acceptable_evidence]))
    explanation = "Key findings: " + "; ".join(gt.key_findings[:3])
    return [
        *investigation,
        {
            "action_type": "submit_final_resolution",
            "final_decision": gt.correct_decision.value,
            "exception_type": gt.correct_exception_type.value,
            "evidence_references": evidence,
            "explanation": explanation,
            "confidence": 0.9,
        },
    ]


def run_format_warmup(
    policy,
    tokenizer,
    optimizer,
    env: InvoiceGuardEnvironment,
    tasks: list[TaskID],
    cfg: TrainConfig,
    device: torch.device,
) -> dict:
    if not cfg.format_warmup or not tasks:
        return {"format_warmup/enabled": 0.0, "format_warmup/n_pairs": 0.0}

    old_lrs = [group["lr"] for group in optimizer.param_groups]
    for group in optimizer.param_groups:
        group["lr"] = cfg.format_warmup_lr

    warmup_trace_lengths = [3, 5, 7]

    policy.train()
    n_pairs = 0
    total_loss = 0.0
    for task_id in tasks[: cfg.format_warmup_tasks]:
        n_inv = warmup_trace_lengths[n_pairs % len(warmup_trace_lengths)]
        obs = env.reset(task_id=task_id.value)
        messages: list[dict] = [{"role": "system", "content": SYSTEM_PROMPT}]
        for action_dict in _format_warmup_actions(env, task_id, max_investigation_steps=n_inv):
            user_msg = build_observation_prompt(obs, is_first=(len(messages) == 1))
            messages.append({"role": "user", "content": user_msg})
            try:
                prompt_text = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True,
                    enable_thinking=False,
                )
            except TypeError:
                prompt_text = tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=True,
                )
            completion_text = json.dumps(action_dict, ensure_ascii=False)
            prompt_ids = tokenizer(
                prompt_text,
                return_tensors="pt",
                add_special_tokens=False,
                truncation=True,
                max_length=cfg.max_prompt_tokens,
            ).input_ids[0]
            comp_enc = tokenizer(
                completion_text,
                return_tensors="pt",
                add_special_tokens=False,
            ).input_ids[0]
            eos_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
            if eos_id is not None and eos_id != tokenizer.unk_token_id:
                completion_ids = torch.cat([comp_enc, torch.tensor([eos_id])])
            else:
                completion_ids = comp_enc
            lp = _completion_logprobs(policy, prompt_ids, completion_ids, device)
            loss = -lp / max(int(completion_ids.shape[0]), 1)
            loss.backward()
            total_loss += float(loss.detach().item())
            n_pairs += 1

            messages.append({"role": "assistant", "content": completion_text})
            obs = env.step(build_action(action_dict))
            if obs.done:
                break

    if n_pairs:
        torch.nn.utils.clip_grad_norm_(
            [p for p in policy.parameters() if p.requires_grad],
            cfg.grad_clip,
        )
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

    for group, lr in zip(optimizer.param_groups, old_lrs):
        group["lr"] = lr

    return {
        "format_warmup/enabled": 1.0,
        "format_warmup/n_pairs": float(n_pairs),
        "format_warmup/loss": total_loss / max(n_pairs, 1),
        "format_warmup/n_tasks": float(min(len(tasks), cfg.format_warmup_tasks)),
    }


def push_adapter_checkpoint(
    policy,
    tokenizer,
    repo_id: str,
    token: str,
    *,
    commit_message: str,
) -> None:
    from huggingface_hub import create_repo

    create_repo(
        repo_id=repo_id,
        repo_type="model",
        exist_ok=True,
        private=False,
        token=token,
    )
    policy.push_to_hub(repo_id, private=False, token=token, commit_message=commit_message)
    tokenizer.push_to_hub(repo_id, private=False, token=token, commit_message=commit_message)


# -----------------------------------------------------------------------------
# 6. Main training loop.
# -----------------------------------------------------------------------------

def train(cfg: TrainConfig) -> None:
    print(f"[setup] code_root={_CODE_ROOT}", flush=True)
    print(f"[setup] base_model={cfg.base_model}", flush=True)
    print(f"[setup] cuda available={torch.cuda.is_available()}", flush=True)

    artifact_dir = Path(cfg.artifact_dir)
    artifact_dir.mkdir(parents=True, exist_ok=True)
    metrics_path = artifact_dir / "metrics.jsonl"
    samples_path = artifact_dir / "rollout_samples.jsonl"
    summary_path = artifact_dir / "training_summary.json"
    metrics_history: list[dict] = []
    eval_history: list[dict] = []
    train_history: list[dict] = []
    sampled_rollouts: list[dict] = []

    run_started_at = datetime.now(timezone.utc).isoformat()

    random.seed(cfg.seed)
    torch.manual_seed(cfg.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(cfg.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    dtype = torch.bfloat16 if cfg.bf16 else torch.float32

    # ----- Tokenizer & policy --------------------------------------------------
    tokenizer = AutoTokenizer.from_pretrained(cfg.base_model, use_fast=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print("[setup] loading base model ...", flush=True)
    quant_cfg = None
    if cfg.use_4bit and torch.cuda.is_available():
        quant_cfg = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_use_double_quant=True,
            bnb_4bit_compute_dtype=dtype,
        )
    base = AutoModelForCausalLM.from_pretrained(
        cfg.base_model,
        torch_dtype=dtype,
        device_map="auto" if torch.cuda.is_available() else None,
        low_cpu_mem_usage=True,
        quantization_config=quant_cfg,
    )
    base.config.pad_token_id = tokenizer.pad_token_id
    base.config.use_cache = False
    if cfg.gradient_checkpointing:
        if cfg.use_4bit:
            base = prepare_model_for_kbit_training(base, use_gradient_checkpointing=True)
        base.gradient_checkpointing_enable()

    lora_cfg = LoraConfig(
        r=cfg.lora_r,
        lora_alpha=cfg.lora_alpha,
        lora_dropout=cfg.lora_dropout,
        target_modules=list(cfg.lora_target_modules),
        bias="none",
        task_type="CAUSAL_LM",
    )
    policy = get_peft_model(base, lora_cfg)
    policy.print_trainable_parameters()
    policy.train()

    # Reference (frozen) = base only, no adapter applied. We use the same
    # PeftModel with adapters disabled (`policy.disable_adapter()`) to compute
    # reference log-probs in-place and avoid loading a second copy of the base.

    optimizer = torch.optim.AdamW(
        [p for p in policy.parameters() if p.requires_grad],
        lr=cfg.lr,
    )

    # ----- Env & task split ----------------------------------------------------
    env = InvoiceGuardEnvironment()
    train_tasks, eval_tasks = split_tasks(cfg)
    print(f"[setup] train_tasks={len(train_tasks)}  eval_tasks={len(eval_tasks)}",
          flush=True)
    print(f"[setup] holdout_eval={[t.value for t in eval_tasks]}", flush=True)

    # ----- Trackio -------------------------------------------------------------
    tracker = None
    try:
        import trackio
        tracker = trackio.init(
            project=cfg.trackio_project,
            name=cfg.trackio_run_name,
            config={
                "base_model": cfg.base_model,
                "num_iterations": cfg.num_iterations,
                "group_size": cfg.group_size,
                "lr": cfg.lr,
                "kl_coef": cfg.kl_coef,
                "ppo_clip": cfg.ppo_clip,
                "grader_bonus": cfg.grader_bonus,
                "lora_r": cfg.lora_r,
                "n_train_tasks": len(train_tasks),
                "n_eval_tasks": len(eval_tasks),
            },
        )
        print("[setup] trackio initialised", flush=True)
    except Exception as e:
        print(f"[setup] trackio disabled: {e}", flush=True)

    def _write_jsonl(path: Path, row: dict) -> None:
        with path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    def _log(metrics: dict, step: int) -> None:
        row = {
            "step": step,
            "time": datetime.now(timezone.utc).isoformat(),
            **metrics,
        }
        metrics_history.append(row)
        _write_jsonl(metrics_path, row)
        msg = " | ".join(f"{k}={v:.4f}" if isinstance(v, float) else f"{k}={v}"
                         for k, v in metrics.items())
        print(f"[step {step}] {msg}", flush=True)
        if tracker is not None:
            try:
                trackio.log(metrics, step=step)
            except Exception:
                pass

    # ----- Eval helper ---------------------------------------------------------
    def evaluate(label: str, step: int) -> dict:
        policy.eval()
        scores, rewards, steps_used = [], [], []
        successes = []
        for tid in eval_tasks:
            traj = rollout_episode(
                policy, tokenizer, env, tid,
                temperature=0.0001,  # near-greedy for eval
                top_p=1.0,
                max_new_tokens=cfg.max_new_tokens,
                max_prompt_tokens=cfg.max_prompt_tokens,
                device=device,
            )
            scores.append(traj.grader_score)
            rewards.append(traj.cumulative_reward)
            steps_used.append(traj.n_steps)
            successes.append(1.0 if traj.success else 0.0)
        policy.train()
        eval_metrics = {
            f"{label}/avg_grader_score": sum(scores) / max(len(scores), 1),
            f"{label}/avg_cum_reward": sum(rewards) / max(len(rewards), 1),
            f"{label}/avg_steps": sum(steps_used) / max(len(steps_used), 1),
            f"{label}/success_rate": sum(successes) / max(len(successes), 1),
            f"{label}/n_tasks": len(eval_tasks),
        }
        _log(eval_metrics, step)
        eval_history.append({"label": label, "step": step, **eval_metrics})
        return eval_metrics

    def _record_rollout_sample(
        *,
        phase: str,
        global_step: int,
        task_id: TaskID,
        trajectories: List[Trajectory],
        advantages: List[float],
    ) -> None:
        # Keep evidence compact: one high-reward and one low-reward trace per task group.
        if not trajectories:
            return
        scored = [
            (trajectory_reward(t, cfg.grader_bonus), adv, t)
            for t, adv in zip(trajectories, advantages)
        ]
        selected = [max(scored, key=lambda x: x[0]), min(scored, key=lambda x: x[0])]
        seen = set()
        for reward_value, advantage, traj in selected:
            key = id(traj)
            if key in seen:
                continue
            seen.add(key)
            row = {
                "phase": phase,
                "step": global_step,
                "task_id": task_id.value,
                "trajectory_reward": reward_value,
                "advantage": advantage,
                "grader_score": traj.grader_score,
                "cumulative_reward": traj.cumulative_reward,
                "success": traj.success,
                "n_steps": traj.n_steps,
                "terminal_decision": traj.terminal_decision,
                "actions": [s.completion_text[:500] for s in traj.steps],
                "step_rewards": [s.reward for s in traj.steps],
            }
            sampled_rollouts.append(row)
            _write_jsonl(samples_path, row)

    def _write_plots() -> None:
        try:
            import matplotlib.pyplot as plt
        except Exception as e:
            print(f"[artifacts] plot generation skipped: {e}", flush=True)
            return

        if train_history:
            xs = [r["step"] for r in train_history]
            fig, ax1 = plt.subplots(figsize=(8, 4.5))
            ax1.plot(xs, [r["train/group_reward_mean"] for r in train_history], label="group reward")
            ax1.plot(xs, [r["train/group_grader_mean"] for r in train_history], label="grader score")
            ax1.set_xlabel("training step")
            ax1.set_ylabel("score")
            ax1.set_title("InvoiceGuard training reward")
            ax1.legend()
            fig.tight_layout()
            fig.savefig(artifact_dir / "training_reward_curve.png", dpi=160)
            plt.close(fig)

            fig, ax = plt.subplots(figsize=(8, 4.5))
            ax.plot(xs, [r["train/loss"] for r in train_history], label="loss")
            ax.plot(xs, [r["train/kl_loss"] for r in train_history], label="kl loss")
            ax.set_xlabel("training step")
            ax.set_ylabel("loss")
            ax.set_title("InvoiceGuard GRPO losses")
            ax.legend()
            fig.tight_layout()
            fig.savefig(artifact_dir / "training_loss_curve.png", dpi=160)
            plt.close(fig)

        eval_rows = [
            r for r in eval_history
            if any(k.endswith("/avg_grader_score") for k in r)
        ]
        if eval_rows:
            xs = [r["step"] for r in eval_rows]
            ys = []
            labels = []
            for r in eval_rows:
                key = next(k for k in r if k.endswith("/avg_grader_score"))
                labels.append(r["label"])
                ys.append(r[key])
            fig, ax = plt.subplots(figsize=(8, 4.5))
            ax.plot(xs, ys, marker="o")
            ax.set_xlabel("training step")
            ax.set_ylabel("holdout grader score")
            ax.set_title("InvoiceGuard holdout eval during training")
            for x, y, label in zip(xs, ys, labels):
                ax.annotate(label.replace("eval/", ""), (x, y), textcoords="offset points", xytext=(0, 6), ha="center")
            fig.tight_layout()
            fig.savefig(artifact_dir / "holdout_eval_curve.png", dpi=160)
            plt.close(fig)

    # ----- Resume from existing adapter or format warm-start ------------------
    global_step = 0
    if cfg.resume_adapter:
        print(f"\n=== resuming from adapter: {cfg.resume_adapter} ===", flush=True)
        from peft import set_peft_model_state_dict
        from safetensors.torch import load_file
        from huggingface_hub import hf_hub_download
        try:
            adapter_path = hf_hub_download(
                cfg.resume_adapter, "adapter_model.safetensors", token=_hf_token()
            )
            adapter_weights = load_file(adapter_path)
            set_peft_model_state_dict(policy, adapter_weights)
            print(f"[resume] loaded {len(adapter_weights)} tensors from {cfg.resume_adapter}", flush=True)
        except Exception as e:
            print(f"[resume] WARNING: could not load adapter: {e}", flush=True)
        cfg.format_warmup = False

    if cfg.format_warmup:
        print("\n=== format warm-start (JSON action behavior) ===", flush=True)
        warmup_metrics = run_format_warmup(
            policy, tokenizer, optimizer, env, train_tasks, cfg, device
        )
        _log(warmup_metrics, global_step)
        if (
            cfg.push_to_hub
            and cfg.hub_username
            and cfg.save_format_warmup_checkpoint
            and warmup_metrics.get("format_warmup/n_pairs", 0.0) > 0
        ):
            token = _hf_token()
            if not token:
                raise RuntimeError(
                    "HF_TOKEN/API_TOKEN_HF is required to save the format warm-start checkpoint."
                )
            warmup_model_id = (
                cfg.format_warmup_model_id
                or f"{cfg.hub_model_id}-format-warmup"
            )
            warmup_repo_id = f"{cfg.hub_username}/{warmup_model_id}"
            print(
                f"[push] saving format warm-start adapter to {warmup_repo_id}",
                flush=True,
            )
            push_adapter_checkpoint(
                policy,
                tokenizer,
                warmup_repo_id,
                token,
                commit_message="Save InvoiceGuard format warm-start adapter",
            )
            print(f"[push] format warm-start saved -> https://huggingface.co/{warmup_repo_id}", flush=True)

    # ----- Initial eval --------------------------------------------------------
    print("\n=== initial eval (after format warm-start) ===", flush=True)
    evaluate("eval/init", global_step)

    # ----- Training loop -------------------------------------------------------
    t_start = time.time()
    for it in range(cfg.num_iterations):
        random.shuffle(train_tasks)

        for ti, task_id in enumerate(train_tasks):
            # 1. Sample G trajectories on the same task (group).
            policy.eval()
            trajectories: List[Trajectory] = []
            for g in range(cfg.group_size):
                traj = rollout_episode(
                    policy, tokenizer, env, task_id,
                    temperature=cfg.sample_temperature,
                    top_p=cfg.sample_top_p,
                    max_new_tokens=cfg.max_new_tokens,
                    max_prompt_tokens=cfg.max_prompt_tokens,
                    device=device,
                )
                trajectories.append(traj)
            policy.train()

            # 2. Group-relative advantages.
            advantages = compute_group_advantages(trajectories, cfg.grader_bonus)

            # 3. PPO-clipped policy gradient on every (prompt, completion) pair,
            #    weighted by that trajectory's advantage, with KL vs. reference.
            optimizer.zero_grad(set_to_none=True)
            total_loss_val = 0.0
            n_pairs = 0
            kl_sum = 0.0
            pg_sum = 0.0

            for traj, adv in zip(trajectories, advantages):
                if abs(adv) < 1e-8 or not traj.steps:
                    continue
                for step in traj.steps:
                    # Current policy log-prob (with adapter active).
                    cur_lp = _completion_logprobs(
                        policy, step.prompt_ids, step.completion_ids, device
                    )
                    # Reference policy log-prob (adapter disabled).
                    with torch.no_grad():
                        with policy.disable_adapter():
                            ref_lp = _completion_logprobs(
                                policy, step.prompt_ids, step.completion_ids, device
                            )

                    # PPO-clipped surrogate. The "old" policy here is the same
                    # snapshot used to sample (we just rolled out moments ago),
                    # so on the first opt step the ratio == 1; the clip becomes
                    # active only across multiple opt steps per batch. We still
                    # apply it for stability when group_size is large.
                    # Stable PPO surrogate: clamp log-ratio before exp to avoid
                    # overflow/underflow from very large policy deltas.
                    log_ratio = (cur_lp - ref_lp.detach()).clamp(-20.0, 20.0)
                    ratio = torch.exp(log_ratio)
                    clipped_ratio = torch.clamp(
                        ratio, 1.0 - cfg.ppo_clip, 1.0 + cfg.ppo_clip
                    )
                    adv_t = torch.tensor(float(adv), device=device, dtype=cur_lp.dtype)
                    pg_term = -torch.min(ratio * adv_t, clipped_ratio * adv_t)

                    kl_term = cfg.kl_coef * (cur_lp - ref_lp.detach()).pow(2)

                    loss = pg_term + kl_term
                    if not torch.isfinite(loss):
                        # Skip pathological pairs instead of poisoning optimizer
                        # state with inf/nan gradients.
                        continue
                    loss.backward()

                    total_loss_val += float(loss.detach().item())
                    pg_sum += float(pg_term.detach().item())
                    kl_sum += float(kl_term.detach().item())
                    n_pairs += 1

            if n_pairs > 0:
                torch.nn.utils.clip_grad_norm_(
                    [p for p in policy.parameters() if p.requires_grad],
                    cfg.grad_clip,
                )
                optimizer.step()

            global_step += 1
            group_rewards = [trajectory_reward(t, cfg.grader_bonus) for t in trajectories]
            group_scores = [t.grader_score for t in trajectories]
            train_metrics = {
                "train/iter": it,
                "train/task_idx": ti,
                "train/task_id": task_id.value,
                "train/group_reward_mean": sum(group_rewards) / len(group_rewards),
                "train/group_reward_std":
                    (sum((r - sum(group_rewards) / len(group_rewards)) ** 2
                         for r in group_rewards) / len(group_rewards)) ** 0.5,
                "train/group_grader_mean": sum(group_scores) / len(group_scores),
                "train/group_success_rate": sum(1.0 if t.success else 0.0 for t in trajectories) / len(trajectories),
                "train/avg_steps": sum(t.n_steps for t in trajectories) / len(trajectories),
                "train/n_pairs": n_pairs,
                "train/loss": total_loss_val / max(n_pairs, 1),
                "train/pg_loss": pg_sum / max(n_pairs, 1),
                "train/kl_loss": kl_sum / max(n_pairs, 1),
            }
            train_history.append({"step": global_step, **train_metrics})
            _record_rollout_sample(
                phase="train",
                global_step=global_step,
                task_id=task_id,
                trajectories=trajectories,
                advantages=advantages,
            )
            _log(
                train_metrics,
                global_step,
            )

        # End-of-iteration eval.
        print(f"\n=== eval after iteration {it + 1}/{cfg.num_iterations} ===",
              flush=True)
        evaluate(f"eval/iter{it+1}", global_step)

    total_wall_clock = time.time() - t_start
    print(f"\n[done] total wall clock: {total_wall_clock:.1f}s", flush=True)

    _write_plots()

    summary = {
        "run_started_at": run_started_at,
        "run_finished_at": datetime.now(timezone.utc).isoformat(),
        "base_model": cfg.base_model,
        "hub_model_id": cfg.hub_model_id,
        "num_iterations": cfg.num_iterations,
        "group_size": cfg.group_size,
        "train_tasks": [t.value for t in train_tasks],
        "eval_tasks": [t.value for t in eval_tasks],
        "wall_clock_s": round(total_wall_clock, 2),
        "n_metric_rows": len(metrics_history),
        "n_rollout_samples": len(sampled_rollouts),
        "artifact_files": [
            p.name for p in sorted(artifact_dir.iterdir()) if p.is_file()
        ],
    }
    summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(f"[artifacts] wrote {artifact_dir}", flush=True)

    # ----- Push LoRA adapter ---------------------------------------------------
    if cfg.push_to_hub and cfg.hub_username:
        from huggingface_hub import HfApi

        token = _hf_token()
        if not token:
            raise RuntimeError(
                "HF_TOKEN/API_TOKEN_HF is required when push_to_hub=True. "
                "Set the token secret or run with --no-push."
            )
        repo_id = f"{cfg.hub_username}/{cfg.hub_model_id}"
        print(f"[push] pushing LoRA adapter to {repo_id}", flush=True)
        push_adapter_checkpoint(
            policy,
            tokenizer,
            repo_id,
            token,
            commit_message="Save InvoiceGuard GRPO adapter",
        )
        print(f"[push] uploading training artifacts to {repo_id}/training_artifacts", flush=True)
        HfApi(token=token).upload_folder(
            folder_path=str(artifact_dir),
            repo_id=repo_id,
            repo_type="model",
            path_in_repo="training_artifacts",
            commit_message="Add InvoiceGuard GRPO training artifacts",
            token=token,
        )
        print(f"[push] done -> https://huggingface.co/{repo_id}", flush=True)
    else:
        out_dir = Path(os.environ.get("OUTPUT_DIR", "/tmp/invoiceguard-grpo"))
        out_dir.mkdir(parents=True, exist_ok=True)
        policy.save_pretrained(out_dir)
        tokenizer.save_pretrained(out_dir)
        print(f"[save] LoRA adapter saved locally -> {out_dir}", flush=True)


# -----------------------------------------------------------------------------
# 7. CLI.
# -----------------------------------------------------------------------------

def _parse_args() -> TrainConfig:
    p = argparse.ArgumentParser()
    p.add_argument("--model-name", dest="base_model", default=None)
    p.add_argument("--num-iterations", type=int, default=None)
    p.add_argument("--group-size", type=int, default=None)
    p.add_argument("--max-train-tasks", type=int, default=None)
    p.add_argument("--lr", type=float, default=None)
    p.add_argument("--no-push", action="store_true")
    p.add_argument("--seed", type=int, default=None)
    p.add_argument("--no-4bit", action="store_true")
    p.add_argument("--no-gradient-checkpointing", action="store_true")
    p.add_argument("--eval-holdout-canonical", type=int, default=None)
    p.add_argument("--eval-holdout-hard", type=int, default=None)
    p.add_argument("--max-new-tokens", type=int, default=None)
    p.add_argument("--max-prompt-tokens", type=int, default=None)
    p.add_argument("--no-format-warmup", action="store_true")
    p.add_argument("--format-warmup-tasks", type=int, default=None)
    p.add_argument("--no-save-format-warmup", action="store_true")
    p.add_argument("--format-warmup-model-id", default=None)
    p.add_argument("--sample-temperature", type=float, default=None)
    p.add_argument("--resume-adapter", default=None)
    args = p.parse_args()

    cfg = TrainConfig()
    if args.base_model:
        cfg.base_model = args.base_model
    if args.num_iterations is not None:
        cfg.num_iterations = args.num_iterations
    if args.group_size is not None:
        cfg.group_size = args.group_size
    if args.max_train_tasks is not None:
        cfg.max_train_tasks = args.max_train_tasks
    if args.lr is not None:
        cfg.lr = args.lr
    if args.seed is not None:
        cfg.seed = args.seed
    if args.eval_holdout_canonical is not None:
        cfg.eval_holdout_canonical = args.eval_holdout_canonical
    if args.eval_holdout_hard is not None:
        cfg.eval_holdout_hard = args.eval_holdout_hard
    if args.max_new_tokens is not None:
        cfg.max_new_tokens = args.max_new_tokens
    if args.max_prompt_tokens is not None:
        cfg.max_prompt_tokens = args.max_prompt_tokens
    if args.no_format_warmup:
        cfg.format_warmup = False
    if args.format_warmup_tasks is not None:
        cfg.format_warmup_tasks = args.format_warmup_tasks
    if args.no_save_format_warmup:
        cfg.save_format_warmup_checkpoint = False
    if args.format_warmup_model_id:
        cfg.format_warmup_model_id = args.format_warmup_model_id
    if args.sample_temperature is not None:
        cfg.sample_temperature = args.sample_temperature
    if args.resume_adapter:
        cfg.resume_adapter = args.resume_adapter
    if args.no_push:
        cfg.push_to_hub = False
    if args.no_4bit:
        cfg.use_4bit = False
    if args.no_gradient_checkpointing:
        cfg.gradient_checkpointing = False
    return cfg


if __name__ == "__main__":
    cfg = _parse_args()
    train(cfg)