physics_dynamics.py

"""
physics_dynamics.py
===================
Physics-informed dynamics model for WeatherForecastEnv.

Architecture
------------
This is a Dyna-style learned dynamics model: trained offline on ERA5 data,
then used during PPO training to generate synthetic rollouts that augment
real environment experience. It does NOT replace the environment — it
supplements it, improving sample efficiency and generalization.

The model predicts how the zone-level forecast state evolves over time,
constrained by an advection-diffusion PDE residual that prevents physically
impossible predictions (e.g. precipitation materialising from nothing,
uncertainty decreasing without new observations).


Classes
-------
ZoneStateTensor         — named container for the three observation arrays
TemporalDynamicsModel   — core learned model (GRU encoder + MLP transition)
PhysicsResidualLoss     — advection-diffusion PDE residual along time axis
DynamicsTrainer         — offline pre-training on ERA5 sequences
EnsembleDynamics        — N models for epistemic uncertainty quantification
DynaRolloutBuffer       — generates synthetic transitions for PPO augmentation

Usage
-----
    # 1. Pre-train on ERA5 sequences
    trainer = DynamicsTrainer(n_zones=4, horizon_days=14)
    trainer.train(era5_sequences)                  # list of ZoneStateTensor
    trainer.save("dynamics_model.pt")

    # 2. Load in training loop and generate synthetic rollouts
    dynamics = TemporalDynamicsModel.load("dynamics_model.pt")
    buffer = DynaRolloutBuffer(dynamics, n_synthetic_per_real=4)
    # Pass buffer to custom PPO callback (see train_curriculum.py notes)
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Data container
# ---------------------------------------------------------------------------

@dataclass
class ZoneStateTensor:
    """
    A single time-step of zone-level state, as tensors.

    Mirrors the WeatherForecastEnv observation space:
        precip:       [batch, n_zones, horizon_days]
        uncertainty:  [batch, n_zones]
        belief:       [batch, n_zones]

    All values float32 in their natural ranges:
        precip        [0, 500]  mm
        uncertainty   [0, 1]
        belief        [0, 1]
    """
    precip: torch.Tensor       # [batch, n_zones, horizon_days]
    uncertainty: torch.Tensor  # [batch, n_zones]
    belief: torch.Tensor       # [batch, n_zones]

    @property
    def batch_size(self) -> int:
        return self.precip.shape[0]

    @property
    def n_zones(self) -> int:
        return self.precip.shape[1]

    @property
    def horizon_days(self) -> int:
        return self.precip.shape[2]

    def to(self, device: torch.device) -> "ZoneStateTensor":
        return ZoneStateTensor(
            precip=self.precip.to(device),
            uncertainty=self.uncertainty.to(device),
            belief=self.belief.to(device),
        )

    def flat(self) -> torch.Tensor:
        """Flatten to [batch, n_zones * (horizon_days + 2)] for MLP input."""
        B, Z, H = self.precip.shape
        precip_flat = self.precip.reshape(B, Z * H)
        return torch.cat([precip_flat, self.uncertainty, self.belief], dim=-1)

    @property
    def flat_dim(self) -> int:
        return self.n_zones * (self.horizon_days + 2)

    @classmethod
    def from_numpy(
        cls,
        precip: np.ndarray,
        uncertainty: np.ndarray,
        belief: np.ndarray,
    ) -> "ZoneStateTensor":
        """Construct from numpy arrays (adds batch dim if missing)."""
        if precip.ndim == 2:
            precip = precip[None]
        if uncertainty.ndim == 1:
            uncertainty = uncertainty[None]
        if belief.ndim == 1:
            belief = belief[None]
        return cls(
            precip=torch.from_numpy(precip.astype(np.float32)),
            uncertainty=torch.from_numpy(uncertainty.astype(np.float32)),
            belief=torch.from_numpy(belief.astype(np.float32)),
        )


# ---------------------------------------------------------------------------
# Physics residual
# ---------------------------------------------------------------------------

class PhysicsResidualLoss(nn.Module):
    """
    Advection-diffusion PDE residual along the TEMPORAL dimension.

    The physical intuition: as time advances by dt days, the precipitation
    forecast at day t in the horizon should approximately equal the forecast
    at day (t - dt) from the previous time step, shifted by advection and
    smoothed by diffusion. This is the atmospheric forecast evolution equation.

    Residual: ∂u/∂t + v ∂u/∂τ - D ∂²u/∂τ² = 0

    where:
        u  = precipitation forecast value
        t  = real time (day-to-day model evolution)
        τ  = forecast lead time (the horizon axis, days 0..H-1)
        v  = advection speed (learnable)
        D  = diffusion coefficient (learnable)

    This is applied per zone independently (zones are not a spatial grid).

    Args:
        weight: Scalar multiplier for the physics loss term.
                Start with 0.01-0.05; increase if predictions are unphysical.
    """

    def __init__(self, weight: float = 0.01):
        super().__init__()
        self.weight = weight
        # Learnable physics parameters — initialised to physically plausible values
        # v: forecast advection ~1 day/day (forecast evolves with real time)
        # D: diffusion smoothing ~0.1 (moderate smoothing of forecast errors)
        self.log_v = nn.Parameter(torch.tensor(0.0))   # exp(0) = 1.0
        self.log_D = nn.Parameter(torch.tensor(-2.3))  # exp(-2.3) ≈ 0.1

    @property
    def v(self) -> torch.Tensor:
        return torch.exp(self.log_v)

    @property
    def D(self) -> torch.Tensor:
        return torch.exp(self.log_D)

    def forward(
        self,
        u_current: torch.Tensor,   # [batch, n_zones, horizon_days]
        u_next: torch.Tensor,      # [batch, n_zones, horizon_days]
        dt: float = 1.0,
    ) -> torch.Tensor:
        """
        Compute mean squared PDE residual.

        u_current: forecast at time t
        u_next:    predicted forecast at time t + dt
        dt:        real-time step in days
        """
        B, Z, H = u_current.shape

        # ∂u/∂t ≈ (u_next - u_current) / dt
        du_dt = (u_next - u_current) / dt

        # ∂u/∂τ — first derivative along horizon axis (central differences)
        # Shape: [batch, n_zones, horizon_days]
        du_dtau = torch.zeros_like(u_current)
        if H > 2:
            du_dtau[:, :, 1:-1] = (u_current[:, :, 2:] - u_current[:, :, :-2]) / 2.0
            du_dtau[:, :, 0]    = u_current[:, :, 1] - u_current[:, :, 0]
            du_dtau[:, :, -1]   = u_current[:, :, -1] - u_current[:, :, -2]

        # ∂²u/∂τ² — second derivative along horizon axis (Laplacian)
        d2u_dtau2 = torch.zeros_like(u_current)
        if H > 2:
            d2u_dtau2[:, :, 1:-1] = (
                u_current[:, :, 2:] - 2 * u_current[:, :, 1:-1] + u_current[:, :, :-2]
            )
            d2u_dtau2[:, :, 0]  = d2u_dtau2[:, :, 1]
            d2u_dtau2[:, :, -1] = d2u_dtau2[:, :, -2]

        # PDE residual: ∂u/∂t + v·∂u/∂τ - D·∂²u/∂τ² = 0
        residual = du_dt + self.v * du_dtau - self.D * d2u_dtau2

        return self.weight * torch.mean(residual ** 2)


# ---------------------------------------------------------------------------
# Core dynamics model
# ---------------------------------------------------------------------------

class TemporalDynamicsModel(nn.Module):
    """
    Learned dynamics model: predicts next zone state from current state.

    Architecture:
        1. Per-zone GRU encoder compresses the forecast horizon sequence
           into a latent zone embedding.
        2. MLP transition model maps current latent → next latent.
        3. MLP decoder reconstructs full next-state from latent.

    Why GRU encoder (not FNO): Your forecast data is [n_zones, horizon_days]
    where n_zones is small (2-4) and horizon_days is short (7-14). FNO is
    designed for large spatial fields (64x64+). A GRU over the horizon axis
    per zone is exact for this scale and directly compatible with your
    existing GRUWeatherFeaturesExtractor architecture.

    Args:
        n_zones:      Number of geographic zones (matches env config)
        horizon_days: Forecast horizon length (matches env config)
        latent_dim:   Dimension of per-zone latent embedding
        hidden_dim:   MLP hidden size for transition and decoder
    """

    def __init__(
        self,
        n_zones: int = 4,
        horizon_days: int = 14,
        latent_dim: int = 32,
        hidden_dim: int = 128,
    ):
        super().__init__()
        self.n_zones = n_zones
        self.horizon_days = horizon_days
        self.latent_dim = latent_dim

        # --- Encoder: horizon sequence → zone latent ---
        # Applied identically to each zone (weight sharing)
        self.precip_encoder = nn.GRU(
            input_size=1,
            hidden_size=latent_dim,
            num_layers=1,
            batch_first=True,
        )

        # Zone metadata (uncertainty + belief) → extra latent dims
        self.meta_encoder = nn.Sequential(
            nn.Linear(2, latent_dim),
            nn.Tanh(),
        )

        zone_latent_dim = latent_dim * 2  # precip latent + meta latent

        # --- Transition: current latent → next latent (all zones jointly) ---
        full_latent_dim = n_zones * zone_latent_dim
        self.transition = nn.Sequential(
            nn.Linear(full_latent_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, full_latent_dim),
        )

        # --- Decoder: latent → next state components ---
        self.precip_decoder = nn.Sequential(
            nn.Linear(zone_latent_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, horizon_days),
            nn.Softplus(),  # precipitation ≥ 0
        )
        self.uncertainty_decoder = nn.Sequential(
            nn.Linear(zone_latent_dim, 32),
            nn.SiLU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),  # uncertainty in [0, 1]
        )
        self.belief_decoder = nn.Sequential(
            nn.Linear(zone_latent_dim, 32),
            nn.SiLU(),
            nn.Linear(32, 1),
            nn.Sigmoid(),  # belief in [0, 1]
        )

        # Physics loss module (parameters learned jointly with model)
        self.physics_loss = PhysicsResidualLoss(weight=0.01)

        logger.info(
            "TemporalDynamicsModel: n_zones=%d horizon=%d latent=%d hidden=%d",
            n_zones, horizon_days, latent_dim, hidden_dim,
        )

    def _encode(self, state: ZoneStateTensor) -> torch.Tensor:
        """Encode zone state → latent. Returns [batch, n_zones, zone_latent_dim]."""
        B, Z, H = state.precip.shape

        # Encode each zone's precipitation forecast sequence with the GRU
        # Reshape to [batch * n_zones, horizon_days, 1] for GRU
        precip_seq = state.precip.reshape(B * Z, H, 1)
        _, h_n = self.precip_encoder(precip_seq)  # h_n: [1, B*Z, latent_dim]
        precip_latent = h_n.squeeze(0).reshape(B, Z, self.latent_dim)

        # Encode per-zone metadata [uncertainty, belief]
        meta = torch.stack([state.uncertainty, state.belief], dim=-1)  # [B, Z, 2]
        meta_flat = meta.reshape(B * Z, 2)
        meta_latent = self.meta_encoder(meta_flat).reshape(B, Z, self.latent_dim)

        return torch.cat([precip_latent, meta_latent], dim=-1)  # [B, Z, 2*latent_dim]

    def forward(
        self,
        current: ZoneStateTensor,
        return_physics_loss: bool = True,
        dt: float = 1.0,
    ) -> Tuple[ZoneStateTensor, Optional[torch.Tensor]]:
        """
        Predict next state from current state.

        Args:
            current: Current zone state
            return_physics_loss: Whether to compute and return the physics residual
            dt: Real-time gap in days between current and next snapshot.
                Must match pairing cadence (e.g. 5.0 for cache step_days=5).

        Returns:
            next_state: Predicted next zone state
            physics_loss: PDE residual loss (None if return_physics_loss=False)
        """
        B, Z, H = current.precip.shape

        # Encode current state
        latent = self._encode(current)              # [B, Z, zone_latent_dim]
        latent_flat = latent.reshape(B, -1)         # [B, Z * zone_latent_dim]

        # Transition in latent space (all zones jointly — captures inter-zone correlations)
        next_latent_flat = self.transition(latent_flat)
        next_latent = next_latent_flat.reshape(B, Z, -1)  # [B, Z, zone_latent_dim]

        # Decode next state per zone
        next_latent_per_zone = next_latent.reshape(B * Z, -1)

        next_precip = self.precip_decoder(next_latent_per_zone).reshape(B, Z, H)
        next_uncertainty = self.uncertainty_decoder(next_latent_per_zone).reshape(B, Z)
        next_belief = self.belief_decoder(next_latent_per_zone).reshape(B, Z)

        next_state = ZoneStateTensor(
            precip=next_precip,
            uncertainty=next_uncertainty,
            belief=next_belief,
        )

        # Physics residual loss on the precipitation forecast evolution
        phys_loss = None
        if return_physics_loss:
            phys_loss = self.physics_loss(
                current.precip, next_precip, dt=float(dt),
            )

        return next_state, phys_loss

    def rollout(
        self,
        initial: ZoneStateTensor,
        steps: int = 5,
    ) -> List[ZoneStateTensor]:
        """
        Generate a multi-step synthetic rollout.

        Used by DynaRolloutBuffer to produce model-imagined transitions
        for PPO augmentation. Gradients are not tracked here (inference only).

        Returns list of states [s_0, s_1, ..., s_steps] where s_0 = initial.
        """
        states = [initial]
        current = initial
        with torch.no_grad():
            for _ in range(steps):
                next_state, _ = self.forward(current, return_physics_loss=False)
                # Clamp to valid ranges
                next_state = ZoneStateTensor(
                    precip=torch.clamp(next_state.precip, 0.0, 500.0),
                    uncertainty=torch.clamp(next_state.uncertainty, 0.0, 1.0),
                    belief=torch.clamp(next_state.belief, 0.0, 1.0),
                )
                states.append(next_state)
                current = next_state
        return states

    def save(self, path: str | Path) -> None:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        torch.save({
            "state_dict": self.state_dict(),
            "config": {
                "n_zones": self.n_zones,
                "horizon_days": self.horizon_days,
                "latent_dim": self.latent_dim,
                "hidden_dim": self.physics_loss.__class__.__name__,  # for verification
            }
        }, path)
        logger.info("Saved dynamics model to %s", path)

    @classmethod
    def load(cls, path: str | Path, device: Optional[torch.device] = None) -> "TemporalDynamicsModel":
        path = Path(path)
        checkpoint = torch.load(path, map_location=device or "cpu")
        cfg = checkpoint["config"]
        model = cls(
            n_zones=cfg["n_zones"],
            horizon_days=cfg["horizon_days"],
            latent_dim=cfg.get("latent_dim", 32),
        )
        model.load_state_dict(checkpoint["state_dict"])
        logger.info("Loaded dynamics model from %s", path)
        return model


# ---------------------------------------------------------------------------
# Offline trainer
# ---------------------------------------------------------------------------

class DynamicsTrainer:
    """
    Pre-trains TemporalDynamicsModel on sequences of zone states.

    Training data format: list of consecutive (current, next) ZoneStateTensor
    pairs extracted from ERA5 reanalysis or from environment rollouts.

    Loss: data_loss + physics_loss
        data_loss = MSE(predicted_next, actual_next) for all three components
        physics_loss = advection-diffusion PDE residual (see PhysicsResidualLoss)

    Args:
        n_zones:         Must match your env config
        horizon_days:    Must match your env config
        physics_weight:  Weight for physics residual in total loss.
                         Start at 0.01, increase to 0.1 if predictions violate physics.
        device:          'cuda' if available, else 'cpu'
    """

    def __init__(
        self,
        n_zones: int = 4,
        horizon_days: int = 14,
        latent_dim: int = 32,
        hidden_dim: int = 128,
        physics_weight: float = 0.01,
        device: Optional[str] = None,
    ):
        self.device = torch.device(
            device or ("cuda" if torch.cuda.is_available() else "cpu")
        )
        self.model = TemporalDynamicsModel(
            n_zones=n_zones,
            horizon_days=horizon_days,
            latent_dim=latent_dim,
            hidden_dim=hidden_dim,
        ).to(self.device)
        self.physics_weight = physics_weight
        logger.info("DynamicsTrainer: device=%s  physics_weight=%.3f", self.device, physics_weight)

    def train(
        self,
        sequence_pairs: List[Tuple[ZoneStateTensor, ZoneStateTensor]],
        epochs: int = 50,
        batch_size: int = 64,
        lr: float = 1e-3,
        val_split: float = 0.1,
        dts: Optional[List[float]] = None,
        default_dt: float = 1.0,
    ) -> dict:
        """
        Train the dynamics model on (current_state, next_state) pairs.

        Args:
            sequence_pairs: List of (current, next) ZoneStateTensor pairs.
            epochs:         Training epochs
            batch_size:     Batch size
            lr:             Learning rate
            val_split:      Fraction of data held out for validation
            dts:            Optional per-pair real-time gaps in days (same
                            length as sequence_pairs). When None, uses
                            default_dt for every pair.
            default_dt:     Fallback dt (days). Use 5.0 for historical
                            cache pairs built at step_days=5.

        Returns:
            Training history dict with 'train_loss' and 'val_loss' lists.
        """
        if not sequence_pairs:
            raise ValueError("sequence_pairs is empty — provide ERA5 data")

        if dts is not None and len(dts) != len(sequence_pairs):
            raise ValueError(
                f"dts length {len(dts)} != sequence_pairs length "
                f"{len(sequence_pairs)}"
            )

        # Build tensor dataset from pairs
        current_precips, current_uncerts, current_beliefs = [], [], []
        next_precips, next_uncerts, next_beliefs = [], [], []
        dt_list: List[float] = []

        for i, (curr, nxt) in enumerate(sequence_pairs):
            current_precips.append(curr.precip)
            current_uncerts.append(curr.uncertainty)
            current_beliefs.append(curr.belief)
            next_precips.append(nxt.precip)
            next_uncerts.append(nxt.uncertainty)
            next_beliefs.append(nxt.belief)
            dt_list.append(float(dts[i]) if dts is not None else float(default_dt))

        # Stack along batch dimension
        cp = torch.cat(current_precips, dim=0)
        cu = torch.cat(current_uncerts, dim=0)
        cb = torch.cat(current_beliefs, dim=0)
        np_ = torch.cat(next_precips, dim=0)
        nu = torch.cat(next_uncerts, dim=0)
        nb = torch.cat(next_beliefs, dim=0)
        dt_t = torch.tensor(dt_list, dtype=torch.float32)

        N = cp.shape[0]
        # Guard: we need at least 1 sample in the training split.
        # When val_split=0 or N is too small, skip validation entirely.
        n_val = int(N * val_split) if val_split > 0 else 0
        if n_val >= N:
            n_val = max(0, N - 1)   # leave at least 1 sample for training
        n_train = N - n_val
        if n_train <= 0:
            raise ValueError(
                f"Dataset too small for the requested val_split: "
                f"N={N}, val_split={val_split} produces n_train={n_train}. "
                f"Reduce val_split or provide more pairs."
            )

        train_ds = TensorDataset(
            cp[:n_train], cu[:n_train], cb[:n_train],
            np_[:n_train], nu[:n_train], nb[:n_train],
            dt_t[:n_train],
        )
        val_ds = TensorDataset(
            cp[n_train:], cu[n_train:], cb[n_train:],
            np_[n_train:], nu[n_train:], nb[n_train:],
            dt_t[n_train:],
        )

        train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
        val_loader   = DataLoader(val_ds,   batch_size=batch_size, shuffle=False)

        optimizer = torch.optim.AdamW(self.model.parameters(), lr=lr, weight_decay=1e-4)
        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)

        history = {"train_loss": [], "val_loss": [], "physics_loss": [], "dt_mean": float(dt_t.mean())}

        for epoch in range(epochs):
            # --- Train ---
            self.model.train()
            epoch_data_loss = 0.0
            epoch_phys_loss = 0.0

            for batch in train_loader:
                cp_b, cu_b, cb_b, np_b, nu_b, nb_b, dt_b = [
                    t.to(self.device) for t in batch
                ]

                current = ZoneStateTensor(precip=cp_b, uncertainty=cu_b, belief=cb_b)
                target  = ZoneStateTensor(precip=np_b, uncertainty=nu_b, belief=nb_b)

                # Batch may mix dts; use batch mean (pairs are homogeneous
                # when built with exact step_days only).
                batch_dt = float(dt_b.mean().item())
                pred, phys_loss = self.model(
                    current, return_physics_loss=True, dt=batch_dt,
                )

                # Data fidelity: MSE on all three components
                # Normalise precip by max scale (500mm) to balance loss magnitudes
                data_loss = (
                    F.mse_loss(pred.precip / 500.0, target.precip / 500.0)
                    + F.mse_loss(pred.uncertainty, target.uncertainty)
                    + F.mse_loss(pred.belief,      target.belief)
                )

                total_loss = data_loss + self.physics_weight * phys_loss

                optimizer.zero_grad()
                total_loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                optimizer.step()

                epoch_data_loss += data_loss.item()
                epoch_phys_loss += phys_loss.item()

            scheduler.step()

            avg_data = epoch_data_loss / len(train_loader)
            avg_phys = epoch_phys_loss / len(train_loader)

            # --- Validate ---
            self.model.eval()
            val_loss = 0.0
            with torch.no_grad():
                for batch in val_loader:
                    cp_b, cu_b, cb_b, np_b, nu_b, nb_b, dt_b = [
                        t.to(self.device) for t in batch
                    ]
                    current = ZoneStateTensor(precip=cp_b, uncertainty=cu_b, belief=cb_b)
                    target  = ZoneStateTensor(precip=np_b, uncertainty=nu_b, belief=nb_b)
                    batch_dt = float(dt_b.mean().item()) if dt_b.numel() else 1.0
                    pred, _ = self.model(current, return_physics_loss=False)
                    val_loss += (
                        F.mse_loss(pred.precip / 500.0, target.precip / 500.0)
                        + F.mse_loss(pred.uncertainty, target.uncertainty)
                        + F.mse_loss(pred.belief,      target.belief)
                    ).item()

            avg_val = val_loss / max(len(val_loader), 1)

            history["train_loss"].append(avg_data)
            history["val_loss"].append(avg_val)
            history["physics_loss"].append(avg_phys)

            if epoch % 10 == 0 or epoch == epochs - 1:
                logger.info(
                    "Epoch %3d/%d  train=%.4f  val=%.4f  physics=%.4f  "
                    "v=%.3f  D=%.3f",
                    epoch + 1, epochs, avg_data, avg_val, avg_phys,
                    self.model.physics_loss.v.item(),
                    self.model.physics_loss.D.item(),
                )

        return history

    def save(self, path: str | Path) -> None:
        self.model.save(path)


# ---------------------------------------------------------------------------
# Ensemble for uncertainty quantification
# ---------------------------------------------------------------------------

class EnsembleDynamics:
    """
    Ensemble of N dynamics models for epistemic uncertainty quantification.

    Each model is trained with different random seed initialization.
    Disagreement between models = epistemic uncertainty = regions where
    the policy should be conservative (important for plasma/fusion applications
    where high uncertainty = potentially dangerous operating regime).

    Usage in PPO training:
        ensemble = EnsembleDynamics(n_models=5, ...)
        mean_next, uncertainty = ensemble.predict(current_state)
        # Add uncertainty penalty to reward: reward -= uncertainty_weight * uncertainty
    """

    def __init__(self, n_models: int = 5, **model_kwargs):
        self.models = [TemporalDynamicsModel(**model_kwargs) for _ in range(n_models)]
        logger.info("EnsembleDynamics: %d models", n_models)

    def predict(
        self,
        current: ZoneStateTensor,
    ) -> Tuple[ZoneStateTensor, torch.Tensor]:
        """
        Return mean prediction and epistemic uncertainty.

        uncertainty is a scalar tensor: mean std across all zones and features.
        Use this to penalise the RL policy for actions that lead to high
        uncertainty states (encourages conservative, well-characterised behaviour).
        """
        all_precips, all_uncerts, all_beliefs = [], [], []

        for model in self.models:
            model.eval()
            with torch.no_grad():
                pred, _ = model(current, return_physics_loss=False)
            all_precips.append(pred.precip)
            all_uncerts.append(pred.uncertainty)
            all_beliefs.append(pred.belief)

        precip_stack = torch.stack(all_precips)    # [N, B, Z, H]
        uncert_stack = torch.stack(all_uncerts)    # [N, B, Z]
        belief_stack = torch.stack(all_beliefs)    # [N, B, Z]

        mean_state = ZoneStateTensor(
            precip=precip_stack.mean(0),
            uncertainty=uncert_stack.mean(0),
            belief=belief_stack.mean(0),
        )

        # Epistemic uncertainty: normalised std across ensemble members
        epistemic = (
            # correction=0 avoids NaN when n_models=1 (Bessel correction
            # would divide by zero with a single sample).
            (precip_stack.std(0, correction=0) / 500.0).mean()
            + uncert_stack.std(0, correction=0).mean()
            + belief_stack.std(0, correction=0).mean()
        ) / 3.0

        return mean_state, epistemic

    def to(self, device: torch.device) -> "EnsembleDynamics":
        for m in self.models:
            m.to(device)
        return self


# ---------------------------------------------------------------------------
# Dyna rollout buffer
# ---------------------------------------------------------------------------

class DynaRolloutBuffer:
    """
    Generates synthetic (s, a, r, s') transitions for PPO augmentation.

    Dyna-style model-based RL: use the learned dynamics model to generate
    additional training transitions from states already in the replay buffer.
    This improves sample efficiency without changing the PPO algorithm.

    Integration with train_curriculum.py:
    -------------------------------------
    Add a DynaCallback to the PPO training loop:

        class DynaCallback(BaseCallback):
            def __init__(self, dynamics: TemporalDynamicsModel, n_synthetic: int = 4):
                super().__init__()
                self.dynamics = dynamics
                self.n_synthetic = n_synthetic

            def _on_rollout_end(self) -> None:
                # After each real rollout, generate synthetic transitions
                # and inject them into the rollout buffer before the update.
                # (Implementation depends on SB3 internals — see notes below.)
                pass

    Note: Direct rollout buffer injection is not officially supported in SB3.
    The practical approach is to use the dynamics model for reward shaping:
    predict next state, measure surprise (|| actual - predicted ||), and add
    a small exploration bonus for high-surprise transitions. This requires
    no SB3 modifications and still improves sample efficiency.

    Args:
        dynamics:           Trained TemporalDynamicsModel
        n_synthetic_steps:  Steps to roll out from each seed state
        uncertainty_weight: Weight for epistemic uncertainty penalty in reward
    """

    def __init__(
        self,
        dynamics: TemporalDynamicsModel,
        n_synthetic_steps: int = 3,
        uncertainty_weight: float = 0.1,
    ):
        self.dynamics = dynamics
        self.n_synthetic_steps = n_synthetic_steps
        self.uncertainty_weight = uncertainty_weight

    def compute_surprise_bonus(
        self,
        obs_current: ZoneStateTensor,
        obs_actual_next: ZoneStateTensor,
    ) -> torch.Tensor:
        """
        Reward bonus for transitions that surprise the dynamics model.

        High surprise = model uncertainty = exploration bonus.
        This is the simplest Dyna integration: no SB3 modifications needed,
        just add this to the reward in a step callback.

        Returns scalar bonus in [0, ~1].
        """
        self.dynamics.eval()
        with torch.no_grad():
            pred_next, _ = self.dynamics(obs_current, return_physics_loss=False)

        # Normalised prediction error
        precip_err = F.mse_loss(
            pred_next.precip / 500.0,
            obs_actual_next.precip / 500.0,
        )
        uncert_err = F.mse_loss(pred_next.uncertainty, obs_actual_next.uncertainty)
        belief_err = F.mse_loss(pred_next.belief, obs_actual_next.belief)

        surprise = (precip_err + uncert_err + belief_err) / 3.0
        return torch.clamp(surprise * self.uncertainty_weight, 0.0, 1.0)

    def generate_rollout(
        self,
        seed_state: ZoneStateTensor,
    ) -> List[ZoneStateTensor]:
        """
        Generate synthetic state sequence from a seed state.

        Returns list of n_synthetic_steps + 1 states starting from seed_state.
        """
        return self.dynamics.rollout(seed_state, steps=self.n_synthetic_steps)