#!/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 class RealDataUsageCallback(BaseCallback): """Confirms at runtime -- not just at startup -- what fraction of episodes actually drew real vs synthetic data. info["context_source"] (set by WeatherForecastEnv.reset(), one of "injected"/"real_sampled"/ "synthetic") is only present on the step where an episode boundary triggered an autoreset, so this counts opportunistically rather than on every step. A silent zero real_sampled count with real_data_pkl_path configured is exactly the kind of thing that goes unnoticed for a long time otherwise -- this exists so it can't. """ def __init__(self, log_every: int = 10_000) -> None: super().__init__() self.log_every = log_every self._counts: Dict[str, int] = {} self._last_log = 0 def _on_step(self) -> bool: infos = self.locals.get("infos") if infos: for info in infos: src = info.get("context_source") if isinstance(info, dict) else None if src is not None: self._counts[src] = self._counts.get(src, 0) + 1 if self.num_timesteps - self._last_log >= self.log_every: total = sum(self._counts.values()) if total > 0: logger.info( "RealDataUsageCallback: step=%d episode context_source " "counts so far: %s", self.num_timesteps, dict(self._counts), ) self._last_log = self.num_timesteps return True def _on_training_end(self) -> None: total = sum(self._counts.values()) if total == 0: logger.warning( "RealDataUsageCallback: never observed a context_source in " "any info dict this run -- could not confirm real-vs-synthetic " "mix (this can happen with some VecEnv wrapping; it does not " "necessarily mean sampling failed, but it also cannot confirm " "it succeeded -- verify with a direct env.reset() smoke test " "if this matters for the run)." ) return logger.info( "RealDataUsageCallback: final episode context_source counts: " "%s (%d total, real fraction=%.3f)", dict(self._counts), total, self._counts.get("real_sampled", 0) / total, ) # --------------------------------------------------------------------------- # 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: # GRUWeatherFeaturesExtractor's forward() reshapes per-zone features # into (batch, n_zones * hidden_size * 2) -- features_dim MUST match # that exactly, or SB3's MlpExtractor (built from the declared # features_dim) gets a differently-shaped tensor and crashes at the # first policy.forward() call. Confirmed empirically: the actual # working checkpoints (final_normal/drought/humidity.zip) all have # features_dim == n_zones * hidden_size * 2 (256/384/512 for # n_zones=2/3/4, hidden_size=64) -- NOT hidden_size * 2, which only # coincides with the correct value when n_zones == 1. features_dim = args.n_zones * args.hidden_size * 2 policy_kwargs = create_gru_weather_policy_kwargs( hidden_size=args.hidden_size, features_dim=features_dim, ) policy = get_equivariant_policy_class() if get_equivariant_policy_class is not None else "MultiInputPolicy" logger.info( "Using GRU policy (hidden_size=%d, n_zones=%d, features_dim=%d)", args.hidden_size, args.n_zones, features_dim, ) 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"]) p.add_argument( "--real-data-pkl", default=None, help=( "Path to a historical trajectory cache (see " "real_episode_sampler.py / RealEpisodeIndex). When set, " "training samples real historical episodes with probability " "--real-data-ratio each reset, respecting RealEpisodeIndex's " "built-in eval-window holdout (see DEFAULT_HOLDOUT_RANGES). " "When NOT set (the default), training is 100%% synthetic -- " "this was true of every checkpoint in this project before " "2026-08-24; see technical_details.md. Any file the cache " "pipeline produces works here regardless of name -- naming " "convention is documentation only, see real_episode_sampler.py." ), ) p.add_argument( "--real-data-ratio", type=float, default=None, help=( "Per-episode probability of sampling real data when " "--real-data-pkl is set. Omit to use " "ForecastConfig.real_data_ratio's own default (0.7)." ), ) p.add_argument( "--real-data-noise-scale", type=float, default=None, help=( "Perturbation magnitude applied to sampled real data (see " "real_episode_sampler._perturb_zone_obs). Omit to use " "ForecastConfig.noise_scale's own default (0.05). Noise " "injection is enabled automatically whenever --real-data-pkl " "is set -- the real cache is thousands of points, training " "draws hundreds of thousands of episodes, and unperturbed " "replay risks memorizing specific real days rather than " "learning a generalizable policy." ), ) 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_kwargs: Dict[str, Any] = dict( n_zones=args.n_zones, max_steps=max_steps, soft_reset=True, seed=args.seed, real_data_pkl_path=args.real_data_pkl, inject_noise=bool(args.real_data_pkl), ) if args.real_data_ratio is not None: config_kwargs["real_data_ratio"] = args.real_data_ratio if args.real_data_noise_scale is not None: config_kwargs["noise_scale"] = args.real_data_noise_scale config = ForecastConfig(**config_kwargs) if config.real_data_pkl_path: logger.info( "=" * 72 + "\n" "REAL-DATA TRAINING ENABLED\n" " pkl = %s\n" " real_data_ratio = %.2f (fraction of episodes drawn from real data)\n" " inject_noise = %s\n" " noise_scale = %.3f\n" " holdout = enforced by RealEpisodeIndex " "(see real_episode_sampler.DEFAULT_HOLDOUT_RANGES)\n" + "=" * 72, config.real_data_pkl_path, config.real_data_ratio, config.inject_noise, config.noise_scale, ) else: logger.warning( "=" * 72 + "\n" "TRAINING IS 100%% SYNTHETIC -- no --real-data-pkl was provided.\n" "The agent will not see a single real historical episode this run.\n" "Pass --real-data-pkl to change that.\n" + "=" * 72 ) env = Monitor(make_weather_env(config)) model = build_model(env, args) callbacks = [RegressionWatch(), RealDataUsageCallback()] 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()