training/train.py

"""
MindFlayer GRPO training script.

Architecture (post-fix):
- TRL generates Flayer completions directly (no custom rollout_func — TRL
  0.23 silently drops that kwarg, which is why every reward used to be 0
  and OpenAI was never billed).
- reward_combined.* reward functions parse each completion into per-round
  Flayer messages and replay them through the live env, which is what
  actually calls the investigators (and therefore OpenAI). The first
  reward fn warms a per-step cache so all 5 reward fns share one episode.
- Episodes within a batch run concurrently via asyncio.gather, capped by
  MINDFLAYER_PARALLEL_EPISODES (set to match the per-call batch size so
  all 8 episodes per microbatch run in parallel, not serially).

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

# Set BEFORE importing reward_combined — module-level constants are read
# once at import. The defaults here target a single OpenAI Tier 1 day
# (200k TPM, 500 RPM, 2M TPD, 10k RPD) for gpt-4o-mini.
#
# With per_device_train_batch_size=2 × num_generations=4 we get 8
# completions per reward call → 8-way parallelism (peak ~96k–192k TPM,
# well under 200k).
os.environ.setdefault("MINDFLAYER_PARALLEL_EPISODES", "8")
# Easy mode: 3 rounds × 1 investigator (eleven only) ≈ 4 calls / 2k tokens
# per episode — ~4× cheaper than normal mode. Required to fit the 2M TPD
# budget. Override to "normal" only if you've upgraded to Tier 2+.
os.environ.setdefault("MINDFLAYER_TASK_ID", "easy")
# SFT is FREE (no OpenAI calls) — use it to absorb wall-clock budget
# without burning token quota. 3 epochs ≈ 25–30 min on Qwen 0.5B.
os.environ.setdefault("MINDFLAYER_SFT_EPOCHS", "3")

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 import run_sft_warmup
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 import run_sft_warmup

_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"
GRPO_OUTPUT_DIR = "./mindflayer-grpo-output"
FINAL_OUTPUT_DIR = "./mindflayer-trained"


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:
    """
    Conversational rows: each prompt is the chat template input that TRL
    passes to model.generate(). The 'scenario' column is forwarded to
    reward functions as a kwarg so the env replay uses the matching
    investigator framing.

    Sized for OpenAI Tier 1 (2M TPD): 1 row per scenario × 2 epochs gives
    ~30 GRPO optimizer steps × 32 episodes ≈ 1.92M tokens in easy mode.
    Bump n_per_scenario only if you've upgraded to Tier 2+.
    """
    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_tier1_budget() -> dict:
    """
    Back-of-envelope cost projection for one GRPO training run, assuming
    easy mode (3 rounds × 1 investigator + 1 ToM judge call ≈ 3.5 calls
    and ~2,000 tokens per episode on average).
    """
    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  # mirrors GRPOConfig below
    grad_accum = 4
    epochs = 2
    eps_per_step = per_device * 4 * grad_accum  # × num_generations
    steps = max(1, (n_rows * epochs) // (per_device * grad_accum))
    eps_total = steps * eps_per_step
    is_easy = os.environ.get("MINDFLAYER_TASK_ID", "easy").startswith("easy")
    tok_per_ep = 2_000 if is_easy else 8_000
    calls_per_ep = 3.5 if is_easy else 12
    return {
        "mode": "easy" if is_easy else "normal",
        "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 / 10_000 * 100,
    }


class ClearRewardCacheCallback(TrainerCallback):
    """
    The reward cache is keyed by completion text and grows unboundedly
    across training. Each completion is unique per step, so we wipe the
    cache after every optimizer step. Per-step intra-batch reuse (5 reward
    fns sharing 1 episode result) is preserved because the cache is only
    cleared at step boundaries, not between reward fn calls.
    """

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


class GenerationLogCallback(TrainerCallback):
    """Logs a sample interactive episode transcript every 50 steps."""

    def on_step_end(self, args, state, control, **kwargs):
        if state.global_step % 50 != 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 // 50) % 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"normal:{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(5):
                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}"
                if getattr(obs, "max_response", ""):
                    inv_text += f"\nmax: {obs.max_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}\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")

    openai_key = os.environ.get("OPENAI_API_KEY")
    if not openai_key:
        raise EnvironmentError("OPENAI_API_KEY environment variable is required")

    check_gpu()

    budget = estimate_tier1_budget()
    print(
        f"\nBudget projection ({budget['mode']} mode): "
        f"{budget['steps']} steps × {budget['episodes'] // budget['steps']} 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 10k RPD)"
    )
    if budget["tpd_pct"] > 100 or budget["rpd_pct"] > 100:
        print("  WARNING: projected over Tier 1 daily quota — training will hit 429s.")

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

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

    dataset = build_dataset()

    from trl import GRPOConfig, GRPOTrainer

    # 8-way parallel episode generation per reward call:
    #   per_device_train_batch_size (2) × num_generations (4) = 8 completions
    # The reward functions run all 8 through the env via asyncio.gather,
    # bounded by MINDFLAYER_PARALLEL_EPISODES (set to 8 above).
    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=1024,
        num_generations=4,
        temperature=0.9,
        logging_steps=10,
        save_steps=50,
        save_total_limit=2,
        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 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("\nTraining complete.")


if __name__ == "__main__":
    main()