Upload 3 files

climatology.py

@@ -0,0 +1,834 @@
+"""
+climatology.py
+==============
+Per-zone day-of-year climatology and anomaly (z-score) computation.
+
+WHY THIS MODULE EXISTS (the gap it closes)
+------------------------------------------
+Every real fetcher in era5_data_pipeline.py (_build_era5_obs,
+_fetch_openmeteo, _fetch_imerg, _fetch_smap) sets precip_anomaly_idx /
+temp_anomaly_idx / soil_moisture_anom to 0.0 with the comment "requires
+climatology -- set in scorer". Nothing ever supplied that climatology, and
+the scorer never set the fields either. The consequences are structural:
+
+  * ZoneObs.drought_signal() = 0.6 * clip(-precip_anomaly/3)
+                             + 0.4 * clip(-soil_anom/3)
+  * ZoneObs.flood_signal()   = 0.4 * clip(+precip_anomaly/3) + ...
+
+With all anomalies pinned at 0.0, drought_signal() is identically 0.0 and
+flood_signal() loses its primary term on every real observation. The whole
+risk-scoring stack (crop_risk_scorer, the env's belief initialisation, the
+hierarchical search gating) was effectively only sensitive on SYNTHETIC
+data, where anomalies are injected directly by the event flags. This module
+is the missing climatology layer: it turns absolute real-world readings
+into the z-score anomalies the rest of the pipeline was designed around.
+
+DATA SOURCES (two tiers, matching the codebase's degrade-safely philosophy)
+----------------------------------------------------------------------------
+  1. Real: Open-Meteo historical archive API (free, no API key -- the same
+     endpoint era5_data_pipeline._fetch_openmeteo already uses). Daily
+     precipitation_sum + temperature_2m_mean are well-established archive
+     variables. soil_moisture_0_to_7cm_mean is requested as documented in
+     the Open-Meteo archive docs at write time; if the API rejects it or
+     returns nothing, soil climatology degrades to the precip-tracked
+     model below, logged -- never silently zeroed.
+  2. Synthetic: a deterministic, latitude-aware maritime-continent monsoon
+     model (see _synthetic_climatology). It is a HEURISTIC, calibrated to
+     the broad shape of the Indonesian wet/dry season (SH monsoon: wet
+     Dec-Mar, dry Jun-Sep; weaker/bimodal near the equator; shifted peak
+     for northern Sumatra). It exists so the pipeline keeps producing
+     sensible anomalies offline and in tests. It is NOT a retrieval --
+     treat its absolute values as plausible shapes, not measurements.
+
+IMPORTANT USAGE NOTE
+--------------------
+apply_climatology_anomalies() SKIPS observations whose source is
+DataSource.SYNTHETIC. The synthetic generator injects its own meaningful
+anomalies via the drought/flood event flags; re-scoring those against a
+climatology would double-transform a deliberately-constructed signal.
+Climatology anomalies are for REAL observations only.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import math
+import os
+from dataclasses import dataclass, field
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import zone_observation as _zo
+
+assert _zo.SCHEMA_VERSION == 3, (
+    f"climatology: zone_observation schema mismatch "
+    f"(expected 3, got {_zo.SCHEMA_VERSION})"
+)
+
+from zone_observation import DataSource, ZoneObs, _clip, _stable_seed
+
+logger = logging.getLogger(__name__)
+
+try:
+    import requests
+    _REQUESTS_AVAILABLE = True
+except ImportError:
+    _REQUESTS_AVAILABLE = False
+    logger.info("requests not installed -- climatology will use the synthetic model")
+
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+# Same archive endpoint era5_data_pipeline.py uses; duplicated here (rather
+# than imported) so era5_data_pipeline can import THIS module lazily without
+# a circular import at module load time.
+_OPENMETEO_ARCHIVE_URL = "https://archive-api.open-meteo.com/v1/archive"
+_TIMEOUT_S = int(os.environ.get("WEATHER_HTTP_TIMEOUT", "60"))
+
+_CACHE_DIR = Path(os.environ.get("WEATHER_CACHE_DIR", ".cache/era5")) / "climatology"
+_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+_CACHE_TTL_DAYS = 90
+
+_DAYS_PER_YEAR = 365.25
+_TABLE_LEN = 366  # DOY table indexed doy-1; DOY 60 = Feb 29 (leap mapping below)
+
+# Std floors: prevent division blow-ups in convectively uniform seasons.
+_PRECIP_STD_FLOOR = 1.5   # mm/day
+_TEMP_STD_FLOOR   = 0.4   # deg C
+_SOIL_STD_FLOOR   = 2.0   # percent
+_RH_STD_FLOOR = 3.0  # percent; prevents z-score blowup when tropical RH variance is naturally small
+
+# Trailing window for the precipitation anomaly. Matches ZoneObs.precip_30d_mm,
+# the longest aggregate every real fetcher populates.
+_PRECIP_WINDOW_DAYS = 30
+
+
+# ---------------------------------------------------------------------------
+# Day-of-year helpers (fixed 366-entry table regardless of leap years)
+# ---------------------------------------------------------------------------
+
+def _is_leap(year: int) -> bool:
+    return year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)
+
+
+def doy_index(dt: datetime) -> int:
+    """Map a date to a 1..366 index in the fixed climatology table.
+
+    In non-leap years, dates after Feb 28 are shifted up by one so that e.g.
+    Mar 1 always maps to the same table entry (61) in every year. Table entry
+    60 (Feb 29) is only ever hit by leap-year dates.
+    """
+    doy = dt.timetuple().tm_yday
+    if not _is_leap(dt.year) and doy >= 60:
+        doy += 1
+    return min(doy, _TABLE_LEN)
+
+
+def _circular_smooth(values: List[float], half_window: int = 7) -> List[float]:
+    """Circular moving average over the DOY table (Dec wraps to Jan)."""
+    n = len(values)
+    out = []
+    for i in range(n):
+        acc = 0.0
+        cnt = 0
+        for j in range(-half_window, half_window + 1):
+            acc += values[(i + j) % n]
+            cnt += 1
+        out.append(acc / cnt)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# ZoneClimatology
+# ---------------------------------------------------------------------------
+
+@dataclass
+class ZoneClimatology:
+    """Day-of-year climatology for one zone.
+
+    All lists have length 366 and are indexed by (doy_index(dt) - 1).
+    precip is a DAILY mean rate (mm/day); window aggregates are computed by
+    summing daily means over the window (see window_precip_stats).
+    """
+    zone_id: str
+    source: str                       # 'openmeteo_archive' | 'synthetic_model' | 'mixed'
+    n_years: int
+    period_start_year: int
+    period_end_year: int
+    precip_mean_mm: List[float] = field(default_factory=list)   # mm/day
+    precip_std_mm:  List[float] = field(default_factory=list)
+    temp_mean_c:    List[float] = field(default_factory=list)
+    temp_std_c:     List[float] = field(default_factory=list)
+    soil_mean_pct:  List[float] = field(default_factory=list)
+    soil_std_pct:   List[float] = field(default_factory=list)
+    # FIX: baseline uses rh_mean (not rh_max -- unreliable in Open-Meteo);
+    # correlates well enough to correct fungi_risk_signal()'s unchanged
+    # max-based threshold.
+    rh_mean_pct:    List[float] = field(default_factory=list)
+    rh_std_pct:     List[float] = field(default_factory=list)
+
+    def __post_init__(self) -> None:
+        for name in ("precip_mean_mm", "precip_std_mm", "temp_mean_c",
+                     "temp_std_c", "soil_mean_pct", "soil_std_pct",
+                     "rh_mean_pct", "rh_std_pct"):
+            v = getattr(self, name)
+            if len(v) != _TABLE_LEN:
+                raise ValueError(
+                    f"ZoneClimatology('{self.zone_id}'): {name} has length "
+                    f"{len(v)}, expected {_TABLE_LEN}"
+                )
+
+    # --- Window statistics ------------------------------------------------
+    def window_precip_stats(self, dt: datetime, window_days: int) -> Tuple[float, float]:
+        """Climatological mean and std of a TRAILING `window_days` precip total
+        ending at dt's day-of-year.
+
+        Mean: sum of daily means (exact under the daily model).
+        Std: sqrt(sum of daily variances) -- assumes day-to-day independence,
+             so it UNDERSTATES true variance during correlated multi-day
+             spells, inflating anomaly magnitude for persistent events.
+             Bounded by the floors and the [-5, 5] clip in ZoneObs.
+        """
+        idx0 = doy_index(dt) - 1
+        mean = 0.0
+        var = 0.0
+        for k in range(window_days):
+            i = (idx0 - k) % _TABLE_LEN
+            mean += self.precip_mean_mm[i]
+            var += self.precip_std_mm[i] ** 2
+        return mean, max(math.sqrt(var), _PRECIP_STD_FLOOR)
+
+    def daily_temp_stats(self, dt: datetime) -> Tuple[float, float]:
+        i = doy_index(dt) - 1
+        return self.temp_mean_c[i], max(self.temp_std_c[i], _TEMP_STD_FLOOR)
+
+    def daily_soil_stats(self, dt: datetime) -> Tuple[float, float]:
+        i = doy_index(dt) - 1
+        return self.soil_mean_pct[i], max(self.soil_std_pct[i], _SOIL_STD_FLOOR)
+
+    def daily_rh_stats(self, dt: datetime) -> Tuple[float, float]:
+        i = doy_index(dt) - 1
+        return self.rh_mean_pct[i], max(self.rh_std_pct[i], _RH_STD_FLOOR)
+
+    # --- Serialisation (JSON cache) ---------------------------------------
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "zone_id": self.zone_id,
+            "source": self.source,
+            "n_years": self.n_years,
+            "period_start_year": self.period_start_year,
+            "period_end_year": self.period_end_year,
+            "precip_mean_mm": self.precip_mean_mm,
+            "precip_std_mm": self.precip_std_mm,
+            "temp_mean_c": self.temp_mean_c,
+            "temp_std_c": self.temp_std_c,
+            "soil_mean_pct": self.soil_mean_pct,
+            "soil_std_pct": self.soil_std_pct,
+            "rh_mean_pct": self.rh_mean_pct,
+            "rh_std_pct": self.rh_std_pct,
+        }
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "ZoneClimatology":
+        # FIX: rh_mean_pct/rh_std_pct are new fields; old cached files
+        # predate them. Default to flat 85% (synthetic-model range) instead
+        # of crashing -- self-heals within _CACHE_TTL_DAYS as real RH is
+        # fetched.
+        rh_mean = d.get("rh_mean_pct")
+        rh_std = d.get("rh_std_pct")
+        if rh_mean is None or len(rh_mean) != _TABLE_LEN:
+            rh_mean = [85.0] * _TABLE_LEN
+        if rh_std is None or len(rh_std) != _TABLE_LEN:
+            rh_std = [_RH_STD_FLOOR] * _TABLE_LEN
+        return cls(
+            zone_id=d["zone_id"],
+            source=d.get("source", "unknown"),
+            n_years=int(d.get("n_years", 0)),
+            period_start_year=int(d.get("period_start_year", 0)),
+            period_end_year=int(d.get("period_end_year", 0)),
+            precip_mean_mm=[float(v) for v in d["precip_mean_mm"]],
+            precip_std_mm=[float(v) for v in d["precip_std_mm"]],
+            temp_mean_c=[float(v) for v in d["temp_mean_c"]],
+            temp_std_c=[float(v) for v in d["temp_std_c"]],
+            soil_mean_pct=[float(v) for v in d["soil_mean_pct"]],
+            soil_std_pct=[float(v) for v in d["soil_std_pct"]],
+            rh_mean_pct=[float(v) for v in rh_mean],
+            rh_std_pct=[float(v) for v in rh_std],
+        )
+
+
+# ---------------------------------------------------------------------------
+# Tier 2 -- deterministic synthetic maritime-continent climatology
+# ---------------------------------------------------------------------------
+
+def _synthetic_climatology(zone_id: str, lat: float, n_years: int = 0) -> ZoneClimatology:
+    """Deterministic heuristic climatology for the Indonesian maritime continent.
+    A documented heuristic, NOT a retrieval -- see module docstring.
+
+      * Precip: single-harmonic wet season, peak ~DOY 30 (late Jan) for the
+        southern archipelago (Java, Bali, Nusa Tenggara, Sulawesi, S. Sumatra);
+        peak shifts earlier (Oct-Dec) moving north past ~1 deg N (N. Sumatra).
+        Amplitude grows with distance from equator; equatorial belt stays wet
+        year-round.
+      * Temp: weak annual cycle (~2.6 deg C peak-to-peak), coolest Jul-Aug in
+        the south (SH dry season), weaker and phase-reversed north of equator.
+      * Soil: precip-tracked with ~20-day lag, scaled to ERA5 swvl1's typical
+        volumetric-% range for the region.
+
+    Per-zone jitter (+/-10% on base/amp, via _stable_seed) keeps neighbouring
+    zones numerically distinct without changing the seasonal shape.
+    """
+    seed = _stable_seed(f"clim_{zone_id}")
+    # Deterministic jitter in [0.9, 1.1] from the seed's low bits.
+    jitter = 0.9 + 0.2 * ((seed % 1000) / 1000.0)
+
+    abs_lat = abs(lat)
+    # Wet-season peak: late Jan in the south, shifting earlier north of ~1N.
+    peak_doy = 30.0 if lat <= 1.0 else max(300.0, 30.0 - 12.0 * lat)
+    amp_scale = _clip(abs_lat / 8.0, 0.35, 1.0)
+
+    precip_base = max(3.0, (7.0 - 0.25 * abs_lat) * jitter)   # mm/day
+    precip_amp  = 4.5 * amp_scale * jitter                    # mm/day
+    temp_base   = 27.0 - 0.30 * abs_lat
+    # SH zones: coolest around DOY ~200 (mid-Jul). NH: weaker, reversed.
+    temp_amp    = 1.3 if lat < 0.0 else -0.5
+
+    precip_mean, precip_std = [], []
+    temp_mean, temp_std = [], []
+    soil_mean, soil_std = [], []
+    rh_mean, rh_std = [], []
+
+    daily_precip_for_soil: List[float] = []
+    for doy in range(1, _TABLE_LEN + 1):
+        phase = 2.0 * math.pi * (doy - peak_doy) / _DAYS_PER_YEAR
+        p = precip_base + precip_amp * math.cos(phase)
+        p = max(0.8, p)
+        daily_precip_for_soil.append(p)
+        precip_mean.append(p)
+        precip_std.append(max(_PRECIP_STD_FLOOR, 0.9 * p))
+
+        t_phase = 2.0 * math.pi * (doy - 200.0) / _DAYS_PER_YEAR
+        t = temp_base - temp_amp * math.cos(t_phase)
+        temp_mean.append(t)
+        temp_std.append(0.7)
+
+        # RH baseline tracks the wet season (in phase with precip), range
+        # 78-94%. This is a baseline for ANOMALY detection only --
+        # fungi_risk_signal() still applies its own absolute threshold to
+        # rh_max_pct separately.
+        r = 86.0 + amp_scale * 6.0 * math.cos(phase)
+        rh_mean.append(_clip(r, 78.0, 94.0))
+        rh_std.append(max(_RH_STD_FLOOR, 3.5))
+
+    # Soil tracks precip with a 20-day lag.
+    for doy in range(1, _TABLE_LEN + 1):
+        lagged = daily_precip_for_soil[(doy - 1 - 20) % _TABLE_LEN]
+        s = _clip(16.0 + 2.4 * lagged, 8.0, 52.0)
+        soil_mean.append(s)
+        soil_std.append(max(_SOIL_STD_FLOOR, 4.0))
+
+    return ZoneClimatology(
+        zone_id=zone_id,
+        source="synthetic_model",
+        n_years=n_years,
+        period_start_year=0,
+        period_end_year=0,
+        precip_mean_mm=_circular_smooth(precip_mean),
+        precip_std_mm=precip_std,
+        temp_mean_c=_circular_smooth(temp_mean),
+        temp_std_c=temp_std,
+        soil_mean_pct=_circular_smooth(soil_mean),
+        soil_std_pct=soil_std,
+        rh_mean_pct=_circular_smooth(rh_mean),
+        rh_std_pct=rh_std,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Tier 1 -- real climatology from the Open-Meteo archive
+# ---------------------------------------------------------------------------
+
+def _fetch_openmeteo_climatology(
+    zone_id: str,
+    lat: float,
+    lon: float,
+    years: int,
+    end_year: Optional[int] = None,
+) -> ZoneClimatology:
+    """Build a DOY climatology from the Open-Meteo historical archive.
+
+    Downloads `years` full calendar years of daily data in one request and
+    pools by day-of-year. Raises on any failure -- the caller
+    (get_zone_climatology) falls back to the synthetic model, matching the
+    pipeline-wide degrade-safely pattern.
+
+    soil_moisture_0_to_7cm_mean is requested per the Open-Meteo archive
+    documentation at write time (VERIFIED LIVE against the archive API:
+    the variable exists and returns daily means). UNITS: Open-Meteo returns
+    soil moisture in m3/m3; this function converts to percent (x100) so the
+    table matches ZoneObs.soil_moisture_pct and ERA5's swvl1 x 100 handling
+    in era5_data_pipeline._build_era5_obs. (Found via a z-score clipped at
+    +5.0 against an 18% observation -- the raw 0.2-0.4 m3/m3 values were
+    being read as ~0.3%.)
+
+    If the key is absent/empty in the response (API change, or variable not
+    in the daily list for this endpoint), the soil tables are derived from
+    the REAL precip series via the same lagged mapping the synthetic model
+    uses -- so a soil-variable outage degrades one field's provenance, not
+    the whole fetch. The result's `source` is then 'mixed' rather than
+    'openmeteo_archive' so downstream auditing can tell.
+    """
+    if not _REQUESTS_AVAILABLE:
+        raise RuntimeError("requests not installed")
+
+    last_full_year = (end_year if end_year is not None
+                      else datetime.now(timezone.utc).year - 1)
+    start_year = last_full_year - years + 1
+
+    params = {
+        "latitude": lat,
+        "longitude": lon,
+        "start_date": f"{start_year}-01-01",
+        "end_date": f"{last_full_year}-12-31",
+        "daily": ",".join([
+            "precipitation_sum",
+            "temperature_2m_mean",
+            "soil_moisture_0_to_7cm_mean",
+            "relative_humidity_2m_mean",
+        ]),
+        "timezone": "UTC",
+    }
+
+    resp = requests.get(_OPENMETEO_ARCHIVE_URL, params=params, timeout=_TIMEOUT_S)
+    resp.raise_for_status()
+    data = resp.json()
+    daily = data.get("daily", {})
+    dates = daily.get("time", [])
+    if not dates:
+        raise RuntimeError(f"Open-Meteo archive returned no daily rows for {zone_id}")
+
+    precip_series = daily.get("precipitation_sum", [])
+    temp_series = daily.get("temperature_2m_mean", [])
+    soil_series = daily.get("soil_moisture_0_to_7cm_mean", [])
+    rh_series = daily.get("relative_humidity_2m_mean", [])
+
+    p_sum = [0.0] * _TABLE_LEN
+    p_sq = [0.0] * _TABLE_LEN
+    p_n = [0] * _TABLE_LEN
+    t_sum = [0.0] * _TABLE_LEN
+    t_sq = [0.0] * _TABLE_LEN
+    t_n = [0] * _TABLE_LEN
+    s_sum = [0.0] * _TABLE_LEN
+    s_sq = [0.0] * _TABLE_LEN
+    s_n = [0] * _TABLE_LEN
+    r_sum = [0.0] * _TABLE_LEN
+    r_sq = [0.0] * _TABLE_LEN
+    r_n = [0] * _TABLE_LEN
+
+    def _val(series: List[Any], i: int) -> Optional[float]:
+        if i >= len(series):
+            return None
+        v = series[i]
+        if v is None:
+            return None
+        try:
+            return float(v)
+        except (TypeError, ValueError):
+            return None
+
+    for i, date_str in enumerate(dates):
+        try:
+            dt = datetime.fromisoformat(date_str).replace(tzinfo=timezone.utc)
+        except ValueError:
+            continue
+        k = doy_index(dt) - 1
+
+        p = _val(precip_series, i)
+        if p is not None:
+            p_sum[k] += p
+            p_sq[k] += p * p
+            p_n[k] += 1
+        t = _val(temp_series, i)
+        if t is not None:
+            t_sum[k] += t
+            t_sq[k] += t * t
+            t_n[k] += 1
+        s = _val(soil_series, i)
+        if s is not None:
+            s_pct = s * 100.0  # m3/m3 -> % (matches ZoneObs.soil_moisture_pct)
+            s_sum[k] += s_pct
+            s_sq[k] += s_pct * s_pct
+            s_n[k] += 1
+        r = _val(rh_series, i)
+        if r is not None:
+            r_sum[k] += r
+            r_sq[k] += r * r
+            r_n[k] += 1
+
+    if sum(p_n) < 300 * years or sum(t_n) < 300 * years:
+        raise RuntimeError(
+            f"Open-Meteo archive coverage too thin for {zone_id}: "
+            f"precip_days={sum(p_n)} temp_days={sum(t_n)} over {years}y"
+        )
+
+    def _mean_std(sums, sqs, ns, floor):
+        means, stds = [], []
+        for k in range(_TABLE_LEN):
+            n = ns[k]
+            if n == 0:
+                # Should not happen with full-year coverage; guard anyway.
+                means.append(0.0)
+                stds.append(floor)
+                continue
+            m = sums[k] / n
+            var = max(0.0, sqs[k] / n - m * m)
+            means.append(m)
+            stds.append(max(floor, math.sqrt(var)))
+        return means, stds
+
+    precip_mean, precip_std = _mean_std(p_sum, p_sq, p_n, _PRECIP_STD_FLOOR)
+    temp_mean, temp_std = _mean_std(t_sum, t_sq, t_n, _TEMP_STD_FLOOR)
+
+    soil_days = sum(s_n)
+    soil_thin = soil_days < 300 * years
+    if not soil_thin:
+        soil_mean, soil_std = _mean_std(s_sum, s_sq, s_n, _SOIL_STD_FLOOR)
+    else:
+        logger.warning(
+            "climatology: soil_moisture_0_to_7cm_mean coverage thin for %s "
+            "(%d days over %dy) -- deriving soil tables from the real precip "
+            "series (lagged mapping). Provenance marked 'mixed'.",
+            zone_id, soil_days, years,
+        )
+        soil_mean, soil_std = [], []
+        for doy in range(1, _TABLE_LEN + 1):
+            lagged = precip_mean[(doy - 1 - 20) % _TABLE_LEN]
+            soil_mean.append(_clip(16.0 + 2.4 * lagged, 8.0, 52.0))
+            soil_std.append(max(_SOIL_STD_FLOOR, 4.0))
+
+    rh_days = sum(r_n)
+    rh_thin = rh_days < 300 * years
+    if not rh_thin:
+        rh_mean, rh_std = _mean_std(r_sum, r_sq, r_n, _RH_STD_FLOOR)
+    else:
+        logger.warning(
+            "climatology: relative_humidity_2m_mean coverage thin for %s "
+            "(%d days over %dy) -- deriving RH tables from the real precip "
+            "series (wet-season correlation, same phase). Provenance marked "
+            "'mixed'.",
+            zone_id, rh_days, years,
+        )
+        rh_mean, rh_std = [], []
+        p_min, p_max = min(precip_mean), max(precip_mean)
+        p_span = max(p_max - p_min, 1e-6)
+        for doy in range(1, _TABLE_LEN + 1):
+            p_frac = (precip_mean[doy - 1] - p_min) / p_span  # 0..1
+            rh_mean.append(_clip(78.0 + 12.0 * p_frac, 78.0, 94.0))
+            rh_std.append(max(_RH_STD_FLOOR, 3.5))
+
+    source = "openmeteo_archive" if not (soil_thin or rh_thin) else "mixed"
+
+    return ZoneClimatology(
+        zone_id=zone_id,
+        source=source,
+        n_years=years,
+        period_start_year=start_year,
+        period_end_year=last_full_year,
+        precip_mean_mm=_circular_smooth(precip_mean),
+        precip_std_mm=precip_std,
+        temp_mean_c=_circular_smooth(temp_mean),
+        temp_std_c=temp_std,
+        soil_mean_pct=_circular_smooth(soil_mean),
+        soil_std_pct=soil_std,
+        rh_mean_pct=_circular_smooth(rh_mean),
+        rh_std_pct=rh_std,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Public API: cached climatology + anomaly application
+# ---------------------------------------------------------------------------
+
+def _cache_path(zone_id: str, years: int, end_year: Optional[int]) -> Path:
+    key = _stable_seed(f"{zone_id}|{years}|{end_year}")
+    return _CACHE_DIR / f"{zone_id}_{key}.json"
+
+
+def get_zone_climatology(
+    zone_id: str,
+    lat: float,
+    lon: float,
+    years: int = 10,
+    end_year: Optional[int] = None,
+    prefer_real: bool = True,
+    use_cache: bool = True,
+) -> ZoneClimatology:
+    """Return the day-of-year climatology for a zone, cached on disk.
+
+    Resolution order:
+      1. Fresh cache hit (same zone/years/end_year, < _CACHE_TTL_DAYS old).
+      2. Real Open-Meteo archive fetch (if prefer_real and requests present).
+      3. Deterministic synthetic monsoon model (never fails).
+
+    Args:
+        end_year:   Last calendar year included in the climatology period.
+                    Default: the most recent COMPLETE year (now.year - 1).
+                    Pin this explicitly for backtests so the climatology
+                    cannot see the period being backtested (look-ahead).
+        prefer_real: Set False to force the synthetic model (offline tests).
+    """
+    path = _cache_path(zone_id, years, end_year)
+    if use_cache and path.exists():
+        age_days = (
+            datetime.now(timezone.utc)
+            - datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc)
+        ).days
+        if age_days < _CACHE_TTL_DAYS:
+            try:
+                with open(path) as f:
+                    return ZoneClimatology.from_dict(json.load(f))
+            except Exception as e:
+                logger.warning("climatology: cache read failed (%s) -- rebuilding", e)
+
+    clim: Optional[ZoneClimatology] = None
+    if prefer_real and _REQUESTS_AVAILABLE:
+        try:
+            clim = _fetch_openmeteo_climatology(zone_id, lat, lon, years, end_year)
+            logger.info(
+                "climatology: built real %dy climatology for %s (%d-%d)",
+                years, zone_id, clim.period_start_year, clim.period_end_year,
+            )
+        except Exception as e:
+            logger.warning(
+                "climatology: real fetch failed for %s (%s) -- synthetic model",
+                zone_id, e,
+            )
+            clim = None
+
+    if clim is None:
+        clim = _synthetic_climatology(zone_id, lat, n_years=years)
+
+    if use_cache:
+        try:
+            with open(path, "w") as f:
+                json.dump(clim.to_dict(), f)
+        except Exception as e:
+            logger.warning("climatology: cache write failed (%s) -- continuing", e)
+
+    return clim
+
+
+def apply_climatology_anomalies(obs: ZoneObs, clim: ZoneClimatology) -> ZoneObs:
+    """Return a NEW ZoneObs with the four anomaly fields populated as
+    z-scores against `clim`. Never mutates the input (to_dict/from_dict
+    round-trip, matching the codebase idiom).
+
+    Skipped/Guarded cases (all deliberate, all logged at debug level):
+      * obs.source == SYNTHETIC: returned unchanged. Synthetic obs carry
+        injected anomalies from the event flags; re-scoring them against a
+        climatology would double-transform the training signal.
+      * precip anomaly only computed when at least one precip aggregate is
+        non-zero (a precip-less fetch like _fetch_smap would otherwise read
+        as a catastrophic false drought: (0 - mean)/std << 0).
+      * temp anomaly only when temp_mean_c != 0.0 (0.0 is the "unset"
+        default, not a real temperature in this pipeline's operating range).
+      * soil anomaly only when soil_moisture_pct > 0.0.
+      * rh anomaly only when rh_mean_pct > 0.0 (0.0 is "unset", not a real
+        humidity reading -- see rh_anomaly_idx's field comment in
+        zone_observation.py for why this exists: fungi_risk_signal()'s pure
+        absolute-RH threshold was flat across ENSO regimes in a tropical
+        climate, so it was masking correctly regime-sensitive drought/flood
+        signals in the actual alert_level output).
+
+    Z-scores are clipped to [-5, 5] by ZoneObs.__post_init__ as usual.
+    """
+    if obs.source == DataSource.SYNTHETIC:
+        logger.debug(
+            "climatology: %s source is SYNTHETIC -- anomalies left as injected",
+            obs.zone_id,
+        )
+        return obs
+
+    d = obs.to_dict()
+    d.pop("_schema_version", None)
+
+    has_precip = (obs.precip_30d_mm > 0.0) or (obs.precip_14d_mm > 0.0) \
+        or (obs.precip_7d_mm > 0.0) or (obs.precip_24h_mm > 0.0)
+    if has_precip:
+        mean_w, std_w = clim.window_precip_stats(obs.valid_time, _PRECIP_WINDOW_DAYS)
+        d["precip_anomaly_idx"] = _clip(
+            (obs.precip_30d_mm - mean_w) / std_w, -5.0, 5.0
+        )
+
+    if obs.temp_mean_c != 0.0:
+        t_mean, t_std = clim.daily_temp_stats(obs.valid_time)
+        d["temp_anomaly_idx"] = _clip((obs.temp_mean_c - t_mean) / t_std, -5.0, 5.0)
+
+    if obs.soil_moisture_pct > 0.0:
+        s_mean, s_std = clim.daily_soil_stats(obs.valid_time)
+        d["soil_moisture_anom"] = _clip(
+            (obs.soil_moisture_pct - s_mean) / s_std, -5.0, 5.0
+        )
+
+    if obs.rh_mean_pct > 0.0:
+        r_mean, r_std = clim.daily_rh_stats(obs.valid_time)
+        d["rh_anomaly_idx"] = _clip(
+            (obs.rh_mean_pct - r_mean) / r_std, -5.0, 5.0
+        )
+
+    return ZoneObs.from_dict(d)
+
+
+def apply_anomalies_by_zone_id(
+    obs: ZoneObs,
+    lat: float,
+    lon: float,
+    years: int = 10,
+    prefer_real: bool = True,
+) -> ZoneObs:
+    """Convenience wrapper: resolve (or build) the cached climatology for
+    obs.zone_id, then apply it. This is the entry point era5_data_pipeline
+    calls; kept separate from apply_climatology_anomalies so callers that
+    already hold a ZoneClimatology (e.g. the backtester looping over days)
+    don't pay the cache lookup per step.
+    """
+    clim = get_zone_climatology(obs.zone_id, lat, lon, years=years,
+                                prefer_real=prefer_real)
+    return apply_climatology_anomalies(obs, clim)
+
+
+# ---------------------------------------------------------------------------
+# Self-test (python climatology.py) -- fully offline
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import sys
+    from datetime import timezone as _tz
+    from zone_observation import make_synthetic_zone_obs
+
+    logging.basicConfig(level=logging.WARNING)
+    print("climatology.py self-test (offline: prefer_real=False)\n")
+    failures: List[str] = []
+
+    def _assert(cond: bool, msg: str) -> None:
+        if not cond:
+            failures.append(msg)
+            print(f"  FAIL: {msg}")
+
+    LAT, LON = -6.3, 107.3  # Karawang, West Java
+
+    # 1. Synthetic climatology: shape, length, round-trip
+    clim = _synthetic_climatology("test_zone", LAT)
+    _assert(len(clim.precip_mean_mm) == _TABLE_LEN, "precip table length")
+    d = clim.to_dict()
+    clim2 = ZoneClimatology.from_dict(d)
+    _assert(clim2.zone_id == clim.zone_id, "ZoneClimatology round-trip zone_id")
+    _assert(abs(clim2.precip_mean_mm[100] - clim.precip_mean_mm[100]) < 1e-12,
+            "ZoneClimatology round-trip values")
+
+    # 2. Seasonality: Java should be much wetter in Jan than in Aug
+    jan_mean = sum(clim.precip_mean_mm[0:31]) / 31.0
+    aug_mean = sum(clim.precip_mean_mm[212:243]) / 31.0
+    _assert(jan_mean > aug_mean * 1.3,
+            f"monsoon shape wrong: Jan={jan_mean:.1f} vs Aug={aug_mean:.1f} mm/day")
+    print(f"  Seasonality OK: Jan {jan_mean:.1f} mm/day vs Aug {aug_mean:.1f} mm/day")
+
+    # 3. doy_index leap mapping: Mar 1 maps to the same entry in every year
+    d1 = doy_index(datetime(2023, 3, 1, tzinfo=_tz.utc))
+    d2 = doy_index(datetime(2024, 3, 1, tzinfo=_tz.utc))
+    _assert(d1 == d2 == 61, f"Mar 1 mapping inconsistent: {d1} vs {d2}")
+    _assert(doy_index(datetime(2024, 2, 29, tzinfo=_tz.utc)) == 60, "Feb 29 mapping")
+    print(f"  doy_index OK (Mar 1 -> {d1}, Feb 29 -> 60)")
+
+    # 4. Window stats: 30-day wet-season aggregate exceeds dry-season
+    wet_dt = datetime(2024, 1, 31, tzinfo=_tz.utc)
+    dry_dt = datetime(2024, 8, 31, tzinfo=_tz.utc)
+    wet_mean, wet_std = clim.window_precip_stats(wet_dt, 30)
+    dry_mean, _ = clim.window_precip_stats(dry_dt, 30)
+    _assert(wet_mean > dry_mean, "window aggregate seasonality wrong")
+    _assert(wet_std >= _PRECIP_STD_FLOOR, "window std floor violated")
+    print(f"  Window stats OK: wet30={wet_mean:.0f}mm dry30={dry_mean:.0f}mm")
+
+    # 5. apply_climatology_anomalies: real-source obs gets anomalies
+    obs = make_synthetic_zone_obs("realish_zone", seed=1)
+    od = obs.to_dict()
+    od.pop("_schema_version", None)
+    od["source"] = DataSource.OPENMETEO_LIVE.value   # pretend real
+    # A real fetcher leaves anomaly fields at 0.0 ("unset") -- mirror that so
+    # this test measures exactly what this module adds.
+    od["precip_anomaly_idx"] = od["temp_anomaly_idx"] = od["soil_moisture_anom"] = 0.0
+    real_obs = ZoneObs.from_dict(od)
+    pre_precip_z = real_obs.precip_anomaly_idx
+    out = apply_climatology_anomalies(real_obs, clim)
+    _assert(out is not real_obs, "apply should return a NEW object")
+    _assert(real_obs.precip_anomaly_idx == pre_precip_z, "input obs was mutated!")
+    # The obs built by make_synthetic_zone_obs has neutral-ish aggregates;
+    # anomaly must be finite and within clip range.
+    _assert(-5.0 <= out.precip_anomaly_idx <= 5.0, "anomaly outside clip")
+    _assert(-5.0 <= out.temp_anomaly_idx <= 5.0, "temp anomaly outside clip")
+    _assert(-5.0 <= out.soil_moisture_anom <= 5.0, "soil anomaly outside clip")
+    print(f"  Anomaly application OK: precip_z={out.precip_anomaly_idx:+.2f} "
+          f"temp_z={out.temp_anomaly_idx:+.2f} soil_z={out.soil_moisture_anom:+.2f}")
+
+    # 6. Drought/wet extremes produce correctly-signed anomalies
+    # (keep aggregates monotonic: 24h <= 7d <= 14d <= 30d)
+    dry_obs_d = dict(od)
+    dry_obs_d["precip_24h_mm"] = 0.0
+    dry_obs_d["precip_7d_mm"] = 0.01 * wet_mean / 4.0
+    dry_obs_d["precip_14d_mm"] = 0.02 * wet_mean / 2.0
+    dry_obs_d["precip_30d_mm"] = 0.05 * wet_mean   # 5% of wet climatology
+    dry_obs_d["valid_time"] = wet_dt.isoformat()
+    dry_out = apply_climatology_anomalies(ZoneObs.from_dict(dry_obs_d), clim)
+    _assert(dry_out.precip_anomaly_idx < -1.0,
+            f"dry obs should get negative anomaly, got {dry_out.precip_anomaly_idx}")
+    wet_obs_d = dict(od)
+    wet_obs_d["precip_24h_mm"] = 2.5 * dry_mean / 30.0
+    wet_obs_d["precip_7d_mm"] = 2.5 * dry_mean / 4.0
+    wet_obs_d["precip_14d_mm"] = 2.5 * dry_mean / 2.0
+    wet_obs_d["precip_30d_mm"] = 2.5 * dry_mean
+    wet_obs_d["valid_time"] = dry_dt.isoformat()
+    wet_out = apply_climatology_anomalies(ZoneObs.from_dict(wet_obs_d), clim)
+    _assert(wet_out.precip_anomaly_idx > 1.0,
+            f"wet obs should get positive anomaly, got {wet_out.precip_anomaly_idx}")
+    print(f"  Sign check OK: dry_z={dry_out.precip_anomaly_idx:+.2f} "
+          f"wet_z={wet_out.precip_anomaly_idx:+.2f}")
+
+    # 7. SYNTHETIC-source obs is skipped unchanged
+    syn = make_synthetic_zone_obs("syn_zone", drought=True, seed=2)
+    syn_out = apply_climatology_anomalies(syn, clim)
+    _assert(syn_out.precip_anomaly_idx == syn.precip_anomaly_idx,
+            "SYNTHETIC obs anomaly was modified (should be skipped)")
+    print("  SYNTHETIC skip OK")
+
+    # 8. Zero-precip obs (smap-style real fetch: anomaly fields unset at 0.0)
+    # must NOT read as a catastrophic false drought.
+    zero_d = dict(od)
+    zero_d["precip_24h_mm"] = zero_d["precip_7d_mm"] = 0.0
+    zero_d["precip_14d_mm"] = zero_d["precip_30d_mm"] = 0.0
+    zero_d["precip_anomaly_idx"] = 0.0   # real fetchers leave this unset
+    zero_out = apply_climatology_anomalies(ZoneObs.from_dict(zero_d), clim)
+    _assert(zero_out.precip_anomaly_idx == 0.0,
+            "precip-less fetch should keep anomaly 0.0 (no false drought)")
+    print("  Zero-precip guard OK")
+
+    # 9. Cache round-trip through get_zone_climatology (offline path)
+    clim_c = get_zone_climatology("cache_zone", LAT, LON, years=5,
+                                  prefer_real=False, use_cache=True)
+    clim_c2 = get_zone_climatology("cache_zone", LAT, LON, years=5,
+                                   prefer_real=False, use_cache=True)
+    _assert(clim_c.precip_mean_mm == clim_c2.precip_mean_mm,
+            "cached climatology not identical")
+    print("  Cache round-trip OK")
+
+    print()
+    if failures:
+        print(f"FAILED  {len(failures)} test(s):")
+        for f in failures:
+            print(f"  - {f}")
+        sys.exit(1)
+    else:
+        print("All 9 test groups passed.")

requirements.txt

@@ -0,0 +1,62 @@
+# Core (needed for the data/scoring/Indonesia stack -- climatology.py,
+# indonesia_zones.py, backtest_indonesia.py, era5_data_pipeline.py,
+# zone_observation.py, crop_risk_scorer.py, hierarchical_search.py)
+numpy>=1.24
+
+# RL training stack (train_curriculum.py, train_kaggle.py,
+# weather_forecast_env.py, gru_weather_policy.py)
+torch>=2.0
+gymnasium>=0.29
+stable-baselines3>=2.0
+sb3-contrib>=2.0
+
+# Hyperparameter sweeps (sweep_reward_shaping.py) -- this file was temporarily
+# created to assess optimal hyperparameters. Train_kaggle.py's best values
+# dict (learning_rate=6.916624987609979e-05, ent_coef=0.08779238696445962,
+# etc.) was "found by the Optuna sweep (trial 6 of the 20-trial run against
+# n_zones=3 / max_steps=250)" and hardcoded as the CLI defaults --
+# consistent with this having been a one-time iteration tool whose winning
+# trial's output was captured inline, rather than a script meant to persist
+# in the repo. optuna is essential if you wish to reexplore optimal parameters.
+optuna>=3.5
+
+# Edge export (mnn_export.py) -- MNN itself has no pip package; build/install
+# per https://github.com/alibaba/MNN, this only covers the ONNX/export side.
+onnx>=1.15
+
+# Real-data fetching (era5_data_pipeline.py, climatology.py) -- all optional;
+# each degrades to synthetic/cached data gracefully without it, but any real
+# (non-synthetic) fetch needs at least `requests`. The comment at the top of
+# this file listing era5_data_pipeline.py under "Core -- numpy>=1.24" is
+# incomplete: numpy alone is enough for the module to import, not for its
+# real-data code paths to work.
+requests>=2.31
+cdsapi>=0.6             # ERA5 reanalysis tier only
+earthengine-api         # `import ee` -- IMERG/CHIRPS/SMAP satellite tier only
+netCDF4>=1.6            # ERA5 NetCDF reads -- tried first
+xarray>=2023.1          # ERA5 NetCDF reads -- fallback if netCDF4 unavailable
+
+# LocalTimesFMBackend only (timesfm_wrapper.py) -- a deliberately opt-in
+# forecast tier gated behind a manually downloaded, SHA256-verified
+# checkpoint (see LocalTimesFMBackend.__post_init__); most users won't hit
+# this path. Note: timesfm_wrapper.py imports pandas without a try/except
+# guard (unlike its `import timesfm` a few lines above, which does have
+# one) -- if pandas is missing, this fails with an unhelpful raw
+# ImportError rather than the graceful message the rest of this codebase
+# uses for optional deps.
+timesfm
+pandas>=2.0
+
+# TensorBoard training logs (train_kaggle.py) -- optional; training runs
+# fine without it, just without tfevents output. Listed here despite being
+# wrapped in a try/except in code because train_kaggle.py's own quickstart
+# docstring tells users to install it, and every verified training run in
+# this project's history had it installed.
+tensorboard>=2.14
+
+# Optional -- only needed if you actually connect to a broker
+# (node_transport.py's MQTTTransport falls back to LocalTransport without it)
+paho-mqtt>=1.6
+
+# Testing
+pytest>=7.0

train_curriculum.py

@@ -0,0 +1,1045 @@
+"""
+train_curriculum.py
+===================
+"""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import zone_observation as _zo
+
+assert _zo.SCHEMA_VERSION == 3, (
+    f"train_curriculum: 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 (graceful degradation)
+# ---------------------------------------------------------------------------
+
+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        # type: ignore
+    MaskablePPO     = None        # type: ignore
+    Monitor         = None        # type: ignore
+    BaseCallback    = object      # type: ignore
+
+# GRU policy is optional — falls back to MlpPolicy if not present
+try:
+    from gru_weather_policy import (
+        create_gru_weather_policy_kwargs,
+        ZoneEquivariantMaskablePolicy,
+    )
+    _GRU_AVAILABLE = True
+except ImportError:
+    _GRU_AVAILABLE = False
+    create_gru_weather_policy_kwargs = None  # type: ignore
+    ZoneEquivariantMaskablePolicy = None  # type: ignore
+
+# Physics dynamics is optional — Dyna augmentation disabled if unavailable.
+# Import failure is silent: train_phase() runs identically to the original
+# when dynamics_config=None or _DYNAMICS_AVAILABLE=False.
+try:
+    from physics_dynamics import TemporalDynamicsModel, DynaRolloutBuffer, ZoneStateTensor
+    _DYNAMICS_AVAILABLE = True
+except ImportError:
+    _DYNAMICS_AVAILABLE = False
+    TemporalDynamicsModel = None  # type: ignore
+    DynaRolloutBuffer     = None  # type: ignore
+    ZoneStateTensor       = None  # type: ignore
+
+
+# ---------------------------------------------------------------------------
+# Logging
+# ---------------------------------------------------------------------------
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    handlers=[
+        logging.FileHandler("training.log"),
+        logging.StreamHandler(),
+    ],
+)
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Device selection
+# ---------------------------------------------------------------------------
+
+def _select_device(requested: str) -> str:
+    """Return 'cuda' if available and requested, else 'cpu'."""
+    if requested == "cuda":
+        if _TORCH_AVAILABLE and torch.cuda.is_available():
+            return "cuda"
+        logger.warning("CUDA requested but not available — falling back to CPU.")
+        return "cpu"
+    return requested
+
+
+# ---------------------------------------------------------------------------
+# Dynamics configuration
+# ---------------------------------------------------------------------------
+
+@dataclass
+class DynamicsConfig:
+    """
+    Configuration for Dyna-style physics dynamics augmentation.
+
+    When dynamics_model_path is set and the model file exists, DynaCallback
+    loads the pre-trained TemporalDynamicsModel and adds a surprise bonus
+    to the PPO reward at each step. When dynamics_model_path is None (default),
+    training is identical to the original curriculum — no overhead, no change.
+
+    Fields
+    ------
+    dynamics_model_path:
+        Path to a pre-trained TemporalDynamicsModel checkpoint (.pt).
+        Produced by DynamicsTrainer.save() in physics_dynamics.py.
+        If None or the file does not exist, DynaCallback is not attached.
+
+    surprise_weight:
+        Scalar multiplier for the surprise bonus added to the PPO reward.
+        Start at 0.05. Increase to 0.1 if the agent is under-exploring;
+        decrease to 0.01 if the dynamics bonus dominates task reward.
+        The bonus is clipped to [0, surprise_weight] before adding, so
+        this value is also the maximum bonus per step.
+
+    update_dynamics_every_n_steps:
+        Fine-tune the dynamics model on transitions collected during RL
+        training every N environment steps. 0 = no fine-tuning (frozen model).
+        Fine-tuning closes the Dyna loop: better policy -> richer data ->
+        better dynamics -> better policy. Start with 0 until baseline training
+        is stable, then enable at 50_000 steps.
+
+    fine_tune_epochs:
+        Number of gradient steps per fine-tuning update. Keep low (3-5)
+        to avoid overfitting to the most recent transitions.
+
+    transition_buffer_size:
+        Maximum number of (current, next) transition pairs stored for
+        fine-tuning. Ring buffer: oldest pairs dropped when full.
+    """
+    dynamics_model_path:           Optional[str] = None
+    surprise_weight:               float         = 0.05
+    update_dynamics_every_n_steps: int           = 0       # 0 = frozen
+    fine_tune_epochs:              int           = 3
+    transition_buffer_size:        int           = 10_000
+
+
+# ---------------------------------------------------------------------------
+# Curriculum definition
+# ---------------------------------------------------------------------------
+
+def resolve_phase_max_steps(n_zones: int, budget_mode: str, episode_length: int) -> int:
+    """
+    Map budget_mode → max_steps for a curriculum phase.
+
+    The visit-once mask makes the structural ceiling n_zones+1. Old phases
+    used episode_length of 150–300, which never forced zone selection.
+    budget_mode overrides that so later phases train allocation skill.
+
+      full    → n_zones + 1
+      scarce  → n_zones
+      triage  → max(1, n_zones - 1)
+      legacy  → keep episode_length (old behaviour)
+    """
+    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
+    if mode == "triage":
+        return max(1, n - 1)
+    raise ValueError(f"Unknown budget_mode {budget_mode!r}")
+
+
+@dataclass
+class CurriculumPhase:
+    name:           str
+    total_steps:    int
+    episode_length: int   # used only when budget_mode="legacy"
+    n_zones:        int
+    risk_weights:   RiskWeights
+    # Default triage: no phase can pass without learning which zones to skip.
+    budget_mode:    str = "triage"  # full | scarce | triage | legacy
+
+    learning_rate:  float = 3e-4
+    n_steps:        int   = 4_096
+    batch_size:     int   = 256
+    n_epochs:       int   = 10
+    gamma:          float = 0.995
+    gae_lambda:     float = 0.95
+    clip_range:     float = 0.2
+    ent_coef:       float = 0.02
+    vf_coef:        float = 0.5
+    max_grad_norm:  float = 0.5
+
+    def resolved_max_steps(self) -> int:
+        return resolve_phase_max_steps(self.n_zones, self.budget_mode, self.episode_length)
+
+
+class WeatherCurriculum:
+    """Five-phase climate curriculum from baseline through stress extremes.
+
+    Budget progression (forces zone differentiation):
+      normal   → full   (learn inspection has value; 2 zones)
+      monsoon  → scarce (start leaving someone out; 3 zones)
+      drought  → triage (must skip ≥1; 3 zones)
+      heatwave → triage (4 zones)
+      humidity → triage (4 zones)
+    """
+
+    PHASES: Dict[str, CurriculumPhase] = {
+
+        "normal": CurriculumPhase(
+            name="normal",
+            total_steps=200_000,
+            episode_length=150,
+            n_zones=2,
+            budget_mode="full",
+            risk_weights=RiskWeights(),
+            n_steps=4_096,
+            ent_coef=0.05,
+        ),
+
+        "monsoon": CurriculumPhase(
+            name="monsoon",
+            total_steps=150_000,
+            episode_length=300,
+            n_zones=3,
+            budget_mode="scarce",
+            risk_weights=RiskWeights(
+                drought_obs_weight=0.40, drought_forecast_weight=0.60,
+                flood_obs_weight=0.70,   flood_forecast_weight=0.30,
+                fungi_obs_weight=0.75,   fungi_forecast_weight=0.25,
+                supply_drought_weight=0.25,
+                supply_flood_weight=0.50,
+                supply_harvest_pressure_weight=0.25,
+            ),
+            n_steps=4_096,
+        ),
+
+        "drought": CurriculumPhase(
+            name="drought",
+            total_steps=120_000,
+            episode_length=250,
+            n_zones=3,
+            budget_mode="triage",
+            risk_weights=RiskWeights(
+                drought_obs_weight=0.80, drought_forecast_weight=0.20,
+                flood_obs_weight=0.30,   flood_forecast_weight=0.70,
+                fungi_obs_weight=0.55,   fungi_forecast_weight=0.45,
+                supply_drought_weight=0.55,
+                supply_flood_weight=0.25,
+                supply_harvest_pressure_weight=0.20,
+            ),
+            n_steps=4_096,
+        ),
+
+        "heatwave": CurriculumPhase(
+            name="heatwave",
+            total_steps=120_000,
+            episode_length=220,
+            n_zones=4,
+            budget_mode="triage",
+            risk_weights=RiskWeights(
+                drought_obs_weight=0.75, drought_forecast_weight=0.25,
+                flood_obs_weight=0.25,   flood_forecast_weight=0.75,
+                fungi_obs_weight=0.50,   fungi_forecast_weight=0.50,
+                supply_drought_weight=0.60,
+                supply_flood_weight=0.15,
+                supply_harvest_pressure_weight=0.25,
+            ),
+            n_steps=4_096,
+        ),
+
+        "humidity": CurriculumPhase(
+            name="humidity",
+            total_steps=100_000,
+            episode_length=200,
+            n_zones=4,
+            budget_mode="triage",
+            risk_weights=RiskWeights(
+                drought_obs_weight=0.30, drought_forecast_weight=0.70,
+                flood_obs_weight=0.50,   flood_forecast_weight=0.50,
+                fungi_obs_weight=0.85,   fungi_forecast_weight=0.15,
+                supply_drought_weight=0.20,
+                supply_flood_weight=0.30,
+                supply_harvest_pressure_weight=0.50,
+                quality_fungi_weight=0.80,
+                quality_delay_weight=0.20,
+            ),
+            n_steps=4_096,
+        ),
+    }
+
+    @classmethod
+    def get_phase(cls, name: str) -> CurriculumPhase:
+        if name not in cls.PHASES:
+            raise ValueError(
+                f"Unknown phase '{name}'. Options: {sorted(cls.PHASES)}"
+            )
+        return cls.PHASES[name]
+
+    @classmethod
+    def phase_order(cls) -> List[str]:
+        return ["normal", "monsoon", "drought", "heatwave", "humidity"]
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint callback
+# ---------------------------------------------------------------------------
+
+class CheckpointCallback(BaseCallback):
+    """Save a checkpoint every `save_freq` timesteps."""
+
+    def __init__(self, output_dir: Path, save_freq: int = 25_000) -> None:
+        super().__init__()
+        self.output_dir = output_dir
+        self.save_freq  = save_freq
+        self._last_save = 0
+
+    def _on_step(self) -> bool:
+        if self.num_timesteps - self._last_save >= self.save_freq:
+            self._last_save = self.num_timesteps
+            path = self.output_dir / f"checkpoint_{self.num_timesteps}.zip"
+            self.model.save(str(path))
+            logger.info("Checkpoint saved: %s", path.name)
+        return True
+
+
+
+# ---------------------------------------------------------------------------
+# Dyna callback
+# ---------------------------------------------------------------------------
+
+class DynaCallback(BaseCallback):
+    """
+    Augments PPO rewards with a physics-dynamics surprise bonus (Dyna-style).
+
+    At each environment step, this callback:
+      1. Extracts the current and next observation as ZoneStateTensor objects.
+      2. Calls DynaRolloutBuffer.compute_surprise_bonus() — the normalised
+         prediction error of the dynamics model for this transition.
+      3. Writes the bonus directly into the PPO rollout buffer at the position
+         that was just written by env.step().
+      4. Optionally fine-tunes the dynamics model on accumulated transitions.
+
+    Reward injection mechanism
+    --------------------------
+    SB3's RolloutBuffer stores rewards at self.model.rollout_buffer.rewards[pos-1]
+    immediately after env.step() returns, where pos is the buffer write pointer.
+    The callback's _on_step() runs after that write, so we can read and modify
+    the reward before any PPO computation sees it.
+
+    The pos pointer advances BEFORE _on_step() is called, so the correct
+    index is (self.model.rollout_buffer.pos - 1) % n_steps.
+
+    This is the same approach used by SB3's RND and curiosity implementations.
+
+    Safe degradation
+    ----------------
+    If the dynamics model is unavailable, or if obs keys are missing (e.g.
+    during the first step of an episode), the callback returns True silently
+    without modifying any reward. It never raises or interrupts training.
+
+    Args:
+        dyna_buffer:    DynaRolloutBuffer wrapping the loaded dynamics model.
+        dynamics_cfg:   DynamicsConfig controlling weights and fine-tuning.
+        n_zones:        Must match the environment's n_zones.
+        horizon_days:   Must match ForecastConfig.horizon_days.
+        device:         Torch device string for tensor ops.
+    """
+
+    _OBS_KEYS = ("forecast_precip", "forecast_uncertainty", "zone_belief")
+
+    def __init__(
+        self,
+        dyna_buffer:  "DynaRolloutBuffer",
+        dynamics_cfg: DynamicsConfig,
+        n_zones:      int,
+        horizon_days: int,
+        device:       str = "cpu",
+    ) -> None:
+        super().__init__()
+        self.dyna_buffer   = dyna_buffer
+        self.dynamics_cfg  = dynamics_cfg
+        self.n_zones       = n_zones
+        self.horizon_days  = horizon_days
+        self.device        = device
+
+        # Ring buffer for fine-tuning transitions
+        # Stored as (current_ZoneStateTensor, next_ZoneStateTensor) pairs
+        self._transition_buffer: list = []
+        self._tb_max = dynamics_cfg.transition_buffer_size
+
+        # Statistics logged every 10k steps
+        self._bonus_sum   = 0.0
+        self._bonus_count = 0
+        self._log_freq    = 10_000
+        self._last_log    = 0
+
+        # Previous obs for transition construction (obs_t → obs_t+1)
+        self._prev_obs: Optional[dict] = None
+
+    def _obs_to_state_tensor(self, obs: dict) -> Optional["ZoneStateTensor"]:
+        """
+        Convert a raw SB3 obs dict to ZoneStateTensor.
+
+        SB3 stores observations as numpy arrays with a leading env-count
+        dimension even for a single env: shape [1, ...]. We squeeze that dim.
+
+        Returns None if any required key is missing (safe degradation).
+        """
+        if not all(k in obs for k in self._OBS_KEYS):
+            return None
+
+        import torch
+        import numpy as np
+
+        try:
+            # SB3 obs shapes: [n_envs, ...] — squeeze env dim (n_envs=1)
+            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)
+
+            # Handle both [1, n_zones, H] and [n_zones, H] shapes gracefully
+            if precip.ndim == 2:
+                precip = precip[np.newaxis]   # [n_zones, H] -> [1, n_zones, H]
+            if uncert.ndim == 1:
+                uncert = uncert[np.newaxis]   # [n_zones]    -> [1, n_zones]
+            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:
+        """
+        Called after every env.step(). Injects surprise bonus into reward buffer.
+        """
+        # --- Extract current and next observations ---
+        # self.locals["obs_tensor"] is the obs BEFORE the step (obs_t).
+        # self.locals["new_obs"] is the obs AFTER the step (obs_t+1).
+        # Both are available in SB3 >= 1.8 on_step locals.
+        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  # safe: missing locals, skip silently
+
+            # Convert to ZoneStateTensor
+            if hasattr(obs_now, "numpy"):
+                # Tensor: convert dict-of-tensors or single tensor
+                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()}
+            else:
+                obs_now_np = obs_now
+
+            if hasattr(obs_next, "items"):
+                obs_next_np = {k: (v.cpu().numpy() if hasattr(v, "cpu") else v)
+                               for k, v in obs_next.items()}
+            else:
+                obs_next_np = obs_next
+
+            curr_state = self._obs_to_state_tensor(obs_now_np)
+            next_state = self._obs_to_state_tensor(obs_next_np)
+
+            if curr_state is None or next_state is None:
+                return True  # safe: obs keys not present yet
+
+            # --- Compute surprise bonus ---
+            bonus = self.dyna_buffer.compute_surprise_bonus(curr_state, next_state)
+            bonus_val = float(bonus.item())
+            bonus_clipped = min(bonus_val, self.dynamics_cfg.surprise_weight)
+
+            # --- Inject into PPO rollout buffer ---
+            # The rollout buffer pos pointer has already advanced; the reward
+            # for the current step is at (pos - 1) % n_steps.
+            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_clipped
+
+            # --- Accumulate for fine-tuning ---
+            if self.dynamics_cfg.update_dynamics_every_n_steps > 0:
+                self._transition_buffer.append((curr_state, next_state))
+                if len(self._transition_buffer) > self._tb_max:
+                    self._transition_buffer.pop(0)  # ring buffer: drop oldest
+
+            # --- Statistics ---
+            self._bonus_sum   += bonus_clipped
+            self._bonus_count += 1
+
+            if self.num_timesteps - self._last_log >= self._log_freq:
+                avg_bonus = (
+                    self._bonus_sum / self._bonus_count
+                    if self._bonus_count > 0 else 0.0
+                )
+                logger.info(
+                    "DynaCallback: step=%d  avg_surprise_bonus=%.4f  "
+                    "buffer_size=%d",
+                    self.num_timesteps, avg_bonus,
+                    len(self._transition_buffer),
+                )
+                self._bonus_sum   = 0.0
+                self._bonus_count = 0
+                self._last_log    = self.num_timesteps
+
+        except Exception as e:
+            # Never interrupt training on callback error — log and continue
+            logger.debug("DynaCallback._on_step error (non-fatal): %s", e)
+
+        return True
+
+    def _on_rollout_end(self) -> None:
+        """
+        Called at the end of each rollout collection. Optionally fine-tunes
+        the dynamics model on accumulated transitions.
+        """
+        if (
+            self.dynamics_cfg.update_dynamics_every_n_steps <= 0
+            or self.num_timesteps % self.dynamics_cfg.update_dynamics_every_n_steps != 0
+            or len(self._transition_buffer) < 16  # need at least one batch
+        ):
+            return
+
+        try:
+            from physics_dynamics import DynamicsTrainer
+            # Access the dynamics model directly from the buffer
+            dynamics_model = self.dyna_buffer.dynamics
+
+            # Minimal fine-tune: a few gradient steps on recent transitions
+            # We construct a temporary DynamicsTrainer around the existing model
+            # rather than creating a new one, to avoid re-initialising weights.
+            import torch
+            import torch.nn.functional as F
+
+            optimizer = torch.optim.AdamW(
+                dynamics_model.parameters(), lr=1e-4, weight_decay=1e-4
+            )
+            dynamics_model.train()
+
+            pairs = list(self._transition_buffer)  # snapshot
+            batch_size = min(32, len(pairs))
+
+            for epoch in range(self.dynamics_cfg.fine_tune_epochs):
+                import random
+                random.shuffle(pairs)
+                total_loss = 0.0
+                n_batches  = 0
+
+                for i in range(0, len(pairs) - batch_size, batch_size):
+                    batch = pairs[i : i + batch_size]
+                    curr_list = [p[0] for p in batch]
+                    next_list = [p[1] for p in batch]
+
+                    import torch as _t
+                    # Stack batch dimension
+                    curr_b = ZoneStateTensor(
+                        precip=_t.cat([s.precip for s in curr_list], dim=0),
+                        uncertainty=_t.cat([s.uncertainty for s in curr_list], dim=0),
+                        belief=_t.cat([s.belief for s in curr_list], dim=0),
+                    )
+                    next_b = ZoneStateTensor(
+                        precip=_t.cat([s.precip for s in next_list], dim=0),
+                        uncertainty=_t.cat([s.uncertainty for s in next_list], dim=0),
+                        belief=_t.cat([s.belief for s in next_list], dim=0),
+                    )
+
+                    pred, phys_loss = dynamics_model(curr_b, return_physics_loss=True)
+                    data_loss = (
+                        F.mse_loss(pred.precip / 500.0, next_b.precip / 500.0)
+                        + F.mse_loss(pred.uncertainty, next_b.uncertainty)
+                        + F.mse_loss(pred.belief,      next_b.belief)
+                    )
+                    loss = data_loss + 0.01 * phys_loss
+
+                    optimizer.zero_grad()
+                    loss.backward()
+                    _t.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 dynamics model at step=%d  "
+                "avg_loss=%.4f  n_transitions=%d",
+                self.num_timesteps,
+                total_loss / max(n_batches, 1),
+                len(self._transition_buffer),
+            )
+
+        except Exception as e:
+            logger.warning(
+                "DynaCallback._on_rollout_end fine-tune failed (non-fatal): %s", e
+            )
+
+
+def _build_dyna_callback(
+    dynamics_cfg: Optional[DynamicsConfig],
+    n_zones:      int,
+    horizon_days: int,
+    device:       str,
+) -> Optional["DynaCallback"]:
+    """
+    Build a DynaCallback if dynamics are configured and available.
+
+    Returns None (no Dyna augmentation) if:
+      - dynamics_cfg is None
+      - dynamics_model_path is not set
+      - the model file does not exist
+      - physics_dynamics module is unavailable
+      - any load/init error occurs
+
+    Callers pass the return value directly to CallbackList — None is ignored.
+    """
+    if dynamics_cfg is None or dynamics_cfg.dynamics_model_path is None:
+        return None
+
+    if not _DYNAMICS_AVAILABLE:
+        logger.warning(
+            "DynamicsConfig provided but physics_dynamics not installed — "
+            "Dyna augmentation disabled."
+        )
+        return None
+
+    model_path = Path(dynamics_cfg.dynamics_model_path)
+    if not model_path.exists():
+        logger.warning(
+            "Dynamics model not found at %s — Dyna augmentation disabled.",
+            model_path,
+        )
+        return None
+
+    try:
+        dynamics_model = TemporalDynamicsModel.load(str(model_path))
+        dynamics_model.eval()
+        dyna_buffer = DynaRolloutBuffer(
+            dynamics=dynamics_model,
+            uncertainty_weight=dynamics_cfg.surprise_weight,
+        )
+        callback = DynaCallback(
+            dyna_buffer=dyna_buffer,
+            dynamics_cfg=dynamics_cfg,
+            n_zones=n_zones,
+            horizon_days=horizon_days,
+            device=device,
+        )
+        logger.info(
+            "DynaCallback loaded: model=%s  surprise_weight=%.3f  "
+            "fine_tune_every=%d",
+            model_path.name,
+            dynamics_cfg.surprise_weight,
+            dynamics_cfg.update_dynamics_every_n_steps,
+        )
+        return callback
+
+    except Exception as e:
+        logger.warning(
+            "Failed to build DynaCallback (%s) — Dyna augmentation disabled.", e
+        )
+        return None
+
+
+
+# ---------------------------------------------------------------------------
+# Training
+# ---------------------------------------------------------------------------
+
+def transfer_curriculum_weights(
+    resume_from: str,
+    model: "MaskablePPO",
+    device: str = "auto",
+) -> "MaskablePPO":
+    """
+    Warm-start `model` (freshly constructed for the CURRENT phase's env/n_zones)
+    from `resume_from`'s checkpoint, transferring every parameter whose shape
+    matches exactly and leaving the rest at fresh random initialization.
+
+    Exists because MaskablePPO.load(path, env=new_env) raises "Observation
+    spaces do not match" whenever n_zones changes between curriculum phases
+    -- a hard SB3-level space-equality check that fires before any weight-
+    shape question is even considered. Loading without `env=` sidesteps that
+    (the checkpoint reconstructs against its own saved spaces); this function
+    then transfers whatever's compatible directly via the two state_dicts.
+
+    With the permutation-invariant GRUWeatherFeaturesExtractor (see
+    gru_weather_policy.py), every parameter except the action_net output
+    layer (shape tied to n_zones+1, the discrete action count) now matches
+    across any n_zones -- verified empirically at 61/63 tensors transferred
+    in a 2-zone -> 3-zone test. value_net transfers too (scalar output,
+    always n_zones-independent); only action_net needs relearning.
+    """
+    old_model = MaskablePPO.load(resume_from, device=device)
+    old_state = old_model.policy.state_dict()
+    new_state = model.policy.state_dict()
+
+    transferred, skipped = [], []
+    merged = {}
+    for key, new_tensor in new_state.items():
+        old_tensor = old_state.get(key)
+        if old_tensor is not None and old_tensor.shape == new_tensor.shape:
+            merged[key] = old_tensor.clone()
+            transferred.append(key)
+        else:
+            merged[key] = new_tensor
+            skipped.append(key)
+
+    model.policy.load_state_dict(merged)
+
+    logger.info(
+        "transfer_curriculum_weights: transferred %d/%d parameter tensors from %s "
+        "(freshly initialized: %s)",
+        len(transferred), len(new_state), resume_from, skipped or "none",
+    )
+    if not transferred:
+        logger.warning(
+            "transfer_curriculum_weights: transferred ZERO parameters -- the "
+            "architectures are likely genuinely incompatible (e.g. resuming "
+            "from a pre-permutation-invariant checkpoint), not just a normal "
+            "n_zones change. Check resume_from's origin before trusting this run."
+        )
+    return model
+
+
+def train_phase(
+    phase_name:     str,
+    output_dir:     Path,
+    resume_from:    Optional[str]           = None,
+    override_steps: Optional[int]           = None,
+    hidden_size:    int                     = 64,
+    device:         str                     = "auto",
+    seed:           int                     = 42,
+    dynamics_cfg:   Optional[DynamicsConfig] = None,
+    precip_scale:      float                = 40.0,
+) -> str:
+    """Train one curriculum phase. Returns path to the saved final model.
+
+    Args:
+        phase_name:    One of the WeatherCurriculum phase names.
+        output_dir:    Root directory for checkpoints and final model.
+        resume_from:   Path to a previous phase checkpoint to resume from.
+        override_steps: Override total_steps (useful for quick tests).
+        hidden_size:   GRU hidden size; must match any checkpoint being resumed.
+        device:        'cpu', 'cuda', or 'auto'.
+        seed:          Random seed.
+        dynamics_cfg:  Optional DynamicsConfig for Dyna surprise-bonus augmentation.
+                       Pass None (default) for standard training with no overhead.
+        precip_scale:  Fixed (non-learned) divisor applied to forecast_precip
+                       before the GRU extractor. See train_kaggle.py's
+                       INPUT NORMALIZATION docstring section for why this
+                       exists. Default 40.0 matches the value that produced
+                       the validated single-dirty selection-accuracy result
+                       (see model card) -- confirmed working, not confirmed
+                       optimal, and not yet validated across curriculum
+                       phase transitions specifically (only within a single
+                       train_kaggle.py run). Must match across resumed
+                       checkpoints the same way hidden_size must.
+    """
+
+    if not _ML_AVAILABLE:
+        raise RuntimeError(
+            f"ML stack not available: {_ML_IMPORT_ERROR}\n"
+            "Install: pip install stable-baselines3 sb3-contrib torch"
+        )
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    models_dir = output_dir / "models"
+    models_dir.mkdir(exist_ok=True)
+
+    device = _select_device(
+        device if device != "auto"
+        else ("cuda" if _TORCH_AVAILABLE and torch.cuda.is_available() else "cpu")
+    )
+
+    phase = WeatherCurriculum.get_phase(phase_name)
+    total_steps = override_steps or phase.total_steps
+    max_steps = phase.resolved_max_steps()
+    full_ceiling = phase.n_zones + 1
+
+    logger.info(
+        "Phase=%s  steps=%d  max_steps=%d (budget_mode=%s, full_ceiling=%d)  "
+        "n_zones=%d  device=%s  must_skip=%s",
+        phase.name, total_steps, max_steps, phase.budget_mode, full_ceiling,
+        phase.n_zones, device,
+        "yes" if max_steps < full_ceiling else "no",
+    )
+    if max_steps >= full_ceiling and phase.budget_mode not in ("full", "legacy"):
+        logger.warning(
+            "Phase %s: max_steps=%d >= full_ceiling=%d despite budget_mode=%s — "
+            "check resolve_phase_max_steps.",
+            phase.name, max_steps, full_ceiling, phase.budget_mode,
+        )
+
+    # --- Environment ---
+    config = ForecastConfig(
+        n_zones=phase.n_zones,
+        seed=seed,
+        soft_reset=True,
+        max_steps=max_steps,
+    )
+    phase.risk_weights.attach_to_config(config)
+
+    # Monitor wraps correctly: get_wrapper_attr('action_masks') walks the
+    # wrapper stack and finds action_masks() on WeatherForecastEnv.
+    env = Monitor(make_weather_env(config))
+
+    # --- Policy kwargs ---
+    if _GRU_AVAILABLE:
+        policy_kwargs = create_gru_weather_policy_kwargs(
+            hidden_size=hidden_size,
+            precip_scale=precip_scale,
+        )
+        # ZoneEquivariantMaskablePolicy, NOT the "MultiInputPolicy" string.
+        # The default policy builds action logits from action_net(latent_pi)
+        # on top of the extractor's pooled (permutation-invariant) feature
+        # vector -- zone identity is erased before the action head ever
+        # sees it, so under triage the policy can only express a static
+        # per-slot bias (empirically: always inspects action index 0,
+        # regardless of which zone's content actually looks risky). This
+        # was an active bug in this function: GRU features were used, but
+        # every prior curriculum-trained checkpoint went through the same
+        # pooled action_net as the plain-MLP fallback below and could not
+        # have learned risk-conditioned zone selection. See
+        # gru_weather_policy.py module docstring and train_kaggle.py's
+        # POLICY section for the full diagnosis.
+        policy = ZoneEquivariantMaskablePolicy
+        logger.info(
+            "Using zone-equivariant GRU policy (hidden_size=%d, "
+            "precip_scale=%.3g)",
+            hidden_size, precip_scale,
+        )
+    else:
+        # dict with 'pi'/'vf' keys is the correct net_arch format for
+        # MultiInputPolicy in SB3 >= 1.8 (validated on SB3 2.8.0)
+        policy_kwargs = dict(net_arch=dict(pi=[128, 64], vf=[128, 64]))
+        policy = "MultiInputPolicy"
+        logger.info("GRU policy unavailable — using MLP policy (net_arch=128,64)")
+
+    # --- Model ---
+    ppo_kwargs = dict(
+        learning_rate=phase.learning_rate,
+        n_steps=phase.n_steps,
+        batch_size=phase.batch_size,
+        n_epochs=phase.n_epochs,
+        gamma=phase.gamma,
+        gae_lambda=phase.gae_lambda,
+        clip_range=phase.clip_range,
+        ent_coef=phase.ent_coef,
+        vf_coef=phase.vf_coef,
+        max_grad_norm=phase.max_grad_norm,
+        device=device,
+        verbose=1,
+        seed=seed,
+    )
+
+    if resume_from:
+        logger.info("Resuming from %s", resume_from)
+        # Build a FRESH model for the CURRENT phase's env/n_zones (correct
+        # observation/action spaces throughout), then warm-start it from the
+        # checkpoint wherever shapes match. MaskablePPO.load(resume_from,
+        # env=env) directly would raise "Observation spaces do not match"
+        # the moment n_zones changes between phases -- see
+        # transfer_curriculum_weights()'s docstring for why, and what
+        # actually transfers (everything except action_net).
+        model = MaskablePPO(
+            policy=policy,
+            env=env,
+            policy_kwargs=policy_kwargs,
+            **ppo_kwargs,
+        )
+        model = transfer_curriculum_weights(resume_from, model, device=device)
+        reset_timesteps = False
+    else:
+        model = MaskablePPO(
+            policy=policy,
+            env=env,
+            policy_kwargs=policy_kwargs,
+            **ppo_kwargs,
+        )
+        reset_timesteps = True
+
+    # --- Callbacks ---
+    from stable_baselines3.common.callbacks import CallbackList
+    callbacks = [CheckpointCallback(output_dir)]
+
+    horizon_days = getattr(config, "horizon_days", 14)
+    dyna_cb = _build_dyna_callback(
+        dynamics_cfg=dynamics_cfg,
+        n_zones=phase.n_zones,
+        horizon_days=horizon_days,
+        device=device,
+    )
+    if dyna_cb is not None:
+        callbacks.append(dyna_cb)
+        logger.info("Dyna augmentation active for phase=%s", phase.name)
+    else:
+        logger.info("Dyna augmentation inactive for phase=%s", phase.name)
+
+    # --- Train ---
+    # use_masking=True is the default; MaskablePPO calls action_masks()
+    # automatically via get_wrapper_attr during rollout collection.
+    # Do NOT pass action_masks= to learn() — it is not a valid parameter.
+    model.learn(
+        total_timesteps=total_steps,
+        callback=CallbackList(callbacks),
+        reset_num_timesteps=reset_timesteps,
+        use_masking=True,
+    )
+
+    final_path = models_dir / f"final_{phase.name}.zip"
+    model.save(str(final_path))
+    logger.info("Saved final model: %s", final_path)
+
+    return str(final_path)
+
+
+def train_full_curriculum(
+    output_dir:   Path,
+    device:       str                      = "auto",
+    seed:         int                      = 42,
+    dynamics_cfg: Optional[DynamicsConfig] = None,
+    precip_scale:      float               = 40.0,
+) -> None:
+    """Run all phases in order, chaining each phase from the previous."""
+    phases = WeatherCurriculum.phase_order()
+    resume = None
+    for phase_name in phases:
+        logger.info("=== Starting phase: %s ===", phase_name)
+        resume = train_phase(
+            phase_name=phase_name,
+            output_dir=output_dir / phase_name,
+            resume_from=resume,
+            device=device,
+            seed=seed,
+            dynamics_cfg=dynamics_cfg,
+            precip_scale=precip_scale,
+        )
+        logger.info("=== Completed phase: %s ===", phase_name)
+
+
+# ---------------------------------------------------------------------------
+# CLI
+# ---------------------------------------------------------------------------
+
+def main() -> None:
+    p = argparse.ArgumentParser(
+        description="MaskablePPO curriculum trainer for WeatherForecastEnv"
+    )
+    p.add_argument(
+        "--phase",
+        choices=list(WeatherCurriculum.PHASES) + ["all"],
+        default="normal",
+        help="Curriculum phase to run, or 'all' to run full curriculum.",
+    )
+    p.add_argument("--output-dir",   default="./run",  help="Root output directory")
+    p.add_argument("--resume-from",  default=None,     help="Path to checkpoint .zip")
+    p.add_argument("--steps",        type=int, default=None, help="Override total_steps")
+    p.add_argument("--hidden-size",  type=int, default=64)
+    p.add_argument("--device",       default="auto",   help="'cpu', 'cuda', or 'auto'")
+    p.add_argument("--seed",                  type=int,   default=42)
+    p.add_argument(
+        "--dynamics-model",
+        default=None,
+        help="Path to pre-trained TemporalDynamicsModel .pt file. Enables Dyna augmentation.",
+    )
+    p.add_argument(
+        "--dynamics-weight",
+        type=float,
+        default=0.05,
+        help="Surprise bonus weight per step (only used with --dynamics-model). Default 0.05.",
+    )
+    p.add_argument(
+        "--dynamics-finetune-every",
+        type=int,
+        default=0,
+        help="Fine-tune dynamics model every N steps. 0=frozen (default).",
+    )
+    p.add_argument(
+        "--precip-scale",
+        type=float,
+        default=40.0,
+        help="Fixed (non-learned) divisor applied to forecast_precip before "
+             "the GRU extractor. See train_kaggle.py's INPUT NORMALIZATION "
+             "docstring section for the full rationale. 40.0 matches the "
+             "value that produced the validated single-dirty "
+             "selection-accuracy result (see model card). Must stay "
+             "consistent across a resumed checkpoint's phases, the same "
+             "way --hidden-size must.",
+    )
+    args = p.parse_args()
+
+    output_dir = Path(args.output_dir)
+
+    dynamics_cfg: Optional[DynamicsConfig] = None
+    if args.dynamics_model is not None:
+        dynamics_cfg = DynamicsConfig(
+            dynamics_model_path=args.dynamics_model,
+            surprise_weight=args.dynamics_weight,
+            update_dynamics_every_n_steps=args.dynamics_finetune_every,
+        )
+        logger.info(
+            "Dyna config: model=%s  weight=%.3f  finetune_every=%d",
+            args.dynamics_model, args.dynamics_weight, args.dynamics_finetune_every,
+        )
+
+    if args.phase == "all":
+        train_full_curriculum(
+            output_dir=output_dir,
+            device=args.device,
+            seed=args.seed,
+            dynamics_cfg=dynamics_cfg,
+            precip_scale=args.precip_scale,
+        )
+    else:
+        train_phase(
+            phase_name=args.phase,
+            output_dir=output_dir,
+            resume_from=args.resume_from,
+            override_steps=args.steps,
+            hidden_size=args.hidden_size,
+            device=args.device,
+            seed=args.seed,
+            dynamics_cfg=dynamics_cfg,
+            precip_scale=args.precip_scale,
+        )
+
+
+if __name__ == "__main__":
+    main()