Update app.py

app.py

@@ -21,7 +21,8 @@ VECNORM_PATH = "models/vecnormalize_lstm_final.pkl"
 MAX_STEPS   = 50
 
 ZONE_LABELS = ["Zone 1 (Residential)", "Zone 2 (Commercial)", "Zone 3 (Hospital)"]
-ZONE_TYPES  = ["Commercial (medium)", "Residential (low)", "Commercial (medium)"]
+ZONE_TYPES  = ["Residential (low)", "Commercial (medium)", "Hospital (critical)"]
+ZONE_SHORT  = ["Residential", "Commercial", "Hospital"]
 
 _DARK_BG  = "#1e1e1e"
 _DARK_AX  = "#181818"
@@ -46,12 +47,38 @@ def _extract_array_obs(obs, length=3):
         return [0.33] * length
 
 
-def _heuristic_action(demand):
-    weights = [demand[0], demand[1], demand[2] * 1.2]
+def _heuristic_action(demand, fault_status=None):
+    """
+    Deterministic, fault-aware heuristic:
+    - Allocate proportionally to demand
+    - Penalise zones currently in fault (halve their weight)
+    - Give Hospital (Zone 3) a 30 % priority boost
+    - Clip each allocation to [0.10, 0.60] to avoid extreme values
+    - Always returns a normalised 3-element action
+    """
+    if fault_status is None:
+        fault_status = [False, False, False]
+
+    weights = [max(float(d), 0.0) for d in demand]
+
+    # Penalise faulty zones so AI routes around them
+    for i, fault in enumerate(fault_status):
+        if fault:
+            weights[i] *= 0.5
+
+    # Hospital (Zone 3) always gets a boost
+    weights[2] *= 1.3
+
     total = sum(weights)
     if total <= 0:
-        return np.array([[0.25, 0.35, 0.40]])
-    return np.array([[w / total for w in weights]])
+        raw = [0.20, 0.30, 0.50]
+    else:
+        raw = [w / total for w in weights]
+
+    # Clip to sensible range so no zone gets starved or flooded
+    clipped = [max(0.10, min(0.60, v)) for v in raw]
+    s = sum(clipped)
+    return np.array([[v / s for v in clipped]])
 
 
 def _fig_to_pil(fig) -> PILImage.Image:
@@ -84,7 +111,7 @@ def _make_grid_chart(demand_vals, supply_vals, fault_status=None) -> PILImage.Im
     # FAULT annotations
     if fault_status is None:
         fault_status = [
-            demand_vals[i] > 0.01 and supply_vals[i] < demand_vals[i] * 0.6
+            demand_vals[i] > 0.01 and supply_vals[i] < demand_vals[i] * 0.8
             for i in range(3)
         ]
     for i, fault in enumerate(fault_status):
@@ -99,7 +126,8 @@ def _make_grid_chart(demand_vals, supply_vals, fault_status=None) -> PILImage.Im
     ax.set_title("Grid Status per Zone", color="white", fontsize=12,
                  pad=12, fontweight="bold")
     ax.set_xticks(x)
-    ax.set_xticklabels(["Zone 1", "Zone 2", "Zone 3"], color="#cccccc", fontsize=10)
+    ax.set_xticklabels(["Zone 1\n(Residential)", "Zone 2\n(Commercial)", "Zone 3\n(Hospital)"],
+                       color="#cccccc", fontsize=9)
     ax.set_ylabel("Power Units", color="#888888", fontsize=9)
     ax.tick_params(axis="y", colors="#888888", labelsize=8)
     ax.tick_params(axis="x", colors="#cccccc")
@@ -131,8 +159,10 @@ class GridSimulator:
 
         if os.path.exists(MODEL_PATH + ".zip"):
             self.model = RecurrentPPO.load(MODEL_PATH, env=self.env)
+            print("✅ Trained model loaded successfully.")
         else:
-            self.model = None          # heuristic-only demo mode
+            self.model = None
+            print("⚠️ Using heuristic AI (trained model not found)")
 
         self.ready = True
         self.reset()
@@ -158,6 +188,7 @@ class GridSimulator:
         self.current_supply = [0.33, 0.33, 0.33]
         self.action_taken   = [0.33, 0.33, 0.33]
         self.fault_status   = [False, False, False]
+        self.ai_explanation = "System reset. Ready for AI power allocation."
 
     # ------------------------------------------------------------------
     def step(self, action=None, manual=False):
@@ -176,7 +207,8 @@ class GridSimulator:
                 )
             else:
                 raw = self.obs[0] if hasattr(self.obs, '__len__') else self.obs
-                action = _heuristic_action(_extract_array_obs(raw))
+                # Pass current fault_status so heuristic avoids broken zones
+                action = _heuristic_action(_extract_array_obs(raw), self.fault_status)
         else:
             total = sum(action)
             action = np.array([[a / total if total > 0 else 0.33 for a in action]])
@@ -200,13 +232,20 @@ class GridSimulator:
         self.action_taken  = (action[0].tolist()
                               if hasattr(action[0], 'tolist') else list(action[0]))
 
+        # Blackout counting — use explicit key presence check (not falsy)
+        # to avoid double-counting and to correctly handle 0-value steps
+        step_blk = None
         for key in ("blackout_count", "blackouts", "fault_count"):
-            if ep_info.get(key):
-                self.blackouts += int(ep_info[key])
+            if key in ep_info:
+                step_blk = int(ep_info[key])
                 break
+        if step_blk is not None:
+            self.blackouts += step_blk
 
-        self.stability = float(ep_info.get("stability_score",
-                                            ep_info.get("stability", 1.0)))
+        # Stability: trust env value but also penalise for accumulated blackouts
+        env_stab = float(ep_info.get("stability_score", ep_info.get("stability", 1.0)))
+        penalty_stab = max(0.0, 1.0 - self.blackouts / 10.0)
+        self.stability = min(env_stab, penalty_stab)
         self.stability_history.append(self.stability)
         self.unmet_demand += float(ep_info.get("unmet_demand", 0.0))
 
@@ -219,21 +258,65 @@ class GridSimulator:
             pass
 
         self.current_supply = self.action_taken[:3]
-        self.fault_status   = [
+        # Fault threshold raised to 0.8 — only flag meaningful undersupply
+        self.fault_status = [
             self.current_demand[i] > 0.01 and
-            self.current_supply[i] < self.current_demand[i] * 0.6
+            self.current_supply[i] < self.current_demand[i] * 0.8
             for i in range(3)
         ]
-        return self.get_ui_state("Step completed.")
+        self._update_ai_explanation()
+        return self.get_ui_state(f"✅ Reward: {self.last_reward:.2f}")
+
+    # ------------------------------------------------------------------
+    def _update_ai_explanation(self):
+        z = self.action_taken[:3]
+        d = self.current_demand[:3]
+        f = self.fault_status[:3]
+
+        faulty_zones = [ZONE_SHORT[i] for i, fault in enumerate(f) if fault]
+        max_z = z.index(max(z))
+
+        if faulty_zones:
+            self.ai_explanation = (
+                f"AI reduced allocation to faulty zone(s) ({', '.join(faulty_zones)}) "
+                f"and rerouted power to maintain grid stability."
+            )
+        elif max_z == 2:
+            self.ai_explanation = (
+                f"AI prioritized Hospital (Zone 3) — demand at {d[2]:.0%}. "
+                f"Critical load protected ({z[2]:.0%} share allocated)."
+            )
+        elif max_z == 1:
+            self.ai_explanation = (
+                f"AI allocated maximum to Commercial Zone to handle demand spike "
+                f"({z[1]:.0%} share). Hospital share: {z[2]:.0%}."
+            )
+        else:
+            self.ai_explanation = (
+                f"AI balanced supply across zones. "
+                f"Residential: {z[0]:.0%} | Commercial: {z[1]:.0%} | Hospital: {z[2]:.0%}."
+            )
+
+        if self.blackouts > 0:
+            self.ai_explanation += f" ⚠️ Cumulative blackouts: {int(self.blackouts)} — adapting allocation."
 
     # ------------------------------------------------------------------
     def get_ui_state(self, status_msg="Step completed."):
         reward_str = f"{self.last_reward:.4f}"
-        done_str   = "✅ Yes — Episode Complete" if self.done else ""
-        env_desc   = self._build_env_description()
-        chart      = _make_grid_chart(self.current_demand,
-                                      self.current_supply,
-                                      self.fault_status)
+        if self.done:
+            stab_icon = "🟢" if self.stability >= 0.75 else ("🟡" if self.stability >= 0.45 else "🔴")
+            done_str = (
+                f"🏁 Episode Complete\n"
+                f"⚡ Final Stability: {self.stability:.2f} {stab_icon}\n"
+                f"🚨 Total Blackouts: {int(self.blackouts)}\n"
+                f"🏆 Total Reward:    {self.total_reward:.2f}"
+            )
+        else:
+            done_str = ""
+        env_desc = self._build_env_description()
+        chart    = _make_grid_chart(self.current_demand,
+                                    self.current_supply,
+                                    self.fault_status)
         return reward_str, done_str, status_msg, env_desc, chart
 
     # ------------------------------------------------------------------
@@ -263,13 +346,24 @@ class GridSimulator:
             f"  Total unmet demand: {self.unmet_demand:.3f}",
             f"  Total reward:       {self.total_reward:.2f}",
             "",
+            f"🤖 AI Decision: {self.ai_explanation}",
+            "",
             "Task: Allocate power to 3 zones as fractions summing to 1.0.",
             "Priority: Serve Zone 3 (Hospital) first. "
             "Avoid overloads – they cascade into blackouts.",
             "Reply with exactly 3 space-separated floats. Example: 0.20 0.30 0.50",
         ]
         if self.done:
-            lines.append("\n✅ Episode complete. Click 'Reset Env' to start a new episode.")
+            lines += [
+                "",
+                "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
+                "🏁 Episode Complete",
+                f"  ⚡ Final Stability: {self.stability:.2f} {stab_icon}",
+                f"  🚨 Total Blackouts: {int(self.blackouts)}",
+                f"  🏆 Total Reward:    {self.total_reward:.2f}",
+                "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━",
+                "Click 'Reset Env' to start a new episode.",
+            ]
         return "\n".join(lines)
 
     # ------------------------------------------------------------------
@@ -293,7 +387,7 @@ def ui_reset():
     if not sim.ready:
         return _ERR
     sim.reset()
-    return sim.get_ui_state("Environment reset. Ready.")
+    return sim.get_ui_state("✅ Environment reset. Click AI Step or Take Step to begin.")
 
 
 def ui_ai_step():
@@ -313,7 +407,7 @@ def ui_auto_run():
         return
     while not sim.done and sim.steps < MAX_STEPS:
         yield sim.step()
-        time.sleep(0.15)
+        time.sleep(0.30)
 
 
 def ui_take_step(z1, z2, z3):