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