DHDRL/monsoon-rl

Agricultural Weather-Risk RL — MaskablePPO + Zone-Equivariant GRU

Budget-constrained multi-zone weather-risk inspection for Indonesian rice zones. A Gymnasium environment tracks per-zone beliefs under a scarce inspection budget. A zone-equivariant GRU policy supports curriculum phases with changing n_zones. Deterministic crop-risk scoring grounds rewards and product alerts. An ONNX → MNN path is provided for edge inference.

Research codebase, not a production alerting service.

What is included

• weather_forecast_env.py — Gymnasium env (belief map, action masks, triage/scarce/full budget modes, zone-order shuffle, real EpisodeContext injection)
• gru_weather_policy.py — Zone-equivariant GRU extractor (ZoneEquivariantMaskablePolicy; shared per-zone weights; precip input scaling, precip_scale=40.0)
• train_curriculum.py / train_kaggle.py — Curriculum and standalone training (--budget-mode triage|scarce|full, --precip-scale)
• crop_risk_scorer.py — Deterministic, economics-calibrated risk scoring
• climatology.py + indonesia_zones.py — Per-zone climatology and Indonesian agricultural zones
• evaluate_checkpoint_real.py — Real-trajectory eval (L1 labels, belief Δ, zero_inspect control)
• mnn_export.py + edge_wrapper.cpp — ONNX/MNN export and C++ edge runtime with external GRU hidden state

Architecture note

Inspect logits are scored per zone before pooling (ZoneEquivariantMaskablePolicy). The value path stays permutation-invariant so most weights transfer across n_zones. Full-tour episode length is not zone-selection skill; allocation under triage is evaluated with max_steps < n_zones and per-episode zone-order shuffle.

How to use

Train (triage, recommended recipe)

python train_kaggle.py \
  --out ./run_nz2_triage \
  --n-zones 2 \
  --budget-mode triage \
  --steps 2400000 \
  --clean-episode-ratio 0.80 \
  --event-spatial-correlation 0.50 \
  --precip-scale 40.0 \
  --device auto

Requires ZoneEquivariantMaskablePolicy from gru_weather_policy.py (not the string "MultiInputPolicy"). --precip-scale is a real CLI flag (default 40.0, matching the validated recipe). Step count matters: 150k steps is enough for drought/flood selection but not fungi; fungi selection was confirmed solved by 2.4M steps on the same recipe — see Validation status.

Load

from sb3_contrib import MaskablePPO

model = MaskablePPO.load("run_nz2_triage/final_model.zip")
action, _ = model.predict(obs, action_masks=masks, deterministic=True)

Edge export

python mnn_export.py \
  --checkpoint run_nz2_triage/final_model.zip \
  --output weather_rl_model.mnn \
  --quantize int8 \
  --n-zones 2

Validation status

Claim	Status
Weight transfer across n_zones changes	Supported by design
Zone-order shuffle avoids fixed-slot lock-in	Observed, confirmed at 250k and 2.4M steps (raw per-slot action counts stay ~50/50)
Risk-conditioned zone choice — drought (synthetic, single-dirty)	Confirmed 100% at 2.4M steps (n=76 combined, two seed blocks + ρ=0.85 transfer)
Risk-conditioned zone choice — flood (synthetic, single-dirty)	Confirmed 100% at 2.4M steps (n=97 combined, two seed blocks + ρ=0.85 transfer)
Risk-conditioned zone choice — fungi/disease pressure (synthetic, single-dirty)	Confirmed 100% at 2.4M steps (n=95 combined, two seed blocks + ρ=0.85 transfer) — lagged at ~57–67% through 250k on the same run; needed more steps, not a different mechanism
Triage agent evaluated against real L1 trajectories	Open — all accuracy figures above are synthetic; evaluate_checkpoint_real.py has not yet been run on a triage-trained checkpoint
Product skill vs curated L1 labels	Scorer/product path available; L1 density limits apply
Multi-seed replication of the triage recipe	Open — the 2.4M result is one training run (--seed 7, confirmed via the training command; this pins the environment's episode generation but not PyTorch's weight-init RNG, which train_kaggle.py doesn't currently seed separately). A second seed reaching the same result has not been tried.
Held-out real product / emit metrics	Open

Limitations

• Research stack only: no on-call, feed SLA, or regulatory certification.
• Training is largely synthetic for the moment.
• Full-tour belief or episode-length metrics are not zone-selection skill.
• All zone-selection accuracy figures to date are from synthetic episodes; the triage-trained policy has not been evaluated against real L1 trajectories (see Validation status).
• The 2.4M-step selection-accuracy result comes from a single training run; not yet replicated on a second seed.
• MNN/C++ path is export-oriented, not a validated production edge deployment.

Dependencies

torch>=2.0
gymnasium>=0.29
stable-baselines3>=2.0
sb3-contrib>=2.0
numpy>=1.24

While GPU is recommended for training, CPUs can be used successfully.

License

AGPL-3.0. See repository for full terms.