""" gru_weather_policy.py ===================== GRU-based features extractor + zone-equivariant MaskablePPO policy for WeatherForecastEnv. Architecture (v3 -- invariant value path + equivariant inspect logits) ---------------------------------------------------------------------- The extractor receives a Dict observation with five inputs (action_mask is excluded — MaskablePPO handles that separately): zone_belief [batch, n_zones] forecast_precip [batch, n_zones, horizon_days] forecast_uncertainty [batch, n_zones] prior_belief [batch, 1] basin_context [batch, 8] -- basin-scale teleconnections + helio (NOT per-zone) Processing pipeline: 0. forecast_precip /= precip_scale (fixed; default 40.0) 1. Per-zone encode (shared weights): forecast_proj(scaled_precip) ++ belief ++ uncertainty -> zone_encoder -> zone_embeds [batch, n_zones, zone_embed_dim] 2. EQUIVARIANT inspect path (v3 — must run BEFORE pooling): zone_action_head(zone_embeds) -> [batch, n_zones] These scores preserve slot identity so logit_i tracks zone i's content. Weight shape [1, zone_embed_dim] is independent of n_zones. 3. INVARIANT pooling (for value / global belief / terminate context): attention-weighted sum ++ max over the zone axis -> [batch, 2*zone_embed_dim] 4. basin_context -> MLP; concat with pooled + prior -> GRU -> features 5. terminate_head(features) -> [batch, 1] (n_zones-independent) 6. ZoneEquivariantMaskablePolicy concatenates: logits = cat([zone_scores, terminate_logit]) # [batch, n_zones+1] MaskablePPO then applies env action_masks on top (unchanged). """ from __future__ import annotations from typing import Any, Dict, Optional, Type import numpy as np import torch import torch.nn as nn from gymnasium import spaces from stable_baselines3.common.torch_layers import BaseFeaturesExtractor try: from sb3_contrib.common.maskable.policies import MaskableMultiInputActorCriticPolicy _MASKABLE_AVAILABLE = True except ImportError: # pragma: no cover MaskableMultiInputActorCriticPolicy = object # type: ignore _MASKABLE_AVAILABLE = False class GRUWeatherFeaturesExtractor(BaseFeaturesExtractor): PRECIP_SCALE: float = 40.0 def __init__( self, observation_space: spaces.Dict, hidden_size: int = 64, spatial_output_size: int = 8, zone_embed_dim: int = 16, features_dim: int = 256, basin_context_hidden: int = 8, precip_scale: float = 40.0, ) -> None: super().__init__(observation_space, features_dim=features_dim) self.hidden_size = hidden_size self.spatial_output_size = spatial_output_size self.zone_embed_dim = zone_embed_dim self.basin_context_hidden = basin_context_hidden self.precip_scale = float(precip_scale) if precip_scale > 0 else self.PRECIP_SCALE obs_spaces = observation_space.spaces precip_shape = obs_spaces["forecast_precip"].shape horizon_days = precip_shape[-1] prior_dim = int(np.prod(obs_spaces["prior_belief"].shape)) if "basin_context" not in obs_spaces: raise KeyError( "GRUWeatherFeaturesExtractor: observation_space is missing " "'basin_context'. This extractor requires " "weather_forecast_env.py's v3+ observation space." ) basin_dim = int(np.prod(obs_spaces["basin_context"].shape)) self.forecast_proj = nn.Sequential( nn.Linear(horizon_days, spatial_output_size * 2), nn.ReLU(), nn.Linear(spatial_output_size * 2, spatial_output_size), ) self.zone_encoder = nn.Sequential( nn.Linear(spatial_output_size + 2, zone_embed_dim * 2), nn.ReLU(), nn.Linear(zone_embed_dim * 2, zone_embed_dim), ) self.zone_action_head = nn.Linear(zone_embed_dim, 1) self.attn_score = nn.Linear(zone_embed_dim, 1, bias=False) self.basin_proj = nn.Sequential( nn.Linear(basin_dim, basin_context_hidden), nn.ReLU(), ) gru_input_size = 2 * zone_embed_dim + prior_dim + basin_context_hidden self.gru = nn.GRU( input_size=gru_input_size, hidden_size=hidden_size, num_layers=1, batch_first=True, ) self.output_proj = nn.Sequential( nn.Linear(hidden_size, features_dim), nn.ReLU(), ) self.terminate_head = nn.Linear(features_dim, 1) self._hidden: Optional[torch.Tensor] = None self._last_zone_scores: Optional[torch.Tensor] = None self._last_terminate_logit: Optional[torch.Tensor] = None def set_hidden(self, hidden: Optional[torch.Tensor]) -> None: self._hidden = hidden def get_hidden(self) -> torch.Tensor: if self._hidden is None: raise RuntimeError("Hidden state requested before forward pass.") return self._hidden def reset_hidden( self, batch_size: int = 1, device: Optional[torch.device] = None, ) -> None: dev = device if device is not None else next(self.parameters()).device self._hidden = torch.zeros(1, batch_size, self.hidden_size, device=dev) def forward(self, observations: Dict[str, torch.Tensor]) -> torch.Tensor: zone_belief = observations["zone_belief"].float() # [B, n_zones] forecast = observations["forecast_precip"].float() # [B, n_zones, H] uncertainty = observations["forecast_uncertainty"].float() # [B, n_zones] prior = observations["prior_belief"].float() basin_context = observations["basin_context"].float() forecast = forecast / self.precip_scale batch_size, n_zones = zone_belief.shape device = zone_belief.device fc_flat = forecast.reshape(batch_size * n_zones, -1) fc_proj = self.forecast_proj(fc_flat) zone_belief_flat = zone_belief.reshape(batch_size * n_zones, 1) uncertainty_flat = uncertainty.reshape(batch_size * n_zones, 1) per_zone_raw = torch.cat([fc_proj, zone_belief_flat, uncertainty_flat], dim=-1) zone_embeds_flat = self.zone_encoder(per_zone_raw) zone_embeds = zone_embeds_flat.reshape(batch_size, n_zones, self.zone_embed_dim) self._last_zone_scores = self.zone_action_head(zone_embeds).squeeze(-1) # [B, n_zones] attn_logits = self.attn_score(zone_embeds).squeeze(-1) attn_weights = torch.softmax(attn_logits, dim=1).unsqueeze(-1) attended = (attn_weights * zone_embeds).sum(dim=1) max_pooled = zone_embeds.max(dim=1).values pooled = torch.cat([attended, max_pooled], dim=-1) basin_embed = self.basin_proj(basin_context) gru_input = torch.cat([pooled, prior, basin_embed], dim=-1).unsqueeze(1) if self._hidden is None or self._hidden.shape[1] != batch_size: self._hidden = torch.zeros(1, batch_size, self.hidden_size, device=device) else: self._hidden = self._hidden.to(device) gru_out, self._hidden = self.gru(gru_input, self._hidden) self._hidden = self._hidden.detach() features = self.output_proj(gru_out.squeeze(1)) self._last_terminate_logit = self.terminate_head(features) # [B, 1] return features if _MASKABLE_AVAILABLE: class ZoneEquivariantMaskablePolicy(MaskableMultiInputActorCriticPolicy): def _get_action_dist_from_latent(self, latent_pi: torch.Tensor): fe = self.features_extractor if not isinstance(fe, GRUWeatherFeaturesExtractor): return super()._get_action_dist_from_latent(latent_pi) zone_scores = fe._last_zone_scores term = fe._last_terminate_logit if zone_scores is None or term is None: raise RuntimeError( "ZoneEquivariantMaskablePolicy: features extractor has not " "stashed zone/terminate scores. extract_features must run " "before _get_action_dist_from_latent." ) logits = torch.cat([zone_scores, term], dim=-1) return self.action_dist.proba_distribution(action_logits=logits) else: # pragma: no cover class ZoneEquivariantMaskablePolicy: # type: ignore def __init__(self, *args, **kwargs): raise ImportError( "ZoneEquivariantMaskablePolicy requires sb3_contrib " "(MaskableMultiInputActorCriticPolicy)." ) def create_gru_weather_policy_kwargs( hidden_size: int = 64, spatial_output_size: int = 8, zone_embed_dim: int = 16, features_dim: int = 256, basin_context_hidden: int = 8, precip_scale: float = 40.0, net_arch: Optional[dict] = None, ) -> dict: if net_arch is None: net_arch = dict(pi=[128, 64], vf=[128, 64]) return dict( features_extractor_class=GRUWeatherFeaturesExtractor, features_extractor_kwargs=dict( hidden_size=hidden_size, spatial_output_size=spatial_output_size, zone_embed_dim=zone_embed_dim, features_dim=features_dim, basin_context_hidden=basin_context_hidden, precip_scale=precip_scale, ), net_arch=net_arch, ) def get_equivariant_policy_class() -> Type[Any]: if not _MASKABLE_AVAILABLE: raise ImportError("sb3_contrib is required for ZoneEquivariantMaskablePolicy") return ZoneEquivariantMaskablePolicy