#!/usr/bin/env python3 """ train_kaggle.py =============== Single-phase MaskablePPO trainer for Kaggle Notebooks / TPU / GPU. Validated hyperparameters from ablation study (see BEST_HYPERPARAMETERS). """ from __future__ import annotations import argparse import json import logging import os import random import sys import time import warnings from collections import deque from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import numpy as np # --------------------------------------------------------------------------- # Repo bootstrap (Kaggle input is read-only) # --------------------------------------------------------------------------- REPO = Path("/kaggle/input/datasets/dhmmmreally/weather-modeller") if str(REPO) not in sys.path: sys.path.insert(0, str(REPO)) import zone_observation as _zo assert _zo.SCHEMA_VERSION == 3, ( f"train_kaggle: zone_observation schema mismatch " f"(expected 3, got {_zo.SCHEMA_VERSION})" ) from zone_observation import ForecastConfig from crop_risk_scorer import RiskWeights # --------------------------------------------------------------------------- # Optional ML imports # --------------------------------------------------------------------------- try: import torch _TORCH_AVAILABLE = True except ImportError: _TORCH_AVAILABLE = False try: from weather_forecast_env import make_weather_env from sb3_contrib import MaskablePPO from stable_baselines3.common.monitor import Monitor from stable_baselines3.common.callbacks import BaseCallback _ML_AVAILABLE = True except ImportError as _e: _ML_AVAILABLE = False _ML_IMPORT_ERROR = str(_e) make_weather_env = None MaskablePPO = None Monitor = None BaseCallback = object try: from gru_weather_policy import create_gru_weather_policy_kwargs, get_equivariant_policy_class _GRU_AVAILABLE = True except ImportError: _GRU_AVAILABLE = False create_gru_weather_policy_kwargs = None get_equivariant_policy_class = None try: from physics_dynamics import TemporalDynamicsModel, DynaRolloutBuffer, ZoneStateTensor _DYNAMICS_AVAILABLE = True except ImportError: _DYNAMICS_AVAILABLE = False TemporalDynamicsModel = None DynaRolloutBuffer = None ZoneStateTensor = None warnings.filterwarnings("ignore", category=UserWarning) logging.basicConfig( level=logging.INFO, format="%(asctime)s | %(levelname)s | %(name)s | %(message)s", ) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Hyperparameters (ablation-validated) # --------------------------------------------------------------------------- BEST_HYPERPARAMETERS: Dict[str, Any] = { "learning_rate": 3e-4, "n_steps": 4096, "batch_size": 256, "n_epochs": 10, "gamma": 0.995, "gae_lambda": 0.95, "clip_range": 0.2, "ent_coef": 0.02, "vf_coef": 0.5, "max_grad_norm": 0.5, } # --------------------------------------------------------------------------- # Regression watch callback # --------------------------------------------------------------------------- class RegressionWatch(BaseCallback): """Abort training if mean reward collapses vs. a rolling baseline.""" def __init__( self, window: int = 20, threshold: float = -0.30, patience: int = 3, ) -> None: super().__init__() self.window = window self.threshold = threshold self.patience = patience self._history: deque = deque(maxlen=window) self._strikes = 0 def _on_step(self) -> bool: if len(self.model.ep_info_buffer) == 0: return True recent = [ep["r"] for ep in self.model.ep_info_buffer][-self.window :] if len(recent) < self.window // 2: return True mean_recent = float(np.mean(recent)) self._history.append(mean_recent) if len(self._history) < self.window: return True baseline = float(np.mean(list(self._history)[: self.window // 2])) drop = (mean_recent - baseline) / max(abs(baseline), 1.0) if drop < self.threshold: self._strikes += 1 logger.warning( "RegressionWatch: mean reward dropped %.1f%% (%d/%d strikes)", 100 * drop, self._strikes, self.patience, ) if self._strikes >= self.patience: logger.error("RegressionWatch: aborting training — reward collapse.") return False else: self._strikes = max(0, self._strikes - 1) return True # --------------------------------------------------------------------------- # Model builder # --------------------------------------------------------------------------- def build_model(env, args: argparse.Namespace): if not _ML_AVAILABLE: raise RuntimeError(f"ML stack missing: {_ML_IMPORT_ERROR}") if _GRU_AVAILABLE and create_gru_weather_policy_kwargs is not None: policy_kwargs = create_gru_weather_policy_kwargs( hidden_size=args.hidden_size, features_dim=args.hidden_size * 2, ) policy = get_equivariant_policy_class() if get_equivariant_policy_class is not None else "MultiInputPolicy" logger.info("Using GRU policy (hidden_size=%d)", args.hidden_size) else: policy_kwargs = dict(net_arch=dict(pi=[128, 64], vf=[128, 64])) policy = "MultiInputPolicy" logger.info("GRU unavailable — using MLP policy") ppo_kwargs = dict( learning_rate=args.learning_rate, n_steps=args.n_steps, batch_size=args.batch_size, n_epochs=args.n_epochs, gamma=args.gamma, gae_lambda=args.gae_lambda, clip_range=args.clip_range, ent_coef=args.ent_coef, vf_coef=args.vf_coef, max_grad_norm=args.max_grad_norm, device=args.device, verbose=1, seed=args.seed, ) model = MaskablePPO( policy=policy, env=env, policy_kwargs=policy_kwargs, **ppo_kwargs, ) return model # --------------------------------------------------------------------------- # Dyna callback (mirrors train_curriculum.py) # --------------------------------------------------------------------------- class DynaCallback(BaseCallback): _OBS_KEYS = ("forecast_precip", "forecast_uncertainty", "zone_belief") def __init__( self, dyna_buffer: "DynaRolloutBuffer", surprise_weight: float = 0.05, update_every: int = 0, fine_tune_epochs: int = 3, device: str = "cpu", ) -> None: super().__init__() self.dyna_buffer = dyna_buffer self.surprise_weight = surprise_weight self.update_every = update_every self.fine_tune_epochs = fine_tune_epochs self.device = device self._transition_buffer: list = [] self._tb_max = 10_000 self._bonus_sum = 0.0 self._bonus_count = 0 self._log_freq = 10_000 self._last_log = 0 def _obs_to_state_tensor(self, obs: dict) -> Optional["ZoneStateTensor"]: if not all(k in obs for k in self._OBS_KEYS): return None import torch try: precip = np.array(obs["forecast_precip"], dtype=np.float32) uncert = np.array(obs["forecast_uncertainty"], dtype=np.float32) belief = np.array(obs["zone_belief"], dtype=np.float32) if precip.ndim == 2: precip = precip[np.newaxis] if uncert.ndim == 1: uncert = uncert[np.newaxis] if belief.ndim == 1: belief = belief[np.newaxis] return ZoneStateTensor( precip=torch.from_numpy(precip).to(self.device), uncertainty=torch.from_numpy(uncert).to(self.device), belief=torch.from_numpy(belief).to(self.device), ) except Exception as e: logger.debug("DynaCallback._obs_to_state_tensor failed: %s", e) return None def _on_step(self) -> bool: try: obs_now = self.locals.get("obs_tensor") or self.locals.get("obs") obs_next = self.locals.get("new_obs") if obs_now is None or obs_next is None: return True obs_now_np = ( {k: v.cpu().numpy() for k, v in obs_now.items()} if hasattr(obs_now, "items") else {"_raw": obs_now.cpu().numpy()} ) if hasattr(obs_now, "cpu") else obs_now obs_next_np = ( {k: (v.cpu().numpy() if hasattr(v, "cpu") else v) for k, v in obs_next.items()} if hasattr(obs_next, "items") else obs_next ) if hasattr(obs_next, "cpu") else obs_next curr = self._obs_to_state_tensor(obs_now_np) nxt = self._obs_to_state_tensor(obs_next_np) if curr is None or nxt is None: return True bonus = self.dyna_buffer.compute_surprise_bonus(curr, nxt) bonus_val = min(float(bonus.item()), self.surprise_weight) rb = self.model.rollout_buffer if rb is not None and hasattr(rb, "rewards") and rb.rewards is not None: idx = (rb.pos - 1) % rb.buffer_size rb.rewards[idx] += bonus_val if self.update_every > 0: self._transition_buffer.append((curr, nxt)) if len(self._transition_buffer) > self._tb_max: self._transition_buffer.pop(0) self._bonus_sum += bonus_val self._bonus_count += 1 if self.num_timesteps - self._last_log >= self._log_freq: avg = self._bonus_sum / max(self._bonus_count, 1) logger.info( "DynaCallback: step=%d avg_bonus=%.4f buffer=%d", self.num_timesteps, avg, len(self._transition_buffer), ) self._bonus_sum = 0.0 self._bonus_count = 0 self._last_log = self.num_timesteps except Exception as e: logger.debug("DynaCallback._on_step error (non-fatal): %s", e) return True def _on_rollout_end(self) -> None: if ( self.update_every <= 0 or self.num_timesteps % self.update_every != 0 or len(self._transition_buffer) < 16 ): return try: import torch import torch.nn.functional as F dynamics_model = self.dyna_buffer.dynamics optimizer = torch.optim.AdamW( dynamics_model.parameters(), lr=1e-4, weight_decay=1e-4 ) dynamics_model.train() pairs = list(self._transition_buffer) batch_size = min(32, len(pairs)) for epoch in range(self.fine_tune_epochs): random.shuffle(pairs) total_loss = 0.0 n_batches = 0 for i in range(0, len(pairs), batch_size): batch = pairs[i : i + batch_size] curr_list = [p[0] for p in batch] nxt_list = [p[1] for p in batch] curr_b = ZoneStateTensor( precip=torch.cat([s.precip for s in curr_list], dim=0), uncertainty=torch.cat([s.uncertainty for s in curr_list], dim=0), belief=torch.cat([s.belief for s in curr_list], dim=0), ) nxt_b = ZoneStateTensor( precip=torch.cat([s.precip for s in nxt_list], dim=0), uncertainty=torch.cat([s.uncertainty for s in nxt_list], dim=0), belief=torch.cat([s.belief for s in nxt_list], dim=0), ) pred, phys_loss = dynamics_model(curr_b, return_physics_loss=True) data_loss = ( F.mse_loss(pred.precip / 500.0, nxt_b.precip / 500.0) + F.mse_loss(pred.uncertainty, nxt_b.uncertainty) + F.mse_loss(pred.belief, nxt_b.belief) ) loss = data_loss + 0.01 * phys_loss optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(dynamics_model.parameters(), 1.0) optimizer.step() total_loss += loss.item() n_batches += 1 dynamics_model.eval() logger.info( "DynaCallback: fine-tuned at step=%d avg_loss=%.4f n=%d", self.num_timesteps, total_loss / max(n_batches, 1), len(self._transition_buffer), ) except Exception as e: logger.warning("DynaCallback fine-tune failed (non-fatal): %s", e) # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def parse_args() -> argparse.Namespace: p = argparse.ArgumentParser(description="MaskablePPO trainer (Kaggle)") p.add_argument("--dataset-dir", required=True, help="Path to repo / dataset root") p.add_argument("--out", default="./run", help="Output directory") p.add_argument("--n-zones", type=int, default=3) p.add_argument("--max-steps", type=int, default=250) p.add_argument("--budget-mode", choices=["triage", "scarce", "full", "legacy"], default="triage") p.add_argument("--steps", type=int, default=150_000, help="Total timesteps") p.add_argument("--hidden-size", type=int, default=128) p.add_argument("--device", default="auto") p.add_argument("--seed", type=int, default=42) p.add_argument("--dynamics-model", default=None) p.add_argument("--dynamics-weight", type=float, default=0.05) p.add_argument("--dynamics-finetune-every", type=int, default=0) p.add_argument("--learning-rate", type=float, default=BEST_HYPERPARAMETERS["learning_rate"]) p.add_argument("--n-steps", type=int, default=BEST_HYPERPARAMETERS["n_steps"]) p.add_argument("--batch-size", type=int, default=BEST_HYPERPARAMETERS["batch_size"]) p.add_argument("--n-epochs", type=int, default=BEST_HYPERPARAMETERS["n_epochs"]) p.add_argument("--gamma", type=float, default=BEST_HYPERPARAMETERS["gamma"]) p.add_argument("--gae-lambda", type=float, default=BEST_HYPERPARAMETERS["gae_lambda"]) p.add_argument("--clip-range", type=float, default=BEST_HYPERPARAMETERS["clip_range"]) p.add_argument("--ent-coef", type=float, default=BEST_HYPERPARAMETERS["ent_coef"]) p.add_argument("--vf-coef", type=float, default=BEST_HYPERPARAMETERS["vf_coef"]) p.add_argument("--max-grad-norm", type=float, default=BEST_HYPERPARAMETERS["max_grad_norm"]) return p.parse_args() def resolve_max_steps(n_zones: int, budget_mode: str, episode_length: int) -> int: n = max(1, int(n_zones)) mode = (budget_mode or "triage").strip().lower() if mode == "legacy": return max(1, int(episode_length)) if mode == "full": return n + 1 if mode == "scarce": return n return max(1, n - 1) def main() -> None: args = parse_args() random.seed(args.seed) np.random.seed(args.seed) if _TORCH_AVAILABLE: torch.manual_seed(args.seed) out_dir = Path(args.out) out_dir.mkdir(parents=True, exist_ok=True) max_steps = resolve_max_steps(args.n_zones, args.budget_mode, args.max_steps) logger.info( "n_zones=%d max_steps=%d budget_mode=%s", args.n_zones, max_steps, args.budget_mode ) config = ForecastConfig( n_zones=args.n_zones, max_steps=max_steps, soft_reset=True, seed=args.seed, ) env = Monitor(make_weather_env(config)) model = build_model(env, args) callbacks = [RegressionWatch()] if args.dynamics_model and _DYNAMICS_AVAILABLE and TemporalDynamicsModel is not None: try: import torch as _torch dyna_model = TemporalDynamicsModel.load( args.dynamics_model, device=_torch.device(args.device) ) dyna_model.eval() dyna_buffer = DynaRolloutBuffer( dynamics=dyna_model, uncertainty_weight=args.dynamics_weight ) callbacks.append( DynaCallback( dyna_buffer=dyna_buffer, surprise_weight=args.dynamics_weight, update_every=args.dynamics_finetune_every, device=args.device, ) ) logger.info("Dyna augmentation active") except Exception as e: logger.warning("Dyna init failed: %s", e) model.learn( total_timesteps=args.steps, callback=callbacks, reset_num_timesteps=True, use_masking=True, ) final_path = out_dir / "final_model.zip" model.save(str(final_path)) logger.info("Saved: %s", final_path) if __name__ == "__main__": main()