Update weather_forecast_env.py

weather_forecast_env.py

@@ -67,6 +67,14 @@ except ImportError:
     DEFAULT_PRODUCT_GATE = None
     _product_gate_fn = None
 
+try:
+    from real_episode_sampler import RealEpisodeIndex
+except ImportError:
+    RealEpisodeIndex = None  # training real-data sampling unavailable;
+                              # ForecastConfig.real_data_pkl_path is then
+                              # ignored and training stays fully synthetic,
+                              # same as before this module existed.
+
 logger = logging.getLogger(__name__)
 
 
@@ -275,6 +283,28 @@ def _per_zone_beliefs(
     return np.clip(np.array(beliefs, dtype=np.float32), 0.0, 1.0)
 
 
+def _resolved_zone_beliefs(
+    zone_obs: List[Any],
+    prior: float,
+    prior_weight: float,
+) -> np.ndarray:
+    """Belief-map values for zone_obs that are already resolved (real data,
+    whether externally injected via options={"context": ...} or sampled by
+    RealEpisodeIndex during training) -- as opposed to _per_zone_beliefs()
+    above, which generates synthetic obs itself from a seed.
+
+    Single source of truth for the prior/signal blend across both
+    real-data code paths in reset(), so they cannot drift apart the way
+    the injected path and the synthetic path's hard-coded 0.70 literal
+    once did.
+    """
+    beliefs = np.zeros(len(zone_obs), dtype=np.float32)
+    for i, zo in enumerate(zone_obs):
+        signal = float(zo.composite_risk())
+        beliefs[i] = prior_weight * prior + (1.0 - prior_weight) * signal
+    return np.clip(beliefs, 0.0, 1.0)
+
+
 def _make_per_zone_obs(
     zone_ids: List[str],
     base_seed: int,
@@ -351,6 +381,34 @@ class WeatherForecastEnv(gym.Env):
         self._episode_count: int = 0
         self._episode_seed: int = 0
 
+        self._real_index = None
+        pkl_path = getattr(self.config, "real_data_pkl_path", None)
+        if pkl_path:
+            if RealEpisodeIndex is None:
+                logger.warning(
+                    "WeatherForecastEnv: config.real_data_pkl_path=%r set, "
+                    "but real_episode_sampler.py is not importable -- "
+                    "training will remain fully synthetic.",
+                    pkl_path,
+                )
+            else:
+                try:
+                    self._real_index = RealEpisodeIndex(pkl_path)
+                    logger.info(
+                        "WeatherForecastEnv: real-data training sampling "
+                        "ENABLED from %s (%d eligible dates, "
+                        "real_data_ratio=%.2f, inject_noise=%s)",
+                        pkl_path, self._real_index.n_eligible_dates,
+                        self.config.real_data_ratio, self.config.inject_noise,
+                    )
+                except Exception as e:
+                    logger.error(
+                        "WeatherForecastEnv: failed to load "
+                        "real_data_pkl_path=%r (%s) -- training will "
+                        "remain fully synthetic.", pkl_path, e,
+                    )
+                    self._real_index = None
+
         logger.info(
             "WeatherForecastEnv: n_zones=%d  horizon=%dd  max_steps=%d  "
             "terminate_action=%d",
@@ -382,8 +440,31 @@ class WeatherForecastEnv(gym.Env):
                     f"got {type(ctx).__name__}"
                 )
             self._context = ctx
+            context_source = "injected"
         else:
-            self._context = _make_multi_zone_context(self.config, effective_seed)
+            self._context = None
+            context_source = "synthetic"
+            if (
+                self._real_index is not None
+                and float(self.config.real_data_ratio) > 0.0
+            ):
+                # Dedicated RNG stream, offset from the streams
+                # _maybe_shuffle_zone_order and other reset-time randomness
+                # use, so enabling real-data sampling doesn't change the
+                # synthetic path's own random draws when it IS taken.
+                real_rng = random.Random(int(effective_seed) + 9001)
+                if real_rng.random() < float(self.config.real_data_ratio):
+                    sampled = self._real_index.sample(
+                        real_rng, self.config.n_zones, self.config,
+                        inject_noise=bool(self.config.inject_noise),
+                        noise_scale=float(self.config.noise_scale),
+                    )
+                    if sampled is not None:
+                        self._context = sampled
+                        context_source = "real_sampled"
+            if self._context is None:
+                self._context = _make_multi_zone_context(self.config, effective_seed)
+                context_source = "synthetic"
 
         self._zone_ids = list(self._context.zone_ids)
         n_active = len(self._zone_ids)
@@ -394,7 +475,12 @@ class WeatherForecastEnv(gym.Env):
                 f"max {self.max_zones}."
             )
 
-        injected = bool(options and "context" in options)
+        # "injected" (explicit real-eval context) and "real_sampled"
+        # (RealEpisodeIndex draw during training) both arrive as a fully
+        # resolved EpisodeContext with real zone_obs/zone_forecasts already
+        # attached -- they share the same validation and belief-computation
+        # path below. Only "synthetic" generates obs/forecasts itself.
+        use_resolved_path = context_source in ("injected", "real_sampled")
 
         self._forecast_arr.fill(0.0)
         self._uncertainty.fill(0.0)
@@ -409,43 +495,40 @@ class WeatherForecastEnv(gym.Env):
         )
 
         # Single source of truth for the prior/signal belief-blend weight --
-        # both the injected (real-eval) and synthetic (training) reset paths
-        # below must read this same value, not their own copies of it.
+        # the injected, real-sampled, and synthetic reset paths below must
+        # all read this same value, not their own copies of it.
         prior_w = float(getattr(self.config, "belief_prior_weight", 0.70))
 
-        if injected:
+        if use_resolved_path:
             z_obs = self._context.resolved_zone_obs()
             z_fc = self._context.resolved_zone_forecasts()
             if len(z_obs) != n_active or len(z_fc) != n_active:
                 raise ValueError(
-                    f"Injected EpisodeContext zone payload length mismatch: "
-                    f"n_active={n_active} resolved_obs={len(z_obs)} "
+                    f"{context_source} EpisodeContext zone payload length "
+                    f"mismatch: n_active={n_active} resolved_obs={len(z_obs)} "
                     f"resolved_fc={len(z_fc)} zone_ids={self._zone_ids}. "
-                    f"For multi-zone real eval supply zone_obs/zone_forecasts "
+                    f"For multi-zone episodes supply zone_obs/zone_forecasts "
                     f"parallel to zone_ids (do not pad one zone into N slots)."
                 )
             if n_active > 1 and not self._context.zone_obs:
                 raise ValueError(
-                    f"Injected EpisodeContext real-eval path refuses n_zones>1 "
+                    f"{context_source} EpisodeContext refuses n_zones>1 "
                     f"without explicit zone_obs lists (got n_active={n_active}, "
                     f"zone_ids={self._zone_ids}). Padding the primary obs into "
-                    f"multiple slots is not multi-zone eval."
+                    f"multiple slots is not multi-zone data."
                 )
             for i, zid in enumerate(self._zone_ids):
                 if z_obs[i].zone_id != zid or z_fc[i].zone_id != zid:
                     raise ValueError(
-                        f"Injected context zone mismatch at slot {i}: "
+                        f"{context_source} context zone mismatch at slot {i}: "
                         f"zone_ids[{i}]={zid!r} obs={z_obs[i].zone_id!r} "
                         f"fc={z_fc[i].zone_id!r}"
                     )
             self._zone_obs = list(z_obs)
             self._zone_forecasts = list(z_fc)
             prior = float(self.config.prior_belief)
-            for i in range(n_active):
-                signal = float(self._zone_obs[i].composite_risk())
-                self._belief_map[i] = float(
-                    np.clip(prior_w * prior + (1.0 - prior_w) * signal, 0.0, 1.0)
-                )
+            risks = _resolved_zone_beliefs(self._zone_obs, prior, prior_w)
+            self._belief_map[:n_active] = risks[:n_active]
             self._belief_map[n_active:] = 0.0
         else:
             rho = float(getattr(self.config, "event_spatial_correlation", 0.85))
@@ -486,6 +569,7 @@ class WeatherForecastEnv(gym.Env):
 
         info = self._get_info()
         info["outcome_code"] = "RESET_OK"
+        info["context_source"] = context_source
         return self._get_obs(), info
 
     def _maybe_shuffle_zone_order(self, effective_seed: int, n_active: int) -> None: