""" gru_weather_policy.py ==================== Custom GRU feature extractor + zone-equivariant policy head for stable-baselines3 MaskablePPO. Architecture ------------ 1. GRUWeatherFeaturesExtractor: - Per-zone GRU over forecast_precip[zone, :] (horizon_days) - Per-zone MLP over a 10-vector: [uncertainty, belief] + 8-dim basin_context broadcast to every zone (ENSO, IOD, ITCZ, MSLP, solar wind, Kp, log10(GOES), helio regime) - Concatenate → zone-level feature vector - Stash zone scores and terminate logit for the policy head 2. ZoneEquivariantMaskablePolicy: - Overrides _get_action_dist_from_latent - Reads stashed zone scores + terminate logit from the extractor - Returns a Categorical distribution directly - This makes the policy permutation-equivariant across zones (inspecting zone 0 then zone 1 is the same as zone 1 then zone 0) WARNING ------- When using ZoneEquivariantMaskablePolicy, the ``net_arch`` pi layers are instantiated by SB3 inside the MLP extractor but are NEVER called at inference time because ``_get_action_dist_from_latent`` bypasses ``latent_pi`` entirely. The policy capacity is entirely in the extractor. The vf head still uses the ``net_arch`` vf layers normally. This extractor's static MLP input width is 10, not 2. Weights from triage_nz2_rho050_c080_s7_2.4M.zip (2-wide static MLP) cannot be loaded into this class. Train a new checkpoint. The 8-dim basin default MUST stay aligned with weather_forecast_env._BASIN_CONTEXT_NEUTRAL. Dependencies ------------ pip install stable-baselines3 sb3-contrib torch """ from __future__ import annotations import logging from typing import Any, Dict, List, Optional, Type import gymnasium as gym import torch import torch.nn as nn logger = logging.getLogger(__name__) # Matches weather_forecast_env._BASIN_CONTEXT_NEUTRAL and basin_context_vector. # [enso_oni, iod_dmi, itcz_lat, mslp_hpa, sw_kms, kp, log10(goes), regime] # regime: quiet=0, active=1, storm=2 BASIN_CONTEXT_DIM = 8 BASIN_CONTEXT_NEUTRAL = (0.0, 0.0, 0.0, 1013.25, 400.0, 2.0, -7.0, 0.0) STATIC_DIM = 2 + BASIN_CONTEXT_DIM # uncertainty, belief, basin[8] # --------------------------------------------------------------------------- # SB3 availability # --------------------------------------------------------------------------- try: from stable_baselines3.common.torch_layers import BaseFeaturesExtractor from stable_baselines3.common.policies import MultiInputActorCriticPolicy _SB3_AVAILABLE = True except ImportError: _SB3_AVAILABLE = False BaseFeaturesExtractor = object # type: ignore[assignment,misc] MultiInputActorCriticPolicy = object # type: ignore[assignment,misc] try: from sb3_contrib.common.maskable.policies import ( MaskableMultiInputActorCriticPolicy, ) _MASKABLE_AVAILABLE = True except ImportError: _MASKABLE_AVAILABLE = False MaskableMultiInputActorCriticPolicy = object # type: ignore[assignment,misc] def _neutral_basin(batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: base = torch.tensor(BASIN_CONTEXT_NEUTRAL, device=device, dtype=dtype) return base.unsqueeze(0).expand(batch_size, BASIN_CONTEXT_DIM).clone() # --------------------------------------------------------------------------- # GRU feature extractor # --------------------------------------------------------------------------- class GRUWeatherFeaturesExtractor(BaseFeaturesExtractor): def __init__( self, observation_space: gym.spaces.Dict, features_dim: int = 128, hidden_size: int = 64, ): super().__init__(observation_space, features_dim=features_dim) self.hidden_size = hidden_size self._observation_space = observation_space precip_space = observation_space.spaces["forecast_precip"] self.n_zones = int(precip_space.shape[0]) self.horizon_days = int(precip_space.shape[1]) if "basin_context" in observation_space.spaces: bshape = observation_space.spaces["basin_context"].shape if len(bshape) != 1 or int(bshape[0]) != BASIN_CONTEXT_DIM: raise ValueError( f"GRUWeatherFeaturesExtractor: basin_context shape {bshape} " f"must be ({BASIN_CONTEXT_DIM},) to match " f"weather_forecast_env.basin_context_vector" ) self.precip_gru = nn.GRU( input_size=1, hidden_size=hidden_size, num_layers=1, batch_first=True, ) self.static_mlp = nn.Sequential( nn.Linear(STATIC_DIM, hidden_size), nn.Tanh(), ) self.zone_score = nn.Sequential( nn.Linear(hidden_size * 2, 64), nn.Tanh(), nn.Linear(64, 1), ) self.terminate_logit = nn.Linear(hidden_size * 2, 1) self.value_head = nn.Sequential( nn.Linear(self.n_zones * hidden_size * 2, 128), nn.Tanh(), nn.Linear(128, 1), ) def _basin_batch(self, observations: Dict[str, torch.Tensor], batch_size: int) -> torch.Tensor: """[batch, 8]. Missing key → Neutral default, never a silent zero vector.""" precip = observations["forecast_precip"] if "basin_context" not in observations: logger.warning( "GRUWeatherFeaturesExtractor: observation has no basin_context; " "using Neutral default. Combined weather+helio training is a " "no-op until the env supplies this key." ) return _neutral_basin(batch_size, precip.device, precip.dtype) basin = observations["basin_context"] if basin.dim() == 1: basin = basin.unsqueeze(0).expand(batch_size, -1) if basin.shape[-1] != BASIN_CONTEXT_DIM: raise ValueError( f"basin_context last dim is {basin.shape[-1]}, " f"expected {BASIN_CONTEXT_DIM}" ) if basin.shape[0] != batch_size: raise ValueError( f"basin_context batch {basin.shape[0]} != precip batch {batch_size}" ) return basin.to(device=precip.device, dtype=precip.dtype) def forward(self, observations: Dict[str, torch.Tensor]) -> torch.Tensor: precip = observations["forecast_precip"] batch_size = precip.shape[0] uncertainty = observations["forecast_uncertainty"] belief = observations["zone_belief"] basin = self._basin_batch(observations, batch_size) precip_reshaped = precip.reshape( batch_size * self.n_zones, self.horizon_days, 1 ) _, gru_hidden = self.precip_gru(precip_reshaped) gru_features = gru_hidden.squeeze(0) # Broadcast episode-global basin to every zone: [batch, n_zones, 8] basin_z = basin.unsqueeze(1).expand(batch_size, self.n_zones, BASIN_CONTEXT_DIM) static_input = torch.cat( [ uncertainty.unsqueeze(-1), belief.unsqueeze(-1), basin_z, ], dim=-1, ) static_input = static_input.reshape(batch_size * self.n_zones, STATIC_DIM) static_features = self.static_mlp(static_input) zone_features = torch.cat([gru_features, static_features], dim=-1) zone_scores = self.zone_score(zone_features).squeeze(-1) self._last_zone_scores = zone_scores.reshape(batch_size, self.n_zones) self._last_terminate_logit = self.terminate_logit(zone_features).squeeze(-1) self._last_terminate_logit = self._last_terminate_logit.reshape( batch_size, self.n_zones )[:, 0] global_features = zone_features.reshape(batch_size, self.n_zones, -1) global_features = global_features.reshape(batch_size, -1) return global_features def get_value(self, latent_vf: torch.Tensor) -> torch.Tensor: return self.value_head(latent_vf) # --------------------------------------------------------------------------- # Zone-equivariant policy head # --------------------------------------------------------------------------- class ZoneEquivariantMaskablePolicy(MaskableMultiInputActorCriticPolicy): def __init__( self, observation_space: gym.spaces.Dict, action_space: gym.spaces.Discrete, lr_schedule, net_arch: Optional[List[int]] = None, activation_fn: Type[nn.Module] = nn.Tanh, *args, **kwargs, ): super().__init__( observation_space, action_space, lr_schedule, net_arch=net_arch, activation_fn=activation_fn, *args, **kwargs, ) def _get_action_dist_from_latent(self, latent_pi: torch.Tensor) -> Any: features_extractor = self.features_extractor assert isinstance(features_extractor, GRUWeatherFeaturesExtractor) zone_scores = features_extractor._last_zone_scores terminate_logit = features_extractor._last_terminate_logit logits = torch.cat([ zone_scores, terminate_logit.unsqueeze(-1), ], dim=-1) return self.action_dist.proba_distribution(action_logits=logits) # --------------------------------------------------------------------------- # Factory # --------------------------------------------------------------------------- def create_gru_weather_policy_kwargs( hidden_size: int = 64, features_dim: int = 128, ) -> Dict[str, Any]: if not _SB3_AVAILABLE: raise ImportError( "stable-baselines3 not installed. " "Run: pip install stable-baselines3" ) return { "features_extractor_class": GRUWeatherFeaturesExtractor, "features_extractor_kwargs": { "features_dim": features_dim, "hidden_size": hidden_size, }, "net_arch": dict(pi=[128, 64], vf=[128, 64]), } def get_equivariant_policy_class() -> Type[MaskableMultiInputActorCriticPolicy]: if not _MASKABLE_AVAILABLE: raise ImportError( "sb3-contrib not installed. " "Run: pip install sb3-contrib" ) return ZoneEquivariantMaskablePolicy # --------------------------------------------------------------------------- # Self-test # --------------------------------------------------------------------------- def _self_test() -> None: import numpy as np print("gru_weather_policy.py self-test") if not _SB3_AVAILABLE: print(" SKIP: stable-baselines3 not installed") return failures: List[str] = [] def _assert(cond: bool, msg: str) -> None: if not cond: failures.append(msg) print(f" FAIL: {msg}") n_zones = 3 horizon_days = 14 obs_space = gym.spaces.Dict({ "forecast_precip": gym.spaces.Box( low=0, high=500, shape=(n_zones, horizon_days), dtype=np.float32 ), "forecast_uncertainty": gym.spaces.Box( low=0, high=1, shape=(n_zones,), dtype=np.float32 ), "zone_belief": gym.spaces.Box( low=0, high=1, shape=(n_zones,), dtype=np.float32 ), "basin_context": gym.spaces.Box( low=np.array([-5.0, -5.0, -30.0, 900.0, 200.0, 0.0, -9.0, 0.0], dtype=np.float32), high=np.array([5.0, 5.0, 30.0, 1100.0, 900.0, 9.0, -3.0, 2.0], dtype=np.float32), dtype=np.float32, ), }) extractor = GRUWeatherFeaturesExtractor( observation_space=obs_space, features_dim=128, hidden_size=64, ) _assert( extractor.static_mlp[0].in_features == STATIC_DIM, f"static_mlp in_features={extractor.static_mlp[0].in_features} != {STATIC_DIM}", ) batch_size = 2 torch.manual_seed(0) quiet = torch.tensor(BASIN_CONTEXT_NEUTRAL, dtype=torch.float32).unsqueeze(0).repeat(batch_size, 1) storm = quiet.clone() storm[:, 5] = 7.0 # Kp storm[:, 6] = -4.3 # log10(GOES) ~ 5e-5 storm[:, 7] = 2.0 # regime = storm precip = torch.randn(batch_size, n_zones, horizon_days) uncert = torch.rand(batch_size, n_zones) belief = torch.rand(batch_size, n_zones) obs_quiet = { "forecast_precip": precip, "forecast_uncertainty": uncert, "zone_belief": belief, "basin_context": quiet, } obs_storm = { "forecast_precip": precip, "forecast_uncertainty": uncert, "zone_belief": belief, "basin_context": storm, } extractor.eval() with torch.no_grad(): feat_quiet = extractor(obs_quiet) feat_storm = extractor(obs_storm) feat_quiet2 = extractor(obs_quiet) _assert(feat_quiet.shape == (batch_size, n_zones * 64 * 2), f"feature shape {tuple(feat_quiet.shape)}") _assert(extractor._last_zone_scores.shape == (batch_size, n_zones), "zone scores shape") _assert(extractor._last_terminate_logit.shape == (batch_size,), "terminate logit shape") _assert(torch.allclose(feat_quiet, feat_quiet2), "same basin should be deterministic in eval mode") _assert(not torch.allclose(feat_quiet, feat_storm, atol=1e-6), "storm vs quiet basin produced identical features — basin is ignored") print(" Feature extraction + basin sensitivity OK") obs_missing = { "forecast_precip": precip, "forecast_uncertainty": uncert, "zone_belief": belief, } with torch.no_grad(): feat_missing = extractor(obs_missing) feat_neutral = extractor(obs_quiet) _assert(torch.allclose(feat_missing, feat_neutral, atol=1e-5), "missing basin_context should match Neutral default") print(" Missing-key Neutral fallback OK") if _MASKABLE_AVAILABLE: policy_class = get_equivariant_policy_class() _assert(policy_class is ZoneEquivariantMaskablePolicy, "policy class") print(" Policy class OK") kwargs = create_gru_weather_policy_kwargs(hidden_size=64) _assert(kwargs["features_extractor_class"] is GRUWeatherFeaturesExtractor, "kwargs extractor class") _assert(kwargs["features_extractor_kwargs"]["hidden_size"] == 64, "kwargs hidden_size") print(" Policy kwargs OK") if failures: print(f"FAILED {len(failures)} test(s):") for f in failures: print(f" - {f}") raise SystemExit(1) print("All gru_weather_policy self-tests passed.") if __name__ == "__main__": _self_test()