Agricultural Weather-Risk RL — Permutation-Invariant GRU Policy with Curriculum Weight Transfer

Core RL components and training infrastructure for a budget-constrained, multi-zone weather-risk inspection task. The system includes a Gymnasium environment, a permutation-invariant GRU features extractor that enables weight transfer across curriculum phases with varying numbers of zones, full multi-phase curriculum training, and an export path to quantized MNN for edge deployment.

What is included

weather_forecast_env.py — Gymnasium environment with per-zone belief tracking, action masking, soft reset, and NaN-safe wrapper.
gru_weather_policy.py — Permutation-invariant GRU features extractor (shared per-zone processing + attention-weighted aggregation + max pooling). Makes the majority of policy parameters independent of n_zones.
train_curriculum.py — 5-phase curriculum trainer (normal → monsoon → drought → heatwave → humidity) with automatic weight transfer between phases.
crop_risk_scorer.py — Deterministic, economics-calibrated risk scoring used for both training rewards and evaluation.
climatology.py + indonesia_zones.py — Real per-zone climatology and 14 grounded Indonesian agricultural zones with rice crop calendars.
backtest_indonesia.py — Historical replay harness with precision/recall and lead-time metrics.
mnn_export.py + edge_wrapper.cpp — ONNX export and C++/Vulkan edge runtime with external GRU hidden-state management.
era5_data_pipeline.py — Multi-source data pipeline (ERA5 / Open-Meteo / synthetic) with graceful degradation.

Key Technical Achievement

The central engineering result is a permutation-invariant, size-independent policy architecture that solves weight transfer across curriculum phases with changing numbers of zones.

Previous designs tied the GRU input dimension (and all downstream weights) to a fixed n_zones, causing observation space mismatches at phase transitions. The new extractor uses:

Shared per-zone GRU processing
Learned attention + max pooling across the variable zone axis
Separate basin-scale context (ENSO, IOD, etc.)

This design transfers 61 of 63 parameter tensors cleanly when n_zones changes (only the final action head is reinitialized). Full 5-phase curriculum runs confirm that phases sharing the same n_zones inherit competence immediately, while phases after an entity-count change only need to relearn the action head.

Validation Results

Full 5-phase curriculum (690k steps) completed successfully on GPU:

Transition	n_zones change	Tensors transferred	Behavior
normal → monsoon	2 → 3	61/63	Action head reinitialized
monsoon → drought	3 → 3	63/63	Full transfer
drought → heatwave	3 → 4	61/63	Action head reinitialized
heatwave → humidity	4 → 4	63/63	Full transfer

All phases converged to their structural performance ceiling (n_zones + 1). Full-transfer phases started at ceiling performance immediately. Explained variance stayed in the healthy 0.6–0.8 range. The entire curriculum finished in ~50 minutes on one T4.

How to Use

Training

python train_curriculum.py --phase all --output-dir ./runs --device cuda --seed 42
Inference (after training)
Pythonfrom sb3_contrib import MaskablePPO
model = MaskablePPO.load("runs/humidity/models/final_humidity.zip")
Edge Export
Bashpython mnn_export.py --checkpoint runs/humidity/models/final_humidity.zip \
                     --output weather_rl_model.mnn --quantize int8 --n-zones 4
Dependencies

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

GPU recommended for training. The edge runtime (edge_wrapper.cpp) targets Vulkan with CPU fallback.

Downloads last month: -

Video Preview

Reinforcement Learning