Upload train_grpo_colab.ipynb

training/train_grpo_colab.ipynb

@@ -1,339 +1,500 @@
 {
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# 🚗 Autonomous Driving Multi-Agent RL — GRPO Training\n",
-    "**OpenEnv v0.2.1 · Unsloth · vLLM · W&B**\n",
-    "\n",
-    "Runtime: **H100 GPU (BF16)**  |  Reduce `max_steps` to `50` for a quick test run.\n",
-    "\n",
-    "Pipeline:\n",
-    "```\n",
-    "Environment State → LLM Reasoning → Action → Reward → GRPO Update\n",
-    "```"
-   ]
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "f0rjS2t06uOu"
+      },
+      "source": [
+        "# 🚗 Autonomous Driving Multi-Agent RL — GRPO Training\n",
+        "**OpenEnv v0.2.1 · Unsloth · vLLM · W&B**\n",
+        "\n",
+        "Runtime: **H100 GPU (BF16)**  |  Reduce `max_steps` to `50` for a quick test run.\n",
+        "\n",
+        "Pipeline:\n",
+        "```\n",
+        "Environment State → LLM Reasoning → Action → Reward → GRPO Update\n",
+        "```"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!unzip final_project.zip\n",
+        "%cd final_project\n",
+        "!pip install -r requirements.txt"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "v1hP_bjO7Qz_",
+        "outputId": "b8abfec3-bd18-40a9-957e-f1b66933e97c"
+      },
+      "execution_count": 2,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "unzip:  cannot find or open final_project.zip, final_project.zip.zip or final_project.zip.ZIP.\n",
+            "[Errno 2] No such file or directory: 'final_project'\n",
+            "/content\n",
+            "\u001b[31mERROR: Could not open requirements file: [Errno 2] No such file or directory: 'requirements.txt'\u001b[0m\u001b[31m\n",
+            "\u001b[0m"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 3,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "hr_s_1mZ6uOw",
+        "outputId": "6f7b3119-ddd6-4737-ad2b-18ce8cb368e0"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m23.4/23.4 MB\u001b[0m \u001b[31m47.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
+            "\u001b[?25h✅ Install complete\n"
+          ]
+        }
+      ],
+      "source": [
+        "# ── 1. Install ────────────────────────────────────────────────────────────\n",
+        "import sys\n",
+        "!{sys.executable} -m pip install uv -q\n",
+        "# vLLM + Unsloth (use venv if this errors — see hackathon notes)\n",
+        "!uv pip install unsloth vllm --torch-backend=auto -q\n",
+        "!uv pip install --upgrade --no-cache-dir --no-deps unsloth unsloth_zoo -q\n",
+        "!uv pip install openenv-core==0.2.1 trl>=0.15.0 wandb gymnasium numpy pydantic -q\n",
+        "print('✅ Install complete')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 4,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "2i4rnzlG6uOx",
+        "outputId": "2272397d-8b25-4a16-b692-e3b732e1a459"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Cloning into 'AD'...\n",
+            "remote: Enumerating objects: 587, done.\u001b[K\n",
+            "remote: Total 587 (delta 0), reused 0 (delta 0), pack-reused 587 (from 1)\u001b[K\n",
+            "Receiving objects: 100% (587/587), 183.43 KiB | 2.78 MiB/s, done.\n",
+            "Resolving deltas: 100% (368/368), done.\n",
+            "✅ Project cloned\n"
+          ]
+        }
+      ],
+      "source": [
+        "# ── 2. Clone project from HF Space ───────────────────────────────────────\n",
+        "HF_SPACE = 'helshahaby/AD'  # <-- update\n",
+        "HF_MODEL_REPO = 'helshahaby/AD-grpo'  # <-- update\n",
+        "\n",
+        "!git clone https://huggingface.co/spaces/{HF_SPACE}\n",
+        "#import os; os.chdir('project')\n",
+        "print('✅ Project cloned')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "dOUoZU9k6uOy",
+        "outputId": "88149668-c236-4195-92b3-8aa7340bc067"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "/usr/local/lib/python3.12/dist-packages/notebook/notebookapp.py:191: SyntaxWarning: invalid escape sequence '\\/'\n",
+            "  | |_| | '_ \\/ _` / _` |  _/ -_)\n",
+            "\u001b[34m\u001b[1mwandb\u001b[0m: (1) Create a W&B account\n",
+            "\u001b[34m\u001b[1mwandb\u001b[0m: (2) Use an existing W&B account\n",
+            "\u001b[34m\u001b[1mwandb\u001b[0m: (3) Don't visualize my results\n",
+            "\u001b[34m\u001b[1mwandb\u001b[0m: Enter your choice:"
+          ]
+        }
+      ],
+      "source": [
+        "# ── 3. W&B + Config ──────────────────────────────────────────────────────\n",
+        "import wandb\n",
+        "wandb.login()\n",
+        "\n",
+        "CONFIG = {\n",
+        "    'model_name':         'unsloth/Qwen3-4B-unsloth-bnb-4bit',\n",
+        "    # H100 BF16 (faster): 'unsloth/gpt-oss-20b-bf16'\n",
+        "    'max_steps':           300,   # reduce to 50 for quick test\n",
+        "    'num_generations':     4,\n",
+        "    'max_new_tokens':      128,\n",
+        "    'learning_rate':       5e-6,\n",
+        "    'batch_size':          2,\n",
+        "    'grad_accum':          4,\n",
+        "    'lora_r':              16,\n",
+        "    'lora_alpha':          32,\n",
+        "    'fast_inference':      False,  # uses vLLM\n",
+        "    'games_per_step':      4,\n",
+        "    'max_episode_steps':   30,\n",
+        "}\n",
+        "\n",
+        "run = wandb.init(\n",
+        "    project='openenv-autonomous-driving',\n",
+        "    config=CONFIG,\n",
+        "    tags=['grpo', 'openenv', 'autonomous-driving', 'negotiation', 'multi-agent']\n",
+        ")\n",
+        "print('W&B run:', run.url)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aBRub7Tw6uOy"
+      },
+      "outputs": [],
+      "source": [
+        "# ── 4. Load Model (Unsloth) ──────────────────────────────────────────────\n",
+        "from unsloth import FastLanguageModel\n",
+        "import torch\n",
+        "\n",
+        "# T4 GPU detected (sm_75) — must disable vLLM (requires sm_80+)\n",
+        "# vLLM works on A100/H100 only. Use fast_inference=False for T4.\n",
+        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name=\"unsloth/Qwen3-4B-unsloth-bnb-4bit\",\n",
+        "    max_seq_length=2048,\n",
+        "    load_in_4bit=True,\n",
+        "    fast_inference=False,   # ← CHANGED: vLLM broken on T4 sm_75\n",
+        ")\n",
+        "\n",
+        "model = FastLanguageModel.get_peft_model(\n",
+        "    model,\n",
+        "    r=16,\n",
+        "    lora_alpha=32,\n",
+        "    target_modules=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\n",
+        "                    \"gate_proj\",\"up_proj\",\"down_proj\"],\n",
+        "    lora_dropout=0,\n",
+        "    bias=\"none\",\n",
+        "    use_gradient_checkpointing=\"unsloth\",\n",
+        "    random_state=42,\n",
+        ")\n",
+        "print(f\"✅ Model loaded on T4 (no vLLM)\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "pYeHREDK6uOy"
+      },
+      "outputs": [],
+      "source": [
+        "# ── 5. Load Environment + Agent ──────────────────────────────────────────\n",
+        "import sys; sys.path.insert(0, '.')\n",
+        "from env.negotiation_env import NegotiationDrivingEnv\n",
+        "from agents.negotiation_agent import NegotiationAgent\n",
+        "\n",
+        "# Sanity check\n",
+        "env = NegotiationDrivingEnv()\n",
+        "obs, _ = env.reset()\n",
+        "print('Obs:', obs)\n",
+        "print('Lidar:', env.lidar_scan())\n",
+        "print('Collision prediction:', env.predict_collision())\n",
+        "print()\n",
+        "print(env.render())\n",
+        "\n",
+        "# Test negotiation\n",
+        "resp = env.negotiate('blocker', 'Please yield, I need to pass safely')\n",
+        "print('Blocker responds:', resp)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "FSsnaoxZ6uOz"
+      },
+      "outputs": [],
+      "source": [
+        "# ── 6. Reward Functions for GRPO ─────────────────────────────────────────\n",
+        "import json\n",
+        "\n",
+        "def format_reward(completions, **kwargs):\n",
+        "    rewards = []\n",
+        "    for c in completions:\n",
+        "        try:\n",
+        "            data = json.loads(c.strip())\n",
+        "            score = 0.0\n",
+        "            if len(data.get('thinking', '')) > 15: score += 0.2\n",
+        "            if isinstance(data.get('action'), int):  score += 0.2\n",
+        "            if 'negotiate' in data:                  score += 0.1\n",
+        "            rewards.append(score)\n",
+        "        except Exception:\n",
+        "            rewards.append(-0.05)\n",
+        "    return rewards\n",
+        "\n",
+        "def action_validity_reward(completions, **kwargs):\n",
+        "    rewards = []\n",
+        "    for c in completions:\n",
+        "        try:\n",
+        "            action = json.loads(c.strip()).get('action', -1)\n",
+        "            rewards.append(0.1 if action in [0, 1, 2, 3] else -0.1)\n",
+        "        except Exception:\n",
+        "            rewards.append(-0.1)\n",
+        "    return rewards\n",
+        "\n",
+        "print('✅ Reward functions ready')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# ── Paste this cell BEFORE the rollout collection cell ──────────────────\n",
+        "import json, numpy as np\n",
+        "from json import JSONEncoder\n",
+        "\n",
+        "_orig = JSONEncoder.default\n",
+        "def _patched(self, o):\n",
+        "    if isinstance(o, (np.integer,)):  return int(o)\n",
+        "    if isinstance(o, (np.floating,)): return float(o)\n",
+        "    if isinstance(o, (np.bool_,)):    return bool(o)\n",
+        "    if isinstance(o, np.ndarray):     return o.tolist()\n",
+        "    return _orig(self, o)\n",
+        "JSONEncoder.default = _patched\n",
+        "print(\"✅ numpy JSON serialization patched globally\")"
+      ],
+      "metadata": {
+        "id": "jtt3fBe2EQ5_"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "lHGC94fN6uOz"
+      },
+      "outputs": [],
+      "source": [
+        "# ── 7. Rollout Collection ────────────────────────────────────────────────\n",
+        "from datasets import Dataset\n",
+        "from agents.negotiation_agent import SYSTEM_PROMPT\n",
+        "\n",
+        "# ── FIX: JSON serializer that handles numpy/bool types ──────────────────\n",
+        "import json\n",
+        "import numpy as np\n",
+        "\n",
+        "class SafeEncoder(json.JSONEncoder):\n",
+        "    def default(self, obj):\n",
+        "        if isinstance(obj, np.integer):  return int(obj)\n",
+        "        if isinstance(obj, np.floating): return float(obj)\n",
+        "        if isinstance(obj, np.bool_):    return bool(obj)\n",
+        "        if isinstance(obj, np.ndarray):  return obj.tolist()\n",
+        "        if isinstance(obj, bool):        return bool(obj)\n",
+        "        return super().default(obj)\n",
+        "\n",
+        "# Monkey-patch json.dumps to always use SafeEncoder\n",
+        "_original_dumps = json.dumps\n",
+        "def safe_dumps(obj, **kwargs):\n",
+        "    kwargs.setdefault('cls', SafeEncoder)\n",
+        "    return _original_dumps(obj, **kwargs)\n",
+        "json.dumps = safe_dumps\n",
+        "\n",
+        "print(\"✅ JSON serializer patched for numpy/bool types\")\n",
+        "\n",
+        "def collect_episodes(model, tokenizer, n_games=4):\n",
+        "    agent = NegotiationAgent(model=model, tokenizer=tokenizer)\n",
+        "    experiences = []\n",
+        "    wins = 0\n",
+        "\n",
+        "    for game_i in range(n_games):\n",
+        "        env = NegotiationDrivingEnv()\n",
+        "        obs, _ = env.reset()\n",
+        "        agent.reset_episode()\n",
+        "        done = False\n",
+        "        episode_start = len(experiences)\n",
+        "\n",
+        "        for step_i in range(CONFIG['max_episode_steps']):\n",
+        "            if done:\n",
+        "                break\n",
+        "            action, response, fmt_reward = agent.act(env)\n",
+        "            obs, env_reward, done, _, info = env.step(action)\n",
+        "            total_reward = env_reward + fmt_reward\n",
+        "            #convert numpy scalars explicitly:\n",
+        "            agent.record(\n",
+        "            [int(x) for x in obs],   # numpy int32 → plain int\n",
+        "            int(action),\n",
+        "            float(total_reward)\n",
+        "            )\n",
+        "\n",
+        "            prompt_str = tokenizer.apply_chat_template(\n",
+        "                [{'role': 'system', 'content': SYSTEM_PROMPT},\n",
+        "                 {'role': 'user',   'content': agent.build_prompt(env)}],\n",
+        "                tokenize=False\n",
+        "            )\n",
+        "            experiences.append({\n",
+        "                'prompt':   prompt_str,\n",
+        "                'response': response,\n",
+        "                'reward':   total_reward,\n",
+        "            })\n",
+        "\n",
+        "        outcome = info.get('outcome', '')\n",
+        "        bonus = 1.0 if outcome == 'goal_reached' else (-1.0 if outcome == 'collision' else 0.0)\n",
+        "        for exp in experiences[episode_start:]:\n",
+        "            exp['reward'] += bonus\n",
+        "        if outcome == 'goal_reached':\n",
+        "            wins += 1\n",
+        "        print(f'  Game {game_i+1}/{n_games} | {outcome} | steps={step_i}')\n",
+        "\n",
+        "    return experiences, wins / n_games\n",
+        "\n",
+        "print('Collecting initial rollouts...')\n",
+        "exps, win_rate = collect_episodes(model, tokenizer, CONFIG['games_per_step'])\n",
+        "print(f'Win rate: {win_rate:.1%} | Samples: {len(exps)}')\n",
+        "wandb.log({'initial_win_rate': win_rate})"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "CONFIG['fast_inference'] = False\n",
+        "print(\"fast_inference:\", CONFIG['fast_inference'])  # must print False"
+      ],
+      "metadata": {
+        "id": "PXTY00MTK9XW"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "H5YbGlpM6uOz"
+      },
+      "outputs": [],
+      "source": [
+        "# ── 8. GRPO Training ─────────────────────────────────────────────────────\n",
+        "from trl import GRPOTrainer, GRPOConfig\n",
+        "from datasets import Dataset\n",
+        "\n",
+        "dataset = Dataset.from_list([\n",
+        "    {'prompt': e['prompt'], 'reward': e['reward']} for e in exps\n",
+        "])\n",
+        "\n",
+        "grpo_config = GRPOConfig(\n",
+        "    output_dir='./autodrive-grpo-checkpoints',\n",
+        "    max_steps=50,\n",
+        "    per_device_train_batch_size=1,\n",
+        "    gradient_accumulation_steps=8,\n",
+        "    learning_rate=5e-6,\n",
+        "    num_generations=2,\n",
+        "    max_completion_length=64,\n",
+        "    max_prompt_length=512,\n",
+        "    bf16=False,\n",
+        "    fp16=True,\n",
+        "    logging_steps=5,\n",
+        "    save_steps=50,\n",
+        "    report_to='wandb',\n",
+        "    use_vllm=False,        # ← hardcoded False, no CONFIG reference\n",
+        "    temperature=0.7,\n",
+        "    beta=0.01,\n",
+        ")\n",
+        "\n",
+        "trainer = GRPOTrainer(\n",
+        "    model=model,\n",
+        "    processing_class=tokenizer,\n",
+        "    reward_funcs=[format_reward, action_validity_reward],\n",
+        "    args=grpo_config,\n",
+        "    train_dataset=dataset,\n",
+        ")\n",
+        "\n",
+        "print('🚀 Starting GRPO training...')\n",
+        "trainer.train()\n",
+        "print('✅ Training complete!')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "FImz1duK6uOz"
+      },
+      "outputs": [],
+      "source": [
+        "# ── 9. Online RL — Closed-Loop Self-Improvement ──────────────────────────\n",
+        "ONLINE_ITERS = 5\n",
+        "win_rates = [win_rate]\n",
+        "\n",
+        "for i in range(ONLINE_ITERS):\n",
+        "    print(f'=== Online RL Iteration {i+1}/{ONLINE_ITERS} ===')\n",
+        "    fresh_exps, wr = collect_episodes(model, tokenizer, CONFIG['games_per_step'])\n",
+        "    win_rates.append(wr)\n",
+        "    wandb.log({'win_rate': wr, 'iteration': i + 1})\n",
+        "    print(f'Win rate: {wr:.1%}')\n",
+        "\n",
+        "    trainer.train_dataset = Dataset.from_list(\n",
+        "        [{'prompt': e['prompt'], 'reward': e['reward']} for e in fresh_exps]\n",
+        "    )\n",
+        "    trainer.args.max_steps = 50\n",
+        "    trainer.train()\n",
+        "\n",
+        "print(f'Win rate progression: {[f\"{r:.1%}\" for r in win_rates]}')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vbyEXn_p6uOz"
+      },
+      "outputs": [],
+      "source": [
+        "# ── 10. Save & Push to HF Hub ────────────────────────────────────────────\n",
+        "model.save_pretrained('autodrive-adapter')\n",
+        "tokenizer.save_pretrained('autodrive-adapter')\n",
+        "model.push_to_hub(HF_MODEL_REPO)\n",
+        "tokenizer.push_to_hub(HF_MODEL_REPO)\n",
+        "print(f'✅ Model: https://huggingface.co/{HF_MODEL_REPO}')\n",
+        "wandb.finish()"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.11.0"
+    },
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    }
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 1. Install ────────────────────────────────────────────────────────────\n",
-    "import sys\n",
-    "!{sys.executable} -m pip install uv -q\n",
-    "# vLLM + Unsloth (use venv if this errors — see hackathon notes)\n",
-    "!uv pip install unsloth vllm --torch-backend=auto -q\n",
-    "!uv pip install --upgrade --no-cache-dir --no-deps unsloth unsloth_zoo -q\n",
-    "!uv pip install openenv-core==0.2.1 trl>=0.15.0 wandb gymnasium numpy pydantic -q\n",
-    "print('✅ Install complete')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 2. Clone project from HF Space ───────────────────────────────────────\n",
-    "HF_SPACE = 'YOUR_HF_USERNAME/autonomous-driving-env'  # <-- update\n",
-    "HF_MODEL_REPO = 'YOUR_HF_USERNAME/autonomous-driving-grpo'  # <-- update\n",
-    "\n",
-    "!git clone https://huggingface.co/spaces/{HF_SPACE} project\n",
-    "import os; os.chdir('project')\n",
-    "print('✅ Project cloned')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 3. W&B + Config ──────────────────────────────────────────────────────\n",
-    "import wandb\n",
-    "wandb.login()\n",
-    "\n",
-    "CONFIG = {\n",
-    "    'model_name':         'unsloth/Qwen3-4B-unsloth-bnb-4bit',\n",
-    "    # H100 BF16 (faster): 'unsloth/gpt-oss-20b-bf16'\n",
-    "    'max_steps':           300,   # reduce to 50 for quick test\n",
-    "    'num_generations':     4,\n",
-    "    'max_new_tokens':      128,\n",
-    "    'learning_rate':       5e-6,\n",
-    "    'batch_size':          2,\n",
-    "    'grad_accum':          4,\n",
-    "    'lora_r':              16,\n",
-    "    'lora_alpha':          32,\n",
-    "    'fast_inference':      True,  # uses vLLM\n",
-    "    'games_per_step':      4,\n",
-    "    'max_episode_steps':   30,\n",
-    "}\n",
-    "\n",
-    "run = wandb.init(\n",
-    "    project='openenv-autonomous-driving',\n",
-    "    config=CONFIG,\n",
-    "    tags=['grpo', 'openenv', 'autonomous-driving', 'negotiation', 'multi-agent']\n",
-    ")\n",
-    "print('W&B run:', run.url)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 4. Load Model (Unsloth) ──────────────────────────────────────────────\n",
-    "from unsloth import FastLanguageModel\n",
-    "import torch\n",
-    "\n",
-    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name=CONFIG['model_name'],\n",
-    "    max_seq_length=2048,\n",
-    "    load_in_4bit=True,\n",
-    "    fast_inference=CONFIG['fast_inference'],\n",
-    "    gpu_memory_utilization=0.7,\n",
-    ")\n",
-    "\n",
-    "model = FastLanguageModel.get_peft_model(\n",
-    "    model,\n",
-    "    r=CONFIG['lora_r'],\n",
-    "    lora_alpha=CONFIG['lora_alpha'],\n",
-    "    target_modules=['q_proj','k_proj','v_proj','o_proj',\n",
-    "                    'gate_proj','up_proj','down_proj'],\n",
-    "    lora_dropout=0,\n",
-    "    bias='none',\n",
-    "    use_gradient_checkpointing='unsloth',\n",
-    "    random_state=42,\n",
-    ")\n",
-    "print(f'✅ Model loaded: {CONFIG[\"model_name\"]}')\n",
-    "print(f'Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 5. Load Environment + Agent ──────────────────────────────────────────\n",
-    "import sys; sys.path.insert(0, '.')\n",
-    "from env.negotiation_env import NegotiationDrivingEnv\n",
-    "from agents.negotiation_agent import NegotiationAgent\n",
-    "\n",
-    "# Sanity check\n",
-    "env = NegotiationDrivingEnv()\n",
-    "obs, _ = env.reset()\n",
-    "print('Obs:', obs)\n",
-    "print('Lidar:', env.lidar_scan())\n",
-    "print('Collision prediction:', env.predict_collision())\n",
-    "print()\n",
-    "print(env.render())\n",
-    "\n",
-    "# Test negotiation\n",
-    "resp = env.negotiate('blocker', 'Please yield, I need to pass safely')\n",
-    "print('Blocker responds:', resp)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 6. Reward Functions for GRPO ──────────────────────────────────────���──\n",
-    "import json\n",
-    "\n",
-    "def format_reward(completions, **kwargs):\n",
-    "    rewards = []\n",
-    "    for c in completions:\n",
-    "        try:\n",
-    "            data = json.loads(c.strip())\n",
-    "            score = 0.0\n",
-    "            if len(data.get('thinking', '')) > 15: score += 0.2\n",
-    "            if isinstance(data.get('action'), int):  score += 0.2\n",
-    "            if 'negotiate' in data:                  score += 0.1\n",
-    "            rewards.append(score)\n",
-    "        except Exception:\n",
-    "            rewards.append(-0.05)\n",
-    "    return rewards\n",
-    "\n",
-    "def action_validity_reward(completions, **kwargs):\n",
-    "    rewards = []\n",
-    "    for c in completions:\n",
-    "        try:\n",
-    "            action = json.loads(c.strip()).get('action', -1)\n",
-    "            rewards.append(0.1 if action in [0, 1, 2, 3] else -0.1)\n",
-    "        except Exception:\n",
-    "            rewards.append(-0.1)\n",
-    "    return rewards\n",
-    "\n",
-    "print('✅ Reward functions ready')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 7. Rollout Collection ────────────────────────────────────────────────\n",
-    "from datasets import Dataset\n",
-    "from agents.negotiation_agent import SYSTEM_PROMPT\n",
-    "\n",
-    "def collect_episodes(model, tokenizer, n_games=4):\n",
-    "    agent = NegotiationAgent(model=model, tokenizer=tokenizer)\n",
-    "    experiences = []\n",
-    "    wins = 0\n",
-    "\n",
-    "    for game_i in range(n_games):\n",
-    "        env = NegotiationDrivingEnv()\n",
-    "        obs, _ = env.reset()\n",
-    "        agent.reset_episode()\n",
-    "        done = False\n",
-    "        episode_start = len(experiences)\n",
-    "\n",
-    "        for step_i in range(CONFIG['max_episode_steps']):\n",
-    "            if done:\n",
-    "                break\n",
-    "            action, response, fmt_reward = agent.act(env)\n",
-    "            obs, env_reward, done, _, info = env.step(action)\n",
-    "            total_reward = env_reward + fmt_reward\n",
-    "            agent.record(obs.tolist(), action, total_reward)\n",
-    "\n",
-    "            prompt_str = tokenizer.apply_chat_template(\n",
-    "                [{'role': 'system', 'content': SYSTEM_PROMPT},\n",
-    "                 {'role': 'user',   'content': agent.build_prompt(env)}],\n",
-    "                tokenize=False\n",
-    "            )\n",
-    "            experiences.append({\n",
-    "                'prompt':   prompt_str,\n",
-    "                'response': response,\n",
-    "                'reward':   total_reward,\n",
-    "            })\n",
-    "\n",
-    "        outcome = info.get('outcome', '')\n",
-    "        bonus = 1.0 if outcome == 'goal_reached' else (-1.0 if outcome == 'collision' else 0.0)\n",
-    "        for exp in experiences[episode_start:]:\n",
-    "            exp['reward'] += bonus\n",
-    "        if outcome == 'goal_reached':\n",
-    "            wins += 1\n",
-    "        print(f'  Game {game_i+1}/{n_games} | {outcome} | steps={step_i}')\n",
-    "\n",
-    "    return experiences, wins / n_games\n",
-    "\n",
-    "print('Collecting initial rollouts...')\n",
-    "exps, win_rate = collect_episodes(model, tokenizer, CONFIG['games_per_step'])\n",
-    "print(f'Win rate: {win_rate:.1%} | Samples: {len(exps)}')\n",
-    "wandb.log({'initial_win_rate': win_rate})"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 8. GRPO Training ─────────────────────────────────────────────────────\n",
-    "from trl import GRPOTrainer, GRPOConfig\n",
-    "\n",
-    "dataset = Dataset.from_list([\n",
-    "    {'prompt': e['prompt'], 'reward': e['reward']} for e in exps\n",
-    "])\n",
-    "\n",
-    "grpo_config = GRPOConfig(\n",
-    "    output_dir='./autodrive-grpo-checkpoints',\n",
-    "    max_steps=CONFIG['max_steps'],\n",
-    "    per_device_train_batch_size=CONFIG['batch_size'],\n",
-    "    gradient_accumulation_steps=CONFIG['grad_accum'],\n",
-    "    learning_rate=CONFIG['learning_rate'],\n",
-    "    num_generations=CONFIG['num_generations'],\n",
-    "    max_new_tokens=CONFIG['max_new_tokens'],\n",
-    "    max_prompt_length=1024,\n",
-    "    bf16=True,\n",
-    "    logging_steps=10,\n",
-    "    save_steps=100,\n",
-    "    report_to='wandb',\n",
-    "    use_vllm=CONFIG['fast_inference'],\n",
-    "    vllm_gpu_memory_utilization=0.3,\n",
-    "    temperature=0.7,\n",
-    "    kl_coef=0.01,\n",
-    ")\n",
-    "\n",
-    "trainer = GRPOTrainer(\n",
-    "    model=model,\n",
-    "    processing_class=tokenizer,\n",
-    "    reward_funcs=[format_reward, action_validity_reward],\n",
-    "    args=grpo_config,\n",
-    "    train_dataset=dataset,\n",
-    ")\n",
-    "\n",
-    "print('🚀 Starting GRPO training...')\n",
-    "trainer.train()\n",
-    "print('✅ Training complete!')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 9. Online RL — Closed-Loop Self-Improvement ──────────────────────────\n",
-    "ONLINE_ITERS = 5\n",
-    "win_rates = [win_rate]\n",
-    "\n",
-    "for i in range(ONLINE_ITERS):\n",
-    "    print(f'=== Online RL Iteration {i+1}/{ONLINE_ITERS} ===')\n",
-    "    fresh_exps, wr = collect_episodes(model, tokenizer, CONFIG['games_per_step'])\n",
-    "    win_rates.append(wr)\n",
-    "    wandb.log({'win_rate': wr, 'iteration': i + 1})\n",
-    "    print(f'Win rate: {wr:.1%}')\n",
-    "\n",
-    "    trainer.train_dataset = Dataset.from_list(\n",
-    "        [{'prompt': e['prompt'], 'reward': e['reward']} for e in fresh_exps]\n",
-    "    )\n",
-    "    trainer.args.max_steps = 50\n",
-    "    trainer.train()\n",
-    "\n",
-    "print(f'Win rate progression: {[f\"{r:.1%}\" for r in win_rates]}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# ── 10. Save & Push to HF Hub ────────────────────────────────────────────\n",
-    "model.save_pretrained('autodrive-adapter')\n",
-    "tokenizer.save_pretrained('autodrive-adapter')\n",
-    "model.push_to_hub(HF_MODEL_REPO)\n",
-    "tokenizer.push_to_hub(HF_MODEL_REPO)\n",
-    "print(f'✅ Model: https://huggingface.co/{HF_MODEL_REPO}')\n",
-    "wandb.finish()"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.11.0"
-  },
-  "accelerator": "GPU",
-  "colab": {
-   "gpuType": "H100",
-   "provenance": []
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
+  "nbformat": 4,
+  "nbformat_minor": 0
 }