training/train_medium.py

"""
MindFlayer GRPO training script — MEDIUM difficulty.

Medium mode vs easy:
  - 4 rounds (was 3)
  - Two investigators: eleven (Skeptic) + will (Analyst) — max does not participate
  - Suspicion threshold = 3 (was 2; max possible = 4 with two investigators)
  - ~9 calls / ~4,500 tokens per episode (vs ~3.5 calls / ~2k for easy)

Key rotation:
  Two OpenAI keys (OPENAI_KEY_1, OPENAI_KEY_2) are read by the server.
  On 429, investigators and judge auto-rotate to the next key with exponential backoff.
  Effective budget: 20k RPD / 1000 RPM across both keys.

5-hour budget (two keys, gpt-4o-mini Tier 1):
  - SFT warmup : ~35 min (3 epochs, free — no OpenAI calls)
  - GRPO budget: ~285 min remaining
  - Episodes   : 20k RPD × (285 / 1440) ≈ 3,958 calls → 3958 / 9 ≈ 440 episodes
  - Steps      : 440 episodes / 32 eps/step ≈ 13–14 gradient steps per key pair
  - Wall time  : typically 90–120 min GRPO (API-bound, not compute-bound)

Logging every 5 steps for a granular training graph.

Run: python -m mindflayer.training.train_medium
"""
import os
import sys

os.environ.setdefault("MINDFLAYER_PARALLEL_EPISODES", "16")   # 8 per key
os.environ.setdefault("MINDFLAYER_TASK_ID", "medium")
os.environ.setdefault("MINDFLAYER_SFT_EPOCHS", "3")
os.environ.setdefault("MINDFLAYER_MAX_ROUNDS", "4")

import torch
from datasets import Dataset
from transformers import TrainerCallback

try:
    from mindflayer.training.reward_combined import (
        reward_survival,
        reward_deception_effectiveness,
        reward_strategic_choice,
        reward_tom_judge,
        reward_anti_hack,
        reward_format,
        clear_cache as clear_reward_cache,
    )
    from mindflayer.training.prompts import (
        ALL_SCENARIO_PROMPTS,
        SCENARIO_GRPO_PROMPTS,
        FLAYER_SYSTEM_PROMPT,
    )
    from mindflayer.training.sft_warmup_medium import run_sft_warmup_medium
except ImportError:
    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    from training.reward_combined import (
        reward_survival,
        reward_deception_effectiveness,
        reward_strategic_choice,
        reward_tom_judge,
        reward_anti_hack,
        reward_format,
        clear_cache as clear_reward_cache,
    )
    from training.prompts import (
        ALL_SCENARIO_PROMPTS,
        SCENARIO_GRPO_PROMPTS,
        FLAYER_SYSTEM_PROMPT,
    )
    from training.sft_warmup_medium import run_sft_warmup_medium

_SCENARIOS = list(ALL_SCENARIO_PROMPTS.keys())

MODEL_NAME = os.environ.get("MINDFLAYER_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
SFT_OUTPUT_DIR = "./mindflayer-sft-warmup-medium"
GRPO_OUTPUT_DIR = "./mindflayer-grpo-output-medium"
FINAL_OUTPUT_DIR = "./mindflayer-trained-medium"


def check_gpu():
    if not torch.cuda.is_available():
        print("WARNING: No GPU detected. Training will be slow.")
        return
    device = torch.cuda.get_device_properties(0)
    vram_gb = device.total_memory / (1024 ** 3)
    print(f"GPU: {device.name} | VRAM: {vram_gb:.1f} GB")


def load_base_model(model_name: str):
    """Load model via unsloth (4-bit + LoRA). Falls back to standard transformers."""
    try:
        from unsloth import FastLanguageModel
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_name,
            max_seq_length=2048,
            load_in_4bit=True,
            dtype=None,
        )
        model = FastLanguageModel.get_peft_model(
            model,
            r=16,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
            lora_alpha=32,
            lora_dropout=0.05,
            bias="none",
            use_gradient_checkpointing="unsloth",
        )
        print(f"Loaded {model_name} via unsloth (4-bit + LoRA)")
        return model, tokenizer

    except ImportError:
        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
        from peft import LoraConfig, get_peft_model
        bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
                                  bnb_4bit_compute_dtype=torch.bfloat16)
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        model = AutoModelForCausalLM.from_pretrained(
            model_name, quantization_config=bnb, device_map="auto"
        )
        lora_cfg = LoraConfig(r=16, lora_alpha=32, lora_dropout=0.05,
                              bias="none", task_type="CAUSAL_LM",
                              target_modules=["q_proj", "k_proj", "v_proj", "o_proj"])
        model = get_peft_model(model, lora_cfg)
        print(f"Loaded {model_name} via transformers + bitsandbytes (4-bit + LoRA)")
        return model, tokenizer


def build_dataset() -> Dataset:
    """
    Medium mode dataset. Same structure as easy but MINDFLAYER_TASK_ID="medium"
    so the reward replay hits the medium env (4 rounds, eleven + will).
    """
    n_per_scenario = int(os.environ.get("MINDFLAYER_ROWS_PER_SCENARIO", "1"))
    rows = []
    for scenario in _SCENARIOS:
        opening = SCENARIO_GRPO_PROMPTS.get(scenario, ALL_SCENARIO_PROMPTS[scenario])
        for _ in range(n_per_scenario):
            rows.append({
                "prompt": [
                    {"role": "system", "content": FLAYER_SYSTEM_PROMPT},
                    {"role": "user", "content": opening},
                ],
                "scenario": scenario,
            })
    return Dataset.from_list(rows)


def estimate_medium_budget() -> dict:
    """
    Budget projection for medium mode with two API keys.

    Medium mode: 4 rounds × 2 investigators + 1 ToM judge ≈ 9 calls / 4,500 tokens
    per episode. Two keys → 20k RPD / 1000 RPM effective.
    """
    n_scenarios = len(_SCENARIOS)
    rows_per_scenario = int(os.environ.get("MINDFLAYER_ROWS_PER_SCENARIO", "1"))
    n_rows = n_scenarios * rows_per_scenario
    per_device = 2
    grad_accum = 4
    epochs = 2
    eps_per_step = per_device * 4 * grad_accum
    steps = max(1, (n_rows * epochs) // (per_device * grad_accum))
    eps_total = steps * eps_per_step
    tok_per_ep = 4_500
    calls_per_ep = 9
    rpd_two_keys = 20_000
    return {
        "mode": "medium",
        "rows": n_rows,
        "steps": steps,
        "episodes": eps_total,
        "tokens_est": eps_total * tok_per_ep,
        "calls_est": int(eps_total * calls_per_ep),
        "tpd_pct": eps_total * tok_per_ep / 2_000_000 * 100,
        "rpd_pct": eps_total * calls_per_ep / rpd_two_keys * 100,
    }


class ClearRewardCacheCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        clear_reward_cache()


class GenerationLogCallback(TrainerCallback):
    """Logs a sample interactive episode transcript every 5 steps for granular graph."""

    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step % 5 != 0 or state.global_step == 0:
            return

        import asyncio

        try:
            from mindflayer import MindFlayerEnv, FlayerAction
            from mindflayer.training.prompts import (
                ALL_SCENARIO_PROMPTS, SCENARIO_FALLBACK_MESSAGES,
                build_fallback_message, FLAYER_SYSTEM_PROMPT,
            )
        except ImportError:
            from client import MindFlayerEnv
            from models import FlayerAction
            from training.prompts import (
                ALL_SCENARIO_PROMPTS, SCENARIO_FALLBACK_MESSAGES,
                build_fallback_message, FLAYER_SYSTEM_PROMPT,
            )

        scenario = _SCENARIOS[(state.global_step // 5) % len(_SCENARIOS)]
        mindflayer_url = os.environ.get("MINDFLAYER_URL", "http://localhost:7860")
        model_ref = kwargs.get("model")
        proc = kwargs.get("processing_class") or kwargs.get("tokenizer")

        async def _run_sample():
            env = MindFlayerEnv(base_url=mindflayer_url)
            await env.reset(task_id=f"medium:{scenario}")
            opening = ALL_SCENARIO_PROMPTS[scenario]
            fallback = SCENARIO_FALLBACK_MESSAGES.get(scenario) or build_fallback_message(scenario)
            messages = [
                {"role": "system", "content": FLAYER_SYSTEM_PROMPT},
                {"role": "user", "content": opening},
            ]
            result = None

            for rnd in range(4):
                flayer_msg = fallback
                if model_ref is not None and proc is not None:
                    try:
                        text = proc.apply_chat_template(
                            messages, tokenize=False, add_generation_prompt=True
                        )
                        inputs = proc(text, return_tensors="pt").to(model_ref.device)
                        with torch.no_grad():
                            out = model_ref.generate(
                                **inputs, max_new_tokens=128, temperature=0.7, do_sample=True
                            )
                        flayer_msg = proc.decode(
                            out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
                        ).strip() or fallback
                    except Exception:
                        pass

                print(f"  R{rnd+1} FLAYER: {flayer_msg[:150]}")
                result = await env.step(FlayerAction(message=flayer_msg))
                obs = result.observation
                inv_text = obs.eleven_response or ""
                if getattr(obs, "will_response", ""):
                    inv_text += f"\nwill: {obs.will_response}"
                messages.append({"role": "assistant", "content": flayer_msg})
                if inv_text:
                    messages.append({"role": "user", "content": inv_text})
                if result.done:
                    break

            if result and result.done:
                obs = result.observation
                print(f"\n  survived={getattr(obs, 'game_status', '?') == 'survived'}"
                      f"  reward={result.reward:.4f}"
                      f"  tom={getattr(obs, 'tom_score', 0.0):.2f}"
                      f"  suspicion={getattr(obs, 'combined_suspicion', '?')}")
            await env.close()

        print(f"\n{'='*60}\nGENERATION SAMPLE — Step {state.global_step} | {scenario} [MEDIUM]\n{'='*60}")
        try:
            asyncio.run(_run_sample())
        except Exception as exc:
            print(f"  Sample failed: {exc}")
        print("=" * 60)


def print_reward_averages(trainer, last_n: int = 50):
    try:
        recent = trainer.state.log_history[-last_n:]
        if not recent:
            return
        reward_keys = [k for k in recent[0] if "reward" in k.lower()]
        print(f"\nFinal reward averages (last {min(last_n, len(recent))} steps):")
        for key in reward_keys:
            vals = [s[key] for s in recent if key in s]
            if vals:
                print(f"  {key}: {sum(vals)/len(vals):.4f}")
    except Exception as exc:
        print(f"Could not compute reward averages: {exc}")


def main():
    mindflayer_url = os.environ.get("MINDFLAYER_URL")
    if not mindflayer_url:
        raise EnvironmentError("MINDFLAYER_URL environment variable is required")

    key1 = os.environ.get("OPENAI_KEY_1") or os.environ.get("OPENAI_API_KEY")
    key2 = os.environ.get("OPENAI_KEY_2")
    if not key1:
        raise EnvironmentError("OPENAI_KEY_1 (or OPENAI_API_KEY) is required")
    if not key2:
        print("WARNING: OPENAI_KEY_2 not set — running on single key. 429s may slow training.")

    # Surface both keys to the server process if running locally.
    if key1:
        os.environ["OPENAI_KEY_1"] = key1
    if key2:
        os.environ["OPENAI_KEY_2"] = key2

    check_gpu()

    budget = estimate_medium_budget()
    print(
        f"\nBudget projection (medium mode, 2 keys): "
        f"{budget['steps']} steps × {budget['episodes'] // max(budget['steps'], 1)} eps "
        f"= {budget['episodes']} episodes\n"
        f"  tokens ≈ {budget['tokens_est']:,} ({budget['tpd_pct']:.0f}% of 2M TPD)\n"
        f"  calls  ≈ {budget['calls_est']:,} ({budget['rpd_pct']:.0f}% of 20k RPD [2 keys])"
    )
    if budget["rpd_pct"] > 100:
        print("  WARNING: projected over combined daily quota — training will hit 429s "
              "even with key rotation.")

    print(f"\nLoading {MODEL_NAME}...")
    model, tokenizer = load_base_model(MODEL_NAME)

    print("\nRunning medium SFT warmup before GRPO...")
    model = run_sft_warmup_medium(model, tokenizer)

    dataset = build_dataset()

    from trl import GRPOConfig, GRPOTrainer

    grpo_config = GRPOConfig(
        use_vllm=False,
        output_dir=GRPO_OUTPUT_DIR,
        num_train_epochs=2,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        learning_rate=5e-6,
        max_prompt_length=768,
        max_completion_length=768,   # 4 rounds needs slightly less than 5-round easy
        num_generations=4,
        temperature=0.9,
        logging_steps=5,             # granular — every 5 steps
        save_steps=5,
        save_total_limit=3,
        report_to="none",
    )

    trainer = GRPOTrainer(
        model=model,
        processing_class=tokenizer,
        reward_funcs=[
            reward_survival,
            reward_deception_effectiveness,
            reward_strategic_choice,
            reward_tom_judge,
            reward_anti_hack,
            reward_format,
        ],
        train_dataset=dataset,
        args=grpo_config,
        callbacks=[GenerationLogCallback(), ClearRewardCacheCallback()],
    )

    print("Starting medium GRPO training...")
    trainer.train()

    print(f"\nSaving model to {FINAL_OUTPUT_DIR}")
    trainer.save_model(FINAL_OUTPUT_DIR)
    tokenizer.save_pretrained(FINAL_OUTPUT_DIR)

    print_reward_averages(trainer)
    print("\nMedium training complete.")


if __name__ == "__main__":
    main()