fix: notebook uses compute_grpo_reward_env, updated hyperparams, no emojis

training/opengrid_grpo_colab.ipynb

@@ -4,17 +4,17 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# 🔋 OpenGrid — GRPO Training Notebook\n",
+    "#  OpenGrid \u2014 GRPO Training Notebook\n",
     "\n",
     "**Multi-Agent RL for Power Grid Operations**\n",
     "\n",
     "This notebook trains an LLM (Qwen 2.5 1.5B) to operate a power grid using GRPO (Group Relative Policy Optimization).\n",
     "\n",
-    "- **Environment**: OpenGrid — multi-agent POMDP with safety layer & oversight agent\n",
+    "- **Environment**: OpenGrid \u2014 multi-agent POMDP with safety layer & oversight agent\n",
     "- **Task**: Maintain 50 Hz frequency, prevent line overloads, avoid blackouts\n",
     "- **Training**: TRL GRPOTrainer + Unsloth 4-bit quantization\n",
     "\n",
-    "⚡ **Runtime**: Select `T4 GPU` from Runtime → Change runtime type"
+    " **Runtime**: Select `T4 GPU` from Runtime \u2192 Change runtime type"
    ]
   },
   {
@@ -51,7 +51,7 @@
    "source": [
     "import os\n",
     "\n",
-    "# ⚠️ UPDATE THIS with your actual repo URL\n",
+    "#  UPDATE THIS with your actual repo URL\n",
     "REPO_URL = \"https://github.com/krishnagoyal099/Opengrid_env.git\"\n",
     "\n",
     "if not os.path.exists(\"opengrid\"):\n",
@@ -84,7 +84,7 @@
     "    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n",
     "    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n",
     "else:\n",
-    "    print(\"⚠️ No GPU detected! Go to Runtime → Change runtime type → T4 GPU\")"
+    "    print(\" No GPU detected! Go to Runtime \u2192 Change runtime type \u2192 T4 GPU\")"
    ]
   },
   {
@@ -177,14 +177,14 @@
     "        \"std_reward\": np.std(rewards),\n",
     "        \"rewards\": rewards\n",
     "    }\n",
-    "    print(f\"[BASELINE] {task_id}: {np.mean(rewards):.2f} ± {np.std(rewards):.2f}\")\n",
+    "    print(f\"[BASELINE] {task_id}: {np.mean(rewards):.2f} \u00b1 {np.std(rewards):.2f}\")\n",
     "\n",
     "# Save baseline for later comparison\n",
     "import pickle\n",
     "os.makedirs(\"training/outputs\", exist_ok=True)\n",
     "with open(\"training/outputs/baseline_results.pkl\", \"wb\") as f:\n",
     "    pickle.dump(baseline_results, f)\n",
-    "print(\"\\n✅ Baseline scores saved.\")"
+    "print(\"\\n Baseline scores saved.\")"
    ]
   },
   {
@@ -222,7 +222,7 @@
     "if tokenizer.pad_token is None:\n",
     "    tokenizer.pad_token = tokenizer.eos_token\n",
     "\n",
-    "print(f\"✅ Model loaded: {MODEL_NAME}\")\n",
+    "print(f\" Model loaded: {MODEL_NAME}\")\n",
     "print(f\"   Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}\")"
    ]
   },
@@ -261,7 +261,7 @@
     "\n",
     "    for t in range(min(10, task_config['max_steps'])):\n",
     "        for agent_id, obs in zone_obs.items():\n",
-    "            # model_dump_json() → json.loads() ensures all keys are strings\n",
+    "            # model_dump_json() \u2192 json.loads() ensures all keys are strings\n",
     "            obs_dict = _json.loads(obs.model_dump_json())\n",
     "            prompt_text = format_observation_prompt(obs_dict, zone_name=obs.zone_name)\n",
     "            messages = [\n",
@@ -272,7 +272,7 @@
     "                messages, tokenize=False, add_generation_prompt=True\n",
     "            )\n",
     "            prompts.append(formatted)\n",
-    "            # Store as JSON string — flat scalar, no schema-inference issues\n",
+    "            # Store as JSON string \u2014 flat scalar, no schema-inference issues\n",
     "            obs_contexts.append(_json.dumps(obs_dict))\n",
     "\n",
     "        # Advance env with diverse random actions (no slack bus)\n",
@@ -293,7 +293,7 @@
     "            break\n",
     "        zone_obs = result.observations\n",
     "\n",
-    "print(f\"✅ Generated {len(prompts)} training prompts\")\n",
+    "print(f\" Generated {len(prompts)} training prompts\")\n",
     "print(f\"\\nSample prompt (first 400 chars):\")\n",
     "print(prompts[0][:400])"
    ]
@@ -312,21 +312,18 @@
    "outputs": [],
    "source": [
     "import json as _json\n",
-    "from training.train_grpo import compute_grpo_reward, extract_action\n",
+    "from training.train_grpo import compute_grpo_reward_env, extract_action\n",
     "\n",
     "def reward_fn(completions, obs_context=None, **kwargs):\n",
-    "    \"\"\"GRPO-compatible reward function for OpenGrid.\n",
-    "    obs_context arrives as JSON strings from the dataset column.\n",
-    "    \"\"\"\n",
+    "    \"\"\"GRPO reward function with env-grounded physics rewards.\"\"\"\n",
     "    texts = []\n",
     "    for c in completions:\n",
     "        if isinstance(c, list):\n",
-    "            text = c[-1]['content'] if c else \"\"\n",
+    "            text = c[-1][\"content\"] if c else \"\"\n",
     "        else:\n",
     "            text = str(c)\n",
     "        texts.append(text)\n",
     "\n",
-    "    # Deserialize JSON strings → dicts for the reward scorer\n",
     "    if obs_context is None:\n",
     "        batch_obs = [None] * len(texts)\n",
     "    else:\n",
@@ -334,23 +331,23 @@
     "            _json.loads(ctx) if isinstance(ctx, str) else ctx\n",
     "            for ctx in obs_context\n",
     "        ]\n",
-    "    return compute_grpo_reward(texts, batch_obs)\n",
+    "    return compute_grpo_reward_env(texts, batch_obs, task_config, horizon=3)\n",
     "\n",
-    "# Quick sanity test\n",
+    "# Sanity test\n",
     "test_rewards = reward_fn([\n",
     "    '{\"bus_adjustments\": [{\"bus_id\": 1, \"delta\": 5.0}], \"topology_actions\": []}',\n",
-    "    'invalid json here',\n",
+    "    \"invalid json here\",\n",
     "])\n",
     "print(f\"Test rewards: {test_rewards}\")\n",
-    "assert len(test_rewards) == 2, \"reward_fn must return one score per completion\"\n",
-    "print(\"✅ reward_fn OK\")"
+    "assert len(test_rewards) == 2\n",
+    "print(\"[OK] reward_fn works\")\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 9. Train with GRPO 🚀"
+    "## 9. Train with GRPO "
    ]
   },
   {
@@ -368,21 +365,21 @@
     "\n",
     "grpo_config = GRPOConfig(\n",
     "    output_dir=\"training/outputs/grpo_checkpoints\",\n",
-    "    num_train_epochs=1,\n",
+    "    num_train_epochs=3,\n",
     "    per_device_train_batch_size=2,\n",
-    "    gradient_accumulation_steps=4,\n",
-    "    learning_rate=5e-6,\n",
+    "    gradient_accumulation_steps=8,\n",
+    "    learning_rate=1e-5,\n",
     "    logging_steps=5,\n",
     "    save_steps=50,\n",
     "    max_completion_length=256,\n",
-    "    num_generations=4,\n",
+    "    num_generations=8,\n",
     "    report_to=\"none\",\n",
     "    remove_unused_columns=False,\n",
     "    bf16=_bf16,\n",
     "    fp16=_fp16,\n",
     ")\n",
     "\n",
-    "# obs_contexts are JSON strings — PyArrow handles flat strings with no issues\n",
+    "# obs_contexts are JSON strings \u2014 PyArrow handles flat strings with no issues\n",
     "train_dataset = Dataset.from_dict({\"prompt\": prompts, \"obs_context\": obs_contexts})\n",
     "print(f\"Dataset: {len(train_dataset)} rows, columns: {train_dataset.column_names}\")\n",
     "\n",
@@ -396,11 +393,11 @@
     "\n",
     "print(f\"Training on {len(prompts)} prompts, {grpo_config.num_train_epochs} epoch(s)\")\n",
     "print(f\"Effective batch size: {grpo_config.per_device_train_batch_size * grpo_config.gradient_accumulation_steps}\")\n",
-    "print(\"\\n🚀 Starting GRPO training...\")\n",
+    "print(\"\\n Starting GRPO training...\")\n",
     "\n",
     "train_result = trainer.train()\n",
     "\n",
-    "print(\"\\n✅ Training complete!\")\n",
+    "print(\"\\n Training complete!\")\n",
     "print(f\"   Total steps: {trainer.state.global_step}\")"
    ]
   },
@@ -420,7 +417,7 @@
     "OUTPUT_PATH = \"training/outputs/trained_model\"\n",
     "trainer.save_model(OUTPUT_PATH)\n",
     "tokenizer.save_pretrained(OUTPUT_PATH)\n",
-    "print(f\"✅ Model saved to {OUTPUT_PATH}\")"
+    "print(f\" Model saved to {OUTPUT_PATH}\")"
    ]
   },
   {
@@ -479,14 +476,14 @@
     "        \"std_reward\": np.std(rewards),\n",
     "        \"rewards\": rewards\n",
     "    }\n",
-    "    print(f\"[TRAINED] {task_id}: {np.mean(rewards):.2f} ± {np.std(rewards):.2f}\\n\")"
+    "    print(f\"[TRAINED] {task_id}: {np.mean(rewards):.2f} \u00b1 {np.std(rewards):.2f}\\n\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 12. Generate Before/After Plots 📊"
+    "## 12. Generate Before/After Plots "
    ]
   },
   {
@@ -502,7 +499,7 @@
     "with open(\"training/outputs/baseline_results.pkl\", \"rb\") as f:\n",
     "    baseline_results = pickle.load(f)\n",
     "\n",
-    "# ── Plot 1: Before vs After Bar Chart ──\n",
+    "# \u2500\u2500 Plot 1: Before vs After Bar Chart \u2500\u2500\n",
     "common_tasks = [t for t in baseline_results if t in trained_results]\n",
     "fig, ax = plt.subplots(figsize=(10, 6))\n",
     "x = np.arange(len(common_tasks))\n",
@@ -516,7 +513,7 @@
     "\n",
     "ax.set_xlabel('Task', fontsize=12)\n",
     "ax.set_ylabel('Average Episode Reward', fontsize=12)\n",
-    "ax.set_title('OpenGrid — GRPO Training: Before vs After', fontsize=14, fontweight='bold')\n",
+    "ax.set_title('OpenGrid \u2014 GRPO Training: Before vs After', fontsize=14, fontweight='bold')\n",
     "ax.set_xticks(x)\n",
     "ax.set_xticklabels([t.replace('task_', '').title() for t in common_tasks])\n",
     "ax.legend(fontsize=11)\n",
@@ -536,7 +533,7 @@
     "plt.tight_layout()\n",
     "plt.savefig('training/outputs/before_after.png', dpi=150)\n",
     "plt.show()\n",
-    "print(\"✅ Saved: training/outputs/before_after.png\")"
+    "print(\" Saved: training/outputs/before_after.png\")"
    ]
   },
   {
@@ -545,7 +542,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# ── Plot 2: Training Reward Curve ──\n",
+    "# \u2500\u2500 Plot 2: Training Reward Curve \u2500\u2500\n",
     "history = trainer.state.log_history\n",
     "\n",
     "steps = [h['step'] for h in history if 'loss' in h]\n",
@@ -560,13 +557,13 @@
     "\n",
     "ax.set_xlabel('Training Step', fontsize=12)\n",
     "ax.set_ylabel('Loss', fontsize=12)\n",
-    "ax.set_title('OpenGrid GRPO — Training Loss', fontsize=14, fontweight='bold')\n",
+    "ax.set_title('OpenGrid GRPO \u2014 Training Loss', fontsize=14, fontweight='bold')\n",
     "ax.legend()\n",
     "ax.grid(True, alpha=0.3)\n",
     "plt.tight_layout()\n",
     "plt.savefig('training/outputs/training_loss.png', dpi=150)\n",
     "plt.show()\n",
-    "print(\"✅ Saved: training/outputs/training_loss.png\")"
+    "print(\" Saved: training/outputs/training_loss.png\")"
    ]
   },
   {
@@ -585,19 +582,19 @@
    "outputs": [],
    "source": [
     "print(\"=\"*60)\n",
-    "print(\"  OpenGrid GRPO Training — Results Summary\")\n",
+    "print(\"  OpenGrid GRPO Training \u2014 Results Summary\")\n",
     "print(\"=\"*60)\n",
     "\n",
     "# Rebuild common_tasks in case Cell 12 was skipped\n",
     "common_tasks = [t for t in baseline_results if t in trained_results]\n",
     "\n",
-    "print(f\"{'Task':<20} {'Baseline':>12} {'Trained':>12} {'Δ':>10}\")\n",
+    "print(f\"{'Task':<20} {'Baseline':>12} {'Trained':>12} {'\u0394':>10}\")\n",
     "print(\"-\"*60)\n",
     "for t in common_tasks:\n",
     "    b = baseline_results[t]['avg_reward']\n",
     "    a = trained_results[t]['avg_reward']\n",
     "    delta = a - b\n",
-    "    arrow = '↑' if delta > 0 else '↓'\n",
+    "    arrow = '\u2191' if delta > 0 else '\u2193'\n",
     "    print(f\"{t:<20} {b:>10.2f}   {a:>10.2f}   {arrow} {abs(delta):.2f}\")\n",
     "print(\"=\"*60)"
    ]
@@ -608,10 +605,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Download plots for your README\n",
-    "from google.colab import files\n",
-    "files.download('training/outputs/before_after.png')\n",
-    "files.download('training/outputs/training_loss.png')"
+    "# Display plots inline\n",
+    "from IPython.display import Image, display\n",
+    "display(Image(\"training/outputs/before_after.png\"))\n",
+    "display(Image(\"training/outputs/training_loss.png\"))\n"
    ]
   }
  ],
@@ -632,4 +629,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
+}