Update evaluate_checkpoint_real.py

evaluate_checkpoint_real.py

@@ -389,6 +389,33 @@ def group_points_by_date(
 # Decision policies
 # ---------------------------------------------------------------------------
 
+def _multi_zone_risk_score(ctx: EpisodeContext) -> RiskScore:
+    """Risk score for a (possibly multi-zone) EpisodeContext, matching
+    WeatherForecastEnv._compute_multi_zone_risk(): the zone with the
+    highest supply_shortfall_prob across ALL zones in the episode, not
+    just ctx.obs/ctx.forecast (which points_to_multi_zone_episode always
+    sets to zone_ids[0] alone).
+
+    ctx.obs/ctx.forecast is only ever a single zone's data by construction
+    (see points_to_multi_zone_episode / point_to_episode). For single-zone
+    episodes this is identical to compute_risk_score(ctx.obs, ctx.forecast,
+    ctx.config); for multi-zone episodes it must aggregate across
+    ctx.resolved_zone_obs()/ctx.resolved_zone_forecasts() the same way the
+    env does, or "rs" silently stops representing the same decision as the
+    env's actual per-step product_actionable flag.
+    """
+    zone_obs = ctx.resolved_zone_obs()
+    zone_fc = ctx.resolved_zone_forecasts()
+    best: Optional[RiskScore] = None
+    for zo, zf in zip(zone_obs, zone_fc):
+        score = compute_risk_score(zo, zf, ctx.config)
+        if best is None or score.supply_shortfall_prob > best.supply_shortfall_prob:
+            best = score
+    return best if best is not None else compute_risk_score(
+        ctx.obs, ctx.forecast, ctx.config
+    )
+
+
 def decide_scorer_oracle(
     obs: ZoneObs,
     fc: ForecastResult,
@@ -463,7 +490,7 @@ def decide_checkpoint(
     loop_product, loop_elevated = product, elevated
     rs = None
     try:
-        rs = compute_risk_score(ctx.obs, ctx.forecast, ctx.config)
+        rs = _multi_zone_risk_score(ctx)
         rs_product = is_product_actionable(rs, gate)
         _PATH_CHECK_N += 1
         if rs_product != loop_product:
@@ -491,7 +518,7 @@ def decide_zero_inspect(
         believed_p = _initial_belief_from_info(info, ctx.config)
     rs = None
     try:
-        rs = compute_risk_score(ctx.obs, ctx.forecast, ctx.config)
+        rs = _multi_zone_risk_score(ctx)
     except Exception:
         pass
     return product, elevated, rs, 1, believed_p, initial_belief