Update train_kaggle.py

train_kaggle.py

@@ -1,7 +1,10 @@
+#!/usr/bin/env python3
 """
 train_kaggle.py
-================
-Standalone training script for WeatherForecastEnv
+===============
+Single-phase MaskablePPO trainer for Kaggle Notebooks / TPU / GPU.
+
+Validated hyperparameters from ablation study (see BEST_HYPERPARAMETERS).
 """
 
 from __future__ import annotations
@@ -9,409 +12,452 @@ 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 Optional, Tuple
-
-logger = logging.getLogger("train_kaggle")
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s | %(levelname)s | %(message)s",
-)
-
+from typing import Any, Dict, List, Optional, Tuple
 
-def set_global_seeds(seed: int) -> None:
-    """Seed Python, NumPy, and PyTorch so that --seed controls weight init
-    as well as environment episode generation.
-    """
-    import random
-    import numpy as np
-    random.seed(seed)
-    np.random.seed(seed)
-    try:
-        import torch
-        torch.manual_seed(seed)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)
-    except ImportError:
-        pass
-
-
-def _add_file_logging(out_dir: Path) -> None:
-    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 = dict(
-    learning_rate=6.916624987609979e-05,
-    ent_coef=0.08779238696445962,
-    hidden_size=128,
-    spatial_size=12,
-    n_steps=4096,
-)
+import numpy as np
 
+# ---------------------------------------------------------------------------
+# Repo bootstrap (Kaggle input is read-only)
+# ---------------------------------------------------------------------------
 
-def resolve_max_steps(
-    n_zones: int,
-    budget_mode: str,
-    max_steps_arg: Optional[int],
-) -> Tuple[int, str]:
-    n = max(1, int(n_zones))
-    full_ceiling = n + 1
+REPO = Path("/kaggle/input/datasets/dhmmmreally/weather-modeller")
+if str(REPO) not in sys.path:
+    sys.path.insert(0, str(REPO))
 
-    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}"
-        )
+import zone_observation as _zo
 
-    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})"
+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
 
-def _add_dataset_to_path(dataset_dir: Optional[str]) -> None:
-    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)
+# ---------------------------------------------------------------------------
+# Optional ML imports
+# ---------------------------------------------------------------------------
 
+try:
+    import torch
+    _TORCH_AVAILABLE = True
+except ImportError:
+    _TORCH_AVAILABLE = False
 
-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):
+try:
     from weather_forecast_env import make_weather_env
-    from zone_observation import ForecastConfig
+    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)
 
-    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
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+)
+logger = logging.getLogger(__name__)
 
 
-def build_model(args: argparse.Namespace, train_env):
-    from sb3_contrib import MaskablePPO
-    from gru_weather_policy import (
-        ZoneEquivariantMaskablePolicy,
-        create_gru_weather_policy_kwargs,
-    )
+# ---------------------------------------------------------------------------
+# Hyperparameters (ablation-validated)
+# ---------------------------------------------------------------------------
 
-    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)
+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,
+}
 
-    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
+# ---------------------------------------------------------------------------
+# Regression watch callback
+# ---------------------------------------------------------------------------
 
-    return MaskablePPO(
-        ZoneEquivariantMaskablePolicy,
-        train_env,
-        learning_rate=args.lr,
-        ent_coef=args.ent_coef,
-        policy_kwargs=policy_kwargs,
+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=batch_size,
-        gamma=0.98,
-        gae_lambda=0.95,
-        clip_range=0.2,
+        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,
-        tensorboard_log=tb_log,
-        seed=args.seed,  # weight-init + SB3 internal RNG; env already seeded via ForecastConfig
+        seed=args.seed,
     )
 
+    model = MaskablePPO(
+        policy=policy,
+        env=env,
+        policy_kwargs=policy_kwargs,
+        **ppo_kwargs,
+    )
+    return model
+
 
-class RegressionWatchCallback:
+# ---------------------------------------------------------------------------
+# Dyna callback (mirrors train_curriculum.py)
+# ---------------------------------------------------------------------------
+
+class DynaCallback(BaseCallback):
+    _OBS_KEYS = ("forecast_precip", "forecast_uncertainty", "zone_belief")
 
     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 []
+        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),
                 )
-                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,
+                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),
                     )
-                return True
+                    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()
 
-        return _Impl()
 
-    @property
-    def instance(self):
-        return self._instance
+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 run_training(args: argparse.Namespace) -> None:
-    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
+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)
-    _add_file_logging(out_dir)
 
+    max_steps = resolve_max_steps(args.n_zones, args.budget_mode, args.max_steps)
     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,
+        "n_zones=%d max_steps=%d budget_mode=%s", args.n_zones, max_steps, args.budget_mode
     )
 
-    regression_watch = RegressionWatchCallback(
-        total_timesteps=args.steps, check_freq=args.regression_check_freq,
+    config = ForecastConfig(
+        n_zones=args.n_zones,
+        max_steps=max_steps,
+        soft_reset=True,
+        seed=args.seed,
     )
 
-    callbacks = [c for c in [checkpoint_cb, eval_cb, regression_watch.instance] if c is not None]
+    env = Monitor(make_weather_env(config))
+    model = build_model(env, args)
 
-    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)
+    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 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()
-    set_global_seeds(args.seed)
-    logger.info("Global seeds set to %s (Python / NumPy / PyTorch)", args.seed)
-    _add_dataset_to_path(args.dataset_dir)
-    run_training(args)
+    logger.info("Saved: %s", final_path)
 
 
 if __name__ == "__main__":