gru_weather_policy.py

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