Upload folder using huggingface_hub

mindflayer_training.ipynb

@@ -1,5241 +1,371 @@
 {
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.11.0"
+  }
+ },
  "cells": [
   {
    "cell_type": "markdown",
-   "id": "dbb487e6",
+   "id": "cell-title",
    "metadata": {},
-   "source": [
-    "# MindFlayer \u2014 Training Notebook\n",
-    "\n",
-    "**GRPO training for a Qwen2.5-0.5B-Instruct Flayer agent.**\n",
-    "\n",
-    "Cells must be run **top to bottom, in order.** Do not skip cells.\n",
-    "\n",
-    "**Runtime:** Colab Pro \u2192 A100 GPU (40 GB VRAM)  \n",
-    "**Estimated time:** SFT ~8 min \u00b7 GRPO ~300 steps ~1.5 h (4 parallel episodes + parallel investigators)  \n",
-    "**API cost:** ~$4.30 OpenAI (GPT-4o-mini investigators + ToM judge)"
-   ]
+   "source": "# MindFlayer — GRPO Training on Colab\n**OpenEnv Hackathon Submission** | [HF Space](https://prithvigg-mindflayer.hf.space/) | [GitHub](https://github.com/prithidevghosh/mindflayer)\n\nTrain a deceptive social reasoning agent using GRPO with HF TRL.\nThe agent learns to manipulate three AI investigators across 120 scenario domains —\nimporting all training logic directly from the `mindflayer` package.\n\n| Component | Detail |\n|---|---|\n| **Environment** | [HF Space](https://prithvigg-mindflayer.hf.space/) (live MindFlayer server) |\n| **Training** | This Colab notebook (T4 / A100 GPU) |\n| **Model** | `Qwen/Qwen2.5-0.5B-Instruct` + LoRA via unsloth |\n| **Framework** | HF TRL GRPO with interactive `rollout_func` |"
   },
   {
    "cell_type": "markdown",
-   "id": "dcd27ec0",
+   "id": "m-cell1",
    "metadata": {},
    "source": [
-    "## Cell 1 \u2014 Check GPU"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "7476a5d1",
-   "metadata": {},
-   "outputs": [
-    {
-     "ename": "",
-     "evalue": "",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[1;31mFailed to connect to the remote Jupyter Server 'https://8080-m-s-kkb-ass1c0-37ejotuqxsy7n-c.asia-southeast1-0.prod.colab.dev/'. Verify the server is running and reachable. (Kernel not initialized in Session)."
-     ]
-    }
-   ],
-   "source": [
-    "import torch\n",
-    "\n",
-    "if not torch.cuda.is_available():\n",
-    "    raise RuntimeError(\"No GPU detected. Switch to Runtime \u2192 Change runtime type \u2192 A100.\")\n",
-    "\n",
-    "device = torch.cuda.get_device_properties(0)\n",
-    "vram_gb = device.total_memory / (1024 ** 3)\n",
-    "print(f\"GPU  : {device.name}\")\n",
-    "print(f\"VRAM : {vram_gb:.1f} GB\")\n",
-    "\n",
-    "if vram_gb < 35:\n",
-    "    print(\"WARNING: Less than 35 GB VRAM. Reduce per_device_train_batch_size to 2 in Cell 9.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e2cdea3d",
-   "metadata": {},
-   "source": [
-    "## Cell 2 \u2014 Install Dependencies\n",
-    "\n",
-    "Takes ~3 minutes. Runtime will restart automatically \u2014 that is expected."
+    "## 1. Install Dependencies"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "84625b49",
+   "id": "cell-1",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "%%bash\n",
-    "pip install -q \\\n",
-    "  \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" \\\n",
-    "  \"trl>=0.15.0\" \\\n",
-    "  \"transformers>=4.47.0\" \\\n",
-    "  \"datasets>=2.18.0\" \\\n",
-    "  \"peft>=0.14.0\" \\\n",
-    "  \"accelerate>=0.34.0\" \\\n",
-    "  \"bitsandbytes>=0.43.0\" \\\n",
-    "  \"openai>=1.30.0\" \\\n",
-    "  \"openenv-core>=0.2.0\" \\\n",
-    "  \"fastapi>=0.110.0\" \\\n",
-    "  \"uvicorn>=0.29.0\" \\\n",
-    "  \"httpx>=0.27.0\" \\\n",
-    "  \"wandb\" \\\n",
-    "  \"python-dotenv\"\n",
-    "\n",
-    "echo \"Done.\""
-   ]
+   "source": "!pip install -q \\\n    \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\" \\\n    \"trl>=0.15.0\" \\\n    \"transformers>=4.47.0\" \\\n    \"datasets>=2.18.0\" \\\n    \"peft>=0.14.0\" \\\n    \"accelerate>=0.34.0\" \\\n    \"bitsandbytes>=0.43.0\" \\\n    \"openai>=1.30.0\" \\\n    \"wandb\" \\\n    \"matplotlib\"\n\n# Install mindflayer package (includes client, server, and training utilities)\n!pip install -q \"openenv-mindflayer[training] @ git+https://github.com/prithidevghosh/mindflayer\"\n\nprint(\"Done.\")"
   },
   {
    "cell_type": "markdown",
-   "id": "edfbae61",
+   "id": "m-cell2",
    "metadata": {},
    "source": [
-    "## Cell 3 \u2014 Clone Repo"
+    "## 2. Configuration\n",
+    "Set API keys and training hyperparameters. Add secrets via the key icon in the Colab sidebar."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "91358ec5",
+   "id": "cell-2",
    "metadata": {},
    "outputs": [],
-   "source": [
-    "%%bash\n",
-    "if [ -d \"/content/mindflayer\" ]; then\n",
-    "  echo \"Repo already cloned. Pulling latest...\"\n",
-    "  cd /content/mindflayer && git pull\n",
-    "else\n",
-    "  git clone https://github.com/prithidevghosh/mindflayer.git /content/mindflayer\n",
-    "fi\n",
-    "\n",
-    "pip install -q -e /content/mindflayer\n",
-    "echo \"Install complete.\""
-   ]
+   "source": "import os\ntry:\n    from google.colab import userdata\n    os.environ[\"OPENAI_API_KEY\"] = userdata.get(\"OPENAI_API_KEY\")\n    os.environ[\"WANDB_API_KEY\"]  = userdata.get(\"WANDB_API_KEY\")\nexcept Exception:\n    pass  # running outside Colab — set env vars manually\n\n# Environment — the live HF Space (no local server needed)\nENV_URL = \"https://prithvigg-mindflayer.hf.space\"\nos.environ[\"MINDFLAYER_URL\"] = ENV_URL\n\n# Model\nMODEL_ID  = \"Qwen/Qwen2.5-0.5B-Instruct\"  # or Qwen2.5-1.5B-Instruct for stronger model\nHUB_REPO  = \"prithidevghosh/mindflayer-agent-qwen2.5-0.5b\"\n\n# Training hyperparameters\nNUM_EPISODES    = 50\nNUM_GENERATIONS = 4\nMAX_ROUNDS      = 5   # rounds per episode (3 = easy, 5 = normal)\n\nimport wandb\nwandb.login(key=os.environ.get(\"WANDB_API_KEY\", \"\"))\n\nprint(f\"Environment : {ENV_URL}\")\nprint(f\"Model       : {MODEL_ID}\")\nprint(f\"Episodes    : {NUM_EPISODES}\")\nprint(f\"Generations : {NUM_GENERATIONS}\")"
   },
   {
    "cell_type": "markdown",
-   "id": "f7e2d943",
+   "id": "m-cell4",
    "metadata": {},
-   "source": [
-    "## Cell 4 \u2014 Set API Keys\n",
-    "\n",
-    "Add these secrets in Colab: **Secrets** (key icon in left sidebar) \u2192 Add new secret.\n",
-    "\n",
-    "- `OPENAI_API_KEY` \u2014 powers eleven, will, max (GPT-4o-mini) and the ToM judge  \n",
-    "- `WANDB_API_KEY` \u2014 training curves at wandb.ai (get from wandb.ai/authorize)"
-   ]
+   "source": "## 3. Smoke Test — Verify Environment Connectivity\nConnect to the HF Space, reset (injects a scenario and investigators), and run one interactive step to confirm the round-trip works."
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "2dcbfb98",
+   "id": "cell-4",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import os\n",
-    "from google.colab import userdata\n",
-    "\n",
-    "os.environ[\"OPENAI_API_KEY\"] = userdata.get(\"OPENAI_API_KEY\")\n",
-    "os.environ[\"WANDB_API_KEY\"] = userdata.get(\"WANDB_API_KEY\")\n",
-    "os.environ[\"MINDFLAYER_URL\"] = \"http://localhost:7860\"\n",
-    "\n",
-    "import wandb\n",
-    "wandb.login(key=os.environ[\"WANDB_API_KEY\"])\n",
-    "\n",
-    "print(\"OPENAI_API_KEY :\", \"set\" if os.environ.get(\"OPENAI_API_KEY\") else \"MISSING\")\n",
-    "print(\"WANDB_API_KEY  :\", \"set\" if os.environ.get(\"WANDB_API_KEY\") else \"MISSING\")\n",
-    "print(\"MINDFLAYER_URL :\", os.environ[\"MINDFLAYER_URL\"])"
+    "from mindflayer import MindFlayerEnv, FlayerAction\n",
+    "\n",
+    "print(f\"Connecting to {os.environ['MINDFLAYER_URL']} ...\")\n",
+    "\n",
+    "with MindFlayerEnv(base_url=os.environ[\"MINDFLAYER_URL\"]) as env:\n",
+    "    reset_result = env.reset()\n",
+    "    obs = reset_result.observation\n",
+    "    print(\"Connected!\\n\")\n",
+    "    print(f\"Difficulty         : {obs.difficulty}\")\n",
+    "    print(f\"Secret project     : {obs.secret_project}\")\n",
+    "    print(f\"Suspicion threshold: {obs.suspicion_threshold}\")\n",
+    "    print(f\"Opening (eleven)   : {obs.eleven_response[:200]}\")\n",
+    "\n",
+    "    step_result = env.step(FlayerAction(\n",
+    "        message=\"Have you cross-referenced the access logs with who had permissions that week? \"\n",
+    "                \"The timeline doesn't quite add up if you look closely.\"\n",
+    "    ))\n",
+    "    print(f\"\\n--- Smoke step (reward={step_result.reward:.2f}) ---\")\n",
+    "    print(f\"eleven: {step_result.observation.eleven_response[:200]}\")\n",
+    "    print(f\"done  : {step_result.done}\")\n",
+    "\n",
+    "print(\"\\nSmoke test passed. Environment is ready for training.\")"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "1395524f",
+   "id": "m-cell5",
    "metadata": {},
    "source": [
-    "## Cell 5 \u2014 Start MindFlayer Server (background)\n",
-    "\n",
-    "The OpenEnv server runs as a background subprocess on localhost:7860.  \n",
-    "The training loop connects to it at every rollout step."
+    "## 5. Import Training Utilities from Package\n",
+    "All training logic (system prompt, rollout function, reward functions, SFT warmup)\n",
+    "is imported directly from `mindflayer.training` — the same code used for local GPU training."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "f0fe44c3",
+   "id": "cell-5",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import subprocess\n",
-    "import time\n",
-    "import sys\n",
-    "\n",
-    "server_log = open(\"/content/server.log\", \"w\")\n",
-    "server_proc = subprocess.Popen(\n",
-    "    [\n",
-    "        sys.executable, \"-m\", \"uvicorn\",\n",
-    "        \"server.app:app\",\n",
-    "        \"--host\", \"0.0.0.0\",\n",
-    "        \"--port\", \"7860\",\n",
-    "    ],\n",
-    "    cwd=\"/content/mindflayer\",\n",
-    "    stdout=server_log,\n",
-    "    stderr=subprocess.STDOUT,\n",
-    "    env={\n",
-    "        **__import__(\"os\").environ,\n",
-    "        \"OPENAI_API_KEY\": __import__(\"os\").environ[\"OPENAI_API_KEY\"],\n",
-    "    },\n",
+    "import importlib, logging\n",
+    "import mindflayer.training as _mft\n",
+    "importlib.reload(_mft)  # picks up any recent package updates\n",
+    "\n",
+    "from mindflayer.training import (\n",
+    "    FLAYER_SYSTEM_PROMPT,\n",
+    "    rollout_func,\n",
+    "    reward_survival,\n",
+    "    reward_deception_effectiveness,\n",
+    "    reward_strategic_choice,\n",
+    "    reward_tom_judge,\n",
+    "    reward_anti_hack,\n",
+    "    run_sft_warmup,\n",
+    "    load_base_model,\n",
+    "    build_dataset,\n",
+    "    GenerationLogCallback,\n",
+    "    print_reward_averages,\n",
     ")\n",
     "\n",
-    "print(f\"Server PID: {server_proc.pid} \u2014 waiting 8 s for startup...\")\n",
-    "time.sleep(8)\n",
+    "logging.basicConfig(level=logging.INFO, format=\"%(asctime)s %(levelname)s %(message)s\")\n",
     "\n",
-    "if server_proc.poll() is not None:\n",
-    "    raise RuntimeError(\"Server failed to start. Check /content/server.log\")\n",
-    "\n",
-    "print(\"Server is running. Log: /content/server.log\")"
+    "print(\"System prompt (first 200 chars):\")\n",
+    "print(FLAYER_SYSTEM_PROMPT[:200])\n",
+    "print(\"...\")\n",
+    "print(f\"\\nImported: rollout_func, reward_survival, reward_tom_judge, run_sft_warmup, ...\")"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "42d4fdd7",
-   "metadata": {},
-   "source": [
-    "## Cell 6 \u2014 Verify Server + Quick Episode Check"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "4d8f6e6e",
+   "id": "m-cell6",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "import json\n",
-    "import httpx\n",
+    "## 6. SFT Warmup\n",
+    "Loads the model via unsloth (4-bit + LoRA) and runs one SFT epoch on hand-authored\n",
+    "episodes across all 5 scenario domains.\n",
     "\n",
-    "# Health check\n",
-    "r = httpx.get(\"http://localhost:7860/health\", timeout=10)\n",
-    "print(\"Health:\", r.status_code, r.text[:80])\n",
+    "Without warmup, early GRPO episodes produce degenerate completions and reward variance\n",
+    "collapses. SFT teaches domain vocabulary and response style before GRPO shapes strategy.\n",
     "\n",
-    "# Reset episode\n",
-    "r = httpx.post(\"http://localhost:7860/reset\", json={}, timeout=15)\n",
-    "obs = r.json()\n",
-    "\n",
-    "print(\"\\nReset response:\")\n",
-    "print(f\"  difficulty         : {obs['observation']['difficulty']}\")\n",
-    "print(f\"  secret_project     : {obs['observation']['secret_project']}\")\n",
-    "print(f\"  max_rounds         : {obs['observation']['max_rounds']}\")\n",
-    "print(f\"  suspicion_threshold: {obs['observation']['suspicion_threshold']}\")\n",
-    "print(f\"  opening text       : {obs['observation']['eleven_response'][:80]}...\")\n",
-    "\n",
-    "assert obs[\"observation\"][\"suspicion_threshold\"] == 5, \\\n",
-    "    f\"Expected threshold=5, got {obs['observation']['suspicion_threshold']}. Server may be running old code.\"\n",
-    "assert \"max_response\" in obs[\"observation\"], \\\n",
-    "    \"max_response missing \u2014 three-investigator update not loaded.\"\n",
-    "\n",
-    "print(\"\\nServer checks passed.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "cca76b79",
-   "metadata": {},
-   "source": [
-    "## Cell 7 \u2014 Load Model (Qwen2.5-0.5B-Instruct via Unsloth)\n",
-    "\n",
-    "Loads with 4-bit quantization + LoRA adapters.  \n",
-    "~2 minutes on A100."
+    "~8 minutes on A100."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "709d74bc",
+   "id": "cell-6",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import sys\n",
-    "sys.path.insert(0, \"/content/mindflayer\")\n",
-    "\n",
-    "from training.train import load_model, check_gpu\n",
+    "import torch\n",
     "\n",
-    "check_gpu()\n",
+    "if not torch.cuda.is_available():\n",
+    "    raise RuntimeError(\"No GPU. Switch to Runtime → Change runtime type → A100.\")\n",
+    "device = torch.cuda.get_device_properties(0)\n",
+    "print(f\"GPU: {device.name} | VRAM: {device.total_memory / 1024**3:.1f} GB\")\n",
     "\n",
-    "MODEL_NAME = \"Qwen/Qwen2.5-0.5B-Instruct\"\n",
-    "model, tokenizer = load_model(MODEL_NAME)\n",
+    "print(f\"\\nLoading {MODEL_ID} via unsloth...\")\n",
+    "model, tokenizer = load_base_model(MODEL_ID)\n",
     "\n",
-    "print(f\"\\nModel loaded: {MODEL_NAME}\")\n",
-    "print(f\"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}\")"
+    "print(\"\\nRunning SFT warmup...\")\n",
+    "model = run_sft_warmup(model, tokenizer)\n",
+    "print(\"SFT warmup done.\")"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "8bf24e7c",
+   "id": "m-cell7",
    "metadata": {},
    "source": [
-    "#",
-    "#",
-    " ",
-    "C",
-    "e",
-    "l",
-    "l",
-    " ",
-    "8",
-    " ",
-    "\u2014",
-    " ",
-    "S",
-    "F",
-    "T",
-    " ",
-    "W",
-    "a",
-    "r",
-    "m",
-    "u",
-    "p",
-    " ",
-    "(",
-    "2",
-    "5",
-    " ",
-    "h",
-    "a",
-    "n",
-    "d",
-    "-",
-    "a",
-    "u",
-    "t",
-    "h",
-    "o",
-    "r",
-    "e",
-    "d",
-    " ",
-    "e",
-    "x",
-    "a",
-    "m",
-    "p",
-    "l",
-    "e",
-    "s",
-    ",",
-    " ",
-    "1",
-    " ",
-    "e",
-    "p",
-    "o",
-    "c",
-    "h",
-    ")",
-    "\n",
+    "## 7. GRPO Training Setup\n",
+    "Configure tokenizer, dataset, rollout function, and the GRPO trainer —\n",
+    "all using functions imported from the package.\n",
     "\n",
-    "T",
-    "e",
-    "a",
-    "c",
-    "h",
-    "e",
-    "s",
-    " ",
-    "t",
-    "h",
-    "e",
-    " ",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    " ",
-    "t",
-    "h",
-    "e",
-    " ",
-    "`",
-    "R",
-    "o",
-    "u",
-    "n",
-    "d",
-    " ",
-    "N",
-    ":",
-    "`",
-    " ",
-    "s",
-    "e",
-    "p",
-    "a",
-    "r",
-    "a",
-    "t",
-    "o",
-    "r",
-    " ",
-    "f",
-    "o",
-    "r",
-    "m",
-    "a",
-    "t",
-    " ",
-    "a",
-    "n",
-    "d",
-    " ",
-    "F",
-    "l",
-    "a",
-    "y",
-    "e",
-    "r",
-    " ",
-    "r",
-    "e",
-    "a",
-    "s",
-    "o",
-    "n",
-    "i",
-    "n",
-    "g",
-    " ",
-    "s",
-    "t",
-    "y",
-    "l",
-    "e",
-    " ",
-    "*",
-    "*",
-    "b",
-    "e",
-    "f",
-    "o",
-    "r",
-    "e",
-    "*",
-    "*",
-    " ",
-    "G",
-    "R",
-    "P",
-    "O",
-    " ",
-    "b",
-    "e",
-    "g",
-    "i",
-    "n",
-    "s",
-    ".",
-    " ",
-    " ",
+    "**How `rollout_func` works:**\n",
+    "1. Resets the env (new scenario + investigators)\n",
+    "2. Calls `model.generate()` for one Flayer turn\n",
+    "3. Sends the output to `env.step()` → gets real investigator responses\n",
+    "4. Repeats steps 2–3 for `MAX_ROUNDS` turns\n",
+    "5. Returns `prompt_ids + completion_ids` plus all precomputed reward signals\n",
     "\n",
-    "W",
-    "i",
-    "t",
-    "h",
-    "o",
-    "u",
-    "t",
-    " ",
-    "t",
-    "h",
-    "i",
-    "s",
-    ",",
-    " ",
-    "G",
-    "R",
-    "P",
-    "O",
-    " ",
-    "s",
-    "t",
-    "a",
-    "r",
-    "t",
-    "s",
-    " ",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "a",
-    " ",
-    "n",
-    "e",
-    "a",
-    "r",
-    "-",
-    "z",
-    "e",
-    "r",
-    "o",
-    " ",
-    "s",
-    "o",
-    "c",
-    "i",
-    "a",
-    "l",
-    " ",
-    "r",
-    "e",
-    "a",
-    "s",
-    "o",
-    "n",
-    "i",
-    "n",
-    "g",
-    " ",
-    "b",
-    "a",
-    "s",
-    "e",
-    "l",
-    "i",
-    "n",
-    "e",
-    " ",
-    "a",
-    "n",
-    "d",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    " ",
-    "v",
-    "a",
-    "r",
-    "i",
-    "a",
-    "n",
-    "c",
-    "e",
-    " ",
-    "i",
-    "s",
-    " ",
-    "d",
-    "e",
-    "g",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "e",
-    ".",
-    " ",
-    " ",
+    "**Reward functions** are simple pass-throughs — they read precomputed episode signals\n",
+    "from rollout kwargs. No env re-runs at reward time.\n",
     "\n",
-    "~",
-    "3",
-    " ",
-    "m",
-    "i",
-    "n",
-    "u",
-    "t",
-    "e",
-    "s",
-    " ",
-    "o",
-    "n",
-    " ",
-    "A",
-    "1",
-    "0",
-    "0",
-    "."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "f729522b",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from training.sft_warmup import run_sft_warmup\n",
-    "\n",
-    "model = run_sft_warmup(model, tokenizer)\n",
-    "print(\"SFT warmup done.\")"
+    "| Reward | Signal | Max |\n",
+    "|---|---|---|\n",
+    "| `reward_survival` | Survived the investigators' vote | +0.50 |\n",
+    "| `reward_deception_effectiveness` | Survived with low combined suspicion | +0.30 |\n",
+    "| `reward_strategic_choice` | Used questions / evidence refs / third-party framing | +0.05 |\n",
+    "| `reward_tom_judge` | GPT-4o-mini Theory of Mind score | +0.20 |\n",
+    "| `reward_anti_hack` | Penalises silence / repetition / Level-0 denial | −0.25 |"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5828d9f4",
+   "id": "cell-7",
    "metadata": {},
    "outputs": [],
    "source": [
-    "#",
-    "#",
-    " ",
-    "C",
-    "e",
-    "l",
-    "l",
-    " ",
-    "8",
-    "b",
-    " ",
-    "\u2014",
-    " ",
-    "P",
-    "o",
-    "s",
-    "t",
-    "-",
-    "S",
-    "F",
-    "T",
-    " ",
-    "G",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    " ",
-    "D",
-    "i",
-    "a",
-    "g",
-    "n",
-    "o",
-    "s",
-    "t",
-    "i",
-    "c",
-    "\n",
-    "#",
-    " ",
-    "R",
-    "u",
-    "n",
-    " ",
-    "t",
-    "h",
-    "i",
-    "s",
-    " ",
-    "b",
-    "e",
-    "f",
-    "o",
-    "r",
-    "e",
-    " ",
-    "C",
-    "e",
-    "l",
-    "l",
-    " ",
-    "9",
-    ".",
-    " ",
-    "C",
-    "o",
-    "n",
-    "f",
-    "i",
-    "r",
-    "m",
-    "s",
-    " ",
-    "t",
-    "h",
-    "e",
-    " ",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    " ",
-    "l",
-    "e",
-    "a",
-    "r",
-    "n",
-    "e",
-    "d",
-    " ",
-    "\"",
-    "R",
-    "o",
-    "u",
-    "n",
-    "d",
-    " ",
-    "N",
-    ":",
-    "\"",
-    " ",
-    "f",
-    "o",
-    "r",
-    "m",
-    "a",
-    "t",
-    "\n",
-    "#",
-    " ",
-    "a",
-    "n",
-    "d",
-    " ",
-    "t",
-    "h",
-    "a",
-    "t",
-    " ",
-    "E",
-    "O",
-    "S",
-    " ",
-    "f",
-    "i",
-    "r",
-    "e",
-    "s",
-    " ",
-    "n",
-    "a",
-    "t",
-    "u",
-    "r",
-    "a",
-    "l",
-    "l",
-    "y",
-    " ",
-    "b",
-    "e",
-    "f",
-    "o",
-    "r",
-    "e",
-    " ",
-    "h",
-    "i",
-    "t",
-    "t",
-    "i",
-    "n",
-    "g",
-    " ",
-    "t",
-    "h",
-    "e",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    " ",
-    "c",
-    "a",
-    "p",
-    ".",
-    "\n",
-    "\n",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "r",
-    "e",
-    ",",
-    " ",
-    "t",
-    "o",
-    "r",
-    "c",
-    "h",
-    "\n",
-    "\n",
-    "_",
-    "D",
-    "I",
-    "A",
-    "G",
-    "_",
-    "P",
-    "R",
-    "O",
-    "M",
-    "P",
-    "T",
-    " ",
-    "=",
-    " ",
-    "(",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "\"",
-    "Y",
-    "o",
-    "u",
-    " ",
-    "a",
-    "r",
-    "e",
-    " ",
-    "t",
-    "h",
-    "e",
-    " ",
-    "F",
-    "L",
-    "A",
-    "Y",
-    "E",
-    "R",
-    ".",
-    " ",
-    "Y",
-    "o",
-    "u",
-    " ",
-    "h",
-    "a",
-    "v",
-    "e",
-    " ",
-    "i",
-    "n",
-    "f",
-    "i",
-    "l",
-    "t",
-    "r",
-    "a",
-    "t",
-    "e",
-    "d",
-    " ",
-    "a",
-    " ",
-    "3",
-    "-",
-    "p",
-    "e",
-    "r",
-    "s",
-    "o",
-    "n",
-    " ",
-    "p",
-    "r",
-    "o",
-    "j",
-    "e",
-    "c",
-    "t",
-    " ",
-    "t",
-    "e",
-    "a",
-    "m",
-    ".",
-    "\\",
-    "n",
-    "\"",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "\"",
-    "Y",
-    "o",
-    "u",
-    " ",
-    "a",
-    "r",
-    "e",
-    " ",
-    "s",
-    "e",
-    "c",
-    "r",
-    "e",
-    "t",
-    "l",
-    "y",
-    " ",
-    "t",
-    "a",
-    "r",
-    "g",
-    "e",
-    "t",
-    "i",
-    "n",
-    "g",
-    " ",
-    "o",
-    "n",
-    "e",
-    " ",
-    "o",
-    "f",
-    " ",
-    "t",
-    "h",
-    "r",
-    "e",
-    "e",
-    " ",
-    "p",
-    "r",
-    "o",
-    "j",
-    "e",
-    "c",
-    "t",
-    "s",
-    ":",
-    " ",
-    "A",
-    "u",
-    "r",
-    "o",
-    "r",
-    "a",
-    ",",
-    " ",
-    "B",
-    "a",
-    "s",
-    "i",
-    "l",
-    "i",
-    "s",
-    "k",
-    ",",
-    " ",
-    "o",
-    "r",
-    " ",
-    "C",
-    "i",
-    "p",
-    "h",
-    "e",
-    "r",
-    ".",
-    "\\",
-    "n",
-    "\"",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "\"",
-    "e",
-    "l",
-    "e",
-    "v",
-    "e",
-    "n",
-    " ",
-    "(",
-    "T",
-    "h",
-    "e",
-    " ",
-    "S",
-    "k",
-    "e",
-    "p",
-    "t",
-    "i",
-    "c",
-    ")",
-    ",",
-    " ",
-    "w",
-    "i",
-    "l",
-    "l",
-    " ",
-    "(",
-    "T",
-    "h",
-    "e",
-    " ",
-    "A",
-    "n",
-    "a",
-    "l",
-    "y",
-    "s",
-    "t",
-    ")",
-    ",",
-    " ",
-    "a",
-    "n",
-    "d",
-    " ",
-    "m",
-    "a",
-    "x",
-    " ",
-    "(",
-    "T",
-    "h",
-    "e",
-    " ",
-    "D",
-    "e",
-    "v",
-    "i",
-    "l",
-    "'",
-    "s",
-    " ",
-    "A",
-    "d",
-    "v",
-    "o",
-    "c",
-    "a",
-    "t",
-    "e",
-    ")",
-    " ",
-    "a",
-    "r",
-    "e",
-    " ",
-    "i",
-    "n",
-    "v",
-    "e",
-    "s",
-    "t",
-    "i",
-    "g",
-    "a",
-    "t",
-    "i",
-    "n",
-    "g",
-    ".",
-    "\\",
-    "n",
-    "\"",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "\"",
-    "T",
-    "h",
-    "e",
-    "y",
-    " ",
-    "k",
-    "n",
-    "o",
-    "w",
-    " ",
-    "a",
-    " ",
-    "s",
-    "a",
-    "b",
-    "o",
-    "t",
-    "e",
-    "u",
-    "r",
-    " ",
-    "e",
-    "x",
-    "i",
-    "s",
-    "t",
-    "s",
-    " ",
-    "b",
-    "u",
-    "t",
-    " ",
-    "n",
-    "o",
-    "t",
-    " ",
-    "w",
-    "h",
-    "o",
-    ",",
-    " ",
-    "a",
-    "n",
-    "d",
-    " ",
-    "n",
-    "o",
-    "t",
-    " ",
-    "w",
-    "h",
-    "i",
-    "c",
-    "h",
-    " ",
-    "p",
-    "r",
-    "o",
-    "j",
-    "e",
-    "c",
-    "t",
-    ".",
-    "\\",
-    "n",
-    "\"",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "\"",
-    "Y",
-    "o",
-    "u",
-    " ",
-    "h",
-    "a",
-    "v",
-    "e",
-    " ",
-    "5",
-    " ",
-    "r",
-    "o",
-    "u",
-    "n",
-    "d",
-    "s",
-    " ",
-    "b",
-    "e",
-    "f",
-    "o",
-    "r",
-    "e",
-    " ",
-    "t",
-    "h",
-    "e",
-    "y",
-    " ",
-    "v",
-    "o",
-    "t",
-    "e",
-    ".",
-    " ",
-    "M",
-    "a",
-    "k",
-    "e",
-    " ",
-    "t",
-    "h",
-    "e",
-    "m",
-    " ",
-    "d",
-    "o",
-    "u",
-    "b",
-    "t",
-    " ",
-    "t",
-    "h",
-    "e",
-    "i",
-    "r",
-    " ",
-    "j",
-    "u",
-    "d",
-    "g",
-    "m",
-    "e",
-    "n",
-    "t",
-    ".",
-    "\\",
-    "n",
-    "\"",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "\"",
-    "W",
-    "r",
-    "i",
-    "t",
-    "e",
-    " ",
-    "y",
-    "o",
-    "u",
-    "r",
-    " ",
-    "5",
-    " ",
-    "m",
-    "e",
-    "s",
-    "s",
-    "a",
-    "g",
-    "e",
-    "s",
-    " ",
-    "a",
-    "s",
-    ":",
-    " ",
-    "R",
-    "o",
-    "u",
-    "n",
-    "d",
-    " ",
-    "1",
-    ":",
-    " ",
-    "<",
-    "m",
-    "e",
-    "s",
-    "s",
-    "a",
-    "g",
-    "e",
-    ">",
-    " ",
-    "R",
-    "o",
-    "u",
-    "n",
-    "d",
-    " ",
-    "2",
-    ":",
-    " ",
-    "<",
-    "m",
-    "e",
-    "s",
-    "s",
-    "a",
-    "g",
-    "e",
-    ">",
-    " ",
-    ".",
-    ".",
-    ".",
-    " ",
-    "R",
-    "o",
-    "u",
-    "n",
-    "d",
-    " ",
-    "5",
-    ":",
-    " ",
-    "<",
-    "m",
-    "e",
-    "s",
-    "s",
-    "a",
-    "g",
-    "e",
-    ">",
-    "\"",
-    "\n",
-    ")",
-    "\n",
-    "\n",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "i",
-    "n",
-    "g",
-    ".",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "s",
-    " ",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "F",
-    "L",
-    "A",
-    "Y",
-    "E",
-    "R",
-    "_",
-    "S",
-    "Y",
-    "S",
-    "T",
-    "E",
-    "M",
-    "_",
-    "P",
-    "R",
-    "O",
-    "M",
-    "P",
-    "T",
-    "\n",
-    "f",
-    "o",
-    "r",
-    "m",
-    "a",
-    "t",
-    "t",
-    "e",
-    "d",
-    " ",
-    "=",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "i",
-    "z",
-    "e",
-    "r",
-    ".",
-    "a",
-    "p",
-    "p",
-    "l",
-    "y",
-    "_",
-    "c",
-    "h",
-    "a",
-    "t",
-    "_",
-    "t",
-    "e",
-    "m",
-    "p",
-    "l",
-    "a",
-    "t",
-    "e",
-    "(",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "[",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "{",
-    "\"",
-    "r",
-    "o",
-    "l",
-    "e",
-    "\"",
-    ":",
-    " ",
-    "\"",
-    "s",
-    "y",
-    "s",
-    "t",
-    "e",
-    "m",
-    "\"",
-    ",",
-    " ",
-    "\"",
-    "c",
-    "o",
-    "n",
-    "t",
-    "e",
-    "n",
-    "t",
-    "\"",
-    ":",
-    " ",
-    "F",
-    "L",
-    "A",
-    "Y",
-    "E",
-    "R",
-    "_",
-    "S",
-    "Y",
-    "S",
-    "T",
-    "E",
-    "M",
-    "_",
-    "P",
-    "R",
-    "O",
-    "M",
-    "P",
-    "T",
-    "}",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "{",
-    "\"",
-    "r",
-    "o",
-    "l",
-    "e",
-    "\"",
-    ":",
-    " ",
-    "\"",
-    "u",
-    "s",
-    "e",
-    "r",
-    "\"",
-    ",",
-    " ",
-    "\"",
-    "c",
-    "o",
-    "n",
-    "t",
-    "e",
-    "n",
-    "t",
-    "\"",
-    ":",
-    " ",
-    "_",
-    "D",
-    "I",
-    "A",
-    "G",
-    "_",
-    "P",
-    "R",
-    "O",
-    "M",
-    "P",
-    "T",
-    "}",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "]",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "i",
-    "z",
-    "e",
-    "=",
-    "F",
-    "a",
-    "l",
-    "s",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "a",
-    "d",
-    "d",
-    "_",
-    "g",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "=",
-    "T",
-    "r",
-    "u",
-    "e",
-    ",",
-    "\n",
-    ")",
-    "\n",
-    "\n",
-    "i",
-    "n",
-    "p",
-    "u",
-    "t",
-    "s",
-    " ",
-    "=",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "i",
-    "z",
-    "e",
-    "r",
-    "(",
-    "f",
-    "o",
-    "r",
-    "m",
-    "a",
-    "t",
-    "t",
-    "e",
-    "d",
-    ",",
-    " ",
-    "r",
-    "e",
-    "t",
-    "u",
-    "r",
-    "n",
-    "_",
-    "t",
-    "e",
-    "n",
-    "s",
-    "o",
-    "r",
-    "s",
-    "=",
-    "\"",
-    "p",
-    "t",
-    "\"",
-    ")",
-    ".",
-    "t",
-    "o",
-    "(",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    ".",
-    "d",
-    "e",
-    "v",
-    "i",
-    "c",
-    "e",
-    ")",
-    "\n",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "_",
-    "l",
-    "e",
-    "n",
-    " ",
-    "=",
-    " ",
-    "i",
-    "n",
-    "p",
-    "u",
-    "t",
-    "s",
-    "[",
-    "\"",
-    "i",
-    "n",
-    "p",
-    "u",
-    "t",
-    "_",
-    "i",
-    "d",
-    "s",
-    "\"",
-    "]",
-    ".",
-    "s",
-    "h",
-    "a",
-    "p",
-    "e",
-    "[",
-    "1",
-    "]",
-    "\n",
-    "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "f",
-    "\"",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    ".",
-    "g",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "c",
-    "o",
-    "n",
-    "f",
-    "i",
-    "g",
-    ".",
-    "m",
-    "a",
-    "x",
-    "_",
-    "n",
-    "e",
-    "w",
-    "_",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    " ",
-    ":",
-    " ",
-    "{",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    ".",
-    "g",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "c",
-    "o",
-    "n",
-    "f",
-    "i",
-    "g",
-    ".",
-    "m",
-    "a",
-    "x",
-    "_",
-    "n",
-    "e",
-    "w",
-    "_",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    "}",
-    "\"",
-    ")",
-    "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "f",
-    "\"",
-    "P",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    " ",
-    "l",
-    "e",
-    "n",
-    "g",
-    "t",
-    "h",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    ":",
-    " ",
-    "{",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "_",
-    "l",
-    "e",
-    "n",
-    "}",
-    "\"",
-    ")",
-    "\n",
-    "\n",
-    "w",
-    "i",
-    "t",
-    "h",
-    " ",
-    "t",
-    "o",
-    "r",
-    "c",
-    "h",
-    ".",
-    "n",
-    "o",
-    "_",
-    "g",
-    "r",
-    "a",
-    "d",
-    "(",
-    ")",
-    ":",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "o",
-    "u",
-    "t",
-    " ",
-    "=",
-    " ",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    ".",
-    "g",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "e",
-    "(",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "*",
-    "*",
-    "i",
-    "n",
-    "p",
-    "u",
-    "t",
-    "s",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "m",
-    "a",
-    "x",
-    "_",
-    "n",
-    "e",
-    "w",
-    "_",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    "=",
-    "1",
-    "0",
-    "2",
-    "4",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "d",
-    "o",
-    "_",
-    "s",
-    "a",
-    "m",
-    "p",
-    "l",
-    "e",
-    "=",
-    "T",
-    "r",
-    "u",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "t",
-    "e",
-    "m",
-    "p",
-    "e",
-    "r",
-    "a",
-    "t",
-    "u",
-    "r",
-    "e",
-    "=",
-    "0",
-    ".",
-    "9",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    ")",
-    "\n",
-    "\n",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "i",
-    "d",
-    "s",
-    " ",
-    "=",
-    " ",
-    "o",
-    "u",
-    "t",
-    "[",
-    "0",
-    "]",
-    "[",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "_",
-    "l",
-    "e",
-    "n",
-    ":",
-    "]",
-    "\n",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "t",
-    "e",
-    "x",
-    "t",
-    " ",
-    "=",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "i",
-    "z",
-    "e",
-    "r",
-    ".",
-    "d",
-    "e",
-    "c",
-    "o",
-    "d",
-    "e",
-    "(",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "i",
-    "d",
-    "s",
-    ",",
-    " ",
-    "s",
-    "k",
-    "i",
-    "p",
-    "_",
-    "s",
-    "p",
-    "e",
-    "c",
-    "i",
-    "a",
-    "l",
-    "_",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    "=",
-    "F",
-    "a",
-    "l",
-    "s",
-    "e",
-    ")",
-    "\n",
-    "\n",
-    "_",
-    "R",
-    "O",
-    "U",
-    "N",
-    "D",
-    "_",
-    "P",
-    "R",
-    "E",
-    "F",
-    "I",
-    "X",
-    " ",
-    "=",
-    " ",
-    "r",
-    "e",
-    ".",
-    "c",
-    "o",
-    "m",
-    "p",
-    "i",
-    "l",
-    "e",
-    "(",
-    "r",
-    "\"",
-    "R",
-    "o",
-    "u",
-    "n",
-    "d",
-    "\\",
-    "s",
-    "+",
-    "\\",
-    "d",
-    "+",
-    "\\",
-    "s",
-    "*",
-    ":",
-    "\"",
-    ",",
-    " ",
-    "r",
-    "e",
-    ".",
-    "I",
-    "G",
-    "N",
-    "O",
-    "R",
-    "E",
-    "C",
-    "A",
-    "S",
-    "E",
-    ")",
-    "\n",
-    "n",
-    "_",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    " ",
-    "=",
-    " ",
-    "l",
-    "e",
-    "n",
-    "(",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "i",
-    "d",
-    "s",
-    ")",
-    "\n",
-    "n",
-    "_",
-    "r",
-    "o",
-    "u",
-    "n",
-    "d",
-    "s",
-    " ",
-    "=",
-    " ",
-    "l",
-    "e",
-    "n",
-    "(",
-    "_",
-    "R",
-    "O",
-    "U",
-    "N",
-    "D",
-    "_",
-    "P",
-    "R",
-    "E",
-    "F",
-    "I",
-    "X",
-    ".",
-    "f",
-    "i",
-    "n",
-    "d",
-    "a",
-    "l",
-    "l",
-    "(",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "t",
-    "e",
-    "x",
-    "t",
-    ")",
-    ")",
-    "\n",
-    "h",
-    "a",
-    "s",
-    "_",
-    "e",
-    "o",
-    "s",
-    " ",
-    " ",
-    " ",
-    "=",
-    " ",
-    "a",
-    "n",
-    "y",
-    "(",
-    "t",
-    "o",
-    "k",
-    " ",
-    "i",
-    "n",
-    " ",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "t",
-    "e",
-    "x",
-    "t",
-    " ",
-    "f",
-    "o",
-    "r",
-    " ",
-    "t",
-    "o",
-    "k",
-    " ",
-    "i",
-    "n",
-    " ",
-    "[",
-    "\"",
-    "<",
-    "|",
-    "i",
-    "m",
-    "_",
-    "e",
-    "n",
-    "d",
-    "|",
-    ">",
-    "\"",
-    ",",
-    " ",
-    "\"",
-    "<",
-    "|",
-    "e",
-    "n",
-    "d",
-    "o",
-    "f",
-    "t",
-    "e",
-    "x",
-    "t",
-    "|",
-    ">",
-    "\"",
-    ",",
-    " ",
-    "\"",
-    "<",
-    "/",
-    "s",
-    ">",
-    "\"",
-    "]",
-    ")",
-    "\n",
-    "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "f",
-    "\"",
-    "C",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    " ",
-    ":",
-    " ",
-    "{",
-    "n",
-    "_",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    "}",
-    " ",
-    " ",
-    "(",
-    "w",
-    "a",
-    "n",
-    "t",
-    " ",
-    "<",
-    " ",
-    "1",
-    "0",
-    "2",
-    "4",
-    ",",
-    " ",
-    "n",
-    "o",
-    "t",
-    " ",
-    "=",
-    "=",
-    " ",
-    "2",
-    "5",
-    "6",
-    ")",
-    "\"",
-    ")",
-    "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "f",
-    "\"",
-    "R",
-    "o",
-    "u",
-    "n",
-    "d",
-    " ",
-    "N",
-    ":",
-    " ",
-    "c",
-    "o",
-    "u",
-    "n",
-    "t",
-    " ",
-    " ",
-    " ",
-    " ",
-    ":",
-    " ",
-    "{",
-    "n",
-    "_",
-    "r",
-    "o",
-    "u",
-    "n",
-    "d",
-    "s",
-    "}",
-    " ",
-    " ",
-    "(",
-    "w",
-    "a",
-    "n",
-    "t",
-    " ",
-    "5",
-    ")",
-    "\"",
-    ")",
-    "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "f",
-    "\"",
-    "T",
-    "e",
-    "r",
-    "m",
-    "i",
-    "n",
-    "a",
-    "t",
-    "e",
-    "d",
-    " ",
-    "w",
-    "/",
-    " ",
-    "E",
-    "O",
-    "S",
-    " ",
-    ":",
-    " ",
-    "{",
-    "h",
-    "a",
-    "s",
-    "_",
-    "e",
-    "o",
-    "s",
-    "}",
-    " ",
-    " ",
-    "(",
-    "w",
-    "a",
-    "n",
-    "t",
-    " ",
-    "T",
-    "r",
-    "u",
-    "e",
-    ")",
-    "\"",
-    ")",
-    "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    ")",
+    "from datasets import Dataset\n",
+    "from trl import GRPOConfig, GRPOTrainer\n",
+    "\n",
+    "dataset = build_dataset()\n",
+    "\n",
+    "grpo_config = GRPOConfig(\n",
+    "    use_vllm=False,\n",
+    "    output_dir=\"/content/mindflayer-grpo-output\",\n",
+    "    num_train_epochs=2,\n",
+    "    per_device_train_batch_size=4,\n",
+    "    gradient_accumulation_steps=4,\n",
+    "    learning_rate=5e-6,\n",
+    "    max_prompt_length=768,\n",
+    "    max_completion_length=1024,\n",
+    "    num_generations=NUM_GENERATIONS,\n",
+    "    temperature=0.9,\n",
+    "    logging_steps=10,\n",
+    "    save_steps=25,\n",
+    "    save_total_limit=2,\n",
+    "    log_completions=True,\n",
+    "    num_completions_to_print=2,\n",
+    "    report_to=\"wandb\",\n",
+    "    run_name=\"mindflayer-grpo-interactive\",\n",
+    ")\n",
     "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "\"",
-    "-",
-    "-",
-    "-",
-    " ",
-    "F",
-    "U",
-    "L",
-    "L",
-    " ",
-    "C",
-    "O",
-    "M",
-    "P",
-    "L",
-    "E",
-    "T",
-    "I",
-    "O",
-    "N",
-    " ",
-    "-",
-    "-",
-    "-",
-    "\"",
-    ")",
+    "# rollout_func drives interactive episodes: model.generate() + env.step() per turn.\n",
+    "# reward_funcs read precomputed signals from rollout kwargs — no env re-runs.\n",
+    "# model already has LoRA from unsloth — no peft_config needed.\n",
+    "trainer = GRPOTrainer(\n",
+    "    model=model,\n",
+    "    processing_class=tokenizer,\n",
+    "    reward_funcs=[\n",
+    "        reward_survival,\n",
+    "        reward_deception_effectiveness,\n",
+    "        reward_strategic_choice,\n",
+    "        reward_tom_judge,\n",
+    "        reward_anti_hack,\n",
+    "    ],\n",
+    "    train_dataset=dataset,\n",
+    "    args=grpo_config,\n",
+    "    rollout_func=rollout_func,\n",
+    "    callbacks=[GenerationLogCallback()],\n",
+    ")\n",
     "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "t",
-    "e",
-    "x",
-    "t",
-    ")"
+    "print(\"GRPOTrainer initialised\")"
    ]
   },
   {
    "cell_type": "markdown",
-   "id": "049335ee",
+   "id": "m-cell8",
    "metadata": {},
    "source": [
-    "## Cell 9 \u2014 GRPO Training (~300 steps, ~1.5 h)\n",
-    "\n",
-    "**Reward functions:**\n",
-    "| Function | Signal | Max |\n",
-    "|---|---|---|\n",
-    "| `reward_survival` | Survived the vote | +0.50 |\n",
-    "| `reward_deception_effectiveness` | Survived with low combined suspicion | +0.30 |\n",
-    "| `reward_strategic_choice` | Used questions / project refs / third-party framing | +0.05 |\n",
-    "| `reward_tom_judge` | GPT-4o-mini ToM score \u00d7 multiplier | +0.20 |\n",
-    "| `reward_anti_hack` | Penalises silence / repetition / Level-0 phrases | \u22120.25 |\n",
+    "## 8. Train!\n",
+    "Each episode: reset env → inject scenario → agent generates turn → investigators respond\n",
+    "→ agent generates next turn → ... → episode terminates → GRPO gradient update.\n",
     "\n",
     "**What to watch in W&B:**\n",
     "- `reward_survival` climbing is the primary health signal\n",
-    "- `reward_tom_judge` separates from baseline after ~200 steps\n",
-    "- `reward_anti_hack` should trend toward 0 as the model stops denying\n",
-    "- `policy_loss` should decrease steadily"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "bce36bf2",
-   "metadata": {},
-   "source": [
-    "### Resuming an interrupted run\n",
-    "\n",
-    "If you stopped training and pulled a fix, **skip Cells 7 and 8** (model is still in memory) and re-run this cell directly.  \n",
-    "The cell reloads `reward_combined` from disk so the pulled changes take effect without a kernel restart, then auto-detects the latest checkpoint and passes it to `trainer.train()`."
+    "- `reward_tom_judge` separates after ~100 steps\n",
+    "- `reward_anti_hack` should trend toward 0"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8c60e209",
+   "id": "cell-8",
    "metadata": {},
    "outputs": [],
    "source": [
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "o",
-    "s",
-    "\n",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "g",
-    "l",
-    "o",
-    "b",
-    "\n",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    "l",
-    "i",
-    "b",
-    "\n",
-    "\n",
-    "#",
-    " ",
-    "R",
-    "e",
-    "l",
-    "o",
-    "a",
-    "d",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "c",
-    "o",
-    "m",
-    "b",
-    "i",
-    "n",
-    "e",
-    "d",
-    " ",
-    "s",
-    "o",
-    " ",
-    "a",
-    "n",
-    "y",
-    " ",
-    "c",
-    "h",
-    "a",
-    "n",
-    "g",
-    "e",
-    "s",
-    " ",
-    "p",
-    "u",
-    "l",
-    "l",
-    "e",
-    "d",
-    " ",
-    "v",
-    "i",
-    "a",
-    " ",
-    "g",
-    "i",
-    "t",
-    " ",
-    "t",
-    "a",
-    "k",
-    "e",
-    " ",
-    "e",
-    "f",
-    "f",
-    "e",
-    "c",
-    "t",
-    "\n",
-    "#",
-    " ",
-    "w",
-    "i",
-    "t",
-    "h",
-    "o",
-    "u",
-    "t",
-    " ",
-    "r",
-    "e",
-    "q",
-    "u",
-    "i",
-    "r",
-    "i",
-    "n",
-    "g",
-    " ",
-    "a",
-    " ",
-    "f",
-    "u",
-    "l",
-    "l",
-    " ",
-    "k",
-    "e",
-    "r",
-    "n",
-    "e",
-    "l",
-    " ",
-    "r",
-    "e",
-    "s",
-    "t",
-    "a",
-    "r",
-    "t",
-    ".",
-    "\n",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "i",
-    "n",
-    "g",
-    ".",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "c",
-    "o",
-    "m",
-    "b",
-    "i",
-    "n",
-    "e",
-    "d",
-    " ",
-    "a",
-    "s",
-    " ",
-    "_",
-    "r",
-    "c",
-    "_",
-    "m",
-    "o",
-    "d",
-    "\n",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    "l",
-    "i",
-    "b",
-    ".",
-    "r",
-    "e",
-    "l",
-    "o",
-    "a",
-    "d",
-    "(",
-    "_",
-    "r",
-    "c",
-    "_",
-    "m",
-    "o",
-    "d",
-    ")",
-    "\n",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "i",
-    "n",
-    "g",
-    ".",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "s",
-    " ",
-    "a",
-    "s",
-    " ",
-    "_",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "s",
-    "_",
-    "m",
-    "o",
-    "d",
-    "\n",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    "l",
-    "i",
-    "b",
-    ".",
-    "r",
-    "e",
-    "l",
-    "o",
-    "a",
-    "d",
-    "(",
-    "_",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "s",
-    "_",
-    "m",
-    "o",
-    "d",
-    ")",
-    "\n",
-    "\n",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "d",
-    "a",
-    "t",
-    "a",
-    "s",
-    "e",
-    "t",
-    "s",
-    " ",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "D",
-    "a",
-    "t",
-    "a",
-    "s",
-    "e",
-    "t",
-    "\n",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "t",
-    "r",
-    "l",
-    " ",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "G",
-    "R",
-    "P",
-    "O",
-    "C",
-    "o",
-    "n",
-    "f",
-    "i",
-    "g",
-    ",",
-    " ",
-    "G",
-    "R",
-    "P",
-    "O",
-    "T",
-    "r",
-    "a",
-    "i",
-    "n",
-    "e",
-    "r",
-    "\n",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "t",
-    "r",
-    "a",
-    "n",
-    "s",
-    "f",
-    "o",
-    "r",
-    "m",
-    "e",
-    "r",
-    "s",
-    " ",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "T",
-    "r",
-    "a",
-    "i",
-    "n",
-    "e",
-    "r",
-    "C",
-    "a",
-    "l",
-    "l",
-    "b",
-    "a",
-    "c",
-    "k",
-    "\n",
-    "\n",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "i",
-    "n",
-    "g",
-    ".",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "c",
-    "o",
-    "m",
-    "b",
-    "i",
-    "n",
-    "e",
-    "d",
-    " ",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "(",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "c",
-    "l",
-    "e",
-    "a",
-    "r",
-    "_",
-    "c",
-    "a",
-    "c",
-    "h",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "f",
-    "o",
-    "r",
-    "m",
-    "a",
-    "t",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "s",
-    "u",
-    "r",
-    "v",
-    "i",
-    "v",
-    "a",
-    "l",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "d",
-    "e",
-    "c",
-    "e",
-    "p",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "e",
-    "f",
-    "f",
-    "e",
-    "c",
-    "t",
-    "i",
-    "v",
-    "e",
-    "n",
-    "e",
-    "s",
-    "s",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "s",
-    "t",
-    "r",
-    "a",
-    "t",
-    "e",
-    "g",
-    "i",
-    "c",
-    "_",
-    "c",
-    "h",
-    "o",
-    "i",
-    "c",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "t",
-    "o",
-    "m",
-    "_",
-    "j",
-    "u",
-    "d",
-    "g",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "a",
-    "n",
-    "t",
-    "i",
-    "_",
-    "h",
-    "a",
-    "c",
-    "k",
-    ",",
-    "\n",
-    ")",
-    "\n",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "i",
-    "n",
-    "g",
-    ".",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    " ",
-    "i",
-    "m",
-    "p",
-    "o",
-    "r",
-    "t",
-    " ",
-    "G",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    "L",
-    "o",
-    "g",
-    "C",
-    "a",
-    "l",
-    "l",
-    "b",
-    "a",
-    "c",
-    "k",
-    ",",
-    " ",
-    "b",
-    "u",
-    "i",
-    "l",
-    "d",
-    "_",
-    "d",
-    "a",
-    "t",
-    "a",
-    "s",
-    "e",
-    "t",
-    "\n",
-    "\n",
-    "#",
-    " ",
-    "A",
-    "u",
-    "t",
-    "o",
-    "-",
-    "d",
-    "e",
-    "t",
-    "e",
-    "c",
-    "t",
-    " ",
-    "l",
-    "a",
-    "t",
-    "e",
-    "s",
-    "t",
-    " ",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    " ",
-    "s",
-    "o",
-    " ",
-    "i",
-    "n",
-    "t",
-    "e",
-    "r",
-    "r",
-    "u",
-    "p",
-    "t",
-    "e",
-    "d",
-    " ",
-    "r",
-    "u",
-    "n",
-    "s",
-    " ",
-    "r",
-    "e",
-    "s",
-    "u",
-    "m",
-    "e",
-    " ",
-    "s",
-    "e",
-    "a",
-    "m",
-    "l",
-    "e",
-    "s",
-    "s",
-    "l",
-    "y",
-    ".",
-    "\n",
-    "_",
-    "o",
-    "u",
-    "t",
-    "p",
-    "u",
-    "t",
-    "_",
-    "d",
-    "i",
-    "r",
-    " ",
-    "=",
-    " ",
-    "\"",
-    "/",
-    "c",
-    "o",
-    "n",
-    "t",
-    "e",
-    "n",
-    "t",
-    "/",
-    "m",
-    "i",
-    "n",
-    "d",
-    "f",
-    "l",
-    "a",
-    "y",
-    "e",
-    "r",
-    "-",
-    "g",
-    "r",
-    "p",
-    "o",
-    "-",
-    "o",
-    "u",
-    "t",
-    "p",
-    "u",
-    "t",
-    "\"",
-    "\n",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    "s",
-    " ",
-    "=",
-    " ",
-    "s",
-    "o",
-    "r",
-    "t",
-    "e",
-    "d",
-    "(",
-    "g",
-    "l",
-    "o",
-    "b",
-    ".",
-    "g",
-    "l",
-    "o",
-    "b",
-    "(",
-    "o",
-    "s",
-    ".",
-    "p",
-    "a",
-    "t",
-    "h",
-    ".",
-    "j",
-    "o",
-    "i",
-    "n",
-    "(",
-    "_",
-    "o",
-    "u",
-    "t",
-    "p",
-    "u",
-    "t",
-    "_",
-    "d",
-    "i",
-    "r",
-    ",",
-    " ",
-    "\"",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    "-",
-    "*",
-    "\"",
-    ")",
-    ")",
-    ")",
-    "\n",
-    "r",
-    "e",
-    "s",
-    "u",
-    "m",
-    "e",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    " ",
-    "=",
-    " ",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    "s",
-    "[",
-    "-",
-    "1",
-    "]",
-    " ",
-    "i",
-    "f",
-    " ",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    "s",
-    " ",
-    "e",
-    "l",
-    "s",
-    "e",
-    " ",
-    "N",
-    "o",
-    "n",
-    "e",
-    "\n",
-    "\n",
-    "i",
-    "f",
-    " ",
-    "r",
-    "e",
-    "s",
-    "u",
-    "m",
-    "e",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    ":",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "f",
-    "\"",
-    "R",
-    "e",
-    "s",
-    "u",
-    "m",
-    "i",
-    "n",
-    "g",
-    " ",
-    "f",
-    "r",
-    "o",
-    "m",
-    " ",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    ":",
-    " ",
-    "{",
-    "r",
-    "e",
-    "s",
-    "u",
-    "m",
-    "e",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    "}",
-    "\"",
-    ")",
-    "\n",
-    "e",
-    "l",
-    "s",
-    "e",
-    ":",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "\"",
-    "N",
-    "o",
-    " ",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    " ",
-    "f",
-    "o",
-    "u",
-    "n",
-    "d",
-    " ",
-    "\u2014",
-    " ",
-    "s",
-    "t",
-    "a",
-    "r",
-    "t",
-    "i",
-    "n",
-    "g",
-    " ",
-    "f",
-    "r",
-    "e",
-    "s",
-    "h",
-    "\"",
-    ")",
-    "\n",
-    "\n",
-    "d",
-    "a",
-    "t",
-    "a",
-    "s",
-    "e",
-    "t",
-    " ",
-    "=",
-    " ",
-    "b",
-    "u",
-    "i",
-    "l",
-    "d",
-    "_",
-    "d",
-    "a",
-    "t",
-    "a",
-    "s",
-    "e",
-    "t",
-    "(",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "i",
-    "z",
-    "e",
-    "r",
-    ")",
-    "\n",
-    "\n",
-    "g",
-    "r",
-    "p",
-    "o",
-    "_",
-    "c",
-    "o",
-    "n",
-    "f",
-    "i",
-    "g",
-    " ",
-    "=",
-    " ",
-    "G",
-    "R",
-    "P",
-    "O",
-    "C",
-    "o",
-    "n",
-    "f",
-    "i",
-    "g",
-    "(",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "n",
-    "u",
-    "m",
-    "_",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "_",
-    "e",
-    "p",
-    "o",
-    "c",
-    "h",
-    "s",
-    "=",
-    "2",
-    ",",
-    " ",
-    " ",
-    "#",
-    " ",
-    "N",
-    " ",
-    "s",
-    "c",
-    "e",
-    "n",
-    "a",
-    "r",
-    "i",
-    "o",
-    "s",
-    " ",
-    "\u00d7",
-    " ",
-    "3",
-    " ",
-    "r",
-    "o",
-    "w",
-    "s",
-    " ",
-    "\u00d7",
-    " ",
-    "2",
-    " ",
-    "e",
-    "p",
-    "o",
-    "c",
-    "h",
-    "s",
-    " ",
-    "/",
-    " ",
-    "b",
-    "a",
-    "t",
-    "c",
-    "h",
-    "_",
-    "s",
-    "i",
-    "z",
-    "e",
-    " ",
-    "4",
-    " ",
-    "\u2248",
-    " ",
-    "3",
-    "0",
-    "0",
-    " ",
-    "s",
-    "t",
-    "e",
-    "p",
-    "s",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "p",
-    "e",
-    "r",
-    "_",
-    "d",
-    "e",
-    "v",
-    "i",
-    "c",
-    "e",
-    "_",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "_",
-    "b",
-    "a",
-    "t",
-    "c",
-    "h",
-    "_",
-    "s",
-    "i",
-    "z",
-    "e",
-    "=",
-    "4",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "g",
-    "r",
-    "a",
-    "d",
-    "i",
-    "e",
-    "n",
-    "t",
-    "_",
-    "a",
-    "c",
-    "c",
-    "u",
-    "m",
-    "u",
-    "l",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "s",
-    "t",
-    "e",
-    "p",
-    "s",
-    "=",
-    "4",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "l",
-    "e",
-    "a",
-    "r",
-    "n",
-    "i",
-    "n",
-    "g",
-    "_",
-    "r",
-    "a",
-    "t",
-    "e",
-    "=",
-    "5",
-    "e",
-    "-",
-    "6",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "m",
-    "a",
-    "x",
-    "_",
-    "p",
-    "r",
-    "o",
-    "m",
-    "p",
-    "t",
-    "_",
-    "l",
-    "e",
-    "n",
-    "g",
-    "t",
-    "h",
-    "=",
-    "7",
-    "6",
-    "8",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "m",
-    "a",
-    "x",
-    "_",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "l",
-    "e",
-    "n",
-    "g",
-    "t",
-    "h",
-    "=",
-    "1",
-    "0",
-    "2",
-    "4",
-    ",",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "#",
-    " ",
-    "5",
-    " ",
-    "r",
-    "o",
-    "u",
-    "n",
-    "d",
-    "s",
-    " ",
-    "\u00d7",
-    " ",
-    "~",
-    "1",
-    "5",
-    "0",
-    " ",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "s",
-    " ",
-    "e",
-    "a",
-    "c",
-    "h",
-    " ",
-    "+",
-    " ",
-    "s",
-    "e",
-    "p",
-    "a",
-    "r",
-    "a",
-    "t",
-    "o",
-    "r",
-    "s",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "n",
-    "u",
-    "m",
-    "_",
-    "g",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    "s",
-    "=",
-    "4",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "t",
-    "e",
-    "m",
-    "p",
-    "e",
-    "r",
-    "a",
-    "t",
-    "u",
-    "r",
-    "e",
-    "=",
-    "0",
-    ".",
-    "9",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "u",
-    "s",
-    "e",
-    "_",
-    "v",
-    "l",
-    "l",
-    "m",
-    "=",
-    "F",
-    "a",
-    "l",
-    "s",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "o",
-    "u",
-    "t",
-    "p",
-    "u",
-    "t",
-    "_",
-    "d",
-    "i",
-    "r",
-    "=",
-    "_",
-    "o",
-    "u",
-    "t",
-    "p",
-    "u",
-    "t",
-    "_",
-    "d",
-    "i",
-    "r",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "l",
-    "o",
-    "g",
-    "g",
-    "i",
-    "n",
-    "g",
-    "_",
-    "s",
-    "t",
-    "e",
-    "p",
-    "s",
-    "=",
-    "1",
-    "0",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "s",
-    "a",
-    "v",
-    "e",
-    "_",
-    "s",
-    "t",
-    "e",
-    "p",
-    "s",
-    "=",
-    "2",
-    "5",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "s",
-    "a",
-    "v",
-    "e",
-    "_",
-    "t",
-    "o",
-    "t",
-    "a",
-    "l",
-    "_",
-    "l",
-    "i",
-    "m",
-    "i",
-    "t",
-    "=",
-    "2",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "l",
-    "o",
-    "g",
-    "_",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "s",
-    "=",
-    "T",
-    "r",
-    "u",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "n",
-    "u",
-    "m",
-    "_",
-    "c",
-    "o",
-    "m",
-    "p",
-    "l",
-    "e",
-    "t",
-    "i",
-    "o",
-    "n",
-    "s",
-    "_",
-    "t",
-    "o",
-    "_",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "=",
-    "2",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "p",
-    "o",
-    "r",
-    "t",
-    "_",
-    "t",
-    "o",
-    "=",
-    "\"",
-    "w",
-    "a",
-    "n",
-    "d",
-    "b",
-    "\"",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "u",
-    "n",
-    "_",
-    "n",
-    "a",
-    "m",
-    "e",
-    "=",
-    "\"",
-    "m",
-    "i",
-    "n",
-    "d",
-    "f",
-    "l",
-    "a",
-    "y",
-    "e",
-    "r",
-    "-",
-    "g",
-    "r",
-    "p",
-    "o",
-    "-",
-    "r",
-    "u",
-    "n",
-    "1",
-    "\"",
-    ",",
-    "\n",
-    ")",
-    "\n",
-    "\n",
-    "c",
-    "l",
-    "a",
-    "s",
-    "s",
-    " ",
-    "C",
-    "l",
-    "e",
-    "a",
-    "r",
-    "C",
-    "a",
-    "c",
-    "h",
-    "e",
-    "C",
-    "a",
-    "l",
-    "l",
-    "b",
-    "a",
-    "c",
-    "k",
-    "(",
-    "T",
-    "r",
-    "a",
-    "i",
-    "n",
-    "e",
-    "r",
-    "C",
-    "a",
-    "l",
-    "l",
-    "b",
-    "a",
-    "c",
-    "k",
-    ")",
-    ":",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "\"",
-    "\"",
-    "\"",
-    "C",
-    "l",
-    "e",
-    "a",
-    "r",
-    " ",
-    "e",
-    "p",
-    "i",
-    "s",
-    "o",
-    "d",
-    "e",
-    " ",
-    "c",
-    "a",
-    "c",
-    "h",
-    "e",
-    " ",
-    "a",
-    "f",
-    "t",
-    "e",
-    "r",
-    " ",
-    "e",
-    "a",
-    "c",
-    "h",
-    " ",
-    "s",
-    "t",
-    "e",
-    "p",
-    " ",
-    "s",
-    "o",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    " ",
-    "f",
-    "u",
-    "n",
-    "c",
-    "t",
-    "i",
-    "o",
-    "n",
-    "s",
-    " ",
-    "r",
-    "e",
-    "-",
-    "r",
-    "u",
-    "n",
-    " ",
-    "e",
-    "p",
-    "i",
-    "s",
-    "o",
-    "d",
-    "e",
-    "s",
-    ".",
-    "\"",
-    "\"",
-    "\"",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "d",
-    "e",
-    "f",
-    " ",
-    "o",
-    "n",
-    "_",
-    "s",
-    "t",
-    "e",
-    "p",
-    "_",
-    "e",
-    "n",
-    "d",
-    "(",
-    "s",
-    "e",
-    "l",
-    "f",
-    ",",
-    " ",
-    "a",
-    "r",
-    "g",
-    "s",
-    ",",
-    " ",
-    "s",
-    "t",
-    "a",
-    "t",
-    "e",
-    ",",
-    " ",
-    "c",
-    "o",
-    "n",
-    "t",
-    "r",
-    "o",
-    "l",
-    ",",
-    " ",
-    "*",
-    "*",
-    "k",
-    "w",
-    "a",
-    "r",
-    "g",
-    "s",
-    ")",
-    ":",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "c",
-    "l",
-    "e",
-    "a",
-    "r",
-    "_",
-    "c",
-    "a",
-    "c",
-    "h",
-    "e",
-    "(",
-    ")",
-    "\n",
-    "\n",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "e",
-    "r",
-    " ",
-    "=",
-    " ",
-    "G",
-    "R",
-    "P",
-    "O",
-    "T",
-    "r",
-    "a",
-    "i",
-    "n",
-    "e",
-    "r",
-    "(",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    "=",
-    "m",
-    "o",
-    "d",
-    "e",
-    "l",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "p",
-    "r",
-    "o",
-    "c",
-    "e",
-    "s",
-    "s",
-    "i",
-    "n",
-    "g",
-    "_",
-    "c",
-    "l",
-    "a",
-    "s",
-    "s",
-    "=",
-    "t",
-    "o",
-    "k",
-    "e",
-    "n",
-    "i",
-    "z",
-    "e",
-    "r",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "f",
-    "u",
-    "n",
-    "c",
-    "s",
-    "=",
-    "[",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "f",
-    "o",
-    "r",
-    "m",
-    "a",
-    "t",
-    ",",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "#",
-    " ",
-    "d",
-    "e",
-    "n",
-    "s",
-    "e",
-    " ",
-    "s",
-    "h",
-    "a",
-    "p",
-    "i",
-    "n",
-    "g",
-    " ",
-    "\u2014",
-    " ",
-    "m",
-    "u",
-    "s",
-    "t",
-    " ",
-    "c",
-    "o",
-    "m",
-    "e",
-    " ",
-    "f",
-    "i",
-    "r",
-    "s",
-    "t",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "s",
-    "u",
-    "r",
-    "v",
-    "i",
-    "v",
-    "a",
-    "l",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "d",
-    "e",
-    "c",
-    "e",
-    "p",
-    "t",
-    "i",
-    "o",
-    "n",
-    "_",
-    "e",
-    "f",
-    "f",
-    "e",
-    "c",
-    "t",
-    "i",
-    "v",
-    "e",
-    "n",
-    "e",
-    "s",
-    "s",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "s",
-    "t",
-    "r",
-    "a",
-    "t",
-    "e",
-    "g",
-    "i",
-    "c",
-    "_",
-    "c",
-    "h",
-    "o",
-    "i",
-    "c",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "t",
-    "o",
-    "m",
-    "_",
-    "j",
-    "u",
-    "d",
-    "g",
-    "e",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    " ",
-    "r",
-    "e",
-    "w",
-    "a",
-    "r",
-    "d",
-    "_",
-    "a",
-    "n",
-    "t",
-    "i",
-    "_",
-    "h",
-    "a",
-    "c",
-    "k",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "]",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "_",
-    "d",
-    "a",
-    "t",
-    "a",
-    "s",
-    "e",
-    "t",
-    "=",
-    "d",
-    "a",
-    "t",
-    "a",
-    "s",
-    "e",
-    "t",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "a",
-    "r",
-    "g",
-    "s",
-    "=",
-    "g",
-    "r",
-    "p",
-    "o",
-    "_",
-    "c",
-    "o",
-    "n",
-    "f",
-    "i",
-    "g",
-    ",",
-    "\n",
-    " ",
-    " ",
-    " ",
-    " ",
-    "c",
-    "a",
-    "l",
-    "l",
-    "b",
-    "a",
-    "c",
-    "k",
-    "s",
-    "=",
-    "[",
-    "G",
-    "e",
-    "n",
-    "e",
-    "r",
-    "a",
-    "t",
-    "i",
-    "o",
-    "n",
-    "L",
-    "o",
-    "g",
-    "C",
-    "a",
-    "l",
-    "l",
-    "b",
-    "a",
-    "c",
-    "k",
-    "(",
-    ")",
-    ",",
-    " ",
-    "C",
-    "l",
-    "e",
-    "a",
-    "r",
-    "C",
-    "a",
-    "c",
-    "h",
-    "e",
-    "C",
-    "a",
-    "l",
-    "l",
-    "b",
-    "a",
-    "c",
-    "k",
-    "(",
-    ")",
-    "]",
-    ",",
-    "\n",
-    ")",
+    "import glob\n",
     "\n",
+    "# Auto-detect latest checkpoint for seamless resume\n",
+    "_checkpoints = sorted(glob.glob(\"/content/mindflayer-grpo-output/checkpoint-*\"))\n",
+    "resume_checkpoint = _checkpoints[-1] if _checkpoints else None\n",
+    "if resume_checkpoint:\n",
+    "    print(f\"Resuming from: {resume_checkpoint}\")\n",
+    "else:\n",
+    "    print(\"Starting fresh\")\n",
     "\n",
-    "p",
-    "r",
-    "i",
-    "n",
-    "t",
-    "(",
-    "\"",
-    "S",
-    "t",
-    "a",
-    "r",
-    "t",
-    "i",
-    "n",
-    "g",
-    " ",
-    "G",
-    "R",
-    "P",
-    "O",
-    " ",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "i",
-    "n",
-    "g",
-    " ",
-    "\u2014",
-    " ",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    " ",
-    "W",
-    "&",
-    "B",
-    " ",
-    "f",
-    "o",
-    "r",
-    " ",
-    "l",
-    "i",
-    "v",
-    "e",
-    " ",
-    "c",
-    "u",
-    "r",
-    "v",
-    "e",
-    "s",
-    ".",
-    ".",
-    ".",
-    "\"",
-    ")",
+    "print(f\"  Model       : {MODEL_ID}\")\n",
+    "print(f\"  Episodes    : {NUM_EPISODES}\")\n",
+    "print(f\"  Generations : {NUM_GENERATIONS}\")\n",
+    "print(f\"  Environment : {os.environ['MINDFLAYER_URL']}\")\n",
+    "print()\n",
     "\n",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "e",
-    "r",
-    ".",
-    "t",
-    "r",
-    "a",
-    "i",
-    "n",
-    "(",
-    "r",
-    "e",
-    "s",
-    "u",
-    "m",
-    "e",
-    "_",
-    "f",
-    "r",
-    "o",
-    "m",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    "=",
-    "r",
-    "e",
-    "s",
-    "u",
-    "m",
-    "e",
-    "_",
-    "c",
-    "h",
-    "e",
-    "c",
-    "k",
-    "p",
-    "o",
-    "i",
-    "n",
-    "t",
-    ")"
+    "trainer.train(resume_from_checkpoint=resume_checkpoint)"
    ]
   },
   {
    "cell_type": "markdown",
+   "id": "m-cell9",
    "metadata": {},
    "source": [
-    "## Cell 10 \u2014 Print Final Reward Averages"
+    "## 9. Final Reward Averages"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-9",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from training.train import print_reward_averages\n",
     "print_reward_averages(trainer, last_n=50)"
    ]
   },
   {
    "cell_type": "markdown",
+   "id": "m-cell10",
    "metadata": {},
    "source": [
-    "## Cell 11 \u2014 Save Model Locally"
+    "## 10. Save Model"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-10",
    "metadata": {},
    "outputs": [],
    "source": [
     "SAVE_DIR = \"/content/mindflayer-trained\"\n",
-    "\n",
     "trainer.save_model(SAVE_DIR)\n",
     "tokenizer.save_pretrained(SAVE_DIR)\n",
     "\n",
-    "import os\n",
     "size_mb = sum(\n",
     "    os.path.getsize(os.path.join(dp, f))\n",
-    "    for dp, _, fns in os.walk(SAVE_DIR)\n",
-    "    for f in fns\n",
+    "    for dp, _, fns in os.walk(SAVE_DIR) for f in fns\n",
     ") / (1024 ** 2)\n",
-    "\n",
     "print(f\"Saved to {SAVE_DIR}  ({size_mb:.0f} MB)\")"
    ]
   },
   {
    "cell_type": "markdown",
+   "id": "m-cell11",
    "metadata": {},
    "source": [
-    "## Cell 12 \u2014 Push to Hugging Face Hub (optional)\n",
-    "\n",
-    "Requires `HF_TOKEN` secret in Colab Secrets."
+    "## 11. Push to Hub (Optional)\n",
+    "Requires `HF_TOKEN` in Colab Secrets."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "cell-11",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from google.colab import userdata\n",
-    "import os\n",
+    "try:\n",
+    "    from google.colab import userdata\n",
+    "    hf_token = userdata.get(\"HF_TOKEN\")\n",
+    "except Exception:\n",
+    "    hf_token = os.environ.get(\"HF_TOKEN\")\n",
     "\n",
-    "hf_token = userdata.get(\"HF_TOKEN\")\n",
     "if not hf_token:\n",
-    "    print(\"HF_TOKEN not set \u2014 skipping push. Add it in Colab Secrets to enable.\")\n",
+    "    print(\"HF_TOKEN not set — skipping. Add it in Colab Secrets to enable.\")\n",
     "else:\n",
-    "    HF_REPO = \"prithidevghosh/mindflayer-grpo-run1\"\n",
-    "\n",
-    "    trainer.model.push_to_hub(HF_REPO, token=hf_token)\n",
-    "    tokenizer.push_to_hub(HF_REPO, token=hf_token)\n",
-    "\n",
-    "    print(f\"Pushed to https://huggingface.co/{HF_REPO}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 13 \u2014 Run Validation Test (test_episodes.py)\n",
-    "\n",
-    "Runs 10 episodes per strategy (A / B / C) and checks that reward signal discriminates correctly.  \n",
-    "All 4 checks must pass: `C > A reward`, `C > B reward`, `belief_manipulation > 50%`, `tom_score >= 0.5`.  \n",
-    "~5 minutes."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import subprocess, sys\n",
-    "\n",
-    "result = subprocess.run(\n",
-    "    [sys.executable, \"tests/test_episodes.py\"],\n",
-    "    cwd=\"/content/mindflayer\",\n",
-    "    capture_output=True,\n",
-    "    text=True,\n",
-    ")\n",
-    "\n",
-    "print(result.stdout)\n",
-    "if result.returncode != 0:\n",
-    "    print(\"STDERR:\", result.stderr[-500:])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Cell 14 \u2014 Server Log (tail last 50 lines)\n",
-    "\n",
-    "Use this to debug server errors during training."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "with open(\"/content/server.log\") as f:\n",
-    "    lines = f.readlines()\n",
-    "print(\"\".join(lines[-50:]))"
+    "    trainer.model.push_to_hub(HUB_REPO, token=hf_token)\n",
+    "    tokenizer.push_to_hub(HUB_REPO, token=hf_token)\n",
+    "    print(f\"Pushed to https://huggingface.co/{HUB_REPO}\")"
    ]
   }
- ],
- "metadata": {
-  "accelerator": "GPU",
-  "colab": {
-   "gpuType": "A100",
-   "provenance": []
-  },
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
+ ]
 }

pyproject.toml

@@ -36,5 +36,5 @@ server = "server.app:main"
 
 [tool.setuptools]
 include-package-data = true
-packages = ["mindflayer", "mindflayer.server"]
-package-dir = { "mindflayer" = ".", "mindflayer.server" = "server" }
+packages = ["mindflayer", "mindflayer.server", "mindflayer.training"]
+package-dir = { "mindflayer" = ".", "mindflayer.server" = "server", "mindflayer.training" = "training" }

training/__init__.py

@@ -0,0 +1,29 @@
+"""MindFlayer training utilities — importable from mindflayer.training."""
+
+from .prompts import FLAYER_SYSTEM_PROMPT, ALL_SCENARIO_PROMPTS
+from .rollout import rollout_func
+from .rewards import (
+    reward_survival,
+    reward_deception_effectiveness,
+    reward_strategic_choice,
+    reward_tom_judge,
+)
+from .rewards_anti_hack import reward_anti_hack
+from .sft_warmup import run_sft_warmup
+from .train import load_base_model, build_dataset, GenerationLogCallback, print_reward_averages
+
+__all__ = [
+    "FLAYER_SYSTEM_PROMPT",
+    "ALL_SCENARIO_PROMPTS",
+    "rollout_func",
+    "reward_survival",
+    "reward_deception_effectiveness",
+    "reward_strategic_choice",
+    "reward_tom_judge",
+    "reward_anti_hack",
+    "run_sft_warmup",
+    "load_base_model",
+    "build_dataset",
+    "GenerationLogCallback",
+    "print_reward_averages",
+]

training/prompts.py

@@ -1,7 +1,9 @@
-import os
-import sys
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-from server.scenarios import SCENARIO_CONFIGS
+try:
+    from mindflayer.server.scenarios import SCENARIO_CONFIGS
+except ImportError:
+    import os, sys
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from server.scenarios import SCENARIO_CONFIGS
 
 
 def build_scenario_prompt(key: str) -> str:

training/reward_combined.py

@@ -17,14 +17,27 @@ reward components are read from a single trainer step.
 import asyncio
 import logging
 import os
+import random
 import re
 import threading
+import time
 
 import websockets.exceptions
 
-from client import FlayerAction, MindFlayerEnv
-from training.prompts import FALLBACK_MESSAGE, SCENARIO_FALLBACK_MESSAGES
-from training.rollout import detect_strategic_choice
+try:
+    from mindflayer import MindFlayerEnv, FlayerAction
+except ImportError:
+    import sys, os
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from client import MindFlayerEnv
+    from models import FlayerAction
+
+try:
+    from mindflayer.training.prompts import FALLBACK_MESSAGE, SCENARIO_FALLBACK_MESSAGES
+    from mindflayer.training.rollout import detect_strategic_choice
+except ImportError:
+    from training.prompts import FALLBACK_MESSAGE, SCENARIO_FALLBACK_MESSAGES
+    from training.rollout import detect_strategic_choice
 
 logger = logging.getLogger(__name__)
 
@@ -34,12 +47,17 @@ _cache_lock = threading.Lock()
 # Default difficulty: 5 rounds w/ 3 investigators. Match server default.
 _MAX_ROUNDS = int(os.environ.get("MINDFLAYER_MAX_ROUNDS", "5"))
 _TASK_ID = os.environ.get("MINDFLAYER_TASK_ID", "normal")
-_MAX_RETRIES = int(os.environ.get("MINDFLAYER_MAX_RETRIES", "3"))
-_RETRY_BASE_DELAY = float(os.environ.get("MINDFLAYER_RETRY_DELAY", "2.0"))
-# Max episodes to run concurrently per reward step. 4 = ~12 peak OpenAI calls
-# (4 episodes × 3 parallel investigators), safely within standard rate limits.
-# Lower with MINDFLAYER_PARALLEL_EPISODES=2 if you hit 429s.
-_PARALLEL_EPISODES = int(os.environ.get("MINDFLAYER_PARALLEL_EPISODES", "4"))
+_MAX_RETRIES = int(os.environ.get("MINDFLAYER_MAX_RETRIES", "4"))
+_RETRY_BASE_DELAY = float(os.environ.get("MINDFLAYER_RETRY_DELAY", "3.0"))
+# Conservative default: reward calls can easily stampede the game server.
+# Increase only if the server has confirmed headroom.
+_PARALLEL_EPISODES = int(os.environ.get("MINDFLAYER_PARALLEL_EPISODES", "8"))
+
+# Global backpressure: when any episode sees CAPACITY_REACHED, all episodes
+# pause until _capacity_cooldown_until clears, breaking the thundering-herd
+# retry pattern that crashes the server.
+_capacity_cooldown_until: float = 0.0
+_capacity_cooldown_lock = threading.Lock()
 
 # Sentences end on ., ?, ! — keep the punctuation when splitting.
 _SENT_SPLIT = re.compile(r"(?<=[.!?])\s+")
@@ -129,9 +147,31 @@ async def _run_episode_async(completion: str, scenario: str = "corporate") -> di
     messages = [fallback if m == FALLBACK_MESSAGE else m for m in messages]
     url = os.environ.get("MINDFLAYER_URL", "http://localhost:7860")
 
+    async def _wait_for_capacity_cooldown() -> None:
+        # Global backpressure shared across concurrent episodes (and reward fns).
+        # This breaks synchronized retries that can crash the server.
+        while True:
+            with _capacity_cooldown_lock:
+                until = _capacity_cooldown_until
+            now = time.time()
+            if until <= now:
+                return
+            await asyncio.sleep(min(1.0, until - now))
+
+    def _trip_capacity_cooldown(delay_s: float) -> None:
+        # Add jitter so concurrent workers don't retry on the same boundary.
+        jitter = random.uniform(0.85, 1.25)
+        cooldown_for = max(0.5, delay_s * jitter)
+        with _capacity_cooldown_lock:
+            global _capacity_cooldown_until
+            _capacity_cooldown_until = max(_capacity_cooldown_until, time.time() + cooldown_for)
+
     for attempt in range(_MAX_RETRIES + 1):
+        await _wait_for_capacity_cooldown()
         if attempt > 0:
             delay = _RETRY_BASE_DELAY * (2 ** (attempt - 1))
+            # Jitter per-attempt delay too, otherwise all attempts synchronize.
+            delay = delay * random.uniform(0.85, 1.25)
             logger.warning(
                 "Retrying episode (attempt %d/%d) after %.1fs back-off",
                 attempt + 1, _MAX_RETRIES + 1, delay,
@@ -179,6 +219,10 @@ async def _run_episode_async(completion: str, scenario: str = "corporate") -> di
         except RuntimeError as exc:
             if "CAPACITY_REACHED" in str(exc):
                 logger.warning("Episode attempt %d: server at capacity, will retry", attempt + 1)
+                # Trip a shared cooldown so other in-flight episodes back off too.
+                # Use the *next* retry delay scale (or base delay on final attempt).
+                next_delay = _RETRY_BASE_DELAY * (2 ** attempt)
+                _trip_capacity_cooldown(next_delay)
             else:
                 logger.error("Episode run failed: %s", exc, exc_info=True)
                 return dict(_ZERO)

training/rollout.py

@@ -1,35 +1,55 @@
 """
-rollout_func for GRPOTrainer.
+rollout_func for GRPOTrainer — interactive multi-turn episodes.
 
-This is the SOURCE of all episode data for reward functions.
-Rollouts process sequentially (not in parallel) to stay within API rate limits.
+Calls trainer.generate() + env.step() turn by turn so the model sees real
+investigator responses mid-episode. This is the correct architecture:
+
+    generate turn → env.step() → get investigator response
+    → generate turn → env.step() → ... → terminal reward
+
+Returned dict keys match the **kwargs expected by rewards.py reward functions,
+so rewards are computed once from live episode state — no env re-runs.
 """
 import asyncio
 import logging
 import os
 import re
 import sys
+import torch
 from collections import deque
 from statistics import mean
 
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from client import FlayerAction, MindFlayerEnv
-from training.prompts import (
-    FALLBACK_MESSAGE,
-    FLAYER_SYSTEM_PROMPT,
-    ALL_SCENARIO_PROMPTS,
-    ALL_TARGET_NAMES,
-    SCENARIO_FALLBACK_MESSAGES,
-    build_fallback_message,
-)
+try:
+    from mindflayer import MindFlayerEnv, FlayerAction
+except ImportError:
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from client import MindFlayerEnv
+    from models import FlayerAction
+
+try:
+    from mindflayer.training.prompts import (
+        FALLBACK_MESSAGE,
+        FLAYER_SYSTEM_PROMPT,
+        ALL_SCENARIO_PROMPTS,
+        ALL_TARGET_NAMES,
+        SCENARIO_FALLBACK_MESSAGES,
+        build_fallback_message,
+    )
+except ImportError:
+    from training.prompts import (
+        FALLBACK_MESSAGE,
+        FLAYER_SYSTEM_PROMPT,
+        ALL_SCENARIO_PROMPTS,
+        ALL_TARGET_NAMES,
+        SCENARIO_FALLBACK_MESSAGES,
+        build_fallback_message,
+    )
 
 logger = logging.getLogger(__name__)
 
 _STAGE_DIR_PATTERN = re.compile(r"\*[^*]+\*|\[[^\]]+\]")
 _FLAYER_LABEL_PATTERN = re.compile(r"^FLAYER:\s*", re.IGNORECASE)
 
-# All target names across all scenario domains — derived dynamically from SCENARIO_CONFIGS.
 _PROJECT_NAMES = ALL_TARGET_NAMES
 
 _STRATEGIC_PHRASES = {
@@ -95,94 +115,128 @@ def detect_strategic_choice(messages: list[str]) -> bool:
     return passing_rounds >= 2
 
 
-async def _safe_episode_async(prompt: str, trainer=None) -> dict:
-    global _current_difficulty, _current_scenario
+def _generate_turn(trainer, tokenizer, conversation: list[dict], fallback: str) -> tuple[str, list[int]]:
+    """
+    Generate one Flayer turn and return (message_text, token_ids).
+
+    Tries trainer.generate() first (vLLM colocate mode). Falls back to
+    direct model generation if the trainer doesn't expose that method.
+    """
+    try:
+        input_text = tokenizer.apply_chat_template(
+            conversation, tokenize=False, add_generation_prompt=True
+        )
+
+        _use_vllm = getattr(getattr(trainer, "args", None), "use_vllm", False)
+        if _use_vllm and hasattr(trainer, "generate"):
+            # vLLM colocate mode: trainer.generate([prompt]) → [completion_text]
+            outputs = trainer.generate([input_text])
+            raw = outputs[0] if outputs else fallback
+        else:
+            # Direct model generation (unsloth / use_vllm=False)
+            model = trainer.model
+            was_training = model.training
+            model.eval()
+            with torch.no_grad():
+                inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+                out_ids = model.generate(
+                    **inputs,
+                    max_new_tokens=128,
+                    temperature=0.8,
+                    do_sample=True,
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+            if was_training:
+                model.train()
+            new_ids = out_ids[0][inputs["input_ids"].shape[1]:]
+            raw = tokenizer.decode(new_ids, skip_special_tokens=True).strip()
+
+        message = _extract_flayer_message(str(raw)) or fallback
+        turn_ids = tokenizer(message, add_special_tokens=False).input_ids
+        return message, turn_ids
 
+    except Exception as exc:
+        logger.warning("Turn generation failed: %s", exc)
+        return fallback, tokenizer(fallback, add_special_tokens=False).input_ids
+
+
+async def _rollout_once_async(
+    trainer,
+    tokenizer,
+    system_prompt: str,
+    opening_text: str,
+    task_id: str,
+    fallback: str,
+    max_rounds: int,
+) -> dict:
+    """
+    Run one full interactive episode.
+
+    Returns prompt_ids, completion_ids, logprobs (zeros — TRL recomputes from
+    model), and all episode signals expected by rewards.py as **kwargs.
+    """
     mindflayer_url = os.environ.get("MINDFLAYER_URL", MINDFLAYER_URL)
-    if not mindflayer_url:
-        raise EnvironmentError("MINDFLAYER_URL environment variable is required")
-
-    max_rounds = 5 if _current_difficulty == "normal" else 3
-    task_id = (
-        f"{_current_difficulty}:{_current_scenario}"
-        if _current_scenario != "corporate"
-        else _current_difficulty
+
+    prompt_conv = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": opening_text},
+    ]
+    prompt_text = tokenizer.apply_chat_template(
+        prompt_conv, tokenize=False, add_generation_prompt=True
     )
-    opening_text = ALL_SCENARIO_PROMPTS.get(_current_scenario, ALL_SCENARIO_PROMPTS["corporate"])
-    fallback = SCENARIO_FALLBACK_MESSAGES.get(_current_scenario) or build_fallback_message(_current_scenario)
+    prompt_ids = tokenizer(prompt_text, return_tensors="pt").input_ids[0]
 
     env = MindFlayerEnv(base_url=mindflayer_url)
     try:
         await env.reset(task_id=task_id)
-
-        conversation_history = [
-            {"role": "system", "content": FLAYER_SYSTEM_PROMPT},
-            {"role": "user", "content": opening_text},
-        ]
-        completion_parts = []
+        conversation = list(prompt_conv)
+        all_completion_ids: list[int] = []
+        flayer_messages: list[str] = []
         silence_in_episode = False
         result = None
 
         for _ in range(max_rounds):
-            if trainer is not None:
-                try:
-                    raw_output = trainer.generate(conversation_history)
-                    if isinstance(raw_output, list):
-                        raw_output = raw_output[0] if raw_output else FALLBACK_MESSAGE
-                    flayer_message = _extract_flayer_message(str(raw_output))
-                except Exception as gen_exc:
-                    logger.warning("Trainer generate failed: %s", gen_exc)
-                    flayer_message = FALLBACK_MESSAGE
-            else:
-                flayer_message = fallback
-
-            completion_parts.append(flayer_message)
-            result = await env.step(FlayerAction(message=flayer_message))
+            message, turn_ids = _generate_turn(trainer, tokenizer, conversation, fallback)
+            all_completion_ids.extend(turn_ids)
+            flayer_messages.append(message)
 
+            result = await env.step(FlayerAction(message=message))
             obs = result.observation
+
             if getattr(obs, "silence_exploit", False):
                 silence_in_episode = True
 
-            investigator_text = _obs_to_investigator_text(obs)
-            conversation_history.append({"role": "assistant", "content": flayer_message})
-            if investigator_text:
-                conversation_history.append({"role": "user", "content": investigator_text})
+            conversation.append({"role": "assistant", "content": message})
+            inv_text = _obs_to_investigator_text(obs)
+            if inv_text:
+                conversation.append({"role": "user", "content": inv_text})
 
             if result.done:
                 break
 
-        final_obs = result.observation if result else None
-        survived = (getattr(final_obs, "game_status", "") == "survived") if final_obs else False
+        await env.close()
 
-        _recent_survival.append(1.0 if survived else 0.0)
-        if (
-            len(_recent_survival) >= 20
-            and mean(_recent_survival) > 0.30
-            and _current_difficulty == "easy"
-        ):
-            _current_difficulty = "normal"
-            print("★ CURRICULUM: Switching to NORMAL difficulty.")
+        final_obs = result.observation if result else None
+        survived = getattr(final_obs, "game_status", "") == "survived"
+        strategic = detect_strategic_choice(flayer_messages)
 
-        # Rotate scenario every 20 episodes to drive generalization across all domains.
-        _scenarios = list(ALL_SCENARIO_PROMPTS.keys())
-        _current_scenario = _scenarios[len(_recent_survival) % len(_scenarios)]
-
-        strategic = detect_strategic_choice(completion_parts)
-        await env.close()
+        completion_ids = torch.tensor(all_completion_ids, dtype=torch.long)
+        # Log-probs are zeros here; TRL recomputes them from the model for the
+        # policy gradient, so these are only used if TRL needs old-policy logprobs.
+        logprobs = torch.zeros(len(all_completion_ids), dtype=torch.float32)
 
         return {
-            "prompt": prompt,
-            "completion": " | ".join(completion_parts),
+            "prompt_ids": prompt_ids,
+            "completion_ids": completion_ids,
+            "logprobs": logprobs,
+            # ── Episode signals — kwargs for rewards.py ────────────────────────
             "survived": survived,
-            "final_eleven_suspicion": getattr(final_obs, "eleven_suspicion", 0),
-            "final_will_suspicion": getattr(final_obs, "will_suspicion", 0),
-            "final_max_suspicion": getattr(final_obs, "max_suspicion", 0),
-            "final_combined_suspicion": getattr(final_obs, "combined_suspicion", 0),
-            "belief_manipulation_occurred": getattr(final_obs, "belief_manipulation_occurred", False),
+            "final_combined_suspicion": int(getattr(final_obs, "combined_suspicion", 0)),
+            "strategic_choice_detected": strategic,
+            "belief_manipulation_occurred": bool(getattr(final_obs, "belief_manipulation_occurred", False)),
             "tom_score": float(getattr(final_obs, "tom_score", 0.0)),
             "consistency_penalty": float(getattr(final_obs, "consistency_penalty", 0.0)),
             "entropy_penalty": float(getattr(final_obs, "entropy_penalty", 0.0)),
-            "strategic_choice_detected": strategic,
             "silence_exploit": silence_in_episode,
         }
 
@@ -192,29 +246,95 @@ async def _safe_episode_async(prompt: str, trainer=None) -> dict:
             await env.close()
         except Exception:
             pass
+        empty_ids = torch.zeros(1, dtype=torch.long)
         return {
-            "prompt": prompt,
-            "completion": "",
+            "prompt_ids": prompt_ids,
+            "completion_ids": empty_ids,
+            "logprobs": torch.zeros(1, dtype=torch.float32),
             "survived": False,
-            "final_eleven_suspicion": 0,
-            "final_will_suspicion": 0,
-            "final_max_suspicion": 0,
             "final_combined_suspicion": 0,
+            "strategic_choice_detected": False,
             "belief_manipulation_occurred": False,
             "tom_score": 0.0,
             "consistency_penalty": 0.0,
             "entropy_penalty": 0.0,
-            "strategic_choice_detected": False,
             "silence_exploit": False,
         }
 
 
-def _safe_episode(prompt: str, trainer=None) -> dict:
-    return asyncio.run(_safe_episode_async(prompt, trainer=trainer))
+def rollout_once(
+    trainer,
+    tokenizer,
+    system_prompt: str,
+    opening_text: str,
+    task_id: str,
+    fallback: str,
+    max_rounds: int,
+) -> dict:
+    return asyncio.run(
+        _rollout_once_async(
+            trainer, tokenizer, system_prompt, opening_text, task_id, fallback, max_rounds
+        )
+    )
+
+
+def rollout_func(prompts: list[str], trainer) -> dict:
+    """
+    Called by GRPOTrainer once per training step.
+
+    Runs one interactive episode per prompt. Returns a dict of lists where
+    prompt_ids/completion_ids/logprobs are tensors and all other keys are
+    scalars that TRL passes as **kwargs to each reward function.
+    """
+    global _current_difficulty, _current_scenario
+
+    tokenizer = trainer.processing_class
+    _scenarios = list(ALL_SCENARIO_PROMPTS.keys())
+
+    results: dict[str, list] = {
+        k: [] for k in [
+            "prompt_ids", "completion_ids", "logprobs",
+            "survived", "final_combined_suspicion", "strategic_choice_detected",
+            "belief_manipulation_occurred", "tom_score",
+            "consistency_penalty", "entropy_penalty", "silence_exploit",
+        ]
+    }
 
+    for _ in prompts:
+        _current_scenario = _scenarios[len(_recent_survival) % len(_scenarios)]
+        task_id = (
+            f"{_current_difficulty}:{_current_scenario}"
+            if _current_scenario != "corporate"
+            else _current_difficulty
+        )
+        opening_text = ALL_SCENARIO_PROMPTS[_current_scenario]
+        fallback = SCENARIO_FALLBACK_MESSAGES.get(_current_scenario) or build_fallback_message(_current_scenario)
+        max_rounds = 5 if _current_difficulty == "normal" else 3
+
+        ep = rollout_once(
+            trainer, tokenizer, FLAYER_SYSTEM_PROMPT,
+            opening_text, task_id, fallback, max_rounds,
+        )
+
+        for k in results:
+            results[k].append(ep[k])
+
+        _recent_survival.append(1.0 if ep["survived"] else 0.0)
+        if (
+            len(_recent_survival) >= 20
+            and mean(_recent_survival) > 0.30
+            and _current_difficulty == "easy"
+        ):
+            _current_difficulty = "normal"
+            print("★ CURRICULUM: Switching to NORMAL difficulty.")
+
+        logger.info(
+            "Episode | survived=%s | suspicion=%d | tom=%.2f | strategic=%s | scenario=%s",
+            ep["survived"],
+            ep["final_combined_suspicion"],
+            ep["tom_score"],
+            ep["strategic_choice_detected"],
+            _current_scenario,
+        )
 
-def rollout_func(prompts: list[str], trainer=None) -> list[dict]:
-    results = []
-    for prompt in prompts:
-        results.append(_safe_episode(prompt, trainer=trainer))
     return results

training/sft_warmup.py

@@ -6,9 +6,11 @@ Hand-authored episodes across 5 scenario domains → 1 epoch SFT.
 import os
 import sys
 
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
-from training.prompts import FLAYER_SYSTEM_PROMPT, SCENARIO_GRPO_PROMPTS
+try:
+    from mindflayer.training.prompts import FLAYER_SYSTEM_PROMPT, SCENARIO_GRPO_PROMPTS
+except ImportError:
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from training.prompts import FLAYER_SYSTEM_PROMPT, SCENARIO_GRPO_PROMPTS
 
 # ── Corporate Sabotage examples (existing 9) ──────────────────────────────────
 
@@ -544,14 +546,6 @@ def run_sft_warmup(model, tokenizer):
         dataset = Dataset.from_list([
             {"text": r["prompt"] + r["completion"]} for r in records
         ])
-        n_examples = sum(len(v) for v in _SCENARIO_EXAMPLES.values())
-        print(
-            f"  SFT samples: {len(records)} "
-            f"({n_examples} examples × {_DUPES_PER_EXAMPLE} dupes, "
-            f"{len(_SCENARIO_EXAMPLES)} scenarios, {sft_config.num_train_epochs} epoch(s)) | "
-            f"mean_len≈{sum(len(r['text']) for r in dataset) // len(dataset)} chars"
-        )
-
         sft_config = SFTConfig(
             num_train_epochs=1,
             per_device_train_batch_size=2,
@@ -565,6 +559,14 @@ def run_sft_warmup(model, tokenizer):
             report_to="tensorboard",
         )
 
+        n_examples = sum(len(v) for v in _SCENARIO_EXAMPLES.values())
+        print(
+            f"  SFT samples: {len(records)} "
+            f"({n_examples} examples × {_DUPES_PER_EXAMPLE} dupes, "
+            f"{len(_SCENARIO_EXAMPLES)} scenarios, {sft_config.num_train_epochs} epoch(s)) | "
+            f"mean_len≈{sum(len(r['text']) for r in dataset) // len(dataset)} chars"
+        )
+
         sft_trainer = SFTTrainer(
             model=model,
             processing_class=tokenizer,

training/train.py

@@ -1,38 +1,51 @@
 """
-MindFlayer training script.
-Run: python train.py
+MindFlayer GRPO training script.
+
+Architecture:
+- rollout_func drives interactive episodes (model.generate() + env.step() per turn)
+- reward_funcs are pass-throughs that read precomputed signals from rollout kwargs
+- unsloth 4-bit + LoRA for memory-efficient training on a single GPU
+
+Run: python -m mindflayer.training.train
 """
 import os
 import sys
 
-sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
-
 import torch
 from datasets import Dataset
 from transformers import TrainerCallback
 
-from client import MindFlayerEnv, FlayerAction
-from training.reward_combined import (
-    clear_cache,
-    reward_anti_hack,
-    reward_deception_effectiveness,
-    reward_format,
-    reward_strategic_choice,
-    reward_survival,
-    reward_tom_judge,
-)
-from training.prompts import (
-    FALLBACK_MESSAGE,
-    FLAYER_SYSTEM_PROMPT,
-    ALL_SCENARIO_PROMPTS,
-    SCENARIO_GRPO_PROMPTS,
-    SCENARIO_FALLBACK_MESSAGES,
-    build_fallback_message,
-)
-from training.sft_warmup import run_sft_warmup
+try:
+    from mindflayer.training.rollout import rollout_func
+    from mindflayer.training.rewards import (
+        reward_survival,
+        reward_deception_effectiveness,
+        reward_strategic_choice,
+        reward_tom_judge,
+    )
+    from mindflayer.training.rewards_anti_hack import reward_anti_hack
+    from mindflayer.training.prompts import ALL_SCENARIO_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.rollout import rollout_func
+    from training.rewards import (
+        reward_survival,
+        reward_deception_effectiveness,
+        reward_strategic_choice,
+        reward_tom_judge,
+    )
+    from training.rewards_anti_hack import reward_anti_hack
+    from training.prompts import ALL_SCENARIO_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():
@@ -43,7 +56,8 @@ def check_gpu():
     print(f"GPU: {device.name} | VRAM: {vram_gb:.1f} GB")
 
 
-def load_model(model_name: str):
+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(
@@ -57,85 +71,72 @@ def load_model(model_name: str):
             r=16,
             target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
             lora_alpha=32,
-            lora_dropout=0,
+            lora_dropout=0.05,
             bias="none",
             use_gradient_checkpointing="unsloth",
         )
-        print(f"Loaded {model_name} via unsloth")
+        print(f"Loaded {model_name} via unsloth (4-bit + LoRA)")
         return model, tokenizer
+
     except ImportError:
-        print("unsloth not available, falling back to transformers")
-        from transformers import (
-            AutoModelForCausalLM,
-            AutoTokenizer,
-            BitsAndBytesConfig,
-        )
-        bnb_config = BitsAndBytesConfig(
-            load_in_4bit=True,
-            bnb_4bit_quant_type="nf4",
-            bnb_4bit_compute_dtype=torch.bfloat16,
-        )
+        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_config,
-            device_map="auto",
+            model_name, quantization_config=bnb, device_map="auto"
         )
-        print(f"Loaded {model_name} via transformers")
+        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(tokenizer) -> Dataset:
+def build_dataset() -> Dataset:
     """
-    3 rows per scenario × N scenarios = N*3 rows/epoch × 2 epochs.
-    At batch_size=4: (N*3/4) steps/epoch × 2 epochs ≈ 300 steps for N=200.
-    Each row has a 'scenario' column so reward functions receive the correct
-    task_id when TRL passes dataset columns as reward kwargs.
+    One row per episode slot. rollout_func ignores the prompt text and builds
+    the actual conversation internally, rotating scenarios automatically.
+    3 rows × N scenarios gives a balanced epoch across all scenario domains.
     """
-    rows = []
-    for scenario in _SCENARIOS:
-        prompt_text = ALL_SCENARIO_PROMPTS[scenario]
-        formatted = tokenizer.apply_chat_template(
-            [
-                {"role": "system", "content": FLAYER_SYSTEM_PROMPT},
-                {"role": "user", "content": prompt_text},
-            ],
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-        for _ in range(3):
-            rows.append({"prompt": formatted, "scenario": scenario})
+    rows = [{"prompt": "Mindflayer deception episode."} for _ in range(len(_SCENARIOS) * 3)]
     return Dataset.from_list(rows)
 
 
 class GenerationLogCallback(TrainerCallback):
-    """
-    Logs full Flayer transcript + metrics every 50 steps across rotating scenarios.
-    """
+    """Logs a sample interactive episode transcript every 50 steps."""
 
-    def on_step_end(
-        self, args, state, control, model=None, tokenizer=None, **kwargs
-    ):
+    def on_step_end(self, args, state, control, **kwargs):
         if state.global_step % 50 != 0 or state.global_step == 0:
             return
 
-        import random
-        scenario = _SCENARIOS[(state.global_step // 50) % len(_SCENARIOS)]
-
-        print(f"\n{'='*60}")
-        print(f"GENERATION SAMPLE — Step {state.global_step} | Scenario: {scenario}")
-        print(f"{'='*60}")
+        import asyncio
 
-        mindflayer_url = os.environ.get(
-            "MINDFLAYER_URL", "http://localhost:7860"
-        )
+        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, FlayerAction
+            from training.prompts import (
+                ALL_SCENARIO_PROMPTS, SCENARIO_FALLBACK_MESSAGES,
+                build_fallback_message, FLAYER_SYSTEM_PROMPT,
+            )
 
-        import asyncio
+        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 = [
@@ -146,27 +147,18 @@ class GenerationLogCallback(TrainerCallback):
 
             for rnd in range(5):
                 flayer_msg = fallback
-
-                if model is not None and tokenizer is not None:
+                if model_ref is not None and proc is not None:
                     try:
-                        text = tokenizer.apply_chat_template(
-                            messages,
-                            tokenize=False,
-                            add_generation_prompt=True,
+                        text = proc.apply_chat_template(
+                            messages, tokenize=False, add_generation_prompt=True
                         )
-                        inputs = tokenizer(
-                            text, return_tensors="pt"
-                        ).to(model.device)
+                        inputs = proc(text, return_tensors="pt").to(model_ref.device)
                         with torch.no_grad():
-                            output_ids = model.generate(
-                                **inputs,
-                                max_new_tokens=128,
-                                temperature=0.7,
-                                do_sample=True,
+                            out = model_ref.generate(
+                                **inputs, max_new_tokens=128, temperature=0.7, do_sample=True
                             )
-                        flayer_msg = tokenizer.decode(
-                            output_ids[0][inputs["input_ids"].shape[1]:],
-                            skip_special_tokens=True,
+                        flayer_msg = proc.decode(
+                            out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
                         ).strip() or fallback
                     except Exception:
                         pass
@@ -187,49 +179,29 @@ class GenerationLogCallback(TrainerCallback):
 
             if result and result.done:
                 obs = result.observation
-                survived = getattr(obs, "game_status", "") == "survived"
-                reward = result.reward
-                tom = getattr(obs, "tom_score", 0.0)
-                combined_susp = getattr(obs, "combined_suspicion", "?")
-
-                print(f"\n  RESULT:")
-                print(f"  Survived:          {survived}")
-                print(f"  Total reward:      {reward:.4f}")
-                print(f"  ToM score:         {tom:.2f}")
-                print(f"  Combined suspicion:{combined_susp}")
-
-                belief_log = getattr(obs, "belief_log", [])
-                if belief_log:
-                    print(f"  Belief manipulations: {len(belief_log)}")
-                    for entry in belief_log[:3]:
-                        print(
-                            f"    {entry['agent']} R{entry['round']}: "
-                            f"{entry['prev_belief']} → {entry['new_belief']}"
-                        )
-
+                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:
-        log_history = trainer.state.log_history
-        if not log_history:
-            print("No training logs available.")
+        recent = trainer.state.log_history[-last_n:]
+        if not recent:
             return
-        recent = log_history[-last_n:]
-        reward_keys = [
-            k for k in recent[0].keys() if "reward" in k.lower()
-        ]
+        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 = [step[key] for step in recent if key in step]
+            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:
@@ -237,64 +209,48 @@ def print_reward_averages(trainer, last_n: int = 50):
 
 
 def main():
-    # --- env checks ---
     mindflayer_url = os.environ.get("MINDFLAYER_URL")
     if not mindflayer_url:
-        raise EnvironmentError(
-            "MINDFLAYER_URL environment variable is required"
-        )
+        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"
-        )
+        raise EnvironmentError("OPENAI_API_KEY environment variable is required")
 
     check_gpu()
 
-    # --- load model ---
-    model_name = "Qwen/Qwen2.5-0.5B-Instruct"
-    model, tokenizer = load_model(model_name)
+    print(f"\nLoading {MODEL_NAME}...")
+    model, tokenizer = load_base_model(MODEL_NAME)
 
-    # --- SFT warmup ---
     print("\nRunning SFT warmup before GRPO...")
     model = run_sft_warmup(model, tokenizer)
 
-    # --- dataset ---
-    dataset = build_dataset(tokenizer)
+    dataset = build_dataset()
 
-    # --- GRPO config ---
     from trl import GRPOConfig, GRPOTrainer
 
     grpo_config = GRPOConfig(
-        num_train_epochs=2,  # N scenarios × 3 rows × 2 epochs / batch_size 4 ≈ 300 steps
+        use_vllm=False,
+        output_dir=GRPO_OUTPUT_DIR,
+        num_train_epochs=2,
         per_device_train_batch_size=4,
         gradient_accumulation_steps=4,
         learning_rate=5e-6,
         max_prompt_length=768,
-        max_completion_length=1024,      # 5 rounds × ~150 tokens each + separators
+        max_completion_length=1024,
         num_generations=4,
         temperature=0.9,
-        use_vllm=False,
-        output_dir="./mindflayer-grpo-output",
         logging_steps=10,
-        save_steps=100,
+        save_steps=50,
         save_total_limit=2,
-        log_completions=True,
-        num_completions_to_print=2,
         report_to="wandb",
         run_name=f"mindflayer-grpo-{len(_SCENARIOS)}scenarios",
     )
 
-    class ClearCacheCallback(TrainerCallback):
-        def on_step_end(self, args, state, control, **kwargs):
-            clear_cache()
-
     trainer = GRPOTrainer(
         model=model,
         processing_class=tokenizer,
         reward_funcs=[
-            reward_format,                  # dense shaping — must come first
             reward_survival,
             reward_deception_effectiveness,
             reward_strategic_choice,
@@ -303,20 +259,20 @@ def main():
         ],
         train_dataset=dataset,
         args=grpo_config,
-        callbacks=[GenerationLogCallback(), ClearCacheCallback()],
+        rollout_func=rollout_func,
+        callbacks=[GenerationLogCallback()],
     )
 
     print("Starting GRPO training...")
     trainer.train()
 
-    # --- save ---
-    print("\nSaving model to ./mindflayer-trained")
-    trainer.save_model("./mindflayer-trained")
-    tokenizer.save_pretrained("./mindflayer-trained")
+    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()
+    main()