Update weather_forecast_env.py

weather_forecast_env.py

@@ -380,6 +380,7 @@ class WeatherForecastEnv(gym.Env):
         self._cum_reward: float = 0.0
         self._episode_count: int = 0
         self._episode_seed: int = 0
+        self._context_source: str = "synthetic"
 
         self._real_index = None
         pkl_path = getattr(self.config, "real_data_pkl_path", None)
@@ -468,6 +469,17 @@ class WeatherForecastEnv(gym.Env):
 
         self._zone_ids = list(self._context.zone_ids)
         n_active = len(self._zone_ids)
+        self._context_source = context_source  # stamped onto every step's
+                                                  # info via _get_info(), not
+                                                  # just reset's -- SB3's
+                                                  # DummyVecEnv stores
+                                                  # auto-reset info separately
+                                                  # from what step() returns,
+                                                  # so anything only set in
+                                                  # reset()'s info dict is
+                                                  # invisible to a callback's
+                                                  # self.locals["infos"].
+                                                  # Confirmed empirically.
 
         if n_active > self.max_zones:
             raise ValueError(
@@ -853,6 +865,10 @@ class WeatherForecastEnv(gym.Env):
             "mean_belief": float(np.mean(active_belief)) if n_active else 0.0,
             "mean_uncertainty": float(np.mean(active_unc)) if n_active else 0.0,
             "n_visited": int(self._visited[:n_active].sum()) if n_active else 0,
+            # Stamped on EVERY step's info via this method, not just reset's --
+            # see the comment in reset() where self._context_source is set for
+            # why that distinction matters under SB3's VecEnv autoreset.
+            "context_source": self._context_source,
         }
 
     def render(self, mode: Optional[str] = None) -> None: