feat: add GridMind GRPO training notebook using Unsloth and HF TRL

scripts/gridmind_grpo_colab.ipynb

@@ -4,11 +4,13 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# ⚡ GridMind-RL: Training an LLM Energy Controller with Unsloth + GRPO\n",
+    "# \u26a1 GridMind-RL: Training an LLM Energy Controller with Unsloth + GRPO\n",
     "\n",
     "This notebook fine-tunes **Qwen2.5-1.5B-Instruct** to manage industrial building energy\n",
     "using Reinforcement Learning via the live **GridMind-RL OpenEnv** environment.\n",
     "\n",
+    "**Hardware:** This notebook is designed to run on a **Hugging Face Space (ZeroGPU/A10G)** or Google Colab (T4/L4).\n",
+    "\n",
     "| | |\n",
     "|---|---|\n",
     "| **Environment** | https://lo-kyu-gridmind.hf.space |\n",
@@ -19,7 +21,7 @@
     "### What does the agent learn?\n",
     "- **Task 1**: Minimize energy cost by charging thermal storage off-peak\n",
     "- **Task 2**: Maintain indoor temperature while minimizing cost\n",
-    "- **Task 3**: Full demand-response — cost + temperature + grid stress + batch scheduling + carbon"
+    "- **Task 3**: Full demand-response \u2014 cost + temperature + grid stress + batch scheduling + carbon"
    ]
   },
   {
@@ -29,7 +31,7 @@
    "outputs": [],
    "source": [
     "%%capture\n",
-    "!pip install unsloth openenv-core\n",
+    "!pip install unsloth requests\n",
     "!pip install --no-deps bitsandbytes accelerate xformers peft trl triton\n",
     "!pip install --no-deps cut_cross_entropy unsloth_zoo\n",
     "!pip install \"datasets>=3.4.1,<4.0.0\" pandas matplotlib nest_asyncio"
@@ -39,7 +41,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 1 — Verify the Live Environment"
+    "## Step 1 \u2014 Verify the Live Environment"
    ]
   },
   {
@@ -48,31 +50,34 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from openenv.core import GenericEnvClient\n",
-    "import asyncio, nest_asyncio\n",
-    "nest_asyncio.apply()\n",
+    "import requests\n",
     "\n",
     "ENV_URL = \"https://lo-kyu-gridmind.hf.space\"\n",
     "\n",
-    "async def verify_env():\n",
-    "    async with GenericEnvClient(base_url=ENV_URL) as env:\n",
-    "        r = await env.reset()\n",
-    "        print(\"✅ Environment live!\")\n",
-    "        print(\"Observation keys:\", list(r.observation.keys()))\n",
-    "        r2 = await env.step({\n",
+    "def verify_env():\n",
+    "    try:\n",
+    "        r = requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": 1, \"seed\": 42})\n",
+    "        r.raise_for_status()\n",
+    "        data = r.json()\n",
+    "        print(\"\u2705 Environment live!\")\n",
+    "        print(\"Observation keys:\", list(data.get(\"observations\", [{}])[0].keys()))\n",
+    "        r2 = requests.post(f\"{ENV_URL}/step\", json=[{\n",
     "            \"hvac_power_level\": 0.5, \"thermal_charge_rate\": 0.0,\n",
     "            \"batch_job_slot\": 0, \"load_shed_fraction\": 0.0, \"building_id\": 0\n",
-    "        })\n",
-    "        print(f\"Step reward: {r2.reward:.3f}, done: {r2.done}\")\n",
+    "        }])\n",
+    "        res = r2.json().get(\"results\", [{}])[0]\n",
+    "        print(f\"Step reward: {res.get('reward', 0):.3f}, done: {res.get('done', False)}\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"\u274c Environment verification failed: {e}\")\n",
     "\n",
-    "asyncio.run(verify_env())"
+    "verify_env()"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 2 — Load Model with Unsloth 4-bit LoRA"
+    "## Step 2 \u2014 Load Model with Unsloth 4-bit LoRA"
    ]
   },
   {
@@ -102,14 +107,14 @@
     "    use_gradient_checkpointing=\"unsloth\",\n",
     "    random_state=42,\n",
     ")\n",
-    "print(\"✅ Model loaded with Unsloth 4-bit LoRA\")"
+    "print(\"\u2705 Model loaded with Unsloth 4-bit LoRA\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 3 — Define Reward Functions\n",
+    "## Step 3 \u2014 Define Reward Functions\n",
     "\n",
     "We use a **composite reward** with three components:\n",
     "\n",
@@ -127,14 +132,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import json, re\n",
+    "import json, re, requests\n",
     "\n",
     "def reward_valid_json(completions, **kwargs):\n",
     "    rewards = []\n",
     "    for completion in completions:\n",
     "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
     "        try:\n",
-    "            match = re.search(r'\\{.*?\\}', text, re.DOTALL)\n",
+    "            match = re.search(r\"\\{.*?\\}\", text, re.DOTALL)\n",
     "            if match:\n",
     "                json.loads(match.group())\n",
     "                rewards.append(0.3)\n",
@@ -150,7 +155,7 @@
     "    for completion in completions:\n",
     "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
     "        try:\n",
-    "            match = re.search(r'\\{.*?\\}', text, re.DOTALL)\n",
+    "            match = re.search(r\"\\{.*?\\}\", text, re.DOTALL)\n",
     "            if match:\n",
     "                action = json.loads(match.group())\n",
     "                rewards.append(0.3 if required.issubset(action.keys()) else 0.1)\n",
@@ -162,9 +167,11 @@
     "\n",
     "def reward_env_interaction(completions, **kwargs):\n",
     "    \"\"\"Reward 0.0-0.4 based on actual environment reward from live GridMind-RL HF Space.\"\"\"\n",
-    "    async def run_step(text):\n",
+    "    rewards = []\n",
+    "    for completion in completions:\n",
+    "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
     "        try:\n",
-    "            match = re.search(r'\\{.*?\\}', text, re.DOTALL)\n",
+    "            match = re.search(r\"\\{.*?\\}\", text, re.DOTALL)\n",
     "            action = json.loads(match.group()) if match else {}\n",
     "            step_action = {\n",
     "                \"hvac_power_level\":    float(max(0, min(1, action.get(\"hvac_power_level\", 0.5)))),\n",
@@ -173,21 +180,23 @@
     "                \"load_shed_fraction\":  float(max(0, min(0.5, action.get(\"load_shed_fraction\", 0.0)))),\n",
     "                \"building_id\": 0\n",
     "            }\n",
-    "            async with GenericEnvClient(base_url=ENV_URL) as env:\n",
-    "                await env.reset()\n",
-    "                result = await env.step(step_action)\n",
-    "                # Normalize raw env reward (~[-2, 3]) → (0.0, 0.4)\n",
-    "                return min(0.4, max(0.0, (result.reward + 2.0) * 0.08))\n",
+    "            r_reset = requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": 1, \"seed\": 42}, timeout=30)\n",
+    "            if r_reset.status_code != 200:\n",
+    "                rewards.append(0.0)\n",
+    "                continue\n",
+    "            r_step = requests.post(f\"{ENV_URL}/step\", json=[step_action], timeout=30)\n",
+    "            if r_step.status_code != 200:\n",
+    "                rewards.append(0.0)\n",
+    "                continue\n",
+    "            res = r_step.json().get(\"results\", [{}])[0]\n",
+    "            step_reward = float(res.get(\"reward\", 0.0))\n",
+    "            val = (step_reward + 2.0) * 0.08\n",
+    "            rewards.append(min(0.4, max(0.0, val)))\n",
     "        except Exception:\n",
-    "            return 0.0\n",
-    "\n",
-    "    rewards = []\n",
-    "    for completion in completions:\n",
-    "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
-    "        rewards.append(asyncio.run(run_step(text)))\n",
+    "            rewards.append(0.0)\n",
     "    return rewards\n",
     "\n",
-    "print(\"✅ Reward functions defined\")\n",
+    "print(\"\u2705 Reward functions defined\")\n",
     "print(\"  Total max reward per step: 1.0\")"
    ]
   },
@@ -195,7 +204,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 4 — Build Training Dataset & Start GRPO Training"
+    "## Step 4 \u2014 Build Training Dataset & Start GRPO Training"
    ]
   },
   {
@@ -204,80 +213,79 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from trl import GRPOTrainer, GRPOConfig\n",
-    "from datasets import Dataset\n",
-    "import pandas as pd, os\n",
-    "from transformers import TrainerCallback\n",
-    "\n",
-    "SYSTEM_PROMPT = \"\"\"You are an expert industrial building energy controller.\n",
-    "Each turn you receive the current building state and must respond with \n",
-    "ONLY a valid JSON action object.\n",
-    "\n",
-    "Action format:\n",
-    "{\"hvac_power_level\": <0.0-1.0>, \"thermal_charge_rate\": <-1.0 to 1.0>, \n",
-    " \"batch_job_slot\": <0-4>, \"load_shed_fraction\": <0.0-0.5>, \"building_id\": 0}\n",
+    "import json, re, requests\n",
     "\n",
-    "Strategy:\n",
-    "- Charge storage when price < $0.08/kWh (positive thermal_charge_rate)\n",
-    "- Discharge storage when price > $0.15/kWh (negative thermal_charge_rate)  \n",
-    "- Shed load 0.3-0.5 when grid_stress_signal > 0.7\n",
-    "- Reduce HVAC during peak hours (8-12, 17-21)\n",
-    "- Keep temperature between 19-23°C\"\"\"\n",
-    "\n",
-    "def make_prompt(i):\n",
-    "    return [{\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n",
-    "            {\"role\": \"user\",\n",
-    "             \"content\": f\"Episode {i+1}: Building simulation starting. Output your first action as JSON.\"}]\n",
-    "\n",
-    "dataset = Dataset.from_dict({\"prompt\": [make_prompt(i) for i in range(300)]})\n",
-    "print(f\"✅ Dataset: {len(dataset)} training prompts\")\n",
-    "\n",
-    "# --- CSV Logger ---\n",
-    "log_history = []\n",
-    "class CSVLogger(TrainerCallback):\n",
-    "    def on_log(self, args, state, control, logs=None, **kwargs):\n",
-    "        if logs and \"loss\" in logs:\n",
-    "            entry = {**logs, \"step\": state.global_step}\n",
-    "            log_history.append(entry)\n",
-    "            os.makedirs(\"results\", exist_ok=True)\n",
-    "            pd.DataFrame(log_history).to_csv(\"results/training_log.csv\", index=False)\n",
+    "def reward_valid_json(completions, **kwargs):\n",
+    "    rewards = []\n",
+    "    for completion in completions:\n",
+    "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
+    "        try:\n",
+    "            match = re.search(r\"\\{.*?\\}\", text, re.DOTALL)\n",
+    "            if match:\n",
+    "                json.loads(match.group())\n",
+    "                rewards.append(0.3)\n",
+    "            else:\n",
+    "                rewards.append(0.0)\n",
+    "        except Exception:\n",
+    "            rewards.append(0.0)\n",
+    "    return rewards\n",
     "\n",
-    "training_args = GRPOConfig(\n",
-    "    output_dir=\"gridmind-grpo-unsloth\",\n",
-    "    num_train_epochs=1,\n",
-    "    per_device_train_batch_size=1,\n",
-    "    gradient_accumulation_steps=4,\n",
-    "    num_generations=4,\n",
-    "    max_prompt_length=256,\n",
-    "    max_completion_length=128,\n",
-    "    learning_rate=5e-6,\n",
-    "    lr_scheduler_type=\"cosine\",\n",
-    "    warmup_ratio=0.1,\n",
-    "    logging_steps=5,\n",
-    "    save_steps=100,\n",
-    "    fp16=True,\n",
-    "    report_to=\"none\",\n",
-    "    seed=42,\n",
-    ")\n",
+    "def reward_has_required_keys(completions, **kwargs):\n",
+    "    required = {\"hvac_power_level\", \"thermal_charge_rate\", \"batch_job_slot\", \"load_shed_fraction\"}\n",
+    "    rewards = []\n",
+    "    for completion in completions:\n",
+    "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
+    "        try:\n",
+    "            match = re.search(r\"\\{.*?\\}\", text, re.DOTALL)\n",
+    "            if match:\n",
+    "                action = json.loads(match.group())\n",
+    "                rewards.append(0.3 if required.issubset(action.keys()) else 0.1)\n",
+    "            else:\n",
+    "                rewards.append(0.0)\n",
+    "        except Exception:\n",
+    "            rewards.append(0.0)\n",
+    "    return rewards\n",
     "\n",
-    "trainer = GRPOTrainer(\n",
-    "    model=model,\n",
-    "    tokenizer=tokenizer,\n",
-    "    args=training_args,\n",
-    "    train_dataset=dataset,\n",
-    "    reward_funcs=[reward_valid_json, reward_has_required_keys, reward_env_interaction],\n",
-    "    callbacks=[CSVLogger()]\n",
-    ")\n",
+    "def reward_env_interaction(completions, **kwargs):\n",
+    "    \"\"\"Reward 0.0-0.4 based on actual environment reward from live GridMind-RL HF Space.\"\"\"\n",
+    "    rewards = []\n",
+    "    for completion in completions:\n",
+    "        text = completion[0][\"content\"] if isinstance(completion, list) else completion\n",
+    "        try:\n",
+    "            match = re.search(r\"\\{.*?\\}\", text, re.DOTALL)\n",
+    "            action = json.loads(match.group()) if match else {}\n",
+    "            step_action = {\n",
+    "                \"hvac_power_level\":    float(max(0, min(1, action.get(\"hvac_power_level\", 0.5)))),\n",
+    "                \"thermal_charge_rate\": float(max(-1, min(1, action.get(\"thermal_charge_rate\", 0.0)))),\n",
+    "                \"batch_job_slot\":      int(max(0, min(4, action.get(\"batch_job_slot\", 0)))),\n",
+    "                \"load_shed_fraction\":  float(max(0, min(0.5, action.get(\"load_shed_fraction\", 0.0)))),\n",
+    "                \"building_id\": 0\n",
+    "            }\n",
+    "            r_reset = requests.post(f\"{ENV_URL}/reset\", json={\"task_id\": 1, \"seed\": 42}, timeout=30)\n",
+    "            if r_reset.status_code != 200:\n",
+    "                rewards.append(0.0)\n",
+    "                continue\n",
+    "            r_step = requests.post(f\"{ENV_URL}/step\", json=[step_action], timeout=30)\n",
+    "            if r_step.status_code != 200:\n",
+    "                rewards.append(0.0)\n",
+    "                continue\n",
+    "            res = r_step.json().get(\"results\", [{}])[0]\n",
+    "            step_reward = float(res.get(\"reward\", 0.0))\n",
+    "            val = (step_reward + 2.0) * 0.08\n",
+    "            rewards.append(min(0.4, max(0.0, val)))\n",
+    "        except Exception:\n",
+    "            rewards.append(0.0)\n",
+    "    return rewards\n",
     "\n",
-    "print(\"🚀 Starting GRPO training...\")\n",
-    "trainer.train()"
+    "print(\"\u2705 Reward functions defined\")\n",
+    "print(\"  Total max reward per step: 1.0\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 5 — Plot Training Curve\n",
+    "## Step 5 \u2014 Plot Training Curve\n",
     "\n",
     "This plot is the key **evidence of learning** for the hackathon judges."
    ]
@@ -312,14 +320,14 @@
     "plt.tight_layout()\n",
     "plt.savefig(\"results/training_curve.png\", dpi=200, bbox_inches='tight')\n",
     "plt.show()\n",
-    "print(\"✅ Training curve saved to results/training_curve.png\")"
+    "print(\"\u2705 Training curve saved to results/training_curve.png\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 6 — Before vs After Comparison\n",
+    "## Step 6 \u2014 Before vs After Comparison\n",
     "\n",
     "Test the same scenario pre-training and post-training to show qualitative improvement."
    ]
@@ -331,7 +339,7 @@
    "outputs": [],
    "source": [
     "test_state = (\n",
-    "    \"Building state: temp=24.5°C (too hot!), price=$0.18/kWh (peak), \"\n",
+    "    \"Building state: temp=24.5\u00b0C (too hot!), price=$0.18/kWh (peak), \"\n",
     "    \"storage=0.7 (charged), grid_stress=0.85 (CRITICAL!), hour=18, step=60/95\\n\"\n",
     "    \"Pending batch job deadlines: [12, 30]\\n\"\n",
     "    \"Cumulative cost so far: $1.24\"\n",
@@ -354,11 +362,11 @@
     "    )\n",
     "\n",
     "response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
-    "print(\"📋 Test Scenario:\")\n",
+    "print(\"\ud83d\udccb Test Scenario:\")\n",
     "print(\" \", test_state.replace(\"\\n\", \"\\n  \"))\n",
-    "print(\"\\n🤖 Fine-tuned Model Response:\")\n",
+    "print(\"\\n\ud83e\udd16 Fine-tuned Model Response:\")\n",
     "print(\" \", response)\n",
-    "print(\"\\n✅ Expected: load_shed_fraction > 0 (grid_stress=0.85), thermal_charge_rate < 0 (discharge at peak price)\")"
+    "print(\"\\n\u2705 Expected: load_shed_fraction > 0 (grid_stress=0.85), thermal_charge_rate < 0 (discharge at peak price)\")"
    ]
   }
  ],