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


# ---------------------------------------------------------------------------
# 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()