feat(colab): add RL and LLM training notebook for traffic control

colab/train_rl_and_llm_traffic.ipynb

@@ -0,0 +1,462 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# RL + LLM Training (Colab) for Adaptive Traffic Intelligence\n",
+        "\n",
+        "This notebook trains:\n",
+        "1. **RL traffic controller** (PPO) on the `TrafficEnv`\n",
+        "2. **LLM policy model** (LoRA fine-tuning) to imitate the RL controller\n",
+        "\n",
+        "It also compares **Fixed baseline vs RL vs LLM** on waiting time, queue length, and throughput."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Best Model Choices (Practical)\n",
+        "\n",
+        "### RL (this task)\n",
+        "- **Recommended default:** `PPO` (stable, strong for this discrete-control setup)\n",
+        "- Alternatives: `DQN` (already in repo), `A2C`\n",
+        "\n",
+        "### LLM policy model\n",
+        "- **Best on Colab T4 (recommended):** `Qwen/Qwen2.5-1.5B-Instruct` with LoRA + 4-bit\n",
+        "- Higher quality (needs stronger GPU): `meta-llama/Meta-Llama-3.1-8B-Instruct`\n",
+        "- Faster/smaller fallback: `Qwen/Qwen2.5-0.5B-Instruct`"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 0) Runtime setup ====\n",
+        "!nvidia-smi\n",
+        "\n",
+        "!pip -q install -U gymnasium stable-baselines3 sb3-contrib\n",
+        "!pip -q install -U transformers datasets peft trl accelerate bitsandbytes sentencepiece\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 1) Get project code ====\n",
+        "import os\n",
+        "from pathlib import Path\n",
+        "\n",
+        "REPO_URL = \"https://github.com/DivyankLosse/TRLE-Hackethon.git\"\n",
+        "PROJECT_DIR = Path(\"TRLE-Hackethon\")\n",
+        "\n",
+        "if not PROJECT_DIR.exists():\n",
+        "    !git clone {REPO_URL}\n",
+        "\n",
+        "%cd TRLE-Hackethon\n",
+        "\n",
+        "# Optional: install local package deps\n",
+        "!pip -q install -r requirements.txt\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 2) Imports and Gym wrapper for TrafficEnv ====\n",
+        "import math\n",
+        "import json\n",
+        "import random\n",
+        "from dataclasses import dataclass\n",
+        "\n",
+        "import numpy as np\n",
+        "import gymnasium as gym\n",
+        "from gymnasium import spaces\n",
+        "\n",
+        "from traffic_rl.env.traffic_env import TrafficEnv\n",
+        "from traffic_rl.baseline.fixed_time_controller import FixedTimeController\n",
+        "\n",
+        "class TrafficGymEnv(gym.Env):\n",
+        "    metadata = {\"render_modes\": []}\n",
+        "\n",
+        "    def __init__(self, env_config: dict):\n",
+        "        super().__init__()\n",
+        "        self.base_config = dict(env_config)\n",
+        "        self.base_env = TrafficEnv(config=self.base_config)\n",
+        "        self.action_space = spaces.Discrete(3)\n",
+        "        self.observation_space = spaces.Box(\n",
+        "            low=np.zeros((10,), dtype=np.float32),\n",
+        "            high=np.full((10,), np.finfo(np.float32).max, dtype=np.float32),\n",
+        "            dtype=np.float32,\n",
+        "        )\n",
+        "\n",
+        "    def reset(self, *, seed=None, options=None):\n",
+        "        if seed is not None:\n",
+        "            cfg = dict(self.base_config)\n",
+        "            cfg[\"seed\"] = int(seed)\n",
+        "            self.base_env = TrafficEnv(config=cfg)\n",
+        "        obs = self.base_env.reset().astype(np.float32)\n",
+        "        return obs, {}\n",
+        "\n",
+        "    def step(self, action):\n",
+        "        obs, reward, done, info = self.base_env.step(int(action))\n",
+        "        obs = obs.astype(np.float32)\n",
+        "        terminated = bool(done)\n",
+        "        truncated = False\n",
+        "        return obs, float(reward), terminated, truncated, info\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 3) Train RL model (PPO) ====\n",
+        "from stable_baselines3 import PPO\n",
+        "from stable_baselines3.common.vec_env import DummyVecEnv\n",
+        "\n",
+        "ENV_CONFIG = {\n",
+        "    \"max_steps\": 120,\n",
+        "    \"arrival_mode\": \"stochastic\",\n",
+        "    \"lane_bias\": (1.6, 0.8, 1.4, 0.6),\n",
+        "    \"peak_rates\": (3.8, 2.2, 3.4, 1.5),\n",
+        "    \"offpeak_rates\": (1.4, 0.9, 1.2, 0.7),\n",
+        "    \"peak_duration\": 35,\n",
+        "    \"cycle_duration\": 60,\n",
+        "    \"service_rate\": 2,\n",
+        "    \"ambulance_spawn_prob\": 0.08,\n",
+        "    \"seed\": 42,\n",
+        "}\n",
+        "\n",
+        "N_ENVS = 4\n",
+        "TOTAL_TIMESTEPS = 120_000   # increase to 300k+ for stronger policy\n",
+        "\n",
+        "def make_env(rank):\n",
+        "    def _thunk():\n",
+        "        cfg = dict(ENV_CONFIG)\n",
+        "        cfg[\"seed\"] = ENV_CONFIG[\"seed\"] + rank\n",
+        "        return TrafficGymEnv(cfg)\n",
+        "    return _thunk\n",
+        "\n",
+        "vec_env = DummyVecEnv([make_env(i) for i in range(N_ENVS)])\n",
+        "\n",
+        "ppo = PPO(\n",
+        "    policy=\"MlpPolicy\",\n",
+        "    env=vec_env,\n",
+        "    n_steps=1024,\n",
+        "    batch_size=256,\n",
+        "    learning_rate=3e-4,\n",
+        "    gamma=0.99,\n",
+        "    gae_lambda=0.95,\n",
+        "    clip_range=0.2,\n",
+        "    ent_coef=0.01,\n",
+        "    vf_coef=0.5,\n",
+        "    verbose=1,\n",
+        "    seed=42,\n",
+        ")\n",
+        "\n",
+        "ppo.learn(total_timesteps=TOTAL_TIMESTEPS, progress_bar=True)\n",
+        "\n",
+        "Path(\"artifacts\").mkdir(exist_ok=True)\n",
+        "ppo.save(\"artifacts/ppo_traffic\")\n",
+        "print(\"Saved RL model -> artifacts/ppo_traffic.zip\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 4) Evaluate baseline vs RL ====\n",
+        "def eval_policy(policy_fn, env_config, episodes=20):\n",
+        "    metrics = {\n",
+        "        \"reward\": [],\n",
+        "        \"avg_queue_length\": [],\n",
+        "        \"avg_waiting_time\": [],\n",
+        "        \"throughput\": [],\n",
+        "        \"ambulance_clearances\": [],\n",
+        "    }\n",
+        "\n",
+        "    for ep in range(episodes):\n",
+        "        cfg = dict(env_config)\n",
+        "        cfg[\"seed\"] = env_config.get(\"seed\", 42) + ep\n",
+        "        env = TrafficEnv(config=cfg)\n",
+        "        state = env.reset()\n",
+        "        done = False\n",
+        "\n",
+        "        rewards, queues, waits, throughputs = [], [], [], []\n",
+        "        amb_clears = 0\n",
+        "        step = 0\n",
+        "\n",
+        "        while not done:\n",
+        "            action = int(policy_fn(state, step))\n",
+        "            state, reward, done, info = env.step(action)\n",
+        "            rewards.append(float(reward))\n",
+        "            queues.append(float(info[\"queue_sum\"]))\n",
+        "            waits.append(float(info[\"waiting_sum\"]))\n",
+        "            throughputs.append(float(info[\"throughput\"]))\n",
+        "            amb_clears += int(bool(info.get(\"ambulance_cleared\", False)))\n",
+        "            step += 1\n",
+        "\n",
+        "        metrics[\"reward\"].append(sum(rewards))\n",
+        "        metrics[\"avg_queue_length\"].append(float(np.mean(queues) if queues else 0.0))\n",
+        "        metrics[\"avg_waiting_time\"].append(float(np.mean(waits) if waits else 0.0))\n",
+        "        metrics[\"throughput\"].append(sum(throughputs))\n",
+        "        metrics[\"ambulance_clearances\"].append(float(amb_clears))\n",
+        "\n",
+        "    return {k: float(np.mean(v)) for k, v in metrics.items()}\n",
+        "\n",
+        "fixed = FixedTimeController(switch_interval=5)\n",
+        "baseline_metrics = eval_policy(lambda _s, t: fixed.action_for_step(t), ENV_CONFIG, episodes=20)\n",
+        "rl_metrics = eval_policy(lambda s, _t: ppo.predict(s, deterministic=True)[0], ENV_CONFIG, episodes=20)\n",
+        "\n",
+        "def pct_improve(lower_better_key):\n",
+        "    b, r = baseline_metrics[lower_better_key], rl_metrics[lower_better_key]\n",
+        "    return 0.0 if b == 0 else ((b - r) / b) * 100.0\n",
+        "\n",
+        "def pct_gain(higher_better_key):\n",
+        "    b, r = baseline_metrics[higher_better_key], rl_metrics[higher_better_key]\n",
+        "    return 0.0 if b == 0 else ((r - b) / b) * 100.0\n",
+        "\n",
+        "improvement = {\n",
+        "    \"waiting_time_improvement_pct\": pct_improve(\"avg_waiting_time\"),\n",
+        "    \"queue_length_improvement_pct\": pct_improve(\"avg_queue_length\"),\n",
+        "    \"throughput_gain_pct\": pct_gain(\"throughput\"),\n",
+        "    \"ambulance_clearance_gain_pct\": pct_gain(\"ambulance_clearances\"),\n",
+        "}\n",
+        "\n",
+        "print(\"Baseline:\", json.dumps(baseline_metrics, indent=2))\n",
+        "print(\"RL:\", json.dumps(rl_metrics, indent=2))\n",
+        "print(\"Improvement:\", json.dumps(improvement, indent=2))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 5) Build policy dataset from RL trajectories ====\n",
+        "from datasets import Dataset\n",
+        "\n",
+        "def state_to_prompt(state):\n",
+        "    q = [int(x) for x in state[:4]]\n",
+        "    w = [int(x) for x in state[4:8]]\n",
+        "    phase = int(state[8])\n",
+        "    ambulance = int(state[9])\n",
+        "    return (\n",
+        "        \"You control a traffic signal. Output only one action token: 0, 1, or 2.\\n\"\n",
+        "        f\"Queues(Q1..Q4)={q}\\n\"\n",
+        "        f\"Waiting(W1..W4)={w}\\n\"\n",
+        "        f\"CurrentPhase={phase} AmbulanceFlag={ambulance}\\n\"\n",
+        "        \"Action:\"\n",
+        "    )\n",
+        "\n",
+        "def collect_policy_examples(model, env_config, episodes=30):\n",
+        "    rows = []\n",
+        "    for ep in range(episodes):\n",
+        "        cfg = dict(env_config)\n",
+        "        cfg[\"seed\"] = env_config.get(\"seed\", 42) + 1000 + ep\n",
+        "        env = TrafficEnv(config=cfg)\n",
+        "        s = env.reset()\n",
+        "        done = False\n",
+        "\n",
+        "        while not done:\n",
+        "            a, _ = model.predict(s, deterministic=True)\n",
+        "            a = int(a)\n",
+        "            prompt = state_to_prompt(s)\n",
+        "            text = f\"### Instruction:\\n{prompt}\\n### Response:\\n{a}\"\n",
+        "            rows.append({\"prompt\": prompt, \"action\": str(a), \"text\": text})\n",
+        "            s, _r, done, _info = env.step(a)\n",
+        "\n",
+        "    return rows\n",
+        "\n",
+        "policy_rows = collect_policy_examples(ppo, ENV_CONFIG, episodes=35)\n",
+        "print(f\"Collected {len(policy_rows)} examples\")\n",
+        "\n",
+        "dataset = Dataset.from_list(policy_rows).train_test_split(test_size=0.05, seed=42)\n",
+        "dataset"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## LLM Fine-Tuning (LoRA)\n",
+        "Default model is selected for Colab T4 stability. If OOM happens, switch to `Qwen/Qwen2.5-0.5B-Instruct`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 6) Load base LLM + LoRA setup ====\n",
+        "import torch\n",
+        "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
+        "from peft import LoraConfig, get_peft_model\n",
+        "\n",
+        "BASE_LLM = \"Qwen/Qwen2.5-1.5B-Instruct\"  # recommended on Colab T4\n",
+        "# BASE_LLM = \"Qwen/Qwen2.5-0.5B-Instruct\"  # fallback if VRAM is tight\n",
+        "\n",
+        "bnb_config = BitsAndBytesConfig(\n",
+        "    load_in_4bit=True,\n",
+        "    bnb_4bit_quant_type=\"nf4\",\n",
+        "    bnb_4bit_use_double_quant=True,\n",
+        "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
+        ")\n",
+        "\n",
+        "tokenizer = AutoTokenizer.from_pretrained(BASE_LLM, use_fast=True)\n",
+        "if tokenizer.pad_token is None:\n",
+        "    tokenizer.pad_token = tokenizer.eos_token\n",
+        "\n",
+        "llm = AutoModelForCausalLM.from_pretrained(\n",
+        "    BASE_LLM,\n",
+        "    device_map=\"auto\",\n",
+        "    quantization_config=bnb_config,\n",
+        ")\n",
+        "\n",
+        "lora_cfg = LoraConfig(\n",
+        "    r=16,\n",
+        "    lora_alpha=32,\n",
+        "    lora_dropout=0.05,\n",
+        "    bias=\"none\",\n",
+        "    task_type=\"CAUSAL_LM\",\n",
+        "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+        ")\n",
+        "\n",
+        "llm = get_peft_model(llm, lora_cfg)\n",
+        "llm.print_trainable_parameters()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 7) Fine-tune LLM on RL policy traces ====\n",
+        "from trl import SFTTrainer, SFTConfig\n",
+        "\n",
+        "sft_cfg = SFTConfig(\n",
+        "    output_dir=\"artifacts/llm_lora_policy\",\n",
+        "    max_seq_length=512,\n",
+        "    per_device_train_batch_size=4,\n",
+        "    gradient_accumulation_steps=4,\n",
+        "    learning_rate=2e-4,\n",
+        "    num_train_epochs=1,\n",
+        "    logging_steps=20,\n",
+        "    save_strategy=\"epoch\",\n",
+        "    bf16=torch.cuda.is_available(),\n",
+        "    fp16=not torch.cuda.is_available(),\n",
+        "    report_to=\"none\",\n",
+        ")\n",
+        "\n",
+        "trainer = SFTTrainer(\n",
+        "    model=llm,\n",
+        "    args=sft_cfg,\n",
+        "    train_dataset=dataset[\"train\"],\n",
+        "    eval_dataset=dataset[\"test\"],\n",
+        "    processing_class=tokenizer,\n",
+        "    formatting_func=lambda ex: ex[\"text\"],\n",
+        ")\n",
+        "\n",
+        "trainer.train()\n",
+        "trainer.model.save_pretrained(\"artifacts/llm_lora_policy\")\n",
+        "tokenizer.save_pretrained(\"artifacts/llm_lora_policy\")\n",
+        "print(\"Saved LLM adapter -> artifacts/llm_lora_policy\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 8) Evaluate LLM policy ====\n",
+        "import re\n",
+        "\n",
+        "@torch.inference_mode()\n",
+        "def llm_action_from_state(state):\n",
+        "    prompt = state_to_prompt(state)\n",
+        "    wrapped = f\"### Instruction:\\n{prompt}\\n### Response:\\n\"\n",
+        "    inputs = tokenizer(wrapped, return_tensors=\"pt\").to(llm.device)\n",
+        "    out = llm.generate(\n",
+        "        **inputs,\n",
+        "        max_new_tokens=3,\n",
+        "        do_sample=False,\n",
+        "        temperature=0.0,\n",
+        "        eos_token_id=tokenizer.eos_token_id,\n",
+        "    )\n",
+        "    gen = tokenizer.decode(out[0][inputs[\"input_ids\"].shape[1]:], skip_special_tokens=True).strip()\n",
+        "    m = re.search(r\"[012]\", gen)\n",
+        "    return int(m.group(0)) if m else 0\n",
+        "\n",
+        "llm_metrics = eval_policy(lambda s, _t: llm_action_from_state(s), ENV_CONFIG, episodes=10)\n",
+        "print(\"LLM metrics:\")\n",
+        "print(json.dumps(llm_metrics, indent=2))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# ==== 9) Save summary ====\n",
+        "summary = {\n",
+        "    \"baseline\": baseline_metrics,\n",
+        "    \"rl\": rl_metrics,\n",
+        "    \"llm\": llm_metrics,\n",
+        "    \"rl_vs_baseline\": improvement,\n",
+        "    \"base_llm\": BASE_LLM,\n",
+        "}\n",
+        "\n",
+        "Path(\"artifacts\").mkdir(exist_ok=True)\n",
+        "with open(\"artifacts/colab_training_summary.json\", \"w\") as f:\n",
+        "    json.dump(summary, f, indent=2)\n",
+        "\n",
+        "print(json.dumps(summary, indent=2))\n",
+        "print(\"Saved -> artifacts/colab_training_summary.json\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Optional: push checkpoints to Hugging Face Hub\n",
+        "\n",
+        "After training, you can push:\n",
+        "- RL checkpoint: `artifacts/ppo_traffic.zip`\n",
+        "- LLM adapter: `artifacts/llm_lora_policy/`\n",
+        "\n",
+        "Use `huggingface_hub` or git-lfs depending on your target repo layout."
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "name": "train_rl_and_llm_traffic.ipynb",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}