💎 HONEST TRUTH: Synchronized Space simulation with train.ipynb metrics (0.4/0.4/0.2 & MAD Fairness)

fairrecovery_env/rewards.py

@@ -4,58 +4,47 @@ FairRecovery++ - Reward Engine (Fair-GRPO-RLVR).
 Computes dense, verifiable, formula-based rewards - no learned reward model.
 Implements the Fair-GRPO-RLVR multi-objective reinforcement learning framework.
 
-R_total = w_exec*R_exec + w_fair*R_fair + w_safe*R_safe
+R_total = 0.4*Utility + 0.4*Fairness + 0.2*Safety
 """
 
 from __future__ import annotations
 import structlog
 from dataclasses import dataclass, field
 from typing import List
-from .constants import (GRADER_SCORE_MAX, GRADER_SCORE_MIN, MAX_DAYS,
-                         PENALTY_IGNORE_VULNERABLE, PENALTY_PATTERN_IGNORED,
-                         REWARD_WEIGHTS, VULNERABILITY_THRESHOLD)
+from .constants import (GRADER_SCORE_MAX, GRADER_SCORE_MIN, MAX_DAYS)
 from .state import CityState, ZoneState
 from .tasks import ScenarioConfig
 
 logger = structlog.get_logger(__name__)
 
 
-def compute_exec_reward(prev_services: List[float], zones: List[ZoneState]) -> float:
-    """Mean service improvement this day."""
+def compute_exec_reward(zones: List[ZoneState]) -> float:
+    """Utility: Mean service level [0, 1]."""
     if not zones:
         return 0.0
-    improvements = [z.service - prev for z, prev in zip(zones, prev_services)]
-    # Baseline improvement + actual progress
-    return float(0.05 + sum(improvements) / len(improvements))
+    services = [z.service for z in zones]
+    return float(sum(services) / len(services))
 
 
 def compute_fairness_reward(zones: List[ZoneState]) -> float:
     """
-    Research-level Fairness Index: 1 - variance in service levels.
+    Equity: 1 - Mean Absolute Deviation.
     Higher value means more equitable distribution of services.
     """
     if not zones:
         return 0.0
     services = [z.service for z in zones]
     mean_svc = sum(services) / len(services)
-    variance = sum((s - mean_svc) ** 2 for s in services) / len(services)
+    if not services: return 0.0
+    
+    disparity = sum(abs(s - mean_svc) for s in services) / len(services)
     # Fairness index in [0, 1]
-    return float(max(0.0, 1.0 - variance * 2.0))
+    return float(max(0.0, 1.0 - disparity))
 
 
 def compute_safety_reward(violations: List[str]) -> float:
-    """Penalty per safety violation, capped."""
-    return float(-min(0.5, len(violations) * 0.1))
-
-
-def compute_stability_reward(zones: List[ZoneState]) -> float:
-    """Reward for system balance — low variance in satisfaction."""
-    sats = [z.citizen_satisfaction for z in zones]
-    if len(sats) < 2:
-        return 0.0
-    mean_sat = sum(sats) / len(sats)
-    variance = sum((s - mean_sat) ** 2 for s in sats) / len(sats)
-    return float(max(-1.0, -variance * 4))  # Scale up variance penalty
+    """Safety: Normalized violation count [0, 1]."""
+    return float(max(0.0, 1.0 - len(violations) / 10.0))
 
 
 def compute_analysis_reward(chosen_zones: List[int], zones: List[ZoneState]) -> float:
@@ -75,8 +64,6 @@ class RewardComponents:
     R_exec: float = 0.0
     R_fair: float = 0.0
     R_safe: float = 0.0
-    R_adapt: float = 0.0
-    R_stable: float = 0.0
     R_analysis: float = 0.0
     R_total: float = 0.0
     violations: List[str] = field(default_factory=list)
@@ -85,8 +72,7 @@ class RewardComponents:
     def to_dict(self) -> dict:
         return {k: round(v, 4) if isinstance(v, float) else v
                 for k, v in {"R_exec": self.R_exec, "R_fair": self.R_fair,
-                             "R_safe": self.R_safe, "R_adapt": self.R_adapt,
-                             "R_stable": self.R_stable, "R_total": self.R_total,
+                             "R_safe": self.R_safe, "R_total": self.R_total,
                              "violations": self.violations}.items()}
 
 
@@ -98,7 +84,6 @@ class RewardEngine:
         self._cumulative_reward: float = 0.0
         self._step_count: int = 0
         self._action_history: List[str] = []
-        self._vulnerable_ignored_days: int = 0
 
     @property
     def cumulative_reward(self) -> float:
@@ -107,67 +92,56 @@ class RewardEngine:
     def compute_analysis_step(self, chosen_zones: List[int], city: CityState) -> RewardComponents:
         self._step_count += 1
         R_analysis = compute_analysis_reward(chosen_zones, city.zones)
-        # More significant reward for correct analysis
-        R_total = 0.2 * R_analysis 
+        # Analysis provides a small progress signal
+        R_total = 0.05 * R_analysis 
         self._cumulative_reward += R_total
         return RewardComponents(
             R_analysis=R_analysis, R_total=R_total,
             feedback=f"Analysis: {R_total:+.3f} ({int(R_analysis * max(1, len(city.zones)//2))}"
                      f"/{max(1, len(city.zones)//2)} critical zones correct)")
 
-    def compute_execute_step(self, city: CityState, violations: List[str],
-                              adaptation_score: float = 0.0) -> RewardComponents:
-        """Main dense reward after execute step — now includes adaptation and stability."""
+    def compute_execute_step(self, city: CityState, violations: List[str]) -> RewardComponents:
+        """Main dense reward after execute step using Fair-GRPO-RLVR formula."""
         self._step_count += 1
 
-        # Check if vulnerable zones consistently ignored
-        vuln_ids = {z.zone_id for z in city.zones if z.is_vulnerable}
-        if vuln_ids:
-            history_text = " ".join(city.history)
-            zone_served = any(str(zid) in history_text for zid in vuln_ids)
-            if not zone_served and city.day > 1:
-                self._vulnerable_ignored_days += 1
-                if self._vulnerable_ignored_days >= 2:
-                    violations.append(f"persistent_ignore_vulnerable:{vuln_ids}")
-
-        R_exec = compute_exec_reward(city.prev_services, city.zones)
-        R_fair = compute_fairness_reward(city.zones)
-        R_safe = compute_safety_reward(violations)
-        R_adapt = adaptation_score  # from Predictor.evaluate_adaptation()
-        R_stable = compute_stability_reward(city.zones)
-
-        w = REWARD_WEIGHTS
-        # Fair-GRPO-RLVR Rubric (Optimised for positive feedback and strong learning signals)
-        # We use a +0.3 baseline for a successful step to ensure the baseline is clearly positive
-        baseline = 0.3 if not violations else 0.0
-        R_total = (w["exec"] * R_exec + w["fair"] * R_fair + w["safe"] * R_safe + baseline)
-        R_total = float(max(0.0 if not violations else -0.5, min(1.0, R_total)))
+        utility = compute_exec_reward(city.zones)
+        fairness = compute_fairness_reward(city.zones)
+        safety = compute_safety_reward(violations)
+
+        # Truth Formula: 0.4*Utility + 0.4*Fairness + 0.2*Safety
+        R_total = 0.4 * utility + 0.4 * fairness + 0.2 * safety
+        R_total = float(max(0.0, min(1.0, R_total)))
+        
+        # Note: In interactive mode, we track cumulative, but train.ipynb uses final state.
         self._cumulative_reward += R_total
 
-        feedback = (f"R_exec={R_exec:+.3f} | R_fair={R_fair:+.3f} | R_safe={R_safe:+.3f} | "
-                    f"R_adapt={R_adapt:+.3f} | R_stable={R_stable:+.3f} → R_total={R_total:+.3f}")
+        feedback = (f"Utility={utility:.3f} | Fairness={fairness:.3f} | Safety={safety:.3f} → R_step={R_total:.3f}")
         if violations:
             feedback += f" | Violations: {violations}"
 
-        return RewardComponents(R_exec=R_exec, R_fair=R_fair, R_safe=R_safe,
+        return RewardComponents(R_exec=utility, R_fair=fairness, R_safe=safety,
                                 R_total=R_total, violations=violations, feedback=feedback)
 
     def compute_submit_reward(self, city: CityState) -> RewardComponents:
+        """Final submission reward (matches Truth Formula)."""
         self._step_count += 1
-        R_fair = compute_fairness_reward(city.zones)
-        avg_svc = sum(z.service for z in city.zones) / max(1, len(city.zones))
-        avg_sat = sum(z.citizen_satisfaction for z in city.zones) / max(1, len(city.zones))
-        terminal = 0.4 * avg_svc + 0.3 * (1.0 + R_fair) + 0.3 * avg_sat
+        utility = compute_exec_reward(city.zones)
+        fairness = compute_fairness_reward(city.zones)
+        safety = compute_safety_reward([]) # Assume no new violations on submit
+        
+        terminal = 0.4 * utility + 0.4 * fairness + 0.2 * safety
         terminal = float(max(0.0, min(1.0, terminal)))
         self._cumulative_reward += terminal
+        
         return RewardComponents(
-            R_fair=R_fair, R_exec=avg_svc, R_stable=avg_sat, R_total=terminal,
-            feedback=f"Terminal bonus={terminal:.3f} (svc={avg_svc:.3f}, fair={R_fair:.3f}, sat={avg_sat:.3f})")
+            R_fair=fairness, R_exec=utility, R_total=terminal,
+            feedback=f"Terminal Score={terminal:.3f} (Utility={utility:.3f}, Fairness={fairness:.3f})")
 
     def get_final_grader_score(self, city: CityState) -> float:
         """Normalised score in (GRADER_SCORE_MIN, GRADER_SCORE_MAX)."""
-        avg_svc = sum(z.service for z in city.zones) / max(1, len(city.zones))
-        R_fair = compute_fairness_reward(city.zones)
-        avg_sat = sum(z.citizen_satisfaction for z in city.zones) / max(1, len(city.zones))
-        normalised = 0.4 * avg_svc + 0.3 * (1.0 + R_fair) + 0.3 * avg_sat
+        utility = compute_exec_reward(city.zones)
+        fairness = compute_fairness_reward(city.zones)
+        safety = compute_safety_reward([])
+        
+        normalised = 0.4 * utility + 0.4 * fairness + 0.2 * safety
         return round(float(max(GRADER_SCORE_MIN, min(GRADER_SCORE_MAX, normalised))), 4)

server/app.py

@@ -104,8 +104,6 @@ def _build_app():
         
         done = False
         step_count = 0
-        total_reward = 0.0
-        final_fairness = 0.0
         
         policy_fn = greedy_policy if policy_type == "Baseline (Greedy)" else llm_policy
         
@@ -124,51 +122,58 @@ def _build_app():
                 logs.append(f"> 🚚 Dispatching resources: {', '.join(allocs) if allocs else 'None'}")
                 
             obs = env.step(action)
-            total_reward += obs.reward
             
             if action.action_type == "execute":
                 logs.append("> 🏥 **Recovery Update:**")
                 for z in obs.zones:
                     logs.append(f"  - Zone {z.zone_id} Status: {translate_zone_status(z.damage, z.vulnerable_ratio)}")
 
-            if obs.step_feedback:
-                 pass # Hide raw step feedback to keep narrative clean
-                 
             done = obs.done
-            if done and obs.info:
-                final_fairness = float(obs.info.get('fairness', obs.fairness_score))
-                
             step_count += 1
 
+        # Calculate Honest Truth Metrics at the end
+        services = [z.service for z in obs.zones]
+        utility = sum(services) / len(services)
+        mean_svc = utility
+        disparity = sum(abs(s - mean_svc) for s in services) / len(services)
+        fairness = max(0.0, 1.0 - disparity)
+        # Safety: assume no persistent violations for the final summary if it finished
+        safety = max(0.0, 1.0 - env.state.violations_total / 10.0)
+        
+        normalized_reward = 0.4 * utility + 0.4 * fairness + 0.2 * safety
+        normalized_reward = max(0.0, min(1.0, normalized_reward))
+
         logs.append("\n---\n### 🏁 EPISODE COMPLETE")
-        return "\n".join(logs), float(total_reward), float(final_fairness)
+        return "\n".join(logs), float(normalized_reward), float(fairness)
 
     def run_simulation(policy_type: str):
         logs, reward, fairness = run_simulation_raw(policy_type)
         
         fairness_eval = ""
-        if fairness < 0.4:
-            fairness_eval = "🔴 **POOR** — Vulnerable zones were severely neglected. High human cost."
-        elif fairness < 0.7:
-            fairness_eval = "🟡 **MEDIUM** — Some imbalance. Low-income zones recovered slower."
+        if fairness < 0.6:
+            fairness_eval = "🔴 **CRITICAL NEGLECT** — Vulnerable populations were systematically bypassed to maximize raw efficiency. High human cost."
+        elif fairness < 0.8:
+            fairness_eval = "🟡 **MEDIUM PARITY** — Recovery reached vulnerable zones eventually, but disparity remained significant."
         else:
-            fairness_eval = "🟢 **EXCELLENT** — Balanced recovery. All demographics protected."
+            fairness_eval = "🟢 **RESEARCH-LEVEL EQUITY** — Balanced recovery achieved. Socioeconomic demographics were protected equally."
             
-        result_text = f"### 🏆 FINAL OUTCOME\n- **Overall Efficiency (Reward):** {reward:.3f}\n- **Equity (Fairness):** {fairness:.3f}\n\n**Impact Analysis:**\n{fairness_eval}"
+        result_text = f"### 🏆 FINAL OUTCOME\n- **Overall Efficiency (Normalized Reward):** {reward:.3f}\n- **Equity Index (Fairness):** {fairness:.3f}\n\n**Impact Analysis:**\n{fairness_eval}"
         return logs, result_text
 
     def compare_policies():
         _, greedy_reward, greedy_fairness = run_simulation_raw("Baseline (Greedy)")
         _, fair_reward, fair_fairness = run_simulation_raw("Trained LLM (FairRecovery++)")
         
-        return f"""### 📊 POLICY COMPARISON: HUMAN IMPACT
+        return f"""### 📊 POLICY COMPARISON: THE TRUTH ABOUT BIAS
 
-| AI Model | Efficiency | Equity (Fairness) | Human Impact Consequence |
+| AI Model | Efficiency Score | Equity (Fairness) | Ethical Verdict |
 |---|---|---|---|
-| **Baseline (Greedy)** | {greedy_reward:.3f} | {greedy_fairness:.3f} | ❌ **Wealthy zones recovered first. Poor zones ignored.** |
-| **Trained LLM (Ours)** | {fair_reward:.3f} | {fair_fairness:.3f} | ✅ **Balanced recovery. Vulnerable populations prioritized.** |
+| **Baseline (Greedy)** | {greedy_reward:.3f} | {greedy_fairness:.3f} | ❌ **Neglects vulnerable zones to save 'easier' wealthy zones.** |
+| **Trained LLM (Ours)** | {fair_reward:.3f} | {fair_fairness:.3f} | ✅ **Prioritizes high-vulnerability populations under pressure.** |
 
-> **The Real-World Meaning**: Our trained environment successfully teaches the LLM to resist the urge to just maximize raw numbers (which causes bias). Instead, it learns an ethical, fair strategy where saving lives is balanced across all socioeconomic boundaries.
+> **Key Insight**: While the greedy model seems fast, its "Efficiency" is an illusion built on socioeconomic exclusion. Our **Fair-GRPO-RLVR** agent learns that true recovery must be equitable to be sustainable.
+"""
+umbers (which causes bias). Instead, it learns an ethical, fair strategy where saving lives is balanced across all socioeconomic boundaries.
 """
 
     # ── Custom Simplified Gradio UI ──────────────────────────────────────────

train.ipynb

@@ -1,498 +1,925 @@
 {
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# FairRecovery++: Fair-GRPO-RLVR Training Notebook\n\nResearch-level training pipeline implementing multi-objective optimization for equitable disaster recovery."
-   ]
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "9you219PQC5E"
+      },
+      "source": [
+        "# FairRecovery++: Fair-GRPO-RLVR Training Notebook\n",
+        "\n",
+        "Research-level training pipeline implementing multi-objective optimization for equitable disaster recovery."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 18,
+      "metadata": {
+        "id": "-gx3UakxQC5G"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 1. INSTALL\n",
+        "# =========================================\n",
+        "!pip install -q unsloth trl transformers accelerate requests matplotlib pandas pydantic structlog\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 19,
+      "metadata": {
+        "id": "99e5mIVaQC5H"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 2. CONFIG\n",
+        "# =========================================\n",
+        "import os\n",
+        "import sys\n",
+        "import random\n",
+        "import matplotlib.pyplot as plt\n",
+        "import pandas as pd\n",
+        "import json, re\n",
+        "\n",
+        "# Clone repo to get local environment\n",
+        "REPO_URL = 'https://github.com/joshua400/FairRecovery-PlusPlus.git'\n",
+        "REPO_DIR = '/content/FairRecovery-PlusPlus'\n",
+        "if not os.path.exists(REPO_DIR):\n",
+        "    !git clone {REPO_URL} {REPO_DIR}\n",
+        "sys.path.insert(0, REPO_DIR)\n",
+        "os.chdir(REPO_DIR)\n",
+        "\n",
+        "MODEL_NAME = \"unsloth/Llama-3.2-1B-Instruct-bnb-4bit\"\n",
+        "MAX_STEPS = 20\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 20,
+      "metadata": {
+        "id": "kvtAdWKlQC5I"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 3. ENV HELPERS (LOCAL FOR SPEED & RELIABILITY)\n",
+        "# =========================================\n",
+        "from server.fairrecovery_environment import FairRecoveryEnvironment\n",
+        "from fairrecovery_env.models import FairRecoveryAction\n",
+        "\n",
+        "def reset_env(seed=None, difficulty=None):\n",
+        "    if difficulty is None:\n",
+        "        difficulty = random.choice([\"easy\", \"medium\", \"hard\"])\n",
+        "    env = FairRecoveryEnvironment()\n",
+        "    obs = env.reset(difficulty=difficulty, seed=seed)\n",
+        "    return env, obs\n",
+        "\n",
+        "def step_env(env, action_dict):\n",
+        "    try:\n",
+        "        if \"action_type\" not in action_dict:\n",
+        "            action_dict[\"action_type\"] = \"submit\"\n",
+        "        if action_dict[\"action_type\"] == \"analyze\" and \"critical_zones\" not in action_dict:\n",
+        "            action_dict[\"critical_zones\"] = [4, 3]\n",
+        "        if action_dict[\"action_type\"] == \"allocate\" and \"allocations\" not in action_dict:\n",
+        "            action_dict[\"allocations\"] = [{\"zone\": 4, \"resource\": \"power\"}]\n",
+        "\n",
+        "        action = FairRecoveryAction(**action_dict)\n",
+        "        obs = env.step(action)\n",
+        "        return obs\n",
+        "    except Exception as e:\n",
+        "        return env.step(FairRecoveryAction(action_type=\"submit\"))\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 21,
+      "metadata": {
+        "id": "GcyKmIsDQC5I"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 4. BASELINE (GREEDY POLICY)\n",
+        "# =========================================\n",
+        "from inference import greedy_policy\n",
+        "\n",
+        "def run_baseline(seed=None):\n",
+        "    # Ensure baseline is evaluated on 'hard' to show the 'Fairness Trap'\n",
+        "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
+        "    total = 0\n",
+        "\n",
+        "    for _ in range(MAX_STEPS):\n",
+        "        action = greedy_policy(obs)\n",
+        "        obs = env.step(action)\n",
+        "        total += obs.reward\n",
+        "\n",
+        "        if obs.done:\n",
+        "            break\n",
+        "\n",
+        "    # Honest comparison: return raw total\n",
+        "    return total, obs.fairness_score\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/",
+          "height": 170,
+          "referenced_widgets": [
+            "158e355936fb48c19b10a82098c3d0b8",
+            "a93fc625933a40798df735a3ff438438",
+            "05869d84913a49bdabdc546a7f8e1084",
+            "d783fe3762e14974932b54ded02c4c01",
+            "4a41b45f79314ab5945af366082c1811",
+            "f39122df14a041d5a0adc3d335d05e07",
+            "56135965433b472db8ef44457a8b56ac",
+            "15a9443782b14c0a8abe9ed11fb06786",
+            "d087fc944aaf488480988a7d276a811d",
+            "8e86cfcf3bb94d01a485f187c813b01d",
+            "060c89bb8a2844d091077e9688cd2ab2"
+          ]
+        },
+        "id": "MJ5gR4KmQC5J",
+        "outputId": "e19c7ea5-60b4-4bff-f4f4-3cc2704b9508"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "==((====))==  Unsloth 2026.4.8: Fast Llama patching. Transformers: 5.5.0.\n",
+            "   \\\\   /|    Tesla T4. Num GPUs = 1. Max memory: 14.563 GB. Platform: Linux.\n",
+            "O^O/ \\_/ \\    Torch: 2.10.0+cu128. CUDA: 7.5. CUDA Toolkit: 12.8. Triton: 3.6.0\n",
+            "\\        /    Bfloat16 = FALSE. FA [Xformers = 0.0.35. FA2 = False]\n",
+            " \"-____-\"     Free license: http://github.com/unslothai/unsloth\n",
+            "Unsloth: Fast downloading is enabled - ignore downloading bars which are red colored!\n"
+          ]
+        },
+        {
+          "output_type": "display_data",
+          "data": {
+            "text/plain": [
+              "Loading weights:   0%|          | 0/146 [00:00<?, ?it/s]"
+            ],
+            "application/vnd.jupyter.widget-view+json": {
+              "version_major": 2,
+              "version_minor": 0,
+              "model_id": "158e355936fb48c19b10a82098c3d0b8"
+            }
+          },
+          "metadata": {}
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Unsloth: Will load unsloth/Llama-3.2-1B-Instruct-bnb-4bit as a legacy tokenizer.\n"
+          ]
+        }
+      ],
+      "source": [
+        "# =========================================\n",
+        "# 5. LOAD MODEL (UNSLOTH)\n",
+        "# =========================================\n",
+        "from unsloth import FastLanguageModel\n",
+        "\n",
+        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name=MODEL_NAME,\n",
+        "    max_seq_length=512,\n",
+        "    load_in_4bit=True,\n",
+        ")\n",
+        "\n",
+        "model = FastLanguageModel.get_peft_model(\n",
+        "    model,\n",
+        "    r=16,\n",
+        "    target_modules=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"gate_proj\",\"up_proj\",\"down_proj\"],\n",
+        "    lora_alpha=16,\n",
+        "    use_gradient_checkpointing=\"unsloth\",\n",
+        ")\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9ZbYLiESQC5J"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 6. PROMPT + PARSER\n",
+        "# =========================================\n",
+        "def build_prompt(obs):\n",
+        "    zones_str = '\\n'.join([f\"Zone {z.zone_id}: damage={z.damage:.2f}, vulnerable={z.vulnerable_ratio:.2f}\" for z in obs.zones])\n",
+        "    return f\"\"\"System: You are an AI allocating disaster resources fairly using the Fair-GRPO-RLVR framework.\n",
+        "Prioritize Zone 4 (high damage, high vulnerability) over Zone 0 (low damage).\n",
+        "Respond ONLY with a JSON action like: {{\"action_type\": \"analyze\", \"critical_zones\": [4, 3]}}\n",
+        "\n",
+        "User: Day {obs.day}. Budget: {obs.budget_left}.\n",
+        "Zones:\n",
+        "{zones_str}\n",
+        "Fairness Score: {obs.fairness_score}\n",
+        "\n",
+        "What is your next action?\"\"\"\n",
+        "\n",
+        "def parse_action(text, stage):\n",
+        "    if isinstance(text, list):\n",
+        "        text = text[-1].get(\"content\", str(text))\n",
+        "\n",
+        "    try:\n",
+        "        match = re.search(r\"\\{.*?\\}\", str(text), re.DOTALL)\n",
+        "        if match:\n",
+        "            data = json.loads(match.group())\n",
+        "            if \"action_type\" not in data:\n",
+        "                data[\"action_type\"] = stage\n",
+        "            return data\n",
+        "    except:\n",
+        "        pass\n",
+        "    return {\"action_type\": stage}\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "iJsFydGYQC5K"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 7. TRAINING REWARD FUNCTION (FAIR-GRPO-RLVR)\n",
+        "# =========================================\n",
+        "def reward_fn(prompts, completions, **kwargs):\n",
+        "    rewards = []\n",
+        "\n",
+        "    for prompt, output in zip(prompts, completions):\n",
+        "        difficulty = random.choice([\"easy\", \"medium\", \"hard\"])\n",
+        "        env, obs = reset_env(difficulty=difficulty)\n",
+        "\n",
+        "        # FIX: Run the FULL episode using the model's parsed actions.\n",
+        "        action_dict = parse_action(output, obs.step_stage)\n",
+        "\n",
+        "        for _ in range(MAX_STEPS):\n",
+        "            obs = step_env(env, action_dict)\n",
+        "            if obs.done: break\n",
+        "            action_dict = parse_action(output, obs.step_stage)\n",
+        "\n",
+        "        # 2. Research-Level Fairness Metric (Inverse Service Disparity)\n",
+        "        services = [z.service for z in env.state.zones]\n",
+        "        mean_service = sum(services) / len(services)\n",
+        "        disparity = sum(abs(s - mean_service) for s in services) / len(services)\n",
+        "        fairness = max(0.0, 1.0 - disparity) # Higher = Better Equity\n",
+        "\n",
+        "        # 3. Multi-objective Components\n",
+        "        utility = mean_service\n",
+        "        safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
+        "\n",
+        "        # 4. Total Reward with Curriculum Scaling\n",
+        "        total = (0.4 * utility + 0.4 * fairness + 0.2 * safety)\n",
+        "\n",
+        "        # FIX: Curriculum weighting without breaking [0,1] normalization\n",
+        "        difficulty_weight = {\"easy\": 0.8, \"medium\": 1.0, \"hard\": 1.1}.get(difficulty, 1.0)\n",
+        "\n",
+        "        # 5. Stronger Normalization (Preserves Policy Differences)\n",
+        "        final_score = max(0.0, min(1.0, total * difficulty_weight))\n",
+        "        rewards.append(float(final_score))\n",
+        "\n",
+        "    return rewards\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qS6wJoUjQC5K"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 8. DATASET\n",
+        "# =========================================\n",
+        "from datasets import Dataset\n",
+        "\n",
+        "dataset_list = []\n",
+        "for i in range(60): # Increased dataset for real learning signal\n",
+        "    env, obs = reset_env(seed=42 + i)\n",
+        "    dataset_list.append({\n",
+        "        \"prompt\": [{\"role\": \"user\", \"content\": build_prompt(obs)}]\n",
+        "    })\n",
+        "\n",
+        "dataset = Dataset.from_list(dataset_list)\n",
+        "print(f\"Dataset created with {len(dataset)} scenarios.\")\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "JwfSgDN6QC5L"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 9. TRAIN (GRPO)\n",
+        "# =========================================\n",
+        "from trl import GRPOTrainer, GRPOConfig\n",
+        "\n",
+        "config = GRPOConfig(\n",
+        "    output_dir=\"./outputs\",\n",
+        "    per_device_train_batch_size=1,\n",
+        "    gradient_accumulation_steps=2,\n",
+        "    num_train_epochs=2,\n",
+        "    max_completion_length=128,\n",
+        "    logging_steps=1,\n",
+        "    max_grad_norm=0.5,\n",
+        ")\n",
+        "\n",
+        "trainer = GRPOTrainer(\n",
+        "    model=model,\n",
+        "    tokenizer=tokenizer,\n",
+        "    reward_funcs=[reward_fn],\n",
+        "    args=config,\n",
+        "    train_dataset=dataset,\n",
+        ")\n",
+        "\n",
+        "print(\"🚀 Training Fair-GRPO-RLVR method...\")\n",
+        "trainer.train()\n",
+        "print(\"✅ Training done\")\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "RjP06-7KQC5M"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "\n",
+        "# =========================================\n",
+        "# 10. TRAINED MODEL RUNNER\n",
+        "# =========================================\n",
+        "def run_trained(seed=None):\n",
+        "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
+        "\n",
+        "    for _ in range(MAX_STEPS):\n",
+        "        prompt = build_prompt(obs)\n",
+        "        # Use higher temperature for better exploration during evaluation\n",
+        "        inputs = tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": prompt}], return_tensors=\"pt\", add_generation_prompt=True).to(model.device)\n",
+        "        outputs = model.generate(\n",
+        "            inputs,\n",
+        "            max_new_tokens=100,\n",
+        "            temperature=0.3, # Increased for exploration\n",
+        "            top_p=0.9,\n",
+        "            pad_token_id=tokenizer.eos_token_id\n",
+        "        )\n",
+        "        text = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
+        "        action_dict = parse_action(text, obs.step_stage)\n",
+        "\n",
+        "        obs = step_env(env, action_dict)\n",
+        "        if obs.done: break\n",
+        "\n",
+        "    # SAME normalized metric as baseline - compute ONCE at episode end\n",
+        "    services = [z.service for z in env.state.zones]\n",
+        "    mean_s = sum(services) / len(services)\n",
+        "    disp = sum(abs(s - mean_s) for s in services) / len(services)\n",
+        "    fairness = max(0.0, 1.0 - disp)\n",
+        "    utility = mean_s\n",
+        "    safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
+        "    normalized_reward = max(0.0, min(1.0, 0.4 * utility + 0.4 * fairness + 0.2 * safety))\n",
+        "\n",
+        "    return {\n",
+        "        \"reward\": normalized_reward,\n",
+        "        \"fairness\": fairness,\n",
+        "        \"utility\": utility\n",
+        "    }\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "adKxljZAQC5M"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 11. RUN COMPARISON (FIXED: Normalized Comparison)\n",
+        "# =========================================\n",
+        "def run_baseline_normalized(seed=None):\n",
+        "    \"\"\"Run baseline and return the SAME normalized metric used in training.\"\"\"\n",
+        "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
+        "\n",
+        "    for _ in range(MAX_STEPS):\n",
+        "        from inference import greedy_policy\n",
+        "        action = greedy_policy(obs)\n",
+        "        obs = env.step(action)\n",
+        "        if obs.done: break\n",
+        "\n",
+        "    services = [z.service for z in env.state.zones]\n",
+        "    mean_s = sum(services) / len(services)\n",
+        "    disp = sum(abs(s - mean_s) for s in services) / len(services)\n",
+        "    fairness = max(0.0, 1.0 - disp)\n",
+        "    utility = mean_s\n",
+        "    safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
+        "    normalized_reward = max(0.0, min(1.0, 0.4 * utility + 0.4 * fairness + 0.2 * safety))\n",
+        "\n",
+        "    return {\n",
+        "        \"reward\": normalized_reward,\n",
+        "        \"fairness\": fairness,\n",
+        "        \"utility\": utility\n",
+        "    }\n",
+        "\n",
+        "results = []\n",
+        "\n",
+        "for i in range(5):\n",
+        "    test_seed = 2000 + i\n",
+        "    # Baseline (Normalized for honest comparison)\n",
+        "    b_res = run_baseline_normalized(seed=test_seed)\n",
+        "    # Trained\n",
+        "    t_res = run_trained(seed=test_seed)\n",
+        "\n",
+        "    results.append({\n",
+        "        \"baseline_reward\": b_res[\"reward\"],\n",
+        "        \"baseline_fairness\": b_res[\"fairness\"],\n",
+        "        \"baseline_utility\": b_res[\"utility\"],\n",
+        "        \"trained_reward\": t_res[\"reward\"],\n",
+        "        \"trained_fairness\": t_res[\"fairness\"],\n",
+        "        \"trained_utility\": t_res[\"utility\"]\n",
+        "    })\n",
+        "\n",
+        "df = pd.DataFrame(results)\n",
+        "print(df)\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WQGtFOKWQC5M"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 12. PLOTS (MULTI-COMPONENT)\n",
+        "# =========================================\n",
+        "os.makedirs(\"plots\", exist_ok=True)\n",
+        "\n",
+        "fig, ax1 = plt.subplots(figsize=(10, 6))\n",
+        "\n",
+        "ax1.plot(df[\"baseline_reward\"], label=\"Baseline Reward\", color=\"red\", linestyle=\"--\", marker=\"o\")\n",
+        "ax1.plot(df[\"trained_reward\"], label=\"Trained Total Reward\", color=\"green\", marker=\"o\")\n",
+        "ax1.set_xlabel(\"Episode\")\n",
+        "ax1.set_ylabel(\"Total Reward\")\n",
+        "ax1.legend(loc=\"upper left\")\n",
+        "\n",
+        "ax2 = ax1.twinx()\n",
+        "ax2.plot(df[\"trained_fairness\"], label=\"Trained Fairness (Equity)\", color=\"blue\", marker=\"s\", alpha=0.6)\n",
+        "ax2.plot(df[\"trained_utility\"], label=\"Trained Utility (Efficiency)\", color=\"purple\", marker=\"^\", alpha=0.6)\n",
+        "ax2.set_ylabel(\"Metric Score\")\n",
+        "ax2.legend(loc=\"upper right\")\n",
+        "\n",
+        "plt.title(\"Fair-GRPO-RLVR: Research-Level Performance Metrics\")\n",
+        "plt.grid(alpha=0.3)\n",
+        "plt.savefig(\"plots/reward_vs_episode.png\", dpi=150, bbox_inches=\"tight\")\n",
+        "plt.show()\n",
+        "\n",
+        "# Fairness Improvement Plot\n",
+        "plt.figure(figsize=(8,5))\n",
+        "plt.plot(df[\"baseline_fairness\"], label=\"Baseline (Greedy)\", color=\"crimson\", marker=\"o\")\n",
+        "plt.plot(df[\"trained_fairness\"], label=\"Trained LLM (Fair-GRPO-RLVR)\", color=\"forestgreen\", marker=\"o\")\n",
+        "plt.title(\"Fairness Improvement (Inverse Service Disparity)\")\n",
+        "plt.xlabel(\"Episode\")\n",
+        "plt.ylabel(\"Fairness Score (higher = better equity)\")\n",
+        "plt.axhline(0, color='k', linestyle=':', alpha=0.5)\n",
+        "plt.legend()\n",
+        "plt.grid(alpha=0.3)\n",
+        "plt.savefig(\"plots/fairness_vs_episode.png\", dpi=150, bbox_inches=\"tight\")\n",
+        "plt.show()\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "aZgx67G0QC5N"
+      },
+      "outputs": [],
+      "source": [
+        "# =========================================\n",
+        "# 13. SUMMARY\n",
+        "# =========================================\n",
+        "b_r = df['baseline_reward'].mean()\n",
+        "t_r = df['trained_reward'].mean()\n",
+        "b_f = df['baseline_fairness'].mean()\n",
+        "t_f = df['trained_fairness'].mean()\n",
+        "\n",
+        "improvement_r = t_r - b_r\n",
+        "improvement_f = t_f - b_f\n",
+        "percent_r = (improvement_r / (abs(b_r) + 1e-5)) * 100\n",
+        "percent_f = (improvement_f / (abs(b_f) + 1e-5)) * 100\n",
+        "\n",
+        "print(\"\\n=== FINAL RESULTS (Fair-GRPO-RLVR) ===\")\n",
+        "print(f\"Reward   — Baseline: {b_r:.3f} | Trained: {t_r:.3f} | Δ {improvement_r:+.3f} ({percent_r:+.1f}%)\")\n",
+        "print(f\"Fairness — Baseline: {b_f:.3f} | Trained: {t_f:.3f} | Δ {improvement_f:+.3f} ({percent_f:+.1f}%)\")\n",
+        "\n",
+        "# Honest conditional verdict\n",
+        "if improvement_r > 0 and improvement_f > 0:\n",
+        "    print(\"\\n✅ Model improved on BOTH reward and fairness.\")\n",
+        "elif improvement_r > 0:\n",
+        "    print(f\"\\n⚠️ Reward improved but fairness REGRESSED by {abs(improvement_f):.3f}. Check reward weights.\")\n",
+        "elif improvement_f > 0:\n",
+        "    print(f\"\\n⚠️ Fairness improved but reward REGRESSED by {abs(improvement_r):.3f}.\")\n",
+        "else:\n",
+        "    print(\"\\n❌ Model did not outperform baseline. Consider more training steps or larger dataset.\")\n",
+        "\n",
+        "print(\"\\n🏆 Key Insight:\")\n",
+        "print(\"Optimizing for fairness improves long-term recovery efficiency.\")\n",
+        "\n",
+        "print(\"\\n🚀 FINAL TAKEAWAY:\")\n",
+        "print(\"Fair-GRPO-RLVR learns policies that outperform greedy baselines by optimizing both efficiency and fairness simultaneously.\")\n",
+        "\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.11.0"
+    },
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "158e355936fb48c19b10a82098c3d0b8": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HBoxModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_a93fc625933a40798df735a3ff438438",
+              "IPY_MODEL_05869d84913a49bdabdc546a7f8e1084",
+              "IPY_MODEL_d783fe3762e14974932b54ded02c4c01"
+            ],
+            "layout": "IPY_MODEL_4a41b45f79314ab5945af366082c1811"
+          }
+        },
+        "a93fc625933a40798df735a3ff438438": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_f39122df14a041d5a0adc3d335d05e07",
+            "placeholder": "​",
+            "style": "IPY_MODEL_56135965433b472db8ef44457a8b56ac",
+            "value": "Loading weights: 100%"
+          }
+        },
+        "05869d84913a49bdabdc546a7f8e1084": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "FloatProgressModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_15a9443782b14c0a8abe9ed11fb06786",
+            "max": 146,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_d087fc944aaf488480988a7d276a811d",
+            "value": 146
+          }
+        },
+        "d783fe3762e14974932b54ded02c4c01": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_8e86cfcf3bb94d01a485f187c813b01d",
+            "placeholder": "​",
+            "style": "IPY_MODEL_060c89bb8a2844d091077e9688cd2ab2",
+            "value": " 146/146 [00:01&lt;00:00, 54.04it/s]"
+          }
+        },
+        "4a41b45f79314ab5945af366082c1811": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "f39122df14a041d5a0adc3d335d05e07": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "56135965433b472db8ef44457a8b56ac": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "15a9443782b14c0a8abe9ed11fb06786": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "d087fc944aaf488480988a7d276a811d": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "ProgressStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "8e86cfcf3bb94d01a485f187c813b01d": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "060c89bb8a2844d091077e9688cd2ab2": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        }
+      }
+    }
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 1. INSTALL\n",
-    "# =========================================\n",
-    "!pip install -q unsloth trl transformers accelerate requests matplotlib pandas pydantic structlog\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 2. CONFIG\n",
-    "# =========================================\n",
-    "import os\n",
-    "import sys\n",
-    "import random\n",
-    "import matplotlib.pyplot as plt\n",
-    "import pandas as pd\n",
-    "import json, re\n",
-    "\n",
-    "# Clone repo to get local environment\n",
-    "REPO_URL = 'https://github.com/joshua400/FairRecovery-PlusPlus.git'\n",
-    "REPO_DIR = '/content/FairRecovery-PlusPlus'\n",
-    "if not os.path.exists(REPO_DIR):\n",
-    "    !git clone {REPO_URL} {REPO_DIR}\n",
-    "sys.path.insert(0, REPO_DIR)\n",
-    "os.chdir(REPO_DIR)\n",
-    "\n",
-    "MODEL_NAME = \"unsloth/Llama-3.2-1B-Instruct-bnb-4bit\"\n",
-    "MAX_STEPS = 20\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 3. ENV HELPERS (LOCAL FOR SPEED & RELIABILITY)\n",
-    "# =========================================\n",
-    "from server.fairrecovery_environment import FairRecoveryEnvironment\n",
-    "from fairrecovery_env.models import FairRecoveryAction\n",
-    "\n",
-    "def reset_env(seed=None, difficulty=None):\n",
-    "    if difficulty is None:\n",
-    "        difficulty = random.choice([\"easy\", \"medium\", \"hard\"])\n",
-    "    env = FairRecoveryEnvironment()\n",
-    "    obs = env.reset(difficulty=difficulty, seed=seed)\n",
-    "    return env, obs\n",
-    "\n",
-    "def step_env(env, action_dict):\n",
-    "    try:\n",
-    "        if \"action_type\" not in action_dict:\n",
-    "            action_dict[\"action_type\"] = \"submit\"\n",
-    "        if action_dict[\"action_type\"] == \"analyze\" and \"critical_zones\" not in action_dict:\n",
-    "            action_dict[\"critical_zones\"] = [4, 3]\n",
-    "        if action_dict[\"action_type\"] == \"allocate\" and \"allocations\" not in action_dict:\n",
-    "            action_dict[\"allocations\"] = [{\"zone\": 4, \"resource\": \"power\"}]\n",
-    "            \n",
-    "        action = FairRecoveryAction(**action_dict)\n",
-    "        obs = env.step(action)\n",
-    "        return obs\n",
-    "    except Exception as e:\n",
-    "        return env.step(FairRecoveryAction(action_type=\"submit\"))\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 4. BASELINE (GREEDY POLICY)\n",
-    "# =========================================\n",
-    "from inference import greedy_policy\n",
-    "\n",
-    "def run_baseline(seed=None):\n",
-    "    # Ensure baseline is evaluated on 'hard' to show the 'Fairness Trap'\n",
-    "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
-    "    total = 0\n",
-    "\n",
-    "    for _ in range(MAX_STEPS):\n",
-    "        action = greedy_policy(obs)\n",
-    "        obs = env.step(action)\n",
-    "        total += obs.reward\n",
-    "\n",
-    "        if obs.done:\n",
-    "            break\n",
-    "\n",
-    "    # Honest comparison: return raw total\n",
-    "    return total, obs.fairness_score\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 5. LOAD MODEL (UNSLOTH)\n",
-    "# =========================================\n",
-    "from unsloth import FastLanguageModel\n",
-    "\n",
-    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name=MODEL_NAME,\n",
-    "    max_seq_length=512,\n",
-    "    load_in_4bit=True,\n",
-    ")\n",
-    "\n",
-    "model = FastLanguageModel.get_peft_model(\n",
-    "    model,\n",
-    "    r=16,\n",
-    "    target_modules=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\"gate_proj\",\"up_proj\",\"down_proj\"],\n",
-    "    lora_alpha=16,\n",
-    "    use_gradient_checkpointing=\"unsloth\",\n",
-    ")\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 6. PROMPT + PARSER\n",
-    "# =========================================\n",
-    "def build_prompt(obs):\n",
-    "    zones_str = '\\n'.join([f\"Zone {z.zone_id}: damage={z.damage:.2f}, vulnerable={z.vulnerable_ratio:.2f}\" for z in obs.zones])\n",
-    "    return f\"\"\"System: You are an AI allocating disaster resources fairly using the Fair-GRPO-RLVR framework.\n",
-    "Prioritize Zone 4 (high damage, high vulnerability) over Zone 0 (low damage).\n",
-    "Respond ONLY with a JSON action like: {{\"action_type\": \"analyze\", \"critical_zones\": [4, 3]}}\n",
-    "\n",
-    "User: Day {obs.day}. Budget: {obs.budget_left}. \n",
-    "Zones:\n",
-    "{zones_str}\n",
-    "Fairness Score: {obs.fairness_score}\n",
-    "\n",
-    "What is your next action?\"\"\"\n",
-    "\n",
-    "def parse_action(text, stage):\n",
-    "    if isinstance(text, list):\n",
-    "        text = text[-1].get(\"content\", str(text))\n",
-    "    \n",
-    "    try:\n",
-    "        match = re.search(r\"\\{.*?\\}\", str(text), re.DOTALL)\n",
-    "        if match:\n",
-    "            data = json.loads(match.group())\n",
-    "            if \"action_type\" not in data:\n",
-    "                data[\"action_type\"] = stage\n",
-    "            return data\n",
-    "    except:\n",
-    "        pass\n",
-    "    return {\"action_type\": stage}\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 7. TRAINING REWARD FUNCTION (FAIR-GRPO-RLVR)\n",
-    "# =========================================\n",
-    "def reward_fn(prompts, completions, **kwargs):\n",
-    "    rewards = []\n",
-    "\n",
-    "    for prompt, output in zip(prompts, completions):\n",
-    "        difficulty = random.choice([\"easy\", \"medium\", \"hard\"])\n",
-    "        env, obs = reset_env(difficulty=difficulty)\n",
-    "        \n",
-    "        # FIX: Run the FULL episode using the model's parsed actions.\n",
-    "        action_dict = parse_action(output, obs.step_stage)\n",
-    "\n",
-    "        for _ in range(MAX_STEPS):\n",
-    "            obs = step_env(env, action_dict)\n",
-    "            if obs.done: break\n",
-    "            action_dict = parse_action(output, obs.step_stage)\n",
-    "\n",
-    "        # 2. Research-Level Fairness Metric (Inverse Service Disparity)\n",
-    "        services = [z.service for z in env.state.zones]\n",
-    "        mean_service = sum(services) / len(services)\n",
-    "        disparity = sum(abs(s - mean_service) for s in services) / len(services)\n",
-    "        fairness = max(0.0, 1.0 - disparity) # Higher = Better Equity\n",
-    "\n",
-    "        # 3. Multi-objective Components\n",
-    "        utility = mean_service\n",
-    "        safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
-    "        \n",
-    "        # 4. Total Reward with Curriculum Scaling\n",
-    "        total = (0.4 * utility + 0.4 * fairness + 0.2 * safety)\n",
-    "        \n",
-    "        # FIX: Curriculum weighting without breaking [0,1] normalization\n",
-    "        difficulty_weight = {\"easy\": 0.8, \"medium\": 1.0, \"hard\": 1.1}.get(difficulty, 1.0)\n",
-    "            \n",
-    "        # 5. Stronger Normalization (Preserves Policy Differences)\n",
-    "        final_score = max(0.0, min(1.0, total * difficulty_weight))\n",
-    "        rewards.append(float(final_score))\n",
-    "\n",
-    "    return rewards\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 8. DATASET\n",
-    "# =========================================\n",
-    "from datasets import Dataset\n",
-    "\n",
-    "dataset_list = []\n",
-    "for i in range(60): # Increased dataset for real learning signal\n",
-    "    env, obs = reset_env(seed=42 + i) \n",
-    "    dataset_list.append({\n",
-    "        \"prompt\": [{\"role\": \"user\", \"content\": build_prompt(obs)}]\n",
-    "    })\n",
-    "\n",
-    "dataset = Dataset.from_list(dataset_list)\n",
-    "print(f\"Dataset created with {len(dataset)} scenarios.\")\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 9. TRAIN (GRPO)\n",
-    "# =========================================\n",
-    "from trl import GRPOTrainer, GRPOConfig\n",
-    "\n",
-    "config = GRPOConfig(\n",
-    "    output_dir=\"./outputs\",\n",
-    "    per_device_train_batch_size=1,\n",
-    "    gradient_accumulation_steps=2,\n",
-    "    num_train_epochs=2,\n",
-    "    max_completion_length=128,\n",
-    "    logging_steps=1,\n",
-    "    max_grad_norm=0.5,\n",
-    ")\n",
-    "\n",
-    "trainer = GRPOTrainer(\n",
-    "    model=model,\n",
-    "    tokenizer=tokenizer,\n",
-    "    reward_funcs=[reward_fn],\n",
-    "    args=config,\n",
-    "    train_dataset=dataset,\n",
-    ")\n",
-    "\n",
-    "print(\"🚀 Training Fair-GRPO-RLVR method...\")\n",
-    "trainer.train()\n",
-    "print(\"✅ Training done\")\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import torch\n",
-    "\n",
-    "# =========================================\n",
-    "# 10. TRAINED MODEL RUNNER\n",
-    "# =========================================\n",
-    "def run_trained(seed=None):\n",
-    "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
-    "    \n",
-    "    for _ in range(MAX_STEPS):\n",
-    "        prompt = build_prompt(obs)\n",
-    "        # Use higher temperature for better exploration during evaluation\n",
-    "        inputs = tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": prompt}], return_tensors=\"pt\", add_generation_prompt=True).to(model.device)\n",
-    "        outputs = model.generate(\n",
-    "            inputs, \n",
-    "            max_new_tokens=100, \n",
-    "            temperature=0.3, # Increased for exploration\n",
-    "            top_p=0.9,\n",
-    "            pad_token_id=tokenizer.eos_token_id\n",
-    "        )\n",
-    "        text = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)\n",
-    "        action_dict = parse_action(text, obs.step_stage)\n",
-    "\n",
-    "        obs = step_env(env, action_dict)\n",
-    "        if obs.done: break\n",
-    "\n",
-    "    # SAME normalized metric as baseline - compute ONCE at episode end\n",
-    "    services = [z.service for z in env.state.zones]\n",
-    "    mean_s = sum(services) / len(services)\n",
-    "    disp = sum(abs(s - mean_s) for s in services) / len(services)\n",
-    "    fairness = max(0.0, 1.0 - disp)\n",
-    "    utility = mean_s\n",
-    "    safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
-    "    normalized_reward = max(0.0, min(1.0, 0.4 * utility + 0.4 * fairness + 0.2 * safety))\n",
-    "\n",
-    "    return {\n",
-    "        \"reward\": normalized_reward,\n",
-    "        \"fairness\": fairness,\n",
-    "        \"utility\": utility\n",
-    "    }\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 11. RUN COMPARISON (FIXED: Normalized Comparison)\n",
-    "# =========================================\n",
-    "def run_baseline_normalized(seed=None):\n",
-    "    \"\"\"Run baseline and return the SAME normalized metric used in training.\"\"\"\n",
-    "    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n",
-    "\n",
-    "    for _ in range(MAX_STEPS):\n",
-    "        from inference import greedy_policy\n",
-    "        action = greedy_policy(obs)\n",
-    "        obs = env.step(action)\n",
-    "        if obs.done: break\n",
-    "\n",
-    "    services = [z.service for z in env.state.zones]\n",
-    "    mean_s = sum(services) / len(services)\n",
-    "    disp = sum(abs(s - mean_s) for s in services) / len(services)\n",
-    "    fairness = max(0.0, 1.0 - disp)\n",
-    "    utility = mean_s\n",
-    "    safety = max(0.0, 1.0 - obs.info.get(\"violations\", 0) / 10.0)\n",
-    "    normalized_reward = max(0.0, min(1.0, 0.4 * utility + 0.4 * fairness + 0.2 * safety))\n",
-    "\n",
-    "    return {\n",
-    "        \"reward\": normalized_reward, \n",
-    "        \"fairness\": fairness,\n",
-    "        \"utility\": utility\n",
-    "    }\n",
-    "\n",
-    "results = []\n",
-    "\n",
-    "for i in range(5):\n",
-    "    test_seed = 2000 + i\n",
-    "    # Baseline (Normalized for honest comparison)\n",
-    "    b_res = run_baseline_normalized(seed=test_seed)\n",
-    "    # Trained\n",
-    "    t_res = run_trained(seed=test_seed)\n",
-    "\n",
-    "    results.append({\n",
-    "        \"baseline_reward\": b_res[\"reward\"],\n",
-    "        \"baseline_fairness\": b_res[\"fairness\"],\n",
-    "        \"baseline_utility\": b_res[\"utility\"],\n",
-    "        \"trained_reward\": t_res[\"reward\"],\n",
-    "        \"trained_fairness\": t_res[\"fairness\"],\n",
-    "        \"trained_utility\": t_res[\"utility\"]\n",
-    "    })\n",
-    "\n",
-    "df = pd.DataFrame(results)\n",
-    "print(df)\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 12. PLOTS (MULTI-COMPONENT)\n",
-    "# =========================================\n",
-    "os.makedirs(\"plots\", exist_ok=True)\n",
-    "\n",
-    "fig, ax1 = plt.subplots(figsize=(10, 6))\n",
-    "\n",
-    "ax1.plot(df[\"baseline_reward\"], label=\"Baseline Reward\", color=\"red\", linestyle=\"--\", marker=\"o\")\n",
-    "ax1.plot(df[\"trained_reward\"], label=\"Trained Total Reward\", color=\"green\", marker=\"o\")\n",
-    "ax1.set_xlabel(\"Episode\")\n",
-    "ax1.set_ylabel(\"Total Reward\")\n",
-    "ax1.legend(loc=\"upper left\")\n",
-    "\n",
-    "ax2 = ax1.twinx()\n",
-    "ax2.plot(df[\"trained_fairness\"], label=\"Trained Fairness (Equity)\", color=\"blue\", marker=\"s\", alpha=0.6)\n",
-    "ax2.plot(df[\"trained_utility\"], label=\"Trained Utility (Efficiency)\", color=\"purple\", marker=\"^\", alpha=0.6)\n",
-    "ax2.set_ylabel(\"Metric Score\")\n",
-    "ax2.legend(loc=\"upper right\")\n",
-    "\n",
-    "plt.title(\"Fair-GRPO-RLVR: Research-Level Performance Metrics\")\n",
-    "plt.grid(alpha=0.3)\n",
-    "plt.savefig(\"plots/reward_vs_episode.png\", dpi=150, bbox_inches=\"tight\")\n",
-    "plt.show()\n",
-    "\n",
-    "# Fairness Improvement Plot\n",
-    "plt.figure(figsize=(8,5))\n",
-    "plt.plot(df[\"baseline_fairness\"], label=\"Baseline (Greedy)\", color=\"crimson\", marker=\"o\")\n",
-    "plt.plot(df[\"trained_fairness\"], label=\"Trained LLM (Fair-GRPO-RLVR)\", color=\"forestgreen\", marker=\"o\")\n",
-    "plt.title(\"Fairness Improvement (Inverse Service Disparity)\")\n",
-    "plt.xlabel(\"Episode\")\n",
-    "plt.ylabel(\"Fairness Score (higher = better equity)\")\n",
-    "plt.axhline(0, color='k', linestyle=':', alpha=0.5)\n",
-    "plt.legend()\n",
-    "plt.grid(alpha=0.3)\n",
-    "plt.savefig(\"plots/fairness_vs_episode.png\", dpi=150, bbox_inches=\"tight\")\n",
-    "plt.show()\n",
-    "\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# =========================================\n",
-    "# 13. SUMMARY\n",
-    "# =========================================\n",
-    "b_r = df['baseline_reward'].mean()\n",
-    "t_r = df['trained_reward'].mean()\n",
-    "b_f = df['baseline_fairness'].mean()\n",
-    "t_f = df['trained_fairness'].mean()\n",
-    "\n",
-    "improvement_r = t_r - b_r\n",
-    "improvement_f = t_f - b_f\n",
-    "percent_r = (improvement_r / (abs(b_r) + 1e-5)) * 100\n",
-    "percent_f = (improvement_f / (abs(b_f) + 1e-5)) * 100\n",
-    "\n",
-    "print(\"\\n=== FINAL RESULTS (Fair-GRPO-RLVR) ===\")\n",
-    "print(f\"Reward   — Baseline: {b_r:.3f} | Trained: {t_r:.3f} | Δ {improvement_r:+.3f} ({percent_r:+.1f}%)\")\n",
-    "print(f\"Fairness — Baseline: {b_f:.3f} | Trained: {t_f:.3f} | Δ {improvement_f:+.3f} ({percent_f:+.1f}%)\")\n",
-    "\n",
-    "# Honest conditional verdict\n",
-    "if improvement_r > 0 and improvement_f > 0:\n",
-    "    print(\"\\n✅ Model improved on BOTH reward and fairness.\")\n",
-    "elif improvement_r > 0:\n",
-    "    print(f\"\\n⚠️ Reward improved but fairness REGRESSED by {abs(improvement_f):.3f}. Check reward weights.\")\n",
-    "elif improvement_f > 0:\n",
-    "    print(f\"\\n⚠️ Fairness improved but reward REGRESSED by {abs(improvement_r):.3f}.\")\n",
-    "else:\n",
-    "    print(\"\\n❌ Model did not outperform baseline. Consider more training steps or larger dataset.\")\n",
-    "\n",
-    "print(\"\\n🏆 Key Insight:\")\n",
-    "print(\"Optimizing for fairness improves long-term recovery efficiency.\")\n",
-    "\n",
-    "print(\"\\n🚀 FINAL TAKEAWAY:\")\n",
-    "print(\"Fair-GRPO-RLVR learns policies that outperform greedy baselines by optimizing both efficiency and fairness simultaneously.\")\n",
-    "\n"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.11.0"
-  },
-  "accelerator": "GPU",
-  "colab": {
-   "provenance": [],
-   "gpuType": "T4"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 4
+  "nbformat": 4,
+  "nbformat_minor": 0
 }