gru_weather_policy.py

"""
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