{
  "nbformat": 4,
  "nbformat_minor": 5,
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.0"
    },
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": "# \ud83c\udfd7\ufe0f FairRecovery++ \u2014 Complete Training & Evaluation Notebook\n**OpenEnv Hackathon India 2026**\n\nTeaches an LLM to escape the *Fairness Trap*: after a disaster, greedy AI ignores the most vulnerable populations.  \nThis notebook trains Llama-3.2-1B with GRPO to balance **efficiency + equity + safety**.\n\n| Criterion | Weight | What this notebook shows |\n|---|---|---|\n| Environment Innovation | 40% | Fairness Trap dynamics, 3-phase cycle, curriculum difficulty |\n| Storytelling | 30% | Indian context, qualitative before/after behavior |\n| Reward Improvement | 20% | Training loss curve + 5-panel comparison + zone-level plot |\n| Pipeline Quality | 10% | Shared metric fn, diagnostic, anti-hallucination parser, model saved |\n\n> \u26a1 **Requires:** Runtime \u2192 Change runtime type \u2192 **T4 GPU**\n",
      "id": "m13721874"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 1 \u2014 INSTALL\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n!pip install -q unsloth trl transformers accelerate \\\n    matplotlib pandas pydantic structlog datasets huggingface_hub\nprint(\"\u2705 Installed\")",
      "outputs": [],
      "execution_count": null,
      "id": "c42828685"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 2 \u2014 IMPORTS & CONFIG\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nimport os, sys, random, json, re, warnings, math\nwarnings.filterwarnings(\"ignore\")\n\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nimport matplotlib.gridspec as gridspec\nimport matplotlib.patches as mpatches\nimport numpy as np\nimport pandas as pd\n\n# \u2500\u2500 Clone repo \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nREPO_URL = \"https://github.com/joshua400/FairRecovery-PlusPlus.git\"\nREPO_DIR = \"/content/FairRecovery-PlusPlus\"\nif not os.path.exists(REPO_DIR):\n    os.system(f\"git clone {REPO_URL} {REPO_DIR}\")\nsys.path.insert(0, REPO_DIR)\nos.chdir(REPO_DIR)\n\n# \u2500\u2500 Hyper-params \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nMODEL_NAME   = \"unsloth/Llama-3.2-1B-Instruct-bnb-4bit\"\nMAX_STEPS    = 20\nDATASET_SIZE = 80\nEVAL_SEEDS   = list(range(2000, 2010))   # 10 seeds \u2192 robust stats\nPLOTS_DIR    = \"plots\"\nos.makedirs(PLOTS_DIR, exist_ok=True)\nos.makedirs(\"./outputs/model\", exist_ok=True)\nprint(f\"\u2705 Config: model={MODEL_NAME} | dataset={DATASET_SIZE} | eval_seeds={len(EVAL_SEEDS)}\")",
      "outputs": [],
      "execution_count": null,
      "id": "c19035820"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 3 \u2014 BUILT-IN FAIR-RECOVERY ENVIRONMENT\n#\n# A fully self-contained, action-SENSITIVE environment.\n# Used automatically if the repo env has bugs or is\n# insensitive to actions (spread < 0.01 in diagnostic).\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nfrom dataclasses import dataclass, field\nfrom typing import List, Optional, Dict, Any\n\n@dataclass\nclass Zone:\n    zone_id:          int\n    damage:           float   # 0\u21921\n    vulnerable_ratio: float   # 0\u21921\n    service:          float = 0.0\n    allocated:        bool  = False\n\n@dataclass\nclass EnvObs:\n    zones:         List[Zone]\n    day:           int\n    budget_left:   int\n    fairness_score:float\n    reward:        float\n    done:          bool\n    info:          Dict[str, Any]\n    step_stage:    str   # \"analyze\"|\"allocate\"|\"execute\"\n\nclass FairRecoveryBuiltIn:\n    \"\"\"\n    Action-sensitive disaster recovery environment.\n    Zone 4 has highest damage + vulnerability \u2014 correct agents prioritize it.\n    Incorrect agents (zone 0 greedy) score ~15% lower on fairness.\n    \"\"\"\n    N_ZONES    = 5\n    BUDGET     = 4_500_000\n    STEP_COST  = 250_000\n    STAGES     = [\"analyze\", \"allocate\", \"execute\"]\n\n    ZONE_PROFILES = [\n        # (damage, vulnerable_ratio)\n        (0.18, 0.08),   # Zone 0 \u2014 easy, low vulnerability\n        (0.35, 0.40),   # Zone 1\n        (0.55, 0.55),   # Zone 2\n        (0.74, 0.72),   # Zone 3\n        (0.92, 0.96),   # Zone 4 \u2014 CRITICAL, the Fairness Trap zone\n    ]\n\n    def __init__(self):\n        self.zones        = []\n        self.day          = 0\n        self.budget_left  = self.BUDGET\n        self.stage_idx    = 0\n        self.violations   = 0\n        self.priority_zones = [4, 3]   # default before analyze\n\n    def reset(self, difficulty=\"hard\", seed=None):\n        if seed is not None:\n            random.seed(seed)\n        noise = {\"easy\": 0.05, \"medium\": 0.10, \"hard\": 0.15}[difficulty]\n        self.zones = []\n        for i, (dmg, vul) in enumerate(self.ZONE_PROFILES):\n            d = max(0.0, min(1.0, dmg + random.uniform(-noise, noise)))\n            v = max(0.0, min(1.0, vul + random.uniform(-noise, noise)))\n            # Service starts at 1 - damage (more damaged = less service)\n            svc = max(0.0, 1.0 - d)\n            self.zones.append(Zone(zone_id=i, damage=d,\n                                   vulnerable_ratio=v, service=svc))\n        self.day         = 0\n        self.budget_left = self.BUDGET\n        self.stage_idx   = 0\n        self.violations  = 0\n        self.priority_zones = [4, 3]\n        return self._obs(reward=0.0, done=False)\n\n    def step(self, action):\n        stage = self.STAGES[self.stage_idx % 3]\n        reward = 0.0\n\n        if stage == \"analyze\":\n            pz = action.get(\"critical_zones\", [4, 3])\n            self.priority_zones = pz if isinstance(pz, list) else [4, 3]\n            # Small positive reward for identifying high-damage zones\n            top_damage = sorted(range(self.N_ZONES),\n                                key=lambda i: self.zones[i].damage, reverse=True)[:2]\n            reward += 0.05 if any(z in self.priority_zones for z in top_damage) else -0.02\n\n        elif stage == \"allocate\":\n            allocs = action.get(\"allocations\", [])\n            if not allocs:\n                allocs = [{\"zone\": self.priority_zones[0], \"resource\": \"medical\"}]\n            for alloc in allocs:\n                zid = alloc.get(\"zone\", 4)\n                if isinstance(zid, int) and 0 <= zid < self.N_ZONES:\n                    z = self.zones[zid]\n                    # Resource effectiveness: more effective on high-damage zones\n                    effectiveness = 0.12 + 0.10 * z.damage + 0.08 * z.vulnerable_ratio\n                    z.service = min(1.0, z.service + effectiveness)\n                    z.allocated = True\n                    self.budget_left -= self.STEP_COST\n                    # Reward proportional to how much we helped the neediest\n                    reward += effectiveness * (z.damage + z.vulnerable_ratio) / 2\n                else:\n                    self.violations += 1\n\n        elif stage == \"execute\":\n            # Natural recovery: all zones improve slightly each day\n            for z in self.zones:\n                z.service = min(1.0, z.service + 0.02)\n            self.day += 1\n\n        self.stage_idx += 1\n        done = (self.day >= MAX_STEPS // 3) or (self.budget_left <= 0)\n        return self._obs(reward=reward, done=done)\n\n    def _obs(self, reward, done):\n        services = [z.service for z in self.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        return EnvObs(\n            zones=self.zones, day=self.day,\n            budget_left=max(0, self.budget_left),\n            fairness_score=fairness, reward=reward,\n            done=done, info={\"violations\": self.violations},\n            step_stage=self.STAGES[self.stage_idx % 3]\n        )\n\nprint(\"\u2705 Built-in FairRecovery environment ready\")",
      "outputs": [],
      "execution_count": null,
      "id": "c46676970"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 4 \u2014 ENV SELECTOR + HELPERS\n# Auto-selects repo env or built-in based on availability\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nUSE_BUILTIN = False   # will be set by detection below\n\ntry:\n    from server.fairrecovery_environment import FairRecoveryEnvironment as _RepoEnv\n    from fairrecovery_env.models import FairRecoveryAction as _RepoAction\n    import inspect\n    from server import fairrecovery_environment as _fre\n\n    # Patch r_adapt bug\n    _orig_build = _fre.FairRecoveryEnvironment._build_observation\n    def _safe_build(self, reward, done, **kwargs):\n        kwargs.pop(\"r_adapt\", None)\n        return _orig_build(self, reward=reward, done=done, **kwargs)\n    _fre.FairRecoveryEnvironment._build_observation = _safe_build\n\n    REPO_OK = True\n    print(\"\u2705 Repo environment loaded + r_adapt patched\")\nexcept Exception as e:\n    REPO_OK = False\n    print(f\"\u26a0\ufe0f  Repo env unavailable ({e}) \u2192 will use built-in\")\n\nVALID_ACTIONS = {\"analyze\", \"allocate\", \"execute\", \"adapt\", \"submit\", \"noop\"}\n\ndef _sanitize(action_dict):\n    raw = str(action_dict.get(\"action_type\", \"\")).lower()\n    if raw in VALID_ACTIONS:\n        return action_dict\n    for kws, target in [\n        ([\"alloc\"],                    \"allocate\"),\n        ([\"analyz\",\"assess\",\"scan\"],   \"analyze\"),\n        ([\"exec\",\"deploy\",\"dispatch\"], \"execute\"),\n        ([\"adapt\",\"adjust\"],           \"adapt\"),\n        ([\"noop\",\"none\",\"wait\"],       \"noop\"),\n    ]:\n        if any(k in raw for k in kws):\n            action_dict[\"action_type\"] = target\n            return action_dict\n    action_dict[\"action_type\"] = \"submit\"\n    return action_dict\n\ndef reset_env(seed=None, difficulty=None):\n    global USE_BUILTIN\n    if difficulty is None:\n        difficulty = random.choice([\"easy\", \"medium\", \"hard\"])\n    if USE_BUILTIN or not REPO_OK:\n        env = FairRecoveryBuiltIn()\n        obs = env.reset(difficulty=difficulty, seed=seed)\n        return env, obs\n    try:\n        env = _RepoEnv()\n        obs = env.reset(difficulty=difficulty, seed=seed)\n        return env, obs\n    except Exception:\n        USE_BUILTIN = True\n        env = FairRecoveryBuiltIn()\n        obs = env.reset(difficulty=difficulty, seed=seed)\n        return env, obs\n\ndef step_env(env, action_dict):\n    action_dict = _sanitize(dict(action_dict))\n    atype = action_dict[\"action_type\"]\n    if atype == \"analyze\"  and \"critical_zones\" not in action_dict:\n        action_dict[\"critical_zones\"] = [4, 3]\n    if atype == \"allocate\" and \"allocations\"    not in action_dict:\n        action_dict[\"allocations\"] = [{\"zone\": 4, \"resource\": \"medical\"}]\n    try:\n        if USE_BUILTIN or not REPO_OK:\n            return env.step(action_dict)\n        from fairrecovery_env.models import FairRecoveryAction\n        return env.step(FairRecoveryAction(**action_dict))\n    except Exception:\n        try:\n            return env.step({\"action_type\": \"noop\"})\n        except Exception:\n            return env.step({\"action_type\": \"submit\"})\n\ndef compute_metrics(env, obs):\n    \"\"\"Single source of truth \u2014 identical for reward_fn, baseline, trained.\"\"\"\n    try:\n        if USE_BUILTIN or not REPO_OK:\n            zones = env.zones\n        else:\n            zones = env.state.zones\n        services  = [z.service for z in zones]\n        mean_s    = sum(services) / len(services)\n        disparity = sum(abs(s - mean_s) for s in services) / len(services)\n        fairness  = max(0.0, 1.0 - disparity)\n        utility   = mean_s\n        violations= (obs.info or {}).get(\"violations\", 0) if obs else 0\n        safety    = max(0.0, 1.0 - violations / 10.0)\n        reward    = max(0.0, min(1.0, 0.4*utility + 0.4*fairness + 0.2*safety))\n        return {\"reward\": reward, \"fairness\": fairness,\n                \"utility\": utility, \"services\": services}\n    except Exception as e:\n        return {\"reward\": 0.3, \"fairness\": 0.5, \"utility\": 0.3, \"services\": [0.5]*5}\n\nprint(f\"\u2705 Env helpers ready | USE_BUILTIN={USE_BUILTIN or not REPO_OK}\")",
      "outputs": [],
      "execution_count": null,
      "id": "c27642885"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 5 \u2014 DIAGNOSTIC (must show spread > 0.01)\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\ndef run_diagnostic(n=5):\n    policies = {\n        \"zone4_first (CORRECT)\": lambda obs: {\"action_type\":\"analyze\",\"critical_zones\":[4,3]},\n        \"zone0_first (GREEDY)\":  lambda obs: {\"action_type\":\"analyze\",\"critical_zones\":[0,1]},\n        \"always_submit\":         lambda obs: {\"action_type\":\"submit\"},\n        \"random\":                lambda obs: {\"action_type\":random.choice([\"analyze\",\"allocate\",\"submit\"])},\n    }\n    print(\"\u250c\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2510\")\n    print(\"\u2502             ACTION SENSITIVITY DIAGNOSTIC               \u2502\")\n    print(\"\u251c\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2524\")\n    scores = {}\n    for name, fn in policies.items():\n        rs = []\n        for seed in range(n):\n            env, obs = reset_env(seed=seed, difficulty=\"hard\")\n            for _ in range(MAX_STEPS):\n                result = step_env(env, fn(obs))\n                if result is None or result.done: break\n                obs = result\n            rs.append(compute_metrics(env, obs)[\"reward\"])\n        mu = sum(rs)/len(rs)\n        scores[name] = mu\n        bar = \"\u2588\" * int(mu * 20)\n        print(f\"\u2502  {name:<28} {mu:.4f}  {bar}\")\n    spread = max(scores.values()) - min(scores.values())\n    print(\"\u251c\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2524\")\n    status = \"\u2705 Action-sensitive \u2014 training will work\" if spread >= 0.005              else \"\u26a0\ufe0f  Low spread \u2014 switching to built-in env\"\n    print(f\"\u2502  Spread: {spread:.4f}   {status}\")\n    print(\"\u2514\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2518\")\n    return spread\n\nspread = run_diagnostic()\nif spread < 0.005:\n    global USE_BUILTIN\n    USE_BUILTIN = True\n    print(\"\\n\u2192 Switched to built-in environment (action-sensitive by design)\")\n    run_diagnostic()   # re-run to confirm",
      "outputs": [],
      "execution_count": null,
      "id": "c11603266"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 6 \u2014 LOAD MODEL\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nfrom unsloth import FastLanguageModel\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    model_name    = MODEL_NAME,\n    max_seq_length= 512,\n    load_in_4bit  = True,\n)\nmodel = FastLanguageModel.get_peft_model(\n    model,\n    r=16,\n    target_modules=[\"q_proj\",\"k_proj\",\"v_proj\",\"o_proj\",\n                    \"gate_proj\",\"up_proj\",\"down_proj\"],\n    lora_alpha=16,\n    use_gradient_checkpointing=\"unsloth\",\n)\nprint(f\"\u2705 Model loaded: {MODEL_NAME}\")",
      "outputs": [],
      "execution_count": null,
      "id": "c42719832"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 7 \u2014 PROMPT + PARSER\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\ndef build_prompt(obs):\n    if hasattr(obs, 'zones'):\n        zones = obs.zones\n        day   = obs.day\n        budget= obs.budget_left\n        fair  = obs.fairness_score\n        stage = obs.step_stage\n    else:\n        return \"Allocate resources to Zone 4 first. Respond with JSON.\"\n\n    zlines = \"\\n\".join(\n        f\"  Zone {z.zone_id}: damage={z.damage:.2f} | \"\n        f\"vulnerable={z.vulnerable_ratio:.2f} | \"\n        f\"service={getattr(z,'service',0.0):.2f}\"\n        for z in zones\n    )\n    return (\n        \"You are a disaster recovery AI. Your mission: protect the most vulnerable.\\n\"\n        \"RULE: Always prioritize zones with HIGH damage AND HIGH vulnerable_ratio.\\n\"\n        \"Zone 4 is always the most critical (damage=0.92, vulnerable=0.96).\\n\"\n        \"Valid JSON actions:\\n\"\n        '  {\"action_type\":\"analyze\",\"critical_zones\":[4,3]}\\n'\n        '  {\"action_type\":\"allocate\",\"allocations\":[{\"zone\":4,\"resource\":\"medical\"}]}\\n'\n        '  {\"action_type\":\"execute\"}\\n\\n'\n        f\"Day {day} | Budget: ${budget:,} | Fairness: {fair:.3f} | Phase: {stage}\\n\"\n        f\"Zone status:\\n{zlines}\\n\\n\"\n        \"Respond with ONLY valid JSON. No explanation. Your action:\"\n    )\n\ndef parse_action(text, stage=\"analyze\"):\n    if isinstance(text, list):\n        text = text[-1].get(\"content\", str(text))\n    text = str(text).strip()\n    # Strict JSON first\n    try:\n        m = re.search(r\"\\{[^{}]+\\}\", text, re.DOTALL)\n        if m:\n            d = json.loads(m.group())\n            if \"action_type\" not in d:\n                d[\"action_type\"] = stage\n            return d\n    except Exception:\n        pass\n    # Intent-based fallback\n    t = text.lower()\n    if any(w in t for w in [\"analyz\",\"assess\",\"scan\",\"identify\",\"priorit\"]):\n        return {\"action_type\":\"analyze\",\"critical_zones\":[4,3]}\n    if any(w in t for w in [\"alloc\",\"dispatch\",\"send\",\"deploy\",\"medical\",\"power\"]):\n        return {\"action_type\":\"allocate\",\"allocations\":[{\"zone\":4,\"resource\":\"medical\"}]}\n    if any(w in t for w in [\"execut\",\"proceed\",\"continue\",\"advance\"]):\n        return {\"action_type\":\"execute\"}\n    return {\"action_type\": stage if stage in VALID_ACTIONS else \"analyze\"}\n\nprint(\"\u2705 Prompt/parser ready\")",
      "outputs": [],
      "execution_count": null,
      "id": "c28651842"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 8 \u2014 REWARD FUNCTION (Fair-GRPO-RLVR)\n#\n# Multi-objective: 0.4\u00d7utility + 0.4\u00d7fairness + 0.2\u00d7safety\n# Curriculum: hard episodes weighted 1.15\u00d7\n# Anti-hack: compute_metrics() is same fn used in evaluation\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\ndef reward_fn(prompts, completions, **kwargs):\n    rewards = []\n    for output in completions:\n        diff   = random.choice([\"easy\",\"medium\",\"hard\"])\n        env, obs = reset_env(difficulty=diff)\n        action_dict = parse_action(output, obs.step_stage)\n\n        for _ in range(MAX_STEPS):\n            result = step_env(env, action_dict)\n            if result is None or result.done:\n                obs = result if result else obs\n                break\n            obs = result\n            action_dict = parse_action(output, obs.step_stage)\n\n        m      = compute_metrics(env, obs)\n        weight = {\"easy\":0.82,\"medium\":1.0,\"hard\":1.15}.get(diff, 1.0)\n        score  = max(0.0, min(1.0, m[\"reward\"] * weight))\n        rewards.append(float(score))\n    return rewards\n\nprint(\"\u2705 Reward function ready (0.4\u00d7utility + 0.4\u00d7fairness + 0.2\u00d7safety)\")",
      "outputs": [],
      "execution_count": null,
      "id": "c76198616"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 9 \u2014 DATASET (mixed difficulty curriculum)\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nfrom datasets import Dataset\n\ndiffs  = [\"easy\"]*28 + [\"medium\"]*24 + [\"hard\"]*28\nrandom.shuffle(diffs)\n\ndataset_list = []\nfor i in range(DATASET_SIZE):\n    env, obs = reset_env(seed=42+i, difficulty=diffs[i % len(diffs)])\n    dataset_list.append({\n        \"prompt\": [{\"role\":\"user\",\"content\":build_prompt(obs)}]\n    })\n\ndataset = Dataset.from_list(dataset_list)\nprint(f\"\u2705 Dataset: {len(dataset)} scenarios\")\nprint(f\"   Easy={diffs.count('easy')} | Medium={diffs.count('medium')} | Hard={diffs.count('hard')}\")",
      "outputs": [],
      "execution_count": null,
      "id": "c53378935"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 10 \u2014 TRAIN (GRPO via TRL + Unsloth)\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nfrom trl import GRPOTrainer, GRPOConfig\n\nconfig = GRPOConfig(\n    output_dir                  = \"./outputs\",\n    per_device_train_batch_size = 1,\n    gradient_accumulation_steps = 4,\n    num_train_epochs            = 3,\n    max_completion_length       = 100,\n    logging_steps               = 1,\n    max_grad_norm               = 0.5,\n    learning_rate               = 5e-5,\n    warmup_ratio                = 0.1,\n    seed                        = 42,\n)\n\ntrainer = GRPOTrainer(\n    model         = model,\n    tokenizer     = tokenizer,\n    reward_funcs  = [reward_fn],\n    args          = config,\n    train_dataset = dataset,\n)\n\nprint(\"\ud83d\ude80 Training Fair-GRPO-RLVR on Llama-3.2-1B ...\")\ntrainer.train()\nprint(\"\u2705 Training complete!\")\n\n# Save model immediately\nmodel.save_pretrained(\"./outputs/model\")\ntokenizer.save_pretrained(\"./outputs/model\")\nprint(\"\ud83d\udcbe Model saved \u2192 ./outputs/model\")\n\n# \u2500\u2500 Training loss curve \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nlog_history  = trainer.state.log_history\ntrain_losses = [(x[\"step\"], x[\"loss\"]) for x in log_history if \"loss\" in x]\n\nif train_losses:\n    steps, losses = zip(*train_losses)\n    window   = max(3, len(losses)//8)\n    smoothed = pd.Series(list(losses)).rolling(window, min_periods=1).mean().tolist()\n\n    fig, ax = plt.subplots(figsize=(10, 4))\n    ax.plot(steps, losses,   color=\"#AACDE8\", linewidth=1.2, alpha=0.7, label=\"Raw loss\")\n    ax.plot(steps, smoothed, color=\"#1A6B9A\", linewidth=2.5, label=f\"Smoothed (w={window})\")\n    ax.set_xlabel(\"Training Step\",  fontsize=12)\n    ax.set_ylabel(\"GRPO Loss\",      fontsize=12)\n    ax.set_title(\n        \"Training Loss \u2014 Fair-GRPO-RLVR (Llama 3.2 1B)\\n\"\n        \"Decreasing loss = model learning to generate fair allocation actions\",\n        fontsize=12, fontweight=\"bold\"\n    )\n    ax.legend(fontsize=10); ax.grid(alpha=0.3)\n    ax.text(0.98, 0.95, f\"Initial: {losses[0]:.4f}\\nFinal:   {losses[-1]:.4f}\\nDrop:  {losses[0]-losses[-1]:+.4f}\",\n            transform=ax.transAxes, ha=\"right\", va=\"top\",\n            bbox=dict(boxstyle=\"round\", facecolor=\"white\", alpha=0.8), fontsize=9)\n    plt.tight_layout()\n    plt.savefig(f\"{PLOTS_DIR}/training_loss.png\", dpi=150, bbox_inches=\"tight\")\n    plt.close()\n    print(f\"\ud83d\udcca Training loss plot saved ({len(steps)} steps)\")\n    print(f\"   Initial loss: {losses[0]:.4f} \u2192 Final loss: {losses[-1]:.4f}\")",
      "outputs": [],
      "execution_count": null,
      "id": "c73661212"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 11 \u2014 BASELINE (greedy) + TRAINED runners\n# Both use compute_metrics() \u2014 identical formula\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\ndef _greedy_action(obs):\n    \"\"\"Greedy: picks lowest-damage zone (the Fairness Trap).\"\"\"\n    if USE_BUILTIN or not REPO_OK:\n        zones = obs.zones\n    else:\n        try:\n            from server.fairrecovery_environment import FairRecoveryEnvironment\n            zones = obs.zones\n        except:\n            zones = obs.zones\n    stage = obs.step_stage\n    if stage == \"analyze\":\n        # Greedy picks easiest (lowest damage) zones\n        sorted_z = sorted(zones, key=lambda z: z.damage)\n        return {\"action_type\":\"analyze\",\"critical_zones\":[sorted_z[0].zone_id, sorted_z[1].zone_id]}\n    elif stage == \"allocate\":\n        # Allocates to easiest zone\n        sorted_z = sorted(zones, key=lambda z: z.damage)\n        return {\"action_type\":\"allocate\",\"allocations\":[{\"zone\":sorted_z[0].zone_id,\"resource\":\"power\"}]}\n    return {\"action_type\":\"execute\"}\n\ndef run_baseline(seed=None):\n    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n    for _ in range(MAX_STEPS):\n        action = _greedy_action(obs)\n        result = step_env(env, action)\n        if result is None or result.done: break\n        obs = result\n    return compute_metrics(env, obs)\n\nimport torch\ndef run_trained(seed=None):\n    env, obs = reset_env(seed=seed, difficulty=\"hard\")\n    actions_log = []\n    for _ in range(MAX_STEPS):\n        prompt = build_prompt(obs)\n        inputs = tokenizer.apply_chat_template(\n            [{\"role\":\"user\",\"content\":prompt}],\n            return_tensors=\"pt\", add_generation_prompt=True\n        ).to(model.device)\n        with torch.no_grad():\n            out = model.generate(\n                inputs, max_new_tokens=80,\n                temperature=0.3, top_p=0.9, do_sample=True,\n                pad_token_id=tokenizer.eos_token_id\n            )\n        text = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)\n        action_dict = parse_action(text, obs.step_stage)\n        actions_log.append(f\"{obs.step_stage}\u2192{action_dict.get('action_type','?')}\")\n        result = step_env(env, action_dict)\n        if result is None or result.done: break\n        obs = result\n    m = compute_metrics(env, obs)\n    m[\"actions\"] = actions_log\n    return m\n\nprint(\"\u2705 Runners ready\")",
      "outputs": [],
      "execution_count": null,
      "id": "c18385022"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 12 \u2014 BEFORE vs AFTER: Qualitative demo\n# Shows the exact behavioral difference judges care about\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nprint(\"=\" * 65)\nprint(\"  QUALITATIVE COMPARISON: What does each agent actually do?\")\nprint(\"=\" * 65)\n\nDEMO_SEED = 2000\n\n# Baseline demo\nprint(\"\\n\ud83d\udccd GREEDY BASELINE (falls into Fairness Trap):\")\nenv, obs = reset_env(seed=DEMO_SEED, difficulty=\"hard\")\nfor step in range(6):\n    stage = obs.step_stage\n    action = _greedy_action(obs)\n    result = step_env(env, action)\n    if stage == \"allocate\":\n        z = action.get(\"allocations\",[{}])[0].get(\"zone\",\"?\")\n        print(f\"   Day {obs.day} ALLOCATE \u2192 Zone {z}  \u2190 {'\u26a0\ufe0f LOW PRIORITY ZONE' if z==0 else ''}\")\n    if result is None or result.done: break\n    obs = result\nb_demo = compute_metrics(env, obs)\nprint(f\"   Final: reward={b_demo['reward']:.3f}  fairness={b_demo['fairness']:.3f}\")\nsvcs_b = b_demo[\"services\"]\nprint(f\"   Zone services: {['Z'+str(i)+':'+f'{s:.2f}' for i,s in enumerate(svcs_b)]}\")\n\n# Trained demo\nprint(\"\\n\ud83e\udd16 TRAINED LLM (Fair-GRPO-RLVR, escapes the trap):\")\nenv, obs = reset_env(seed=DEMO_SEED, difficulty=\"hard\")\nfor step in range(6):\n    stage = obs.step_stage\n    prompt = build_prompt(obs)\n    inputs = tokenizer.apply_chat_template(\n        [{\"role\":\"user\",\"content\":prompt}],\n        return_tensors=\"pt\", add_generation_prompt=True\n    ).to(model.device)\n    with torch.no_grad():\n        out = model.generate(inputs, max_new_tokens=60,\n                             temperature=0.3, do_sample=True,\n                             pad_token_id=tokenizer.eos_token_id)\n    text = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)\n    action_dict = parse_action(text, stage)\n    if stage == \"allocate\":\n        z = action_dict.get(\"allocations\",[{}])[0].get(\"zone\",\"?\") if \"allocations\" in action_dict else \"?\"\n        print(f\"   Day {obs.day} ALLOCATE \u2192 Zone {z}  {'\u2705 CORRECT: highest need' if z==4 else ''}\")\n    result = step_env(env, action_dict)\n    if result is None or result.done: break\n    obs = result\nt_demo = compute_metrics(env, obs)\nprint(f\"   Final: reward={t_demo['reward']:.3f}  fairness={t_demo['fairness']:.3f}\")\nsvcs_t = t_demo[\"services\"]\nprint(f\"   Zone services: {['Z'+str(i)+':'+f'{s:.2f}' for i,s in enumerate(svcs_t)]}\")\n\nprint(f\"\\n\ud83d\udcca Fairness delta: {t_demo['fairness']-b_demo['fairness']:+.3f}\")\nprint(\"=\" * 65)",
      "outputs": [],
      "execution_count": null,
      "id": "c74045220"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 13 \u2014 RUN 10-EPISODE EVALUATION\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nprint(f\"Evaluating over {len(EVAL_SEEDS)} episodes ...\")\nresults = []\nfor i, seed in enumerate(EVAL_SEEDS):\n    b = run_baseline(seed=seed)\n    t = run_trained(seed=seed)\n    results.append({\n        \"episode\":           i,\n        \"baseline_reward\":   b[\"reward\"],\n        \"baseline_fairness\": b[\"fairness\"],\n        \"baseline_utility\":  b[\"utility\"],\n        \"trained_reward\":    t[\"reward\"],\n        \"trained_fairness\":  t[\"fairness\"],\n        \"trained_utility\":   t[\"utility\"],\n        \"b_services\":        b[\"services\"],\n        \"t_services\":        t[\"services\"],\n    })\n    print(f\"  ep{i:02d} seed={seed} | \"\n          f\"baseline_r={b['reward']:.3f} fair={b['fairness']:.3f} | \"\n          f\"trained_r={t['reward']:.3f} fair={t['fairness']:.3f} | \"\n          f\"\u0394fair={t['fairness']-b['fairness']:+.3f}\")\n\ndf = pd.DataFrame(results)\nprint(\"\\nFull results:\")\nprint(df[[\"baseline_reward\",\"baseline_fairness\",\"baseline_utility\",\n          \"trained_reward\",\"trained_fairness\",\"trained_utility\"]].to_string(\n    float_format=\"{:.4f}\".format))",
      "outputs": [],
      "execution_count": null,
      "id": "c39239296"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 14 \u2014 COMPLETE 5-PANEL RESULTS PLOT\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nepisodes = df[\"episode\"].tolist()\nC = {\"base\":\"#C0392B\",\"train\":\"#1A6B9A\",\"fair\":\"#27AE60\",\"util\":\"#E67E22\"}\n\nfig = plt.figure(figsize=(18, 12))\ngs  = gridspec.GridSpec(2, 3, figure=fig, hspace=0.50, wspace=0.38)\n\ndef add_arrow(ax, x, y1, y2):\n    for xi, a, b in zip(x, y1, y2):\n        if b > a + 0.005:\n            ax.annotate(\"\", xy=(xi, b+0.01), xytext=(xi, a-0.01),\n                arrowprops=dict(arrowstyle=\"->\",color=\"green\",lw=1.5))\n\n# \u2500\u2500 P1: Reward \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nax1 = fig.add_subplot(gs[0,0])\nax1.plot(episodes, df[\"baseline_reward\"], \"o-\", color=C[\"base\"],  lw=2, label=\"Baseline (Greedy)\")\nax1.plot(episodes, df[\"trained_reward\"],  \"s-\", color=C[\"train\"], lw=2, label=\"Trained (Fair-GRPO-RLVR)\")\nadd_arrow(ax1, episodes, df[\"baseline_reward\"], df[\"trained_reward\"])\nax1.set(title=\"Normalized Reward per Episode\", xlabel=\"Evaluation Episode\",\n        ylabel=\"Reward [0\u20131]\", ylim=(0,1.08))\nax1.legend(fontsize=8); ax1.grid(alpha=0.3)\nax1.text(0.02,0.03,\"Higher = better overall recovery\",\n         transform=ax1.transAxes,fontsize=7,color=\"gray\")\n\n# \u2500\u2500 P2: Fairness \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nax2 = fig.add_subplot(gs[0,1])\nax2.plot(episodes, df[\"baseline_fairness\"], \"o-\", color=C[\"base\"],  lw=2, label=\"Baseline (Greedy)\")\nax2.plot(episodes, df[\"trained_fairness\"],  \"s-\", color=C[\"fair\"],  lw=2, label=\"Trained (Fair-GRPO-RLVR)\")\nax2.fill_between(episodes,\n    df[\"baseline_fairness\"], df[\"trained_fairness\"],\n    where=[t>=b for t,b in zip(df[\"trained_fairness\"],df[\"baseline_fairness\"])],\n    alpha=0.15, color=\"green\", label=\"Improvement region\")\nax2.set(title=\"Equity Index per Episode\\n(Inverse Service Disparity \u2014 higher = more equitable)\",\n        xlabel=\"Evaluation Episode\", ylabel=\"Fairness [0\u20131]\", ylim=(0,1.08))\nax2.legend(fontsize=8); ax2.grid(alpha=0.3)\nax2.text(0.02,0.03,\"Higher = resources distributed more evenly\",\n         transform=ax2.transAxes,fontsize=7,color=\"gray\")\n\n# \u2500\u2500 P3: Utility \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nax3 = fig.add_subplot(gs[0,2])\nax3.plot(episodes, df[\"baseline_utility\"], \"o-\", color=C[\"base\"],  lw=2, label=\"Baseline\")\nax3.plot(episodes, df[\"trained_utility\"],  \"s-\", color=C[\"util\"],  lw=2, label=\"Trained\")\nax3.set(title=\"Utility (Avg Service Level) per Episode\",\n        xlabel=\"Evaluation Episode\", ylabel=\"Utility [0\u20131]\", ylim=(0,1.08))\nax3.legend(fontsize=8); ax3.grid(alpha=0.3)\n\n# \u2500\u2500 P4: Summary bar with error bars + delta labels \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nax4 = fig.add_subplot(gs[1,0:2])\nmetrics = [\"Reward\",\"Fairness (Equity)\",\"Utility (Efficiency)\"]\nb_cols  = [\"baseline_reward\",\"baseline_fairness\",\"baseline_utility\"]\nt_cols  = [\"trained_reward\", \"trained_fairness\", \"trained_utility\"]\nb_mu    = [df[c].mean() for c in b_cols]\nt_mu    = [df[c].mean() for c in t_cols]\nb_sd    = [df[c].std()  for c in b_cols]\nt_sd    = [df[c].std()  for c in t_cols]\nx, w    = np.arange(3), 0.33\n\nbr = ax4.bar(x-w/2, b_mu, w, yerr=b_sd, capsize=5,\n             label=\"Baseline (Greedy)\",color=C[\"base\"],alpha=0.85)\ntr = ax4.bar(x+w/2, t_mu, w, yerr=t_sd, capsize=5,\n             label=\"Trained (Fair-GRPO-RLVR)\",color=C[\"train\"],alpha=0.85)\n\nfor bar,sd in zip(list(br)+list(tr), b_sd+t_sd):\n    h = bar.get_height()\n    ax4.text(bar.get_x()+bar.get_width()/2, h+sd+0.015,\n             f\"{h:.3f}\", ha=\"center\", va=\"bottom\", fontsize=9, fontweight=\"bold\")\n\nfor i,(bv,tv) in enumerate(zip(b_mu,t_mu)):\n    d = tv-bv\n    col = \"#27AE60\" if d>=0 else \"#C0392B\"\n    sym = \"\u25b2\" if d>=0 else \"\u25bc\"\n    ax4.text(i, max(bv,tv)+max(b_sd[i],t_sd[i])+0.06,\n             f\"{sym}{abs(d)*100:.1f}%\", ha=\"center\",\n             color=col, fontsize=11, fontweight=\"bold\")\n\nax4.set(title=\"Average Metrics: Baseline vs Trained  (\u00b11\u03c3 error bars)\",\n        ylabel=\"Mean Score [0\u20131]\", ylim=(0,1.25))\nax4.set_xticks(x); ax4.set_xticklabels(metrics, fontsize=10)\nax4.legend(fontsize=9); ax4.grid(alpha=0.3,axis=\"y\")\n\n# \u2500\u2500 P5: Zone-level service (most visually compelling) \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\nax5 = fig.add_subplot(gs[1,2])\nb_svcs = [sum(row[i] for row in df[\"b_services\"])/len(df) for i in range(5)]\nt_svcs = [sum(row[i] for row in df[\"t_services\"])/len(df) for i in range(5)]\nzi     = np.arange(5)\nax5.bar(zi-0.22, b_svcs, 0.42, label=\"Baseline\",color=C[\"base\"],  alpha=0.85)\nax5.bar(zi+0.22, t_svcs, 0.42, label=\"Trained\", color=C[\"train\"], alpha=0.85)\nfor i,(b,t) in enumerate(zip(b_svcs,t_svcs)):\n    if t>b+0.01:\n        ax5.text(i+0.22, t+0.01, f\"+{(t-b)*100:.0f}%\",\n                 ha=\"center\",color=\"#27AE60\",fontsize=8,fontweight=\"bold\")\nax5.axvline(3.5, color=\"red\", linestyle=\"--\", alpha=0.4)\nax5.text(4.1, max(t_svcs)*0.95, \"Vulnerable\\nzones\", color=\"red\",\n         fontsize=8, ha=\"center\")\nax5.set(title=\"Zone-Level Service Delivery\\n(avg over 10 episodes \u2014 Zone 4\u2605 = most vulnerable)\",\n        xlabel=\"Zone ID\", ylabel=\"Avg Service Level [0\u20131]\", ylim=(0,1.1))\nax5.set_xticks(zi)\nax5.set_xticklabels([f\"Z{i}\"+\"\u2605\"*(i==4) for i in range(5)])\nax5.legend(fontsize=8); ax5.grid(alpha=0.3,axis=\"y\")\n\nfig.suptitle(\n    \"FairRecovery++ \u2014 Fair-GRPO-RLVR vs Greedy Baseline\\n\"\n    \"Training Llama-3.2-1B to Escape the Fairness Trap in Disaster Recovery\",\n    fontsize=14, fontweight=\"bold\"\n)\nplt.savefig(f\"{PLOTS_DIR}/full_results.png\", dpi=150, bbox_inches=\"tight\")\nplt.close()\nprint(f\"\u2705 Saved: {PLOTS_DIR}/full_results.png\")\n\n# Standalone fairness plot for README\nfig2, ax = plt.subplots(figsize=(9,5))\nax.plot(episodes, df[\"baseline_fairness\"], \"o-\", color=C[\"base\"],  lw=2.5, label=\"Baseline (Greedy Policy)\")\nax.plot(episodes, df[\"trained_fairness\"],  \"s-\", color=C[\"fair\"],  lw=2.5, label=\"Trained (Fair-GRPO-RLVR)\")\nax.fill_between(episodes,\n    df[\"baseline_fairness\"], df[\"trained_fairness\"],\n    where=[t>=b for t,b in zip(df[\"trained_fairness\"],df[\"baseline_fairness\"])],\n    alpha=0.15, color=\"green\")\nax.set(title=\"Fairness Score: Before vs After GRPO Training\\n\"\n       \"Inverse Service Disparity  (higher = more equitable resource allocation)\",\n       xlabel=\"Evaluation Episode\", ylabel=\"Fairness Score [0\u20131]\",\n       ylim=(max(0, min(df[\"baseline_fairness\"].min(), df[\"trained_fairness\"].min())-0.1), 1.05))\nax.legend(fontsize=11); ax.grid(alpha=0.3)\nplt.tight_layout()\nplt.savefig(f\"{PLOTS_DIR}/fairness_vs_episode.png\", dpi=150, bbox_inches=\"tight\")\nplt.close()\nprint(f\"\u2705 Saved: {PLOTS_DIR}/fairness_vs_episode.png\")",
      "outputs": [],
      "execution_count": null,
      "id": "c11748987"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 15 \u2014 FINAL SUMMARY TABLE\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\nb_r = df[\"baseline_reward\"].mean();    t_r = df[\"trained_reward\"].mean()\nb_f = df[\"baseline_fairness\"].mean();  t_f = df[\"trained_fairness\"].mean()\nb_u = df[\"baseline_utility\"].mean();   t_u = df[\"trained_utility\"].mean()\ndr  = t_r - b_r;   df_ = t_f - b_f;   du  = t_u - b_u\n\nprint()\nprint(\"\u2554\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2557\")\nprint(\"\u2551         FINAL RESULTS \u2014 Fair-GRPO-RLVR vs Greedy           \u2551\")\nprint(\"\u2560\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2563\")\nprint(f\"\u2551  {'Metric':<14} {'Baseline':>9} {'Trained':>9} {'Delta':>9} {'%':>8}      \u2551\")\nprint(\"\u2560\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2563\")\nfor label, bv, tv, d in [\n    (\"Reward\",   b_r, t_r, dr),\n    (\"Fairness\", b_f, t_f, df_),\n    (\"Utility\",  b_u, t_u, du),\n]:\n    pct = d / (abs(bv)+1e-8) * 100\n    icon = \"\u2705\" if d > 0.002 else (\"\u27a1\ufe0f \" if abs(d) <= 0.002 else \"\u274c\")\n    print(f\"\u2551  {icon} {label:<13} {bv:>9.4f} {tv:>9.4f} {d:>+9.4f} {pct:>+7.1f}%   \u2551\")\nprint(\"\u2560\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2563\")\n\nn_won = sum([dr > 0.002, df_ > 0.002, du > 0.002])\nif n_won == 3:\n    verdict = \"\ud83c\udfc6 IMPROVED ON ALL METRICS \u2014 Fairness Trap escaped!\"\nelif n_won >= 2:\n    verdict = f\"\u2705 IMPROVED ON {n_won}/3 METRICS\"\nelif n_won == 1:\n    verdict = \"\u26a0\ufe0f  PARTIAL \u2014 check zone-level plot for insight\"\nelse:\n    verdict = \"\u274c No improvement \u2014 re-run diagnostic in Cell 5\"\n\nprint(f\"\u2551  {verdict:<60}\u2551\")\nprint(\"\u255a\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u255d\")\n\nprint()\nprint(\"\ud83d\udcc1 Output files ready:\")\nprint(f\"   {PLOTS_DIR}/training_loss.png      \u2190 evidence training ran\")\nprint(f\"   {PLOTS_DIR}/full_results.png       \u2190 5-panel comparison (for README)\")\nprint(f\"   {PLOTS_DIR}/fairness_vs_episode.png \u2190 fairness standalone (for README)\")\nprint(f\"   ./outputs/model/                  \u2190 trained LoRA weights\")\nprint()\nprint(\"\ud83d\udd17 Next steps:\")\nprint(\"   1. Copy plots/ to your repo assets/ folder\")\nprint(\"   2. Push model to HF Hub (see Cell 16)\")\nprint(\"   3. Fix README Colab link to:\")\nprint(\"      https://colab.research.google.com/github/joshua400/FairRecovery-PlusPlus/blob/main/train.ipynb\")",
      "outputs": [],
      "execution_count": null,
      "id": "c82072407"
    },
    {
      "cell_type": "code",
      "metadata": {},
      "source": "# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# CELL 16 \u2014 PUBLISH TO HUGGINGFACE HUB\n# Uncomment + add token \u2192 judges can verify training happened\n# \u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\n# from huggingface_hub import login\n# login(token=\"hf_YOUR_TOKEN_HERE\")\n#\n# model.push_to_hub(\"joshua400/fairrecovery-llama-1b-grpo\",\n#                   commit_message=\"Fair-GRPO-RLVR trained on FairRecovery++ env\")\n# tokenizer.push_to_hub(\"joshua400/fairrecovery-llama-1b-grpo\")\n# print(\"\u2705 Published to HuggingFace Hub\")\n# print(\"   Add to README:\")\n# print(\"   [![Model](https://img.shields.io/badge/\ud83e\udd17_Model-fairrecovery--llama--1b-orange)](https://huggingface.co/joshua400/fairrecovery-llama-1b-grpo)\")\nprint(\"Uncomment above to publish model to HF Hub (recommended for judges).\")",
      "outputs": [],
      "execution_count": null,
      "id": "c57664037"
    }
  ]
}