train_colab.py

#!/usr/bin/env python3
"""Component 6 — Minimal Training Script (Colab-Ready).

Trains the detective agent using PPO (HuggingFace TRL) while other agents
use fixed prompt-based LLM calls. Designed for free-tier Colab GPU.
"""

import json
import os
import inspect
import sys
import time
import warnings
from typing import Optional

import matplotlib.pyplot as plt
import numpy as np
import torch
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline,
)
from peft import LoraConfig, TaskType
from trl import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead

# 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.agent_prompts import build_system_prompt


# ── 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 to a stronger model while keeping env override support.
TEST_MODE = _env_bool("TEST_MODE", False)
MODEL_NAME = os.environ.get("MODEL_NAME", "Qwen/Qwen2.5-3B-Instruct")
NPC_MODEL_NAME = os.environ.get(
    "NPC_MODEL_NAME",
    "Qwen/Qwen2.5-0.5B-Instruct" if TEST_MODE else MODEL_NAME,
)
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"
OUTPUT_DIR = "./ppo_detective"
REWARDS_FILE = "rewards.json"
TRANSCRIPTS_FILE = "episode_transcripts.json"
NPC_MAX_NEW_TOKENS = _env_int("NPC_MAX_NEW_TOKENS", 64 if TEST_MODE else 150)
DETECTIVE_MAX_NEW_TOKENS = _env_int(
    "DETECTIVE_MAX_NEW_TOKENS", 48 if TEST_MODE else 80
)
DETECTIVE_RETRY_MAX_NEW_TOKENS = _env_int(
    "DETECTIVE_RETRY_MAX_NEW_TOKENS", 32 if TEST_MODE else 64
)
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)


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


def load_models():
    """Load the detective (trainable) and NPC (fixed) models."""
    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

    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)

    # Quantization for memory efficiency on free Colab
    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 quantization")
    elif DEVICE == "cuda":
        print("Skipping 4-bit quantization for small hidden-size model")

    # LoRA config — enables PEFT so PPOTrainer reuses base weights as the
    # reference model instead of deepcopy-ing the quantized model (OOM fix).
    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

    # Detective model (trainable with value head)
    model_dtype = torch.float16 if DEVICE == "cuda" else torch.float32
    detective_load_kwargs = {
        "quantization_config": quant_config,
        "peft_config": lora_config,
        "dtype": model_dtype,
        "device_map": "auto" if DEVICE == "cuda" else None,
        "trust_remote_code": True,
    }

    try:
        detective_model = AutoModelForCausalLMWithValueHead.from_pretrained(
            MODEL_NAME,
            **detective_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 enabled for quantized layers...")
            offload_quant_config = 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,
            )
            detective_load_kwargs["quantization_config"] = offload_quant_config
            detective_load_kwargs["device_map"] = "auto"
            detective_model = AutoModelForCausalLMWithValueHead.from_pretrained(
                MODEL_NAME,
                **detective_load_kwargs,
            )
        else:
            raise

    # NPC pipeline — load a SEPARATE frozen copy of the base model.
    # Bug 9 fix: Using detective_model.pretrained_model would cause PPO
    # gradient updates to drift the NPC's behavior every step.
    npc_base_model = None
    try:
        npc_base_model = AutoModelForCausalLM.from_pretrained(
            NPC_MODEL_NAME,
            quantization_config=quant_config,
            dtype=model_dtype,
            device_map="auto" if DEVICE == "cuda" else None,
            trust_remote_code=True,
        )
        npc_base_model.eval()  # Freeze: no gradient tracking
        for param in npc_base_model.parameters():
            param.requires_grad = False
    except Exception as e:
        print(f"NPC model load warning: {e}")
        print("Falling back to detective base model for NPC responses to keep training running.")
        npc_base_model = detective_model.pretrained_model

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

    return detective_model, npc_pipeline, tokenizer


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


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

    def llm_call(system_prompt: str, conversation_history: list[dict]) -> str:
        # Build a bounded prompt to avoid model context overflows.
        messages = f"System: {system_prompt[:800]}\n\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", "")
            messages += f"{speaker}: {content[:180]}\n"

        # Hard cap prompt size for small-context models.
        messages = messages[-1800:]

        messages += "\nYour response:"

        try:
            result = npc_pipeline(messages, 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,
    env=None,
    strict_action_format: bool = False,
) -> tuple[str, torch.Tensor, torch.Tensor, bool]:
    """Generate a detective action using the trainable model.

    Returns:
        (action_string, query_tensor, response_tensor, used_fallback)
    """
    # 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']}/{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:
"""

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

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

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

    # Generate
    with torch.no_grad():
        output = detective_model.pretrained_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):]
    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_inputs = tokenizer(
            strict_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.pretrained_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):]
        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:
            action_text = "ACTION: ask_question | TARGET: Witness_1 | CONTENT: What did you see?"

    # Never train PPO on fallback actions: they are scripted and 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


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


def train():
    """Main PPO training loop."""
    print("=" * 60)
    print("  AI Crime Investigation World — PPO 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("=" * 60)

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

    warnings.filterwarnings(
        "ignore",
        message="No dataset is provided.*",
        category=UserWarning,
    )

    # Compatibility guard: this script targets the TRL 0.9 PPOTrainer API.
    ppo_init_params = inspect.signature(PPOTrainer.__init__).parameters
    if "config" not in ppo_init_params:
        raise RuntimeError(
            "Incompatible TRL version detected. "
            "Use `trl>=0.9.0,<0.10.0` for this training script."
        )

    # PPO config (support TRL variants with either ppo_epochs or num_ppo_epochs)
    ppo_cfg_kwargs = {
        "learning_rate": 1e-5,
        "batch_size": 1,
        "mini_batch_size": 1,
        "gradient_accumulation_steps": 1,
    }
    ppo_cfg_params = inspect.signature(PPOConfig.__init__).parameters
    if "num_ppo_epochs" in ppo_cfg_params:
        ppo_cfg_kwargs["num_ppo_epochs"] = 1  # single-sample stability
    elif "ppo_epochs" in ppo_cfg_params:
        ppo_cfg_kwargs["ppo_epochs"] = 1  # single-sample stability

    ppo_config = PPOConfig(**ppo_cfg_kwargs)

    # PPO Trainer
    ppo_trainer = PPOTrainer(
        config=ppo_config,
        model=detective_model,
        tokenizer=tokenizer,
    )

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

    # Tracking
    reward_log = []
    results_log = []
    transcript_log = []

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

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

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

        # Reset environment
        obs = env.reset()
        done = False
        fallback_steps = 0
        total_steps = 0
        ppo_updates = 0

        while not done:
            # Generate detective action
            action_text, query_tensor, response_tensor, used_fallback = (
                generate_detective_action(
                    detective_model,
                    tokenizer,
                    obs,
                    env=env,
                    strict_action_format=strict_action_format,
                )
            )

            # Step environment
            obs, reward, done, info = env.step(action_text)
            total_steps += 1

            if not used_fallback:
                # PPOConfig uses batch_size=1, so update on each valid step.
                try:
                    reward_tensor = torch.tensor([reward], dtype=torch.float32)
                    ppo_trainer.step(
                        [query_tensor],
                        [response_tensor],
                        [reward_tensor],
                    )
                    ppo_updates += 1
                except Exception as e:
                    print(f"  PPO update error (episode {episode}, step {total_steps}): {e}")
            else:
                fallback_steps += 1

        # Get final rewards
        final_rewards = env.reward_calc.get_rewards()
        detective_reward = final_rewards["detective"]
        reward_log.append(detective_reward)

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

        if (episode + 1) in TRANSCRIPT_EPISODES:
            transcript_log.append(
                {
                    "episode": episode + 1,
                    "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
        print(
            f"Episode {episode + 1:>3}/{NUM_EPISODES} | "
            f"Result: {result:<7} | "
            f"Reward: {detective_reward:>+7.2f} | "
            f"Turns: {env.turn:>2} | "
            f"PPO updates: {ppo_updates:>2} | "
            f"Fallback: {fallback_steps:>2}/{max(1, total_steps):<2} "
            f"({fallback_rate*100:>5.1f}%, avg={rolling_fallback_rate*100:>5.1f}%) | "
            f"Time: {elapsed:.1f}s"
        )

    # ── Save results ────────────────────────────────────────────────────

    # Save reward log
    with open(REWARDS_FILE, "w") as f:
        json.dump(
            {
                "rewards": reward_log,
                "results": results_log,
                "num_episodes": NUM_EPISODES,
                "model": MODEL_NAME,
            },
            f,
            indent=2,
        )
    print(f"\nReward log saved to {REWARDS_FILE}")

    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,
        )
    print(f"Episode transcripts saved to {TRANSCRIPTS_FILE}")

    # Plot reward curve
    _plot_reward_curve(reward_log)

    # Summary statistics
    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")

    # Smoothed curve (rolling average)
    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.show()


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

if __name__ == "__main__":
    train()