train_kaggle.py

#!/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 <path-to-cache.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()