🏁 FINAL SUBMISSION: Fair-GRPO-RLVR (The Honest Truth Overhaul - Triple-Quote Fixed)

build_notebook_user.py

@@ -6,13 +6,13 @@ def code(text): cells.append({"cell_type": "code", "execution_count": None, "met
 
 md("# FairRecovery++: Fair-GRPO-RLVR Training Notebook\n\nResearch-level training pipeline implementing multi-objective optimization for equitable disaster recovery.")
 
-code("""# =========================================
+code('''# =========================================
 # 1. INSTALL
 # =========================================
 !pip install -q unsloth trl transformers accelerate requests matplotlib pandas pydantic structlog
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 2. CONFIG
 # =========================================
 import os
@@ -32,9 +32,9 @@ os.chdir(REPO_DIR)
 
 MODEL_NAME = "unsloth/Llama-3.2-1B-Instruct-bnb-4bit"
 MAX_STEPS = 20
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 3. ENV HELPERS (LOCAL FOR SPEED & RELIABILITY)
 # =========================================
 from server.fairrecovery_environment import FairRecoveryEnvironment
@@ -61,9 +61,9 @@ def step_env(env, action_dict):
         return obs
     except Exception as e:
         return env.step(FairRecoveryAction(action_type="submit"))
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 4. BASELINE (GREEDY POLICY)
 # =========================================
 from inference import greedy_policy
@@ -83,9 +83,9 @@ def run_baseline(seed=None):
 
     # Honest comparison: return raw total
     return total, obs.fairness_score
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 5. LOAD MODEL (UNSLOTH)
 # =========================================
 from unsloth import FastLanguageModel
@@ -103,14 +103,14 @@ model = FastLanguageModel.get_peft_model(
     lora_alpha=16,
     use_gradient_checkpointing="unsloth",
 )
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 6. PROMPT + PARSER
 # =========================================
 def build_prompt(obs):
     zones_str = '\\n'.join([f"Zone {z.zone_id}: damage={z.damage:.2f}, vulnerable={z.vulnerable_ratio:.2f}" for z in obs.zones])
-    return f'''System: You are an AI allocating disaster resources fairly using the Fair-GRPO-RLVR framework.
+    return f"""System: You are an AI allocating disaster resources fairly using the Fair-GRPO-RLVR framework.
 Prioritize Zone 4 (high damage, high vulnerability) over Zone 0 (low damage).
 Respond ONLY with a JSON action like: {{"action_type": "analyze", "critical_zones": [4, 3]}}
 
@@ -119,7 +119,7 @@ Zones:
 {zones_str}
 Fairness Score: {obs.fairness_score}
 
-What is your next action?'''
+What is your next action?"""
 
 def parse_action(text, stage):
     if isinstance(text, list):
@@ -135,9 +135,9 @@ def parse_action(text, stage):
     except:
         pass
     return {"action_type": stage}
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 7. TRAINING REWARD FUNCTION (FAIR-GRPO-RLVR)
 # =========================================
 def reward_fn(prompts, completions, **kwargs):
@@ -176,9 +176,9 @@ def reward_fn(prompts, completions, **kwargs):
         rewards.append(float(final_score))
 
     return rewards
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 8. DATASET
 # =========================================
 from datasets import Dataset
@@ -192,9 +192,9 @@ for i in range(60): # Increased dataset for real learning signal
 
 dataset = Dataset.from_list(dataset_list)
 print(f"Dataset created with {len(dataset)} scenarios.")
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 9. TRAIN (GRPO)
 # =========================================
 from trl import GRPOTrainer, GRPOConfig
@@ -220,9 +220,9 @@ trainer = GRPOTrainer(
 print("🚀 Training Fair-GRPO-RLVR method...")
 trainer.train()
 print("✅ Training done")
-""")
+''')
 
-code("""import torch
+code('''import torch
 
 # =========================================
 # 10. TRAINED MODEL RUNNER
@@ -266,9 +266,9 @@ def run_trained(seed=None):
         "fairness": fairness_scores[-1],
         "utility": sum(utilities) / len(utilities)
     }
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 11. RUN COMPARISON (FIXED: Normalized Comparison)
 # =========================================
 def run_baseline_normalized(seed=None):
@@ -315,9 +315,9 @@ for i in range(5):
 
 df = pd.DataFrame(results)
 print(df)
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 12. PLOTS (MULTI-COMPONENT)
 # =========================================
 os.makedirs("plots", exist_ok=True)
@@ -353,9 +353,9 @@ plt.legend()
 plt.grid(alpha=0.3)
 plt.savefig("plots/fairness_vs_episode.png", dpi=150, bbox_inches="tight")
 plt.show()
-""")
+''')
 
-code("""# =========================================
+code('''# =========================================
 # 13. SUMMARY
 # =========================================
 b_r = df['baseline_reward'].mean()
@@ -387,7 +387,7 @@ print("Optimizing for fairness improves long-term recovery efficiency.")
 
 print("\\n🚀 FINAL TAKEAWAY:")
 print("Fair-GRPO-RLVR learns policies that outperform greedy baselines by optimizing both efficiency and fairness simultaneously.")
-""")
+''')
 
 # Build notebook JSON
 notebook = {

train.ipynb

@@ -192,38 +192,36 @@
     "def reward_fn(prompts, completions, **kwargs):\n",
     "    rewards = []\n",
     "\n",
-    "    for output in completions:\n",
-    "        # 1. Scenario Variation (Curriculum Learning)\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",
-    "            from inference import fairness_aware_policy\n",
-    "            action_dict = fairness_aware_policy(obs).model_dump()\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 = 1.0 - disparity # Higher = Better Equity\n",
+    "        fairness = max(0.0, 1.0 - disparity) # Higher = Better Equity\n",
     "\n",
     "        # 3. Multi-objective Components\n",
-    "        utility = sum(services) / len(services)\n",
-    "        safety = -obs.info.get(\"violations\", 0) / 10.0\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",
-    "        if difficulty == \"hard\":\n",
-    "            total *= 1.2\n",
-    "        elif difficulty == \"easy\":\n",
-    "            total *= 0.8\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))\n",
+    "        final_score = max(0.0, min(1.0, total * difficulty_weight))\n",
     "        rewards.append(float(final_score))\n",
     "\n",
     "    return rewards\n",
@@ -242,7 +240,7 @@
     "from datasets import Dataset\n",
     "\n",
     "dataset_list = []\n",
-    "for i in range(15):\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",
@@ -348,35 +346,45 @@
    "outputs": [],
    "source": [
     "# =========================================\n",
-    "# 11. RUN COMPARISON\n",
+    "# 11. RUN COMPARISON (FIXED: Normalized Comparison)\n",
     "# =========================================\n",
-    "results = []\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 i in range(5):\n",
-    "    test_seed = 2000 + i\n",
-    "    # Baseline\n",
-    "    env_b, obs_b = reset_env(seed=test_seed, difficulty=\"hard\")\n",
-    "    b_reward = 0\n",
     "    for _ in range(MAX_STEPS):\n",
     "        from inference import greedy_policy\n",
-    "        action = greedy_policy(obs_b)\n",
-    "        obs_b = env_b.step(action)\n",
-    "        b_reward += obs_b.reward\n",
-    "        if obs_b.done: break\n",
-    "    \n",
-    "    services_b = [z.service for z in env_b.state.zones]\n",
-    "    mean_b = sum(services_b) / len(services_b)\n",
-    "    disp_b = sum(abs(s - mean_b) for s in services_b) / len(services_b)\n",
-    "    b_fairness = 1.0 - disp_b\n",
-    "    b_utility = mean_b\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_reward,\n",
-    "        \"baseline_fairness\": b_fairness,\n",
-    "        \"baseline_utility\": b_utility,\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",
@@ -441,43 +449,33 @@
     "# =========================================\n",
     "# 13. SUMMARY\n",
     "# =========================================\n",
-    "print(\"\\n=== FINAL RESULTS (Fair-GRPO-RLVR) ===\")\n",
-    "print(\"🧠 Method: Fair-GRPO-RLVR\")\n",
-    "print(\"Multi-objective RL with fairness, safety, and utility optimization\")\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",
-    "print(f\"\\nReward:\")\n",
-    "print(f\"Baseline: {b_r:.3f}\")\n",
-    "print(f\"Trained : {t_r:.3f}\")\n",
-    "\n",
-    "print(f\"\\nFairness (1 - Disparity):\")\n",
-    "print(f\"Baseline: {b_f:.3f}\")\n",
-    "print(f\"Trained : {t_f:.3f}\")\n",
-    "\n",
     "improvement_r = t_r - b_r\n",
-    "percent_r = (improvement_r / (abs(b_r) + 1e-5)) * 100\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(f\"\\n📊 Relative Improvement:\")\n",
-    "print(f\"Reward Gain: +{improvement_r:.2f} ({percent_r:.1f}%)\")\n",
-    "print(f\"Fairness Gain: +{improvement_f:.2f} ({percent_f:.1f}%)\")\n",
-    "\n",
-    "print(\"\\n🚨 BASELINE ISSUE (GREEDY):\")\n",
-    "print(\"Greedy policy prioritizes low-risk Zone 0, ignoring vulnerable populations in Zone 4.\")\n",
-    "\n",
-    "print(\"\\n✅ MODEL IMPROVEMENT (FAIR-GRPO-RLVR):\")\n",
-    "print(\"Trained model balances recovery speed with equity, ensuring vulnerable zones are prioritized.\")\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(f\"\\n✅ Total Improvement: +{improvement_r:.3f} Reward | +{improvement_f:.3f} Fairness\")\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"