--- license: mit base_model: Qwen/Qwen3-4B-Instruct-2507 language: - en library_name: transformers pipeline_tag: text-generation tags: - hydrology - agent - tool-use - grpo - reinforcement-learning - qwen3 - ef5 - crest - function-calling datasets: - anonymousOwl/HydroAgent-dataset --- # HydroAgent — Qwen3-4B-Instruct fine-tuned for hydrologic model calibration **HydroAgent** is a tool-using language model that calibrates the [EF5/CREST](https://github.com/HyDROSLab/EF5) distributed hydrologic model. Given a USGS streamflow gage and a precipitation-driven simulation, the agent iteratively proposes physically plausible parameter sets, runs the simulator, inspects the resulting NSE / peak / volume metrics, and revises until the model fits the observations. This release is the **GRPO step-100 checkpoint** of the SFT + RL pipeline described in [chrimerss/HydroLLM](https://github.com/chrimerss/HydroLLM). - **Base model:** [`Qwen/Qwen3-4B-Instruct-2507`](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507) - **Training:** full fine-tuning, BF16, FSDP, no LoRA - **RL framework:** [verl 0.5](https://github.com/volcengine/verl) GRPO with [SGLang](https://github.com/sgl-project/sglang) rollouts - **Tool format:** Hermes-style `` JSON (Qwen3-Instruct native) - **Hardware:** 4× H100, ~30 min/step, K=6 rollouts × max 50 multi-turn calls ## How the agent works The model has access to three tools and runs a multi-turn calibration loop: | Tool | Purpose | |---|---| | `set_parameters` | Set 11 tunable CREST multipliers: `wm`, `b`, `im`, `ke`, `fc`, `under`, `leaki`, `alpha`, `beta`, `alpha0`, `iwu` | | `run_simulation` | Execute EF5 with the current parameters and produce a hydrograph | | `evaluate` | Score the latest run vs. observations: NSE, CC, KGE, peak ratio, lag | Each rollout typically follows: `set_parameters → run_simulation → evaluate → set_parameters → …` until NSE plateaus or the agent runs out of turns. Inputs to the agent are a short system prompt describing the calibration task and a per-gage user message with watershed metadata (basin area, lat/lon, time window). ## Training data Training calibrates the agent on **10 CONUS USGS gages** (basin areas 539 – 2401 km²), each driven by **MRMS 1 km hourly precipitation** and **hourly USGS streamflow observations** from 60-day windows selected to contain a clear flood event (rising + receding limbs, edge-buffered). | Gage ID | Basin (km²) | Lat | Lon | Window (UTC) | |---|---:|---:|---:|---| | 11383500 | 539 | 40.0140 | -121.9483 | 2018-05-19 → 2018-07-17 | | 11043000 | 575 | 33.4798 | -117.1439 | 2019-03-15 → 2019-05-13 | | 11152000 | 632 | 36.2805 | -121.3227 | 2018-05-29 → 2018-07-27 | | 02294781 | 1064 | 27.8245 | -81.8017 | 2018-04-29 → 2018-06-27 | | 02312000 | 1476 | 28.4800 | -82.1776 | 2018-11-15 → 2019-01-13 | | 07195430 | 1489 | 36.1086 | -94.5333 | 2018-01-04 → 2018-03-04 | | 11179000 | 1639 | 37.5871 | -121.9608 | 2018-06-03 → 2018-08-01 | | 14301000 | 1727 | 45.7040 | -123.7554 | 2018-09-11 → 2018-11-09 | | 14207500 | 1828 | 45.3507 | -122.6762 | 2018-04-09 → 2018-06-07 | | 11376000 | 2401 | 40.3871 | -122.2386 | 2018-09-21 → 2018-11-19 | **Held-out evaluation gages** (never seen during training): | Gage ID | Basin (km²) | Lat | Lon | Window (UTC) | |---|---:|---:|---:|---| | 02338660 | 329 | 33.2357 | -84.9876 | 2018-07-01 → 2018-08-31 | | 01403060 | 2033 | 40.5511 | -74.5483 | 2018-11-11 → 2019-01-09 | | 06279500 | 40792 | 44.7585 | -108.1816 | 2018-06-13 → 2018-08-11 | | 07144100 | 3209 | 37.8831 | -97.4245 | 2019-03-30 → 2019-05-28 | The full training dataset — CONUS terrain rasters, per-gage MRMS hourly precipitation clips, USGS hourly streamflow observations, daily PET, the EF5 control template, and the 73 GPT-4o calibration trajectories that seed the SFT phase — is published as [**anonymousOwl/HydroAgent-dataset**](https://huggingface.co/datasets/anonymousOwl/HydroAgent-dataset). See that repo's README for the per-folder layout and provenance. ## Reward Two reward layers shape the policy: **Per-turn (returned by tools):** | Tool call | Reward | |---|---| | `set_parameters` (valid) | `+0.02` | | `run_simulation` (valid) | `+0.05` | | `evaluate` (valid) | `ΔNSE` (this turn − previous best) | | Any tool (invalid) | `−0.5` | **Terminal (returned at end of trajectory):** | Component | Value | |---|---| | Best NSE (clipped) | `[−1, 1]` | | Target-met bonus | `+0.5` if best NSE > gage target | | Iteration bonus | `+0.02 × n_evaluates` | | Improvement bonus | `+0.10 × max(0, n_improvements − 1)` | | Empty-trajectory penalty | `−1.0` | ## GRPO settings | Setting | Value | |---|---| | Algorithm | GRPO (group-relative advantages) | | K (rollouts per prompt) | 6 | | Train batch size | 4 prompts (24 trajectories per step) | | Max assistant turns | 50 | | Learning rate | 1e-6 with 5% warmup | | Entropy coefficient | 0.01 | | KL loss coefficient | 0.05 (anchored to base policy) | | Sampling | `temperature=1.0`, `top_p=0.95` | | Steps in this checkpoint | **100** | ## Quick start ```python from transformers import AutoModelForCausalLM, AutoTokenizer repo = "anonymousOwl/HydroAgent" tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True) model = AutoModelForCausalLM.from_pretrained(repo, torch_dtype="bfloat16", device_map="auto") ``` The model emits Hermes-style tool calls, e.g.: ``` {"name": "set_parameters", "arguments": {"wm": 1.0, "b": 1.0, "im": 0.5, ...}} ``` Parse with `tokenizer.apply_chat_template(..., tools=HYDRO_TOOLS)` and dispatch each call to your EF5 sandbox. See [`modal_app/eval.py`](https://github.com/chrimerss/HydroLLM/blob/main/modal_app/eval.py) for a reference SGLang loop with retry-on-parse-failure logic. For full reproduction (image, EF5 binary, multi-turn rollout, reward computation), use the [HydroLLM repository](https://github.com/chrimerss/HydroLLM). ## Limitations - Trained on **10 small/medium CONUS basins** (≤ 2401 km²) over short flood windows. Generalization to large basins (> 3000 km²), arid catchments, or out-of-CONUS regions is unverified. - Calibrates **CREST parameter multipliers only** — does not modify routing, initial conditions, or sub-basin structure. - The agent depends on a working EF5 toolchain; the weights alone do not perform calibration without the simulation environment in the loop. - This is a research checkpoint, not a production tool. NSE on held-out gages varies substantially with basin and event. ## License MIT — same as the upstream [HydroLLM repository](https://github.com/chrimerss/HydroLLM) and the base [Qwen3-4B-Instruct-2507](https://huggingface.co/Qwen/Qwen3-4B-Instruct-2507). ## Citation ```bibtex @software{hydrollm2026, title = {HydroLLM: Reinforcement Learning Fine-Tuning of LLMs with Hydrologic Simulation Feedback}, year = {2026}, url = {https://github.com/chrimerss/HydroLLM} } ``` ## Acknowledgement Compute for this research was sponsored by [Modal](https://modal.com).