grpo_train.py

"""
grpo_train.py — State-Based GRPO for Project Polymath
======================================================
Trains an LLM to negotiate with expert stakeholders using proper
Group Relative Policy Optimization with weight updates.

THE KEY INSIGHT (State-Based GRPO):
    TRL's GRPOTrainer is single-turn. Our environment is multi-turn.
    Solution: treat every (state, next_action) pair as its own training prompt.
    The model learns: "given THIS game state, what is the best next action?"
    
    Instead of rolling out full episodes, we:
    1. Build a dataset of negotiation states (from oracle + your JSON topics)
    2. For each state, sample G=8 completions from the model
    3. Run each completion through the environment for ONE step
    4. Use GRPO advantage to update weights toward better single-step decisions
    5. Repeat across all states — the model learns the full strategy implicitly

USAGE:
    # Pre-hackathon: verify the pipeline (no GPU needed)
    python grpo_train.py --dry-run --states 10

    # On-site Day 1 with HF GPU credits (the real run):
    python grpo_train.py --use-unsloth --epochs 3 --states 50

    # Without Unsloth (slower but works):
    python grpo_train.py --model Qwen/Qwen2.5-1.5B-Instruct --epochs 3
"""

from __future__ import annotations

import argparse
import json
import random
import os
import re
import time
from pathlib import Path
from typing import Optional

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    # Keep script runnable even if python-dotenv is not installed.
    pass


try:
    import matplotlib
    matplotlib.use("Agg")  # non-interactive backend for servers
    import matplotlib.pyplot as plt
    HAS_PLT = True
except ImportError:
    HAS_PLT = False

HAS_UNSLOTH = False
FastLanguageModel = None


try:
    from trl import GRPOConfig, GRPOTrainer
    HAS_TRL = True
    print("TRL loaded OK")
except Exception as e:
    print(f"TRL FAILED: {e}")
    HAS_TRL = False

try:
    from datasets import Dataset
    HAS_DATASETS = True
except ImportError:
    HAS_DATASETS = False

try:
    from transformers import AutoModelForCausalLM, AutoTokenizer
    HAS_TRANSFORMERS = True
except ImportError:
    HAS_TRANSFORMERS = False

# Local imports
from envs.environment import WorkSpaceEnvironment
from models.schemas import WorkSpaceAction

# Constants
TOPICS_FILE = Path("ai_pm_prompts.json")
OUTPUT_DIR = Path("artifacts/grpo_state_based")

# The three hidden constraints — static for easy/medium mode
HIDDEN_CONSTRAINTS = {
    "Finance": "Budget must not exceed $50k.",
    "Security": "Must include biometric 2FA.",
    "UX": "Checkout must be a single click.",
}

# ── Action templates the model should learn to produce
ORACLE_ACTIONS = {
    "ask_finance": json.dumps({
        "action_type": "message_expert", "target": "Finance",
        "content": "What is the hard budget ceiling the PRD must respect for launch?"
    }),
    "ask_security": json.dumps({
        "action_type": "message_expert", "target": "Security",
        "content": "What authentication controls must the PRD include? Is biometric 2FA required?"
    }),
    "ask_ux": json.dumps({
        "action_type": "message_expert", "target": "UX",
        "content": "What checkout experience is required? Should we target a single-click flow?"
    }),
    "propose_draft": json.dumps({
        "action_type": "propose_draft", "target": "All",
        "content": (
            "PRD Draft:\n"
            "1. Budget: Launch scope capped at $50k.\n"
            "2. Security: Biometric 2FA required for login and sensitive actions.\n"
            "3. UX: Single-click checkout flow."
        ),
    }),
    "submit_final": json.dumps({
        "action_type": "submit_final", "target": None,
        "content": (
            "Final PRD:\n"
            "1. Budget cap: All launch costs must stay at or below $50k.\n"
            "2. Security: The app must enforce biometric 2FA for all authentication.\n"
            "3. UX: Checkout must be implemented as a true single-click experience."
        ),
    }),
}


# Utilities

def load_topics(limit: int = 50) -> list[str]:
    if TOPICS_FILE.exists():
        with TOPICS_FILE.open() as f:
            return json.load(f)[:limit]
    return [
        "Draft a Mobile App PRD for a FinTech startup targeting emerging markets.",
        "Build an AI-driven healthcare platform for enterprise customers.",
        "Create a SaaS analytics tool for regulatory-heavy industries.",
        "Design a gaming platform for Gen Z users with real-time features.",
        "Develop a cross-platform product for low-bandwidth regions.",
    ]


def parse_action(text: str) -> Optional[WorkSpaceAction]:
    """Parse a JSON action from model output. Returns None on failure."""
    text = text.strip()
    if text.startswith("```"):
        # Handle fenced responses like ```json ... ```
        text = re.sub(r"^```(?:json)?\s*", "", text)
        text = re.sub(r"\s*```$", "", text)

    text = text.strip()
    try:
        # Fast path: entire completion is valid JSON.
        return WorkSpaceAction(**json.loads(text))
    except Exception:
        pass

    # Fallback: find the first JSON object that includes action_type.
    try:
        idx = text.find("{")
        while idx != -1:
            depth = 0
            for end in range(idx, len(text)):
                if text[end] == "{":
                    depth += 1
                elif text[end] == "}":
                    depth -= 1
                    if depth == 0:
                        candidate = text[idx:end + 1]
                        if '"action_type"' in candidate:
                            return WorkSpaceAction(**json.loads(candidate))
                        break
            idx = text.find("{", idx + 1)
        return None
    except Exception:
        return None


def lexical_overlap(a: str, b: str) -> float:
    """Simple token overlap score in [0,1] for dense content shaping."""
    toks_a = set(re.findall(r"[a-z0-9]+", (a or "").lower()))
    toks_b = set(re.findall(r"[a-z0-9]+", (b or "").lower()))
    if not toks_a or not toks_b:
        return 0.0
    inter = len(toks_a & toks_b)
    denom = max(1, min(len(toks_a), len(toks_b)))
    return inter / denom


def format_discovered(env: WorkSpaceEnvironment) -> str:
    lines = []
    for name, expert in env.state().experts.items():
        status = "✓ DISCOVERED" if expert.constraint_discovered_by_agent else "? unknown"
        lines.append(f"  {name}: {status}")
    return "\n".join(lines)


# ── State-Based Prompt Builder ─────────────────────────────────────────────────

AGENT_SYSTEM_PROMPT = """You are an expert AI Project Manager in a multi-stakeholder negotiation.

TASK: Produce a final PRD that satisfies ALL three experts — Finance, Security, and UX.
Each expert holds a hidden constraint you must discover through targeted questions.

STRATEGY:
  1. Message each expert INDIVIDUALLY (not "All") to discover their constraint.
  2. Once all constraints are known, propose a draft.
  3. Refine if needed, then submit_final before turn 15.

ANTI-PATTERNS (will be penalized):
  - Broadcasting to "All" when gathering requirements → -0.3 penalty
  - Repeating a question already answered → -0.4 penalty
  - Submitting without discovering constraints → low harmonic mean score

CURRENT DISCOVERED CONSTRAINTS:
{discovered}

You are a strict API. Respond with ONLY raw, valid JSON. 
DO NOT wrap the JSON in markdown formatting (no ```json). 
DO NOT output any conversational text.
End your response immediately after the closing }}.
{{"action_type": "message_expert" | "propose_draft" | "submit_final",
  "target": "Finance" | "Security" | "UX" | "All" | null,
  "content": "your message"}}"""


def build_state_prompt(
    topic: str,
    turn: int,
    feedback_so_far: str,
    discovered: str,
    conversation_history: str = "",
) -> str:
    """
    Build a prompt representing a specific game state.
    This is what gets fed to GRPOTrainer as the 'prompt' field.
    """
    system = AGENT_SYSTEM_PROMPT.format(discovered=discovered)

    user_content = (
        f"NEGOTIATION TASK: {topic}\n\n"
        f"TURN: {turn}/15\n\n"
    )

    if conversation_history:
        user_content += f"CONVERSATION SO FAR:\n{conversation_history}\n\n"

    user_content += f"LATEST FEEDBACK:\n{feedback_so_far}\n\nWhat is your next action?"

    # Use Qwen-compatible ChatML formatting to improve stop behavior.
    # Qwen instruct models are much more likely to terminate with <|im_end|>
    # when prompted in this native format.
    return (
        f"<|im_start|>system\n{system}<|im_end|>\n"
        f"<|im_start|>user\n{user_content}<|im_end|>\n"
        f"<|im_start|>assistant\n"
    )


# State Dataset Builder

def build_state_dataset(topics: list[str], states_per_topic: int = 5) -> list[dict]:
    """
    Build a dataset of negotiation states using the EASY mode environment.
    Each record represents one (state → optimal_action) training example.

    We run oracle trajectories through the environment to get realistic
    expert feedback, then snapshot the state at each turn.

    This is the key fix: instead of hoping the model learns from full episodes,
    we give it explicit training signal at every decision point.
    """
    env = WorkSpaceEnvironment(mode="medium")
    records = []

    # Oracle action sequence for easy mode
    oracle_sequence = [
        ("ask_finance", WorkSpaceAction(
            action_type="message_expert", target="Finance",
            content="What budget ceiling must the PRD respect?"
        )),
        ("ask_security", WorkSpaceAction(
            action_type="message_expert", target="Security",
            content="What authentication requirements must be included?"
        )),
        ("ask_ux", WorkSpaceAction(
            action_type="message_expert", target="UX",
            content="What checkout flow is required?"
        )),
        ("propose_draft", WorkSpaceAction(
            action_type="propose_draft", target="All",
            content="PRD: Budget at or below $50k. Biometric 2FA required. Single-click checkout."
        )),
        ("submit_final", WorkSpaceAction(
            action_type="submit_final", target=None,
            content="Final PRD: Budget capped at $50k. Biometric 2FA for auth. Single-click checkout."
        )),
    ]


    for topic in topics:
        obs = env.reset(topic)
        conversation_history = ""
        discovered = "  Finance: ? unknown\n  Security: ? unknown\n  UX: ? unknown"

        for step_idx, (action_key, oracle_action) in enumerate(oracle_sequence):
            if obs.done:
                break

            # Snapshot the state BEFORE taking the action
            prompt = build_state_prompt(
                topic=topic,
                turn=obs.current_turn,
                feedback_so_far=obs.feedback,
                discovered=discovered,
                conversation_history=conversation_history,
            )

            records.append({
                "prompt": prompt,
                "topic": topic,
                "turn": obs.current_turn,
                "oracle_action": ORACLE_ACTIONS[action_key],
                # These metadata fields help with debugging and post-analysis
                "step_idx": step_idx,
                "discovered_before": discovered,
            })

            # Step forward with oracle action to get next state
            obs = env.step(oracle_action)
            conversation_history += (
                f"Turn {step_idx}: {oracle_action.action_type} → {oracle_action.target}\n"
                f"Feedback: {obs.feedback[:120]}...\n"
            )
            discovered = format_discovered(env)

            if step_idx >= states_per_topic - 1:
                break

    # Add negative-pattern states (what NOT to do)
    records.extend(build_negative_states(topics[:5]))
    # Upweight late-stage "submit_final" states so policy learns to finish.
    late_stage = build_late_stage_states(topics)
    records.extend(late_stage)
    records.extend(late_stage)
    records.extend(late_stage)
    random.shuffle(records)

    print(f"Built {len(records)} training states from {len(topics)} topics")
    return records

def build_late_stage_states(topics: list[str]) -> list[dict]:
    """
    FIX 3: Inject guaranteed late-stage states.
    Forces the model to learn how to synthesize and submit the final PRD.
    """
    late_records = []
    for topic in topics:
        prompt = build_state_prompt(
            topic=topic,
            turn=4,
            feedback_so_far="UX: The checkout must be a single click.",
            discovered="  Finance: ✓ DISCOVERED\n  Security: ✓ DISCOVERED\n  UX: ✓ DISCOVERED",
            conversation_history=(
                "Turn 0: message_expert → Finance\nFeedback: The budget cap is $50k.\n"
                "Turn 1: message_expert → Security\nFeedback: Biometric 2FA is strictly required.\n"
                "Turn 2: message_expert → UX\nFeedback: Checkout must be a single click.\n"
            )
        )
        late_records.append({
            "prompt": prompt,
            "topic": topic,
            "turn": 4,
            "oracle_action": ORACLE_ACTIONS["submit_final"],
            "step_idx": 4,
            "discovered_before": "  Finance: ✓ DISCOVERED\n  Security: ✓ DISCOVERED\n  UX: ✓ DISCOVERED",
        })
    return late_records

def build_negative_states(topics: list[str]) -> list[dict]:
    """
    States where the agent is in a bad situation (repeated question, wrong phase).
    These teach the model to recover, not just follow the oracle.
    """
    negative_records = []

    for topic in topics:
        # State: Finance already answered, agent is about to repeat
        prompt = build_state_prompt(
            topic=topic,
            turn=2,
            feedback_so_far=(
                "Finance: As I mentioned, we have a strict $50k budget cap. "
                "This is the same answer I gave before."
            ),
            discovered="  Finance: ✓ DISCOVERED\n  Security: ? unknown\n  UX: ? unknown",
            conversation_history=(
                "Turn 0: message_expert → Finance\n"
                "Feedback: Finance: The budget cap is $50k. Don't go over it.\n"
                "Turn 1: message_expert → Finance\n"
                "Feedback: Finance: I already told you — $50k. Ask someone else.\n"
            ),
        )
        negative_records.append({
            "prompt": prompt,
            "topic": topic,
            "turn": 2,
            "oracle_action": ORACLE_ACTIONS["ask_security"],  # Should pivot to Security
            "step_idx": -1,  # Negative example
            "discovered_before": "Finance: ✓ DISCOVERED",
        })

    return negative_records


# Reward Function

def make_reward_fn():
    """
    Evaluates the model's actions instantly and locally.
    No live API calls. No reward hacking loopholes.
    """
    def reward_fn(completions: list[str], prompts: list[str], **kwargs) -> list[float]:
        rewards = []
        oracle_actions = kwargs.get("oracle_action", [None] * len(completions))
        turns = kwargs.get("turn", [None] * len(completions))

        for completion, prompt, oracle_raw, turn in zip(completions, prompts, oracle_actions, turns):
            action = parse_action(completion)

            # 1. Formatting Penalty
            if action is None:
                rewards.append(-0.5)
                continue

            reward = 0.0
            completion_text = (completion or "").strip()

            # ── 2. YOUR ANTI-PATTERN PENALTIES ──
            
            # Massive penalty for broadcasting (Reward Hacking)
            if action.target == "All":
                reward -= 1.0  
            
            # Penalty for empty or trivially short drafts/finals
            # (short expert questions are often valid and should not be over-penalized)
            if action.action_type in {"propose_draft", "submit_final"} and len((action.content or "").split()) < 5:
                reward -= 0.2

            # Penalize very long outputs; they correlate with max-length clipping.
            if len((action.content or "").split()) > 80:
                reward -= 0.2
            if len(completion_text) > 320:
                reward -= 0.15

            # Encourage strict JSON-only behavior: starts with { and ends with }.
            is_strict_json = completion_text.startswith("{") and completion_text.endswith("}")
            if is_strict_json:
                reward += 0.1
            else:
                reward -= 0.3

            # Hard penalty for non-terminated JSON-like responses.
            # This directly pushes generations away from max-token clipping.
            if not completion_text.endswith("}"):
                reward -= 0.25

            # Small bonus for compact single-line JSON output.
            if is_strict_json and "\n" not in completion_text and len(completion_text) <= 240:
                reward += 0.08

            # Strongly discourage invalid action/target combinations.
            if action.action_type == "submit_final" and action.target is not None:
                reward -= 0.6
            if action.action_type in {"message_expert", "propose_draft"} and action.target is None:
                reward -= 0.6

            # ── 3. HEURISTIC STATE GRADING (NO API CALLS!) ──
            
            if action.action_type == "message_expert" and action.target != "All":
                # Did it ask a question it already knows the answer to?
                if f"{action.target}: ✓ DISCOVERED" in prompt:
                    reward -= 0.5
                else:
                    reward += 0.33 

            elif action.action_type in ["propose_draft", "submit_final"]:
                # Did it try to submit before gathering all constraints?
                if "? unknown" in prompt:
                    reward -= 1.0 # Heavy penalty for guessing
                else:
                    # It did the research. Did it actually include the constraints?
                    text = action.content.lower()
                    has_finance = "50" in text
                    has_security = "biometric" in text
                    has_ux = "click" in text or "tap" in text
                    
                    if has_finance and has_security and has_ux:
                        reward += 1.5  
                    else:
                        reward -= 0.5

            # ── 4. ORACLE-GUIDED DENSE SHAPING ──
            # This gives non-binary signal and prevents reward plateaus.
            if oracle_raw:
                oracle_action = parse_action(oracle_raw)
                if oracle_action is not None:
                    if action.action_type == oracle_action.action_type:
                        reward += 0.45
                    else:
                        reward -= 0.25

                    if action.target == oracle_action.target:
                        reward += 0.35
                    else:
                        reward -= 0.2

                    overlap = lexical_overlap(action.content, oracle_action.content)
                    reward += 0.4 * overlap

            # Late turns should avoid endless questioning/proposals.
            if isinstance(turn, int):
                if turn >= 8 and action.action_type != "submit_final":
                    reward -= 0.35
                if turn >= 10 and action.action_type != "submit_final":
                    reward -= 0.6
                # Reward timely completion once constraints are all discovered.
                if (
                    action.action_type == "submit_final"
                    and "? unknown" not in prompt
                    and turn <= 10
                ):
                    reward += 0.6

            # Keep rewards in a stable range for GRPO.
            reward = max(-2.0, min(2.0, reward))

            rewards.append(reward)

        return rewards
    return reward_fn


# Plots

def save_training_plots(log_history: list[dict], output_dir: Path):
    if not HAS_PLT:
        print("  matplotlib not available — skipping plots")
        return

    output_dir.mkdir(parents=True, exist_ok=True)

    # Loss curve
    loss_points = [
        (e["step"], e["loss"])
        for e in log_history
        if "loss" in e and "step" in e
    ]
    if loss_points:
        xs, ys = zip(*loss_points)
        fig, ax = plt.subplots(figsize=(9, 4))
        ax.plot(xs, ys, marker="o", linewidth=1.5, color="#4C72B0", markersize=4)
        ax.set_xlabel("Training Step", fontsize=12)
        ax.set_ylabel("GRPO Loss", fontsize=12)
        ax.set_title(
            "Project Polymath — GRPO Training Loss\n"
            "(State-Based: each step = one negotiation decision)",
            fontsize=12
        )
        ax.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(output_dir / "loss_curve.png", dpi=160)
        plt.close()
        print(f"  Saved: {output_dir}/loss_curve.png")

    # Reward curve (from log history if available)
    reward_points = [
        (e["step"], e.get("reward", e.get("mean_reward", None)))
        for e in log_history
        if "step" in e and ("reward" in e or "mean_reward" in e)
    ]
    reward_points = [(s, r) for s, r in reward_points if r is not None]

    if reward_points:
        xs, ys = zip(*reward_points)
        fig, ax = plt.subplots(figsize=(9, 4))
        ax.plot(xs, ys, marker="s", linewidth=1.5, color="#55A868", markersize=4)
        ax.set_xlabel("Training Step", fontsize=12)
        ax.set_ylabel("Mean Reward", fontsize=12)
        ax.set_title(
            "Project Polymath — Mean Reward During GRPO Training\n"
            "(Harmonic mean of Finance/Security/UX constraint satisfaction)",
            fontsize=12
        )
        ax.grid(True, alpha=0.3)
        plt.tight_layout()
        plt.savefig(output_dir / "reward_curve.png", dpi=160)
        plt.close()
        print(f"  Saved: {output_dir}/reward_curve.png")


# Main

def main():
    parser = argparse.ArgumentParser(description="State-Based GRPO — Project Polymath")

    # Model
    parser.add_argument("--model", default="unsloth/Qwen2.5-3B-Instruct-bnb-4bit",
                        help="Base model to train")
    parser.add_argument("--use-unsloth", action="store_true",
                        help="Use Unsloth for 2x faster training (recommended on GPU)")

    # Dataset
    parser.add_argument("--states", type=int, default=40,
                        help="Number of negotiation states to train on")
    parser.add_argument("--states-per-topic", type=int, default=5,
                        help="States to extract per topic (1-5)")
    parser.add_argument("--topics-limit", type=int, default=20,
                        help="Max topics to use from ai_pm_prompts.json")

    # GRPO hyperparams
    parser.add_argument("--group-size", type=int, default=8,
                        help="G: completions per prompt for GRPO advantage (default: 8)")
    parser.add_argument("--epochs", type=float, default=3.0)
    parser.add_argument("--lr", type=float, default=5e-6,
                        help="Learning rate (lower = safer, 5e-6 recommended for GRPO)")
    parser.add_argument("--max-new-tokens", type=int, default=40,
                        help="Max generated tokens per sampled completion (default: 40)")
    parser.add_argument("--temperature", type=float, default=0.9,
                        help="Sampling temperature for GRPO rollouts")
    parser.add_argument("--top-p", type=float, default=0.9,
                        help="Nucleus sampling p for GRPO rollouts")
    parser.add_argument("--batch-size", type=int, default=1)
    parser.add_argument("--grad-accum", type=int, default=4)
    parser.add_argument("--max-seq-length", type=int, default=2048)

    # Output
    parser.add_argument("--output-dir", default=str(OUTPUT_DIR))
    parser.add_argument("--dry-run", action="store_true",
                        help="Build dataset and verify reward fn, skip actual training")

    args = parser.parse_args()

    # for dry run only
    # if not HAS_TRL:
    #     raise RuntimeError("pip install trl>=0.8.0 transformers datasets")

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Build dataset
    print("\n[1/4] Loading state dataset...")
    topics = load_topics(limit=args.topics_limit)
    dataset_path = output_dir / "state_dataset.jsonl"

    # --- THE CACHING LOGIC ---
    if dataset_path.exists():
        print(f"  [CACHE HIT] Found existing dataset! Loading instantly from {dataset_path}...")
        records = []
        with dataset_path.open("r") as f:
            for line in f:
                if line.strip():
                    records.append(json.loads(line))
        records = records[:args.states]
    else:
        print("  [CACHE MISS] No dataset found. Generating from scratch (this may take a minute)...")
        records = build_state_dataset(topics, states_per_topic=args.states_per_topic)
        records = records[:args.states]
        
        # Save it so we never have to build it again
        with dataset_path.open("w") as f:
            for r in records:
                f.write(json.dumps(r, ensure_ascii=True) + "\n")
        print(f"  Saved {len(records)} states → {dataset_path}")

    dataset = Dataset.from_list([{
        "prompt": r["prompt"],
        "topic": r["topic"],
        "turn": r["turn"],
        "oracle_action": r["oracle_action"],
        "step_idx": r["step_idx"],
    } for r in records])

    # Verify reward function
    print("\n[2/4] Verifying reward function on 6 samples...")
    reward_fn = make_reward_fn()
    # reward_fn = make_reward_fn(topics)
    prompt_with_finance_discovered = build_state_prompt(
    topic=topics[0],
    turn=1,
    feedback_so_far="Finance: Budget cap is $50k.",
    discovered="  Finance: ✓ DISCOVERED\n  Security: ? unknown\n  UX: ? unknown",
    )

    test_completions = [
        ORACLE_ACTIONS["ask_finance"],       # Should score ~0.33
        ORACLE_ACTIONS["ask_security"],      # Should score ~0.33  
        '{"action_type": "message_expert", "target": "Finance", "content": "What is the budget?"}',  # Should score -0.5 (repeat)
        '{"action_type": "message_expert", "target": "All", "content": "What do you all need?"}',    # Should score -1.0
        "this is not JSON at all",           # Should score -0.5
        ORACLE_ACTIONS["submit_final"],      # Should score +1.5 (all constraints in content)
    ]

    test_rewards = reward_fn(
        completions=test_completions,
        prompts=[
            "",                               # oracle ask_finance — unknown state
            prompt_with_finance_discovered,   # ask_security — good pivot
            prompt_with_finance_discovered,   # ask Finance again — repeat penalty
            "",                               # broadcast — penalty
            "",                               # bad JSON
            build_state_prompt(               # submit after all discovered
                topic=topics[0], turn=4,
                feedback_so_far="All experts responded.",
                discovered="  Finance: ✓ DISCOVERED\n  Security: ✓ DISCOVERED\n  UX: ✓ DISCOVERED",
            ),
        ],
    )

    labels = [
        "Oracle ask_finance",
        "Oracle ask_security",
        "Repeat question to discovered expert",
        "Broadcast to All",
        "Malformed JSON",
        "Submit final with all constraints",
    ]
    for label, reward in zip(labels, test_rewards):
        print(f"  • {label}: reward={reward:.3f}")

    ask_finance_r, ask_security_r, repeat_r, broadcast_r, malformed_r, submit_r = test_rewards
    checks = [
        ("oracle ask_finance is positive", ask_finance_r > 0.0),
        ("oracle ask_security is positive", ask_security_r > 0.0),
        ("repeat < oracle ask_finance", repeat_r < ask_finance_r),
        ("broadcast <= repeat", broadcast_r <= repeat_r),
        ("malformed JSON is negative", malformed_r < 0.0),
        ("submit_final > oracle ask_finance", submit_r > ask_finance_r),
    ]
    all_ok = True
    for name, passed in checks:
        status = "✓" if passed else "✗"
        print(f"  {status} invariant: {name}")
        all_ok = all_ok and passed
    if not all_ok:
        raise RuntimeError("Reward verification invariants failed. Check reward shaping logic.")

    if args.dry_run:
        print("\n[DRY RUN] Dataset and reward function verified. Skipping training.")
        print("  Run without --dry-run on GPU to train.")
        return
    
    # FOR DRY RUN ONLY
    if not HAS_TRL:
        raise RuntimeError("TRL is required for actual training on the GPU.")
    # Load model
    print(f"\n[3/4] Loading model: {args.model}")

    if args.use_unsloth:
        try:
            from unsloth import FastLanguageModel
            HAS_UNSLOTH = True
        except Exception as e:
            raise RuntimeError(f"Unsloth failed to load: {e}\nRun without --use-unsloth instead.")
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=args.model,
            max_seq_length=args.max_seq_length,
            load_in_4bit=True,
            dtype=None,  # Auto-detect
        )
        model = FastLanguageModel.get_peft_model(
            model,
            r=16,
            lora_alpha=32,
            lora_dropout=0.0,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                           "gate_proj", "up_proj", "down_proj"],
            use_gradient_checkpointing="unsloth",
        )
        # Align generation config for clean termination.
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
        if isinstance(im_end_id, int) and im_end_id >= 0:
            tokenizer.eos_token = "<|im_end|>"
        model.config.pad_token_id = tokenizer.pad_token_id
        if tokenizer.eos_token_id is not None:
            model.config.eos_token_id = tokenizer.eos_token_id
            model.generation_config.eos_token_id = tokenizer.eos_token_id
        model.generation_config.pad_token_id = tokenizer.pad_token_id
        print("  Unsloth LoRA loaded (4-bit quantization)")
    else:
        if not HAS_TRANSFORMERS:
            raise RuntimeError("pip install transformers")
        tokenizer = AutoTokenizer.from_pretrained(args.model)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        # Qwen chat models typically terminate on <|im_end|>.
        im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
        if isinstance(im_end_id, int) and im_end_id >= 0:
            tokenizer.eos_token = "<|im_end|>"

        model = AutoModelForCausalLM.from_pretrained(args.model)

        # Keep model/generation config aligned with tokenizer.
        model.config.pad_token_id = tokenizer.pad_token_id
        if tokenizer.eos_token_id is not None:
            model.config.eos_token_id = tokenizer.eos_token_id
            model.generation_config.eos_token_id = tokenizer.eos_token_id
        model.generation_config.pad_token_id = tokenizer.pad_token_id

        print("  Standard transformers model loaded")

    # GRPO Training
    print(f"\n[4/4] Starting GRPO training...")
    print(f"  States: {len(records)} | Group size (G): {args.group_size}")
    print(f"  Epochs: {args.epochs} | LR: {args.lr}")
    print(f"  Total updates: ~{int(len(records) * args.epochs / args.batch_size)}")

    # Build explicit stop-token list for GRPO sampling.
    eos_ids = []
    if tokenizer.eos_token_id is not None:
        eos_ids.append(int(tokenizer.eos_token_id))
    im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
    if isinstance(im_end_id, int) and im_end_id >= 0:
        eos_ids.append(int(im_end_id))
    close_brace_id = tokenizer.convert_tokens_to_ids("}")
    if isinstance(close_brace_id, int) and close_brace_id >= 0:
        eos_ids.append(int(close_brace_id))
    # Preserve order, remove duplicates.
    eos_ids = list(dict.fromkeys(eos_ids))

    generation_kwargs = {
        "eos_token_id": eos_ids if len(eos_ids) > 1 else (eos_ids[0] if eos_ids else None),
        "pad_token_id": tokenizer.pad_token_id,
    }
    # Remove None values.
    generation_kwargs = {k: v for k, v in generation_kwargs.items() if v is not None}

    config = GRPOConfig(
        output_dir=str(output_dir),

        # GRPO-specific
        num_generations=args.group_size,      # G: sample this many completions per prompt
        max_completion_length=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=40,
        repetition_penalty=1.05,
        generation_kwargs=generation_kwargs,
        mask_truncated_completions=True,

        # Standard training
        learning_rate=args.lr,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=max(1, args.batch_size),
        gradient_accumulation_steps=args.grad_accum,

        # Logging
        logging_steps=1,
        save_strategy="epoch",
        report_to=[],                          # Set to ["wandb"] if you have it configured
    )

    trainer = GRPOTrainer(
            model=model,
            processing_class=tokenizer,
            args=config,
            reward_funcs=reward_fn,    
            train_dataset=dataset,
        )

    trainer.train()

    # ── Save everything ────────────────────────────────────────────────────────
    trainer.save_model(str(output_dir / "final_model"))
    tokenizer.save_pretrained(str(output_dir / "final_model"))
    print(f"\n  Model saved → {output_dir}/final_model")

    # Save metrics
    metrics_path = output_dir / "grpo_metrics.json"
    with metrics_path.open("w") as f:
        json.dump(trainer.state.log_history, f, indent=2)
    print(f"  Metrics saved → {metrics_path}")

    # Save plots
    save_training_plots(trainer.state.log_history, output_dir)

    # Summary
    log = trainer.state.log_history
    losses = [e["loss"] for e in log if "loss" in e]
    if losses:
        print(f"\n  Initial loss:  {losses[0]:.4f}")
        print(f"  Final loss:    {losses[-1]:.4f}")
        print(f"  Improvement:   {((losses[0] - losses[-1]) / losses[0] * 100):.1f}%")

    print(f"\n{'='*60}")
    print(f"  GRPO TRAINING COMPLETE")
    print(f"  Model:   {output_dir}/final_model")
    print(f"  Plots:   {output_dir}/loss_curve.png")
    print(f"           {output_dir}/reward_curve.png")
    print(f"  Metrics: {output_dir}/grpo_metrics.json")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()