train_kaggle.py

"""
train_kaggle.py
================
Standalone training script for WeatherForecastEnv, meant to run on Kaggle
(free GPU, no session-length hyperparameter-search restrictions like Colab).

WHAT THIS IS FOR
----------------
This trains a single fixed-n_zones snapshot of the environment (matching
what the Optuna sweep tuned against), with a configurable step count -- run
it short first (e.g. 100k-300k steps) to validate that the environment
fixes (reward scale, belief grounding, event injection, penalty
normalization, **budget pressure**) hold up, then re-run with a larger
--steps for a real training run.

BUDGET MODES (critical for multi-zone skill)
--------------------------------------------
The visit-once action mask makes the structural episode ceiling
``n_zones + 1`` (inspect each zone once, then terminate). Setting
``max_steps`` much larger than that (the old default of 250) never forces
the agent to *choose which zone to skip*. Full-tour is then both feasible
and reward-optimal (unvisited_zone_penalty → 0 after visiting everyone).

Under that regime policies learn a fixed inspect order (empirically: always
Indramayu under triage eval) and never learn risk-conditioned allocation.

Budget modes derive max_steps from n_zones unless --max-steps is explicit:

  full    max_steps = n_zones + 1   # can visit everyone (legacy behaviour)
  scarce  max_steps = n_zones      # can visit all only by skipping terminate
  triage  max_steps = max(1, n_zones - 1)  # MUST leave ≥1 zone unvisited

Default is **triage**. If you force max_steps >= n_zones+1, a hard WARNING
is logged: triage skill will not be trained.

This is intentionally NOT the 5-phase curriculum in train_curriculum.py --
it trains one fixed configuration end to end. Once you're confident in the
environment at this scale, train_curriculum.py's phase progression is the
next step (and now also carries budget_mode per phase).

POLICY (v3 equivariant)
-----------------------
Always use ZoneEquivariantMaskablePolicy from gru_weather_policy.py — NOT
the string "MultiInputPolicy". The latter builds logits from a pooled
feature vector and cannot express risk-conditioned zone choice (always
slot-0 degeneracy under triage). ZoneEquivariantMaskablePolicy scores each
zone before pooling and concatenates a terminate logit.

INPUT NORMALIZATION (--precip-scale)
-------------------------------------
forecast_precip runs roughly [0, 80] while zone_belief runs roughly
[0, 0.3], with no normalization layer between them and the GRU extractor.
Left unscaled, precip's larger raw magnitude can suppress the smaller but
more reliable belief signal during optimization, independent of which
feature actually carries more information. --precip-scale (default 40.0)
divides forecast_precip before it reaches the extractor. 40.0 is the value
that produced the validated single-dirty selection-accuracy results (see
the model card); it is confirmed working, not confirmed optimal --
recalibrate against your own separability probe if your event-injection
magnitudes differ.

KAGGLE SETUP
------------
1. Upload project files as a Kaggle Dataset (flat folder is fine).
2. Create a notebook, attach the dataset, enable GPU if available.
3. Example (triage, 2 zones, the recipe validated to produce single-dirty
   selection accuracy well above chance -- see model card)::

       !pip install -q gymnasium stable-baselines3 sb3-contrib tensorboard
       !PYTHONPATH=/kaggle/input/datasets/dhmmmreally/weather-modeller:$PYTHONPATH \\
       python train_kaggle.py \\
           --out /kaggle/working/run_nz2_triage \\
           --n-zones 2 \\
           --budget-mode triage \\
           --steps 150000 \\
           --clean-episode-ratio 0.80 \\
           --event-spatial-correlation 0.50 \\
           --precip-scale 40.0 \\
           --device auto

4. Resume mid-run with --resume-from pointing at a checkpoint under --out.

WHAT TO LOOK AT AFTERWARDS
---------------------------
  - ep_len_mean: under triage should sit near max_steps (not n_zones+1).
  - entropy_loss: should decline gradually, not collapse to ~0 immediately.
  - eval/mean_reward: should trend upward.
  - After training, run the zone-selection probe under max_steps=1: inspected
    pattern must NOT be 100% one action index; dirty-day shares should move
    off ~50/50 if allocation skill was learned.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path
from typing import Optional, Tuple

logger = logging.getLogger("train_kaggle")
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)


def _add_file_logging(out_dir: Path) -> None:
    """Attach a FileHandler rooted at out_dir/training.log."""
    formatter = logging.Formatter("%(asctime)s | %(levelname)s | %(message)s")
    file_handler = logging.FileHandler(str(out_dir / "training.log"))
    file_handler.setFormatter(formatter)
    logging.getLogger().addHandler(file_handler)
    logger.info("File logging enabled: %s", out_dir / "training.log")


# ---------------------------------------------------------------------------
# Best hyperparameters found by the Optuna sweep (trial 6 of the 20-trial
# run against n_zones=3 / max_steps=250). Used as defaults below -- override
# any of them from the CLI if you want to explore further.
# ---------------------------------------------------------------------------
BEST_HYPERPARAMETERS = dict(
    learning_rate=6.916624987609979e-05,
    ent_coef=0.08779238696445962,
    hidden_size=128,
    spatial_size=12,
    n_steps=4096,
)


def resolve_max_steps(
    n_zones: int,
    budget_mode: str,
    max_steps_arg: Optional[int],
) -> Tuple[int, str]:
    """
    Derive episode max_steps from budget mode, or honour an explicit override.

    Returns (max_steps, resolved_mode_note).
    """
    n = max(1, int(n_zones))
    full_ceiling = n + 1

    mode = (budget_mode or "triage").strip().lower()
    if mode not in ("full", "scarce", "triage"):
        raise ValueError(
            f"budget_mode must be one of full|scarce|triage, got {budget_mode!r}"
        )

    if mode == "full":
        derived = full_ceiling
    elif mode == "scarce":
        derived = n
    else:  # triage
        derived = max(1, n - 1)

    if max_steps_arg is not None and int(max_steps_arg) > 0:
        ms = int(max_steps_arg)
        note = f"explicit --max-steps={ms} (budget-mode={mode} would have been {derived})"
        if ms >= full_ceiling:
            logger.warning(
                "BUDGET WARNING: max_steps=%d >= n_zones+1=%d. Full tour is "
                "feasible and reward-optimal (unvisited_zone_penalty → 0). "
                "The policy is NOT forced to differentiate which zone to "
                "inspect. For allocation skill use --budget-mode triage "
                "(or scarce) without overriding --max-steps, or set "
                "--max-steps < %d.",
                ms, full_ceiling, full_ceiling,
            )
        return ms, note

    return derived, f"budget-mode={mode} → max_steps={derived} (full ceiling={full_ceiling})"


def _add_dataset_to_path(dataset_dir: Optional[str]) -> None:
    """Make the uploaded project files importable."""
    if dataset_dir:
        p = str(Path(dataset_dir).resolve())
        if p not in sys.path:
            sys.path.insert(0, p)
            logger.info("Added to sys.path: %s", p)
    here = str(Path(__file__).resolve().parent)
    if here not in sys.path:
        sys.path.insert(0, here)


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Single-config MaskablePPO trainer for WeatherForecastEnv (Kaggle)."
    )
    p.add_argument("--dataset-dir", default=None,
                    help="Path to the uploaded Kaggle dataset directory containing the .py files.")
    p.add_argument("--out", default="./run", help="Output directory for checkpoints/logs/final model.")
    p.add_argument("--resume-from", default=None, help="Path to a checkpoint .zip to resume from.")

    p.add_argument("--steps", type=int, default=50_000,
                    help="Total training timesteps. Start small (e.g. 100_000-300_000) to validate, "
                         "then scale up for a real run.")
    p.add_argument("--n-zones", type=int, default=2,
                    help="Number of zones. Default 2 (matches real-eval design window).")
    p.add_argument(
        "--budget-mode",
        choices=("full", "scarce", "triage"),
        default="triage",
        help="How tight the inspection budget is relative to n_zones. "
             "full=n_zones+1, scarce=n_zones, triage=max(1,n_zones-1). "
             "Default triage forces leaving ≥1 zone unvisited.",
    )
    p.add_argument(
        "--max-steps",
        type=int,
        default=None,
        help="Explicit episode length cap. If omitted, derived from --budget-mode. "
             "Setting this >= n_zones+1 disables triage pressure (WARNING logged).",
    )
    p.add_argument("--seed", type=int, default=42)

    p.add_argument("--lr", type=float, default=BEST_HYPERPARAMETERS["learning_rate"])
    p.add_argument("--ent-coef", type=float, default=BEST_HYPERPARAMETERS["ent_coef"])
    p.add_argument("--hidden-size", type=int, default=BEST_HYPERPARAMETERS["hidden_size"])
    p.add_argument("--spatial-size", type=int, default=BEST_HYPERPARAMETERS["spatial_size"])
    p.add_argument("--n-steps", type=int, default=BEST_HYPERPARAMETERS["n_steps"],
                    help="PPO rollout buffer size. If you change max_steps a lot, consider "
                         "resizing this to roughly 15-25x max_steps.")
    p.add_argument("--batch-size", type=int, default=None,
                    help="Defaults to max(32, n_steps // 8) if not given.")

    p.add_argument("--eval-freq", type=int, default=10_000)
    p.add_argument("--eval-episodes", type=int, default=20)
    p.add_argument("--checkpoint-freq", type=int, default=25_000)
    p.add_argument("--regression-check-freq", type=int, default=5_000)

    p.add_argument("--device", default="auto", help="'cpu', 'cuda', or 'auto'.")
    p.add_argument(
        "--clean-episode-ratio",
        type=float,
        default=0.90,
        help="Fraction of synthetic episodes with no regional hazard (match eval).",
    )
    p.add_argument(
        "--event-spatial-correlation",
        type=float,
        default=0.85,
        help="P(zone dirty | regional event). High → El Niño-style joint risk.",
    )
    p.add_argument(
        "--precip-scale",
        type=float,
        default=40.0,
        help="Fixed (non-learned) divisor applied to forecast_precip before "
             "it reaches the GRU extractor's forecast_proj/zone_encoder. "
             "forecast_precip runs roughly [0, 80] while zone_belief runs "
             "roughly [0, 0.3] with no normalization layer between them; "
             "left unscaled, precip's raw magnitude can suppress the "
             "smaller-but-more-reliable belief signal during optimization, "
             "independent of which feature is actually more informative. "
             "40.0 is the value that produced the validated single-dirty "
             "selection-accuracy results (see model card) -- confirmed "
             "working, not confirmed optimal. Recalibrate against your own "
             "separability probe's dirty-zone precip levels if your event "
             "injection magnitudes differ from the defaults.",
    )
    return p.parse_args()


def build_envs(args: argparse.Namespace):
    from weather_forecast_env import make_weather_env
    from zone_observation import ForecastConfig
    from stable_baselines3.common.monitor import Monitor

    train_config = ForecastConfig(
        n_zones=args.n_zones,
        max_steps=args.max_steps,
        seed=args.seed,
        clean_episode_ratio=args.clean_episode_ratio,
        event_spatial_correlation=args.event_spatial_correlation,
    )
    eval_config = ForecastConfig(
        n_zones=args.n_zones,
        max_steps=args.max_steps,
        seed=args.seed + 10_000,
        clean_episode_ratio=args.clean_episode_ratio,
        event_spatial_correlation=args.event_spatial_correlation,
    )
    train_env = Monitor(make_weather_env(train_config))
    eval_env = Monitor(make_weather_env(eval_config))
    return train_env, eval_env


def build_model(args: argparse.Namespace, train_env):
    from sb3_contrib import MaskablePPO
    from gru_weather_policy import (
        ZoneEquivariantMaskablePolicy,
        create_gru_weather_policy_kwargs,
    )

    if args.resume_from:
        logger.info("Resuming from checkpoint: %s", args.resume_from)
        return MaskablePPO.load(args.resume_from, env=train_env, device=args.device)

    policy_kwargs = create_gru_weather_policy_kwargs(
        hidden_size=args.hidden_size,
        spatial_output_size=args.spatial_size,
        features_dim=args.hidden_size * 2,
        basin_context_hidden=12,
        precip_scale=args.precip_scale,
    )
    batch_size = args.batch_size or max(32, args.n_steps // 8)

    try:
        import tensorboard  # noqa: F401
        tb_log = str(Path(args.out) / "tensorboard")
    except ImportError:
        logger.warning(
            "tensorboard not installed -- continuing without TensorBoard logs "
            "(install with `pip install tensorboard` if you want them)."
        )
        tb_log = None

    # ZoneEquivariantMaskablePolicy: inspect logits from per-zone scores
    # (pre-pool); terminate from pooled features. Do not use "MultiInputPolicy"
    # — that routes through the pooled action_net and erases slot identity.
    return MaskablePPO(
        ZoneEquivariantMaskablePolicy,
        train_env,
        learning_rate=args.lr,
        ent_coef=args.ent_coef,
        policy_kwargs=policy_kwargs,
        n_steps=args.n_steps,
        batch_size=batch_size,
        gamma=0.98,
        gae_lambda=0.95,
        clip_range=0.2,
        device=args.device,
        verbose=1,
        tensorboard_log=tb_log,
    )


class RegressionWatchCallback:
    """Flags signs of the old degenerate 'terminate immediately' collapse.

    Under triage, ep_len pinned near 1.0 with no belief movement is still a
    failure mode (never inspect). Under full budget, ep_len pinned at
    n_zones+1 is expected and not a regression by itself.
    """

    def __init__(
        self,
        total_timesteps: int,
        check_freq: int = 5_000,
        after_frac: float = 0.2,
        ep_len_threshold: float = 1.5,
    ):
        from stable_baselines3.common.callbacks import BaseCallback
        import numpy as np

        self._np = np
        self._BaseCallback = BaseCallback
        self.total_timesteps = total_timesteps
        self.check_freq = check_freq
        self.after_frac = after_frac
        self.ep_len_threshold = ep_len_threshold
        self._instance = self._build_instance()

    def _build_instance(self):
        np = self._np
        outer = self

        class _Impl(self._BaseCallback):
            def __init__(self):
                super().__init__()
                self._last_check = 0
                self.history = []  # (timestep, ep_len_mean, entropy_loss)

            def _on_step(self) -> bool:
                if self.num_timesteps - self._last_check < outer.check_freq:
                    return True
                self._last_check = self.num_timesteps

                ep_lens = (
                    [ep["l"] for ep in self.model.ep_info_buffer]
                    if self.model.ep_info_buffer
                    else []
                )
                ep_len_mean = float(np.mean(ep_lens)) if ep_lens else float("nan")

                entropy = None
                if self.model.logger is not None:
                    entropy = self.model.logger.name_to_value.get("train/entropy_loss")

                self.history.append((self.num_timesteps, ep_len_mean, entropy))

                frac = self.num_timesteps / max(1, outer.total_timesteps)
                if frac >= outer.after_frac and ep_lens and ep_len_mean < outer.ep_len_threshold:
                    logger.warning(
                        "REGRESSION WARNING at step %d: ep_len_mean=%.2f after %.0f%% "
                        "of training. This matches the original 'terminate immediately' "
                        "collapse pattern -- worth stopping to check config/reward before "
                        "trusting the rest of this run.",
                        self.num_timesteps, ep_len_mean, frac * 100,
                    )
                return True

        return _Impl()

    @property
    def instance(self):
        return self._instance


def run_training(args: argparse.Namespace) -> None:
    # Resolve budget before anything else so logs and envs agree.
    resolved_ms, budget_note = resolve_max_steps(
        args.n_zones, args.budget_mode, args.max_steps
    )
    args.max_steps = resolved_ms  # mutate so build_envs / summary see the real value

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    _add_file_logging(out_dir)

    logger.info(
        "Config: n_zones=%d max_steps=%d (%s) steps=%d clean=%.3f rho=%.3f "
        "lr=%.3g ent_coef=%.3g hidden_size=%d spatial_size=%d n_steps=%d "
        "precip_scale=%.3g device=%s",
        args.n_zones, args.max_steps, budget_note, args.steps,
        args.clean_episode_ratio, args.event_spatial_correlation,
        args.lr, args.ent_coef,
        args.hidden_size, args.spatial_size, args.n_steps,
        args.precip_scale,
        args.device,
    )
    logger.info(
        "Budget pressure: full_ceiling=%d  resolved_max_steps=%d  "
        "must_skip_zones=%s",
        args.n_zones + 1,
        args.max_steps,
        "yes" if args.max_steps < args.n_zones + 1 else "no (full tour allowed)",
    )

    train_env, eval_env = build_envs(args)
    model = build_model(args, train_env)

    from sb3_contrib.common.maskable.callbacks import MaskableEvalCallback
    try:
        from train_curriculum import CheckpointCallback as _ProjectCheckpointCallback
        checkpoint_cb = _ProjectCheckpointCallback(out_dir, save_freq=args.checkpoint_freq)
    except Exception as e:
        logger.warning(
            "Could not import CheckpointCallback from train_curriculum.py (%s); "
            "continuing without periodic checkpoints -- only the final model will "
            "be saved.", e
        )
        checkpoint_cb = None

    eval_cb = MaskableEvalCallback(
        eval_env,
        n_eval_episodes=args.eval_episodes,
        eval_freq=args.eval_freq,
        deterministic=True,
        best_model_save_path=str(out_dir / "best_model"),
        verbose=1,
    )

    regression_watch = RegressionWatchCallback(
        total_timesteps=args.steps, check_freq=args.regression_check_freq,
    )

    callbacks = [c for c in [checkpoint_cb, eval_cb, regression_watch.instance] if c is not None]

    t0 = time.time()
    model.learn(total_timesteps=args.steps, callback=callbacks, progress_bar=False)
    elapsed = time.time() - t0
    logger.info("Training finished in %.1f minutes.", elapsed / 60.0)

    final_path = out_dir / "final_model.zip"
    model.save(str(final_path))
    logger.info("Saved final model: %s", final_path)

    summary = {
        "args": vars(args),
        "budget_note": budget_note,
        "full_tour_ceiling": args.n_zones + 1,
        "must_skip_zones": args.max_steps < args.n_zones + 1,
        "elapsed_minutes": elapsed / 60.0,
        "best_mean_reward": eval_cb.best_mean_reward,
        "regression_check_history": regression_watch.instance.history,
    }
    summary_path = out_dir / "run_summary.json"
    with open(summary_path, "w") as f:
        json.dump(summary, f, indent=2, default=str)
    logger.info("Saved run summary: %s", summary_path)
    logger.info("Best mean eval reward this run: %s", eval_cb.best_mean_reward)


def main() -> None:
    args = _parse_args()
    _add_dataset_to_path(args.dataset_dir)
    run_training(args)


if __name__ == "__main__":
    main()