train_colab.py

#!/usr/bin/env python3
"""Component 6 — Training Script (Colab-Ready, GRPO + Unsloth).

Trains the detective agent using GRPO (Group Relative Policy Optimisation)
while other agents use fixed prompt-based LLM calls.  Supports curriculum
difficulty staging per §6 of the hackathon guide.
"""

# Unsloth is imported at the start of load_models() so its class-level patches
# are applied before any model is instantiated.  Both detective and NPC are
# loaded via FastLanguageModel, ensuring Qwen2RotaryEmbedding instances get
# extend_rope_embedding set at construction time (required by Unsloth ≥2026.1).
# lora_dropout=0.0 and use_gradient_checkpointing=False are required to avoid
# in-place buffer conflicts with manual per-completion .backward() in GRPO.
USE_UNSLOTH = False  # updated inside load_models() if package is available


import copy
import json
import os
import sys
import time
import warnings
from typing import Optional
import random

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from torch.optim import AdamW
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from peft import LoraConfig, TaskType, get_peft_model

# Unsloth (optional — 2x faster training on supported GPUs)

# Add project root to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from crime_env.environment import CrimeInvestigationEnv
from crime_env.case_generator import generate_case


# ── Configuration ───────────────────────────────────────────────────────────


def _env_bool(name: str, default: bool = False) -> bool:
    value = os.environ.get(name)
    if value is None:
        return default
    return value.strip().lower() in {"1", "true", "yes", "on"}


def _env_int(name: str, default: int) -> int:
    value = os.environ.get(name)
    if value is None:
        return default
    try:
        return int(value)
    except ValueError:
        return default

# Default profile tuned for 6GB VRAM reliability.
TEST_MODE = _env_bool("TEST_MODE", False)
MODEL_NAME = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-1.5B-Instruct")
NPC_MODEL_NAME = os.environ.get("NPC_MODEL_NAME", "Qwen/Qwen2.5-0.5B-Instruct")
NUM_EPISODES = _env_int("NUM_EPISODES", 5 if TEST_MODE else 300)
MAX_TURNS = _env_int("MAX_TURNS", 8 if TEST_MODE else 15)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
REWARDS_FILE = "rewards.json"
TRANSCRIPTS_FILE = "episode_transcripts.json"
NPC_MAX_NEW_TOKENS = _env_int("NPC_MAX_NEW_TOKENS", 48 if TEST_MODE else 80)
DETECTIVE_MAX_NEW_TOKENS = _env_int(
    "DETECTIVE_MAX_NEW_TOKENS", 40 if TEST_MODE else 56
)
DETECTIVE_RETRY_MAX_NEW_TOKENS = _env_int(
    "DETECTIVE_RETRY_MAX_NEW_TOKENS", 24 if TEST_MODE else 40
)
TRANSCRIPT_EPISODES = {
    int(x.strip())
    for x in os.environ.get("TRANSCRIPT_EPISODES", "1,25,50").split(",")
    if x.strip().isdigit()
}
STRICT_FORMAT_FALLBACK_THRESHOLD = float(
    os.environ.get("STRICT_FORMAT_FALLBACK_THRESHOLD", "0.35")
)
STRICT_FORMAT_WINDOW = _env_int("STRICT_FORMAT_WINDOW", 5)

# GRPO hyperparameters
NUM_GENERATIONS = _env_int("NUM_GENERATIONS", 1 if TEST_MODE else 4)
GRPO_BETA = float(os.environ.get("GRPO_BETA", "0.04"))     # KL coefficient
GRPO_LR = float(os.environ.get("GRPO_LR", "1e-5"))         # learning rate
GRPO_MAX_GRAD_NORM = 0.5

# Curriculum thresholds (§6)
CURRICULUM_ADVANCE_RATE = 0.6   # rolling accuracy required to advance tier
CURRICULUM_WINDOW = 10          # episodes over which accuracy is averaged


def _format_chat_prompt(
    tokenizer,
    system_prompt: str,
    user_prompt: str,
) -> str:
    """Format prompts using the tokenizer chat template when available."""
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    if hasattr(tokenizer, "apply_chat_template"):
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )

    formatted = f"System: {system_prompt}\n\nUser: {user_prompt}\n\nAssistant:"
    return formatted


# ── Model Loading ───────────────────────────────────────────────────────────


def _load_detective_standard(quant_config, use_quantization: bool):
    """Load detective model via standard transformers + PEFT (LoRA)."""
    model_dtype = torch.float16 if DEVICE == "cuda" else torch.float32
    lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules=["q_proj", "v_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    ) if DEVICE == "cuda" else None

    load_kwargs = {
        "quantization_config": quant_config,
        "torch_dtype": model_dtype,
        "device_map": "auto" if DEVICE == "cuda" else None,
        "trust_remote_code": True,
    }

    try:
        base_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, **load_kwargs)
    except ValueError as e:
        error_text = str(e)
        if DEVICE == "cuda" and "dispatched on the CPU or the disk" in error_text:
            print("Low VRAM detected, retrying with CPU offload...")
            offload_quant = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                llm_int8_enable_fp32_cpu_offload=True,
            )
            load_kwargs["quantization_config"] = offload_quant
            base_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, **load_kwargs)
        else:
            raise

    if lora_config is not None:
        try:
            detective_model = get_peft_model(base_model, lora_config)
            print("Detective model loaded via transformers + LoRA.")
        except ValueError as e:
            if "Target modules" in str(e) and "not found" in str(e):
                print("LoRA target modules not found, loading without LoRA.")
                detective_model = base_model
            else:
                raise
    else:
        detective_model = base_model

    return detective_model


def load_models():
    """Load the detective (trainable), reference (frozen), and NPC (frozen) models.

    Returns:
        (detective_model, ref_model, npc_pipeline, tokenizer, npc_tokenizer)
        - detective_model: FastLanguageModel + LoRA adapters (or standard PEFT fallback)
        - ref_model: Frozen base model on CPU for GRPO KL penalty (None when Unsloth active)
        - npc_pipeline: text-generation pipeline for NPC agents
        - tokenizer: detective tokenizer
        - npc_tokenizer: NPC tokenizer
    """
    global USE_UNSLOTH

    # ── Import Unsloth first so patches are applied before any instantiation ─
    # Using FastLanguageModel for the detective ensures Qwen2RotaryEmbedding
    # instances get extend_rope_embedding set at construction time, which
    # Unsloth ≥2026.1 requires in its patched LlamaModel_fast_forward.
    _FastLanguageModel = None
    if DEVICE == "cuda":
        try:
            from unsloth import FastLanguageModel as _FastLanguageModel
            USE_UNSLOTH = True
            print("Unsloth: enabled")
        except ImportError:
            print("Unsloth: not available, using standard transformers")

    print(f"Loading model: {MODEL_NAME}")
    print(f"Device: {DEVICE}")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    npc_tokenizer = AutoTokenizer.from_pretrained(
        NPC_MODEL_NAME, trust_remote_code=True,
    )
    if npc_tokenizer.pad_token is None:
        npc_tokenizer.pad_token = npc_tokenizer.eos_token

    model_config = AutoConfig.from_pretrained(MODEL_NAME, trust_remote_code=True)
    hidden_size = getattr(model_config, "hidden_size", None)
    use_quantization = DEVICE == "cuda" and (hidden_size is None or hidden_size >= 64)

    quant_config = None
    if use_quantization:
        try:
            quant_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
            )
        except Exception:
            print("BitsAndBytes not available, loading without quantisation")
    elif DEVICE == "cuda":
        print("Skipping 4-bit quantisation for small hidden-size model")

    # ── Detective model (trainable) ──────────────────────────────────────
    # Load via FastLanguageModel when Unsloth is available so rotary_emb
    # instances get extend_rope_embedding correctly initialised.
    # lora_dropout=0.0 is required for Unsloth's fast kernel path.
    # use_gradient_checkpointing=False avoids in-place buffer reuse that
    # breaks per-completion .backward() calls in grpo_update.
    if _FastLanguageModel is not None:
        try:
            detective_model, tokenizer = _FastLanguageModel.from_pretrained(
                model_name=MODEL_NAME,
                max_seq_length=1024,
                dtype=torch.float16,
                load_in_4bit=use_quantization,
                trust_remote_code=True,
            )
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
            detective_model = _FastLanguageModel.get_peft_model(
                detective_model,
                r=8,
                target_modules=["q_proj", "v_proj"],
                lora_alpha=16,
                lora_dropout=0.0,
                bias="none",
                use_gradient_checkpointing=False,
                random_state=42,
            )
            print("Detective model loaded via Unsloth + LoRA.")
        except Exception as e:
            print(f"Unsloth detective load failed ({e}), falling back to transformers.")
            detective_model = _load_detective_standard(quant_config, use_quantization)
    else:
        detective_model = _load_detective_standard(quant_config, use_quantization)

    # ── Reference model (frozen, CPU) — used for GRPO KL penalty ────────
    # Skipped when Unsloth is active: Unsloth's Triton kernels require CUDA
    # tensors and crash on CPU-resident models.  grpo_update handles
    # ref_model=None gracefully (KL term is omitted).
    ref_model = None
    if TEST_MODE:
        print("TEST_MODE: skipping reference model load (KL penalty disabled).")
    elif USE_UNSLOTH:
        print("Unsloth active: skipping reference model (KL penalty disabled to avoid Triton/CPU clash).")
    else:
        print("Loading reference model on CPU (for GRPO KL penalty) ...")
        try:
            ref_model = AutoModelForCausalLM.from_pretrained(
                MODEL_NAME,
                torch_dtype=torch.float16,
                device_map="cpu",
                trust_remote_code=True,
            )
            ref_model.eval()
            for param in ref_model.parameters():
                param.requires_grad = False
            print("Reference model loaded on CPU.")
        except Exception as e:
            print(f"Warning: reference model failed to load ({e}). KL penalty will be skipped.")

    # ── NPC pipeline — separate frozen copy ─────────────────────────────
    npc_base_model = None
    try:
        if _FastLanguageModel is not None:
            npc_base_model, _npc_tok = _FastLanguageModel.from_pretrained(
                model_name=NPC_MODEL_NAME,
                max_seq_length=512,
                dtype=torch.float16,
                load_in_4bit=use_quantization,
                trust_remote_code=True,
            )
        else:
            npc_base_model = AutoModelForCausalLM.from_pretrained(
                NPC_MODEL_NAME,
                quantization_config=quant_config,
                torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
                device_map="auto" if DEVICE == "cuda" else None,
                trust_remote_code=True,
            )
        npc_base_model.eval()
        for param in npc_base_model.parameters():
            param.requires_grad = False
    except Exception as e:
        raise RuntimeError(f"NPC model failed to load: {e}") from e

    npc_pipeline = pipeline(
        "text-generation",
        model=npc_base_model,
        tokenizer=npc_tokenizer,
        max_new_tokens=NPC_MAX_NEW_TOKENS,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )

    return detective_model, ref_model, npc_pipeline, tokenizer, npc_tokenizer


# ── NPC LLM Call ────────────────────────────────────────────────────────────


def make_npc_call(npc_pipeline, npc_tokenizer):
    """Create a callable for NPC agent responses."""

    def llm_call(system_prompt: str, conversation_history: list[dict]) -> str:
        # Build a bounded user prompt to avoid model context overflows.
        user_prompt = "Conversation so far:\n"

        # Keep only recent turns and constrain prompt length.
        recent_history = conversation_history[-10:]
        for entry in recent_history:
            speaker = entry.get("speaker", "Unknown")
            content = entry.get("content", "")
            user_prompt += f"{speaker}: {content[:180]}\n"

        user_prompt = user_prompt[-1800:]
        prompt_text = _format_chat_prompt(
            npc_tokenizer,
            system_prompt[:800],
            user_prompt + "\nRespond in 1-2 short sentences in character. No action format.",
        )

        try:
            result = npc_pipeline(prompt_text, return_full_text=False)
            response = result[0]["generated_text"].strip()
            # Clean up: take first sentence/paragraph
            if "\n" in response:
                response = response.split("\n")[0]
            return response[:300] if response else "I have nothing to add."
        except Exception as e:
            print(f"  NPC call error: {e}")
            return "I don't recall anything specific about that."

    return llm_call


# ── Detective Action Generation ─────────────────────────────────────────────


def generate_detective_action(
    detective_model,
    tokenizer,
    observation: dict,
    strict_action_format: bool = False,
    episode_max_turns: int = MAX_TURNS,
) -> tuple[str, torch.Tensor, torch.Tensor, bool]:
    """Generate a detective action using the trainable model.

    Returns:
        (action_string, query_tensor, response_tensor, used_fallback)
    """
    # Force-accuse at 70 % of episode length so it scales with difficulty tier.
    force_accuse_at = max(1, int(episode_max_turns * 0.7))

    # Build prompt from observation
    prompt = f"""You are a detective. Based on the conversation so far, choose your next action.

Briefing: {observation['briefing'][:300]}

Turn: {observation['turn']}/{episode_max_turns}

Recent conversation:
"""
    history = observation.get("conversation_history", [])
    for entry in history[-6:]:
        prompt += f"  {entry['speaker']}: {entry['content'][:100]}\n"

    prompt += """
Choose ONE action using EXACTLY this format:
ACTION: ask_question | TARGET: Suspect_A | CONTENT: <question>
ACTION: request_evidence | ITEM: keycard_log
ACTION: accuse | TARGET: Suspect_A

Your action:
"""

    turn_now = observation.get("turn", 0)
    if turn_now >= force_accuse_at:
        prompt += (
            f"\nWARNING: Turn {turn_now}/{episode_max_turns}. "
            "You have gathered enough information. You MUST accuse a suspect NOW. "
            "Use: ACTION: accuse | TARGET: Suspect_A  or  ACTION: accuse | TARGET: Suspect_B\n"
        )

    if strict_action_format:
        prompt += (
            "\nIMPORTANT: Output ONLY a single valid ACTION line. "
            "No explanations, no extra text.\n"
        )

    formatted_prompt = _format_chat_prompt(
        tokenizer,
        "You are a detective. Choose your next action.",
        prompt,
    )

    # Tokenize
    inputs = tokenizer(
        formatted_prompt,
        return_tensors="pt",
        truncation=True,
        max_length=1024,
    )
    query_tensor = inputs["input_ids"].squeeze().detach().cpu()

    if DEVICE == "cuda":
        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}

    # Generate
    with torch.no_grad():
        output = detective_model.generate(
            **inputs,
            max_new_tokens=DETECTIVE_MAX_NEW_TOKENS,
            do_sample=True,
            temperature=0.4 if strict_action_format else 0.8,
            top_p=0.8 if strict_action_format else 0.9,
            pad_token_id=tokenizer.pad_token_id,
        )

    response_tensor = output.squeeze()[len(query_tensor):].detach().cpu()
    action_text = tokenizer.decode(response_tensor, skip_special_tokens=True).strip()

    def _is_valid_action(text: str) -> bool:
        normalized = text.strip().upper()
        return normalized.startswith("ACTION:")

    # Retry once in strict mode before falling back to scripted actions.
    if not _is_valid_action(action_text):
        strict_prompt = prompt + (
            "\nFINAL REMINDER: Reply with one line starting with 'ACTION:' only.\n"
        )
        strict_formatted_prompt = _format_chat_prompt(
            tokenizer,
            "You are a detective. Choose your next action.",
            strict_prompt,
        )
        strict_inputs = tokenizer(
            strict_formatted_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=1024,
        )
        if DEVICE == "cuda":
            strict_inputs = {k: v.to(DEVICE) for k, v in strict_inputs.items()}

        with torch.no_grad():
            retry_output = detective_model.generate(
                **strict_inputs,
                max_new_tokens=DETECTIVE_RETRY_MAX_NEW_TOKENS,
                do_sample=True,
                temperature=0.25,
                top_p=0.75,
                pad_token_id=tokenizer.pad_token_id,
            )

        strict_query_tensor = strict_inputs["input_ids"].squeeze().detach().cpu()
        retry_response_tensor = retry_output.squeeze()[len(strict_query_tensor):].detach().cpu()
        retry_action_text = tokenizer.decode(
            retry_response_tensor, skip_special_tokens=True
        ).strip()
        if _is_valid_action(retry_action_text):
            query_tensor = strict_query_tensor
            response_tensor = retry_response_tensor
            action_text = retry_action_text

    used_fallback = False

    # If model output doesn't parse, generate a deterministic fallback action.
    if not _is_valid_action(action_text):
        used_fallback = True
        turn = observation.get("turn", 0)
        if turn < 4:
            action_text = "ACTION: ask_question | TARGET: Suspect_A | CONTENT: Where were you at the time of the crime?"
        elif turn < 8:
            action_text = "ACTION: ask_question | TARGET: Suspect_B | CONTENT: Can you describe your alibi?"
        elif turn < 10:
            action_text = "ACTION: request_evidence | ITEM: keycard_log"
        else:
            # Late turns: accuse with a randomized suspect to avoid systemic bias.
            fallback_target = random.choice(["Suspect_A", "Suspect_B"])
            action_text = f"ACTION: accuse | TARGET: {fallback_target}"

    # Never train on scripted fallback actions (they are off-policy).
    if used_fallback:
        response_tensor = tokenizer(
            action_text,
            return_tensors="pt",
            add_special_tokens=False,
        )["input_ids"].squeeze()

    return action_text, query_tensor, response_tensor, used_fallback


# ── GRPO Core ────────────────────────────────────────────────────────────────


def score_candidates(
    env: CrimeInvestigationEnv,
    actions: list[str],
) -> list[float]:
    """Score each candidate action by stepping isolated env copies.

    Returns the *incremental* reward delta for each action (not cumulative),
    so advantages reflect only what each candidate action contributes at this
    step — not the accumulated reward from previous turns.
    """
    llm_call_ref = env.llm_call  # callable; restore after deepcopy
    base_reward  = float(env.reward_calc.get_rewards()["detective"])
    rewards: list[float] = []

    for action in actions:
        try:
            env_copy          = copy.deepcopy(env)
            env_copy.llm_call = llm_call_ref
            env_copy.step(action)
            delta = float(env_copy.reward_calc.get_rewards()["detective"]) - base_reward
        except Exception as exc:
            print(f"  score_candidates error: {exc}")
            delta = 0.0
        rewards.append(delta)

    return rewards


def grpo_update(
    detective_model,
    ref_model,
    optimizer: AdamW,
    query_tensors: list,
    response_tensors: list,
    rewards: list[float],
    beta: float = GRPO_BETA,
    max_grad_norm: float = GRPO_MAX_GRAD_NORM,
) -> dict:
    """One GRPO gradient step over a group of G completions.

    Group-relative advantage:  A_i = (r_i − μ_r) / (σ_r + ε)
    Per-completion loss:       −A_i · mean_log_prob(response) + β · KL(π ∥ π_ref)
    Final loss:                mean over the group.

    Args:
        detective_model:  Trainable policy model.
        ref_model:        Frozen reference model on CPU (None ⇒ skip KL).
        optimizer:        AdamW for detective_model.
        query_tensors:    List of G query tensors (CPU int64).
        response_tensors: List of G response tensors (CPU int64).
        rewards:          List of G float rewards.
        beta:             KL coefficient.
        max_grad_norm:    Gradient clipping threshold.

    Returns:
        Dict with keys "loss" and "kl".
    """
    G = len(rewards)
    if G < 2:
        # Need ≥ 2 completions for a non-trivial group-relative advantage.
        return {"loss": 0.0, "kl": 0.0, "skipped": True}

    mean_r = float(np.mean(rewards))
    std_r  = float(np.std(rewards)) + 1e-8
    advantages = [(float(r) - mean_r) / std_r for r in rewards]

    losses: list[torch.Tensor] = []
    kl_vals: list[float] = []

    # Zero grad once before all forward passes so gradient accumulation is
    # correct and Unsloth's in-place gradient-checkpointing ops don't corrupt
    # tensors that are still live in the computation graph.
    optimizer.zero_grad()

    for q_cpu, r_cpu, adv in zip(query_tensors, response_tensors, advantages):
        if r_cpu is None or r_cpu.numel() == 0:
            continue

        q_len = q_cpu.numel()
        r_len = r_cpu.numel()

        # Full sequence on compute device
        input_ids = torch.cat([q_cpu.long(), r_cpu.long()]).unsqueeze(0).to(DEVICE)

        # ── Policy forward (with grad) ────────────────────────────────
        logits = detective_model(input_ids).logits          # [1, seq, vocab]
        resp_logits = logits[0, q_len - 1 : q_len + r_len - 1, :]  # [r_len, vocab]
        resp_ids    = input_ids[0, q_len:]                          # [r_len]

        log_probs     = F.log_softmax(resp_logits, dim=-1)          # has grad
        token_lp      = log_probs.gather(1, resp_ids.unsqueeze(1)).squeeze(1)
        mean_log_prob = token_lp.mean()                             # scalar, has grad

        # ── Reference forward (no grad, CPU) ─────────────────────────
        kl_loss = torch.tensor(0.0, device=DEVICE)
        if ref_model is not None:
            with torch.no_grad():
                ref_out         = ref_model(input_ids.cpu())
                ref_resp_logits = ref_out.logits[0, q_len - 1 : q_len + r_len - 1, :]
                ref_log_probs   = F.log_softmax(ref_resp_logits, dim=-1)   # CPU, no grad

            ref_log_probs_dev = ref_log_probs.to(DEVICE)
            # KL(π ∥ π_ref): E_v[π(v) · (log π(v) − log π_ref(v))]
            pi_detached = log_probs.detach().exp()
            kl_per_token = (pi_detached * (log_probs - ref_log_probs_dev)).sum(dim=-1)
            kl_loss = kl_per_token.mean()
            kl_vals.append(kl_loss.item())

        # ── GRPO step loss — backward immediately to avoid stale graph ─
        step_loss = -adv * mean_log_prob + beta * kl_loss
        # Divide by G so the accumulated gradient equals the group mean.
        (step_loss / G).backward()
        losses.append(step_loss.detach())

    if not losses:
        return {"loss": 0.0, "kl": 0.0, "skipped": True}

    total_loss_val = float(torch.stack(losses).mean().item())
    torch.nn.utils.clip_grad_norm_(
        [p for p in detective_model.parameters() if p.requires_grad],
        max_grad_norm,
    )
    optimizer.step()

    if DEVICE == "cuda":
        torch.cuda.empty_cache()

    return {
        "loss": total_loss_val,
        "kl":   float(np.mean(kl_vals)) if kl_vals else 0.0,
    }


# ── Training Loop ───────────────────────────────────────────────────────────


def train():
    """Main GRPO training loop with curriculum difficulty staging."""
    print("=" * 60)
    print("  AI Crime Investigation World — GRPO Training")
    print("=" * 60)
    if TEST_MODE:
        print("  Running in TEST_MODE (smoke-test settings enabled)")
        print(
            f"  MODEL_NAME={MODEL_NAME} | NPC_MODEL_NAME={NPC_MODEL_NAME} | "
            f"NUM_EPISODES={NUM_EPISODES} | MAX_TURNS={MAX_TURNS}"
        )
    print(
        f"  GRPO: num_generations={NUM_GENERATIONS} | beta={GRPO_BETA} | lr={GRPO_LR}"
    )
    print("=" * 60)

    warnings.filterwarnings("ignore", message="No dataset is provided.*",     category=UserWarning)
    warnings.filterwarnings("ignore", message=r"std\(\).*degrees of freedom", category=UserWarning)
    warnings.filterwarnings("ignore", message=".*pipelines sequentially.*",    category=UserWarning)

    # Load models
    detective_model, ref_model, npc_pipeline, tokenizer, npc_tokenizer = load_models()

    # Optimizer — only trains LoRA parameters
    trainable_params = [p for p in detective_model.parameters() if p.requires_grad]
    optimizer = AdamW(trainable_params, lr=GRPO_LR)

    # Environment
    npc_call = make_npc_call(npc_pipeline, npc_tokenizer)
    env = CrimeInvestigationEnv(llm_call=npc_call)

    # Tracking
    reward_log:     list[float] = []
    results_log:    list[str]   = []
    transcript_log: list[dict]  = []
    difficulty_log: list[str]   = []

    # Curriculum state
    correct_history: list[int] = []   # 1 = correct accusation, 0 = otherwise

    print(f"\nStarting training for {NUM_EPISODES} episodes...\n")

    strict_action_format   = False
    fallback_rate_history: list[float] = []

    def _get_difficulty(episode: int) -> str:
        """Pick difficulty tier with time-based floor + rolling-accuracy auto-advance."""
        floor = "easy" if episode < 50 else ("medium" if episode < 150 else "hard")

        if len(correct_history) >= CURRICULUM_WINDOW:
            rolling = sum(correct_history[-CURRICULUM_WINDOW:]) / CURRICULUM_WINDOW
            if rolling >= CURRICULUM_ADVANCE_RATE:
                if floor == "easy":
                    return "medium"
                if floor == "medium":
                    return "hard"
        return floor

    def save_training_artifacts(checkpoint_dir: Optional[str] = None) -> None:
        with open(REWARDS_FILE, "w") as f:
            json.dump(
                {
                    "rewards":      reward_log,
                    "results":      results_log,
                    "difficulty":   difficulty_log,
                    "num_episodes": NUM_EPISODES,
                    "model":        MODEL_NAME,
                },
                f, indent=2,
            )
        with open(TRANSCRIPTS_FILE, "w") as f:
            json.dump(
                {
                    "episodes_captured": sorted([t["episode"] for t in transcript_log]),
                    "transcripts":       transcript_log,
                },
                f, indent=2,
            )
        _plot_reward_curve(reward_log)
        if checkpoint_dir is not None:
            os.makedirs(checkpoint_dir, exist_ok=True)
            detective_model.save_pretrained(checkpoint_dir)
            tokenizer.save_pretrained(checkpoint_dir)

    for episode in range(NUM_EPISODES):
        t0 = time.time()

        # ── Curriculum: pick difficulty and generate matching case ────
        difficulty = _get_difficulty(episode)
        difficulty_log.append(difficulty)

        case          = generate_case(difficulty=difficulty)
        episode_turns = case.get("max_turns", MAX_TURNS)
        env.MAX_TURNS = episode_turns
        obs  = env.reset(case_data=case)
        done = False

        fallback_steps = 0
        total_steps    = 0
        grpo_updates   = 0
        total_loss_sum = 0.0

        while not done:
            # ── Generate G candidate actions ──────────────────────────
            candidates: list[tuple[str, torch.Tensor, torch.Tensor, bool]] = []
            for _ in range(NUM_GENERATIONS):
                action, q_t, r_t, used_fb = generate_detective_action(
                    detective_model, tokenizer, obs,
                    strict_action_format=strict_action_format,
                    episode_max_turns=episode_turns,
                )
                candidates.append((action, q_t, r_t, used_fb))

            # ── Score all candidates via env copies (reward_fn) ───────
            group_rewards = score_candidates(env, [c[0] for c in candidates])

            # ── GRPO update on non-fallback candidates ────────────────
            valid = [(c, r) for c, r in zip(candidates, group_rewards) if not c[3]]
            if len(valid) >= 2:
                stats = grpo_update(
                    detective_model, ref_model, optimizer,
                    [c[1] for c, _ in valid],
                    [c[2] for c, _ in valid],
                    [r    for _, r in valid],
                )
                if not stats.get("skipped"):
                    grpo_updates   += 1
                    total_loss_sum += stats.get("loss", 0.0)

            # ── Advance main env with highest-reward action ───────────
            best_idx          = int(np.argmax(group_rewards))
            best_action, _, _, best_fb = candidates[best_idx]
            obs, reward, done, info    = env.step(best_action)
            total_steps += 1
            if best_fb:
                fallback_steps += 1

        # ── Episode end ───────────────────────────────────────────────
        final_rewards   = env.reward_calc.get_rewards()
        detective_reward = final_rewards["detective"]
        reward_log.append(detective_reward)

        if info.get("action") == "accuse":
            result = "correct" if info.get("correct") else "wrong"
        else:
            result = "timeout"
        results_log.append(result)

        correct_history.append(1 if result == "correct" else 0)
        if len(correct_history) > CURRICULUM_WINDOW * 3:
            correct_history.pop(0)

        if (episode + 1) in TRANSCRIPT_EPISODES:
            transcript_log.append(
                {
                    "episode":              episode + 1,
                    "difficulty":           difficulty,
                    "result":               result,
                    "detective_reward":     round(detective_reward, 4),
                    "turns":                env.turn,
                    "criminal":             env.case.get("criminal") if env.case else None,
                    "crime":                env.case.get("crime")    if env.case else None,
                    "location":             env.case.get("location") if env.case else None,
                    "conversation_history": list(env.conversation_history),
                    "evidence_log":         list(env.evidence_log),
                }
            )

        fallback_rate = fallback_steps / max(1, total_steps)
        fallback_rate_history.append(fallback_rate)
        if len(fallback_rate_history) > max(1, STRICT_FORMAT_WINDOW):
            fallback_rate_history.pop(0)
        rolling_fallback_rate = sum(fallback_rate_history) / len(fallback_rate_history)
        strict_action_format  = rolling_fallback_rate > STRICT_FORMAT_FALLBACK_THRESHOLD

        elapsed  = time.time() - t0
        avg_loss = total_loss_sum / max(1, grpo_updates)
        print(
            f"Episode {episode + 1:>3}/{NUM_EPISODES} | "
            f"Diff: {difficulty:<6} | "
            f"Result: {result:<7} | "
            f"Reward: {detective_reward:>+7.2f} | "
            f"Turns: {env.turn:>2} | "
            f"GRPO: {grpo_updates:>2} steps | "
            f"Loss: {avg_loss:>6.4f} | "
            f"Fallback: {fallback_steps}/{max(1, total_steps)} "
            f"({fallback_rate*100:>5.1f}%) | "
            f"Time: {elapsed:.1f}s"
        )

        checkpoint_dir = f"./checkpoint_ep{episode + 1}" if (episode + 1) % 25 == 0 else None
        save_training_artifacts(checkpoint_dir=checkpoint_dir)
        if checkpoint_dir:
            print(f"Checkpoint saved: {checkpoint_dir}")

    # ── Final save and summary ────────────────────────────────────────────
    save_training_artifacts()
    print(f"\nReward log saved to {REWARDS_FILE}")
    print(f"Episode transcripts saved to {TRANSCRIPTS_FILE}")
    print("\n" + "=" * 60)
    print("  TRAINING SUMMARY")
    print("=" * 60)
    print(f"  Episodes:             {NUM_EPISODES}")
    print(f"  Correct accusations:  {results_log.count('correct')}")
    print(f"  Wrong accusations:    {results_log.count('wrong')}")
    print(f"  Timeouts:             {results_log.count('timeout')}")
    print(f"  Mean reward (last 50):  {np.mean(reward_log[-50:]):.2f}")
    print(f"  Mean reward (first 50): {np.mean(reward_log[:50]):.2f}")
    print("=" * 60)


def _plot_reward_curve(reward_log: list[float]) -> None:
    """Plot and save the reward curve."""
    fig, ax = plt.subplots(figsize=(12, 5))
    episodes = np.arange(1, len(reward_log) + 1)
    ax.plot(episodes, reward_log, alpha=0.3, color="#4a90d9", label="Per-episode")

    window = min(20, len(reward_log) // 4)
    if window > 1:
        smoothed = np.convolve(reward_log, np.ones(window) / window, mode="valid")
        ax.plot(
            episodes[window - 1:], smoothed,
            color="#e74c3c", linewidth=2, label=f"Rolling avg ({window} ep)",
        )

    ax.set_xlabel("Episode", fontsize=12)
    ax.set_ylabel("Detective Reward", fontsize=12)
    ax.set_title("AI Crime Investigation — Detective Reward Over Training", fontsize=14)
    ax.legend()
    ax.grid(True, alpha=0.3)
    ax.axhline(y=0, color="grey", linestyle="--", alpha=0.5)
    plt.tight_layout()
    plt.savefig("reward_curve.png", dpi=150)
    print("Reward curve saved to reward_curve.png")
    plt.close(fig)


# ── Entry Point ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    train()