reservoir-water-env

💧 Reservoir Water Management Environment

A real-world AI agent benchmark for water utility operations.

Water authority operators make daily decisions about how much water to transfer between a network of reservoirs — balancing overflow prevention, demand fulfilment, and pumping costs, all under uncertain seasonal rainfall and demand. This environment lets you train and benchmark any RL or LLM-based agent on that task.

---

🌍 Motivation

Every day, water utility operators sit at dashboards and manually decide how much water to transfer between reservoirs, tanks, and distribution zones. They must account for:

• Stochastic rainfall — highly seasonal, varying widely by location
• Stochastic demand — spikes in summer, drops in winter
• Asymmetric risk — overflow (flooding, water loss) is 10× more costly than a temporary shortage
• Transfer costs — pumping water long distances is expensive
• Drought events — sudden supply shocks that reward strategic pre-planning

This is exactly the kind of high-stakes, repetitive, partially-observable scheduling task where AI agents should outperform human intuition — but no public benchmark has existed for it. Until now.

---

📦 Installation

git clone https://huggingface.co/spaces/YOUR_USERNAME/reservoir-water-env
cd reservoir-water-env
pip install -r requirements.txt

Docker

docker build -t reservoir-env .
docker run -p 7860:7860 reservoir-env

Then open http://localhost:7860 in your browser.

---

🚀 Quick Start

from reservoir_env.environment import ReservoirEnv
from reservoir_env.tasks.task1_easy import get_task1_config, TASK_HORIZON
from reservoir_env.models import Action

# Load a task
config = get_task1_config()
env    = ReservoirEnv(config, horizon=TASK_HORIZON, task_id="task1_easy")

# OpenEnv interface
obs  = env.reset()
done = False

while not done:
    # Your agent decides here
    action = Action.zeros(config.n_reservoirs)   # no-op baseline

    obs, reward, done, info = env.step(action)
    print(f"Day {obs.day} | Reward: {reward.total:.3f} | {reward.step_info}")

print("Episode complete. Cumulative reward:", env.state()["cumulative_reward"])

---

🧠 Observation & Action Space

Observation

class Observation(BaseModel):
    day:                int           # Current episode day (1-indexed)
    month:              int           # Calendar month 1–12 (seasonal awareness)
    storage:            List[float]   # Current storage per reservoir (units)
    capacity:           List[float]   # Max capacity per reservoir (units)
    storage_pct:        List[float]   # Storage as fraction of capacity [0, 1]
    recent_rain:        List[float]   # Rain received at previous step
    forecast_demand:    List[float]   # Expected demand this step (dist. mean)
    overflow_last_step: List[float]   # Overflow at previous step
    shortage_last_step: List[float]   # Unmet demand at previous step
    episode_step:       int           # Steps elapsed
    drought_active:     bool          # True during drought events

Observations are also available as a natural-language prompt for LLM agents:

print(obs.to_prompt())
# === Water Utility Status — Day 12 (Month 1) ===
#
# Reservoir status:
#   Reservoir 1: 6.82 / 10.0 units  (68%)  🟢 OK
#   Reservoir 2: 3.21 / 10.0 units  (32%)  🟡 LOW
# ...

Action

class Action(BaseModel):
    transfers: List[List[float]]
    # n×n matrix.  transfers[i][j] = units moved FROM reservoir i TO j.
    # Diagonal must be 0.  Invalid values are clipped (not rejected).

Feasibility constraints (enforced by environment via clipping):

Reward

class Reward(BaseModel):
    total:         float   # Scalar reward = -(all costs). Higher = better.
    transfer_cost: float   # Cost from pumping: Σ M_ij · a_ij
    shortage_cost: float   # Penalty: Σ P_short_i · shortage_i  (mild)
    overflow_cost: float   # Penalty: Σ P_over_i  · overflow_i  (severe, 10×)
    step_info:     str     # Human-readable step summary

The dense reward breakdown tells agents exactly which cost is hurting them at each step — not just a single scalar.

---

📋 Tasks

Task 1 — Easy: Two-Reservoir Balancing

Objective: Prevent overflow and meet demand using basic transfer decisions. No seasonal pattern to learn — pure reactive control.

Grader:

score = 0.7 × (1 − overflow_rate) + 0.3 × (1 − shortage_rate)

---

Task 2 — Medium: Three-Reservoir Seasonal Management

Objective: Navigate a dry → wet → dry cycle. Agents must learn to conserve water during dry months and redistribute wet-season surplus to under-served reservoirs.

Grader:

score = 0.5 × (1 − overflow_day_fraction) + 0.5 × (1 − shortage_day_fraction)

---

Task 3 — Hard: Five-Reservoir Annual Operations with Drought

Objective: Manage a full year with complete seasonal variation. A sudden drought on days 200–230 cuts rainfall to 10% of normal and raises demand by 50%. Agents that don't build strategic reserves before day 200 will fail this task.

Grader:

score = 0.3 × (1 − overflow_rate)
      + 0.3 × (1 − shortage_rate)       ← non-drought days
      + 0.4 × drought_survival_score    ← shortage during drought

The 40% weight on drought survival rewards strategic foresight, not just reactive control.

---

📊 Baseline Scores

Evaluated over 5 random seeds (0–4). GPT-4o score requires API key.

Key observations:

• Task 1 is solvable by both agents at ~0.90 — simple reactive control works.
• Tasks 2 and 3 score near 0.50–0.31 even for Greedy — massive headroom for learned agents (RL, LLM) to improve.
• The drought in Task 3 particularly punishes reactive agents; a score > 0.70 requires genuine strategic pre-planning.

To reproduce:

python baseline/baseline_agent.py
# With GPT-4o:
OPENAI_API_KEY=sk-... python baseline/baseline_agent.py

---

⚙️ Full Configuration API

Every aspect of the environment is configurable via ReservoirConfig:

from reservoir_env.config import ReservoirConfig, build_monthly_params

# Define your own seasonal rainfall distributions
rain_params = build_monthly_params(
    base_params=[(1.0, 0.5), (0.8, 0.4), (1.2, 0.6)],  # (mu, sigma) per reservoir
    seasonal_multipliers={
        5: 2.5, 6: 3.0, 7: 2.8,   # monsoon months (3× baseline rain)
        11: 0.3, 12: 0.2, 1: 0.2, # dry season
    },
)

config = ReservoirConfig(
    n_reservoirs=3,
    capacities=[10.0, 15.0, 8.0],
    transfer_cost_matrix=[
        [0.0, 0.3, 0.8],
        [0.3, 0.0, 0.5],
        [0.8, 0.5, 0.0],
    ],
    penalty_shortage=[1.0, 1.0, 2.0],   # Reservoir 3 is critical supply
    penalty_overflow=[10.0, 10.0, 10.0],
    rain_monthly_params=rain_params,
    demand_monthly_params=...,           # same pattern
    initial_storage=[5.0, 7.5, 4.0],
    max_transfer_per_pair=5.0,
    drought_start_day=150,               # optional drought event
    drought_end_day=180,
    drought_rain_scale=0.1,
    drought_demand_scale=1.5,
    seed=42,
)

env = ReservoirEnv(config, horizon=180)

Monthly Distribution Format

# rain_monthly_params: Dict[month, List[(mu, sigma)]]
# One (mu, sigma) tuple per reservoir, for each of the 12 months.
# Sampled from Normal(mu, sigma) truncated at 0.

rain_monthly_params = {
    1:  [(0.3, 0.4), (0.5, 0.4), (0.2, 0.3)],   # January (dry)
    ...
    6:  [(2.5, 1.0), (2.8, 1.2), (2.0, 0.9)],   # June (monsoon peak)
    ...
    12: [(0.4, 0.4), (0.5, 0.4), (0.3, 0.3)],   # December (dry)
}

---

🤖 Writing Your Own Agent

Any callable with this signature works:

def my_agent(obs: Observation, config: ReservoirConfig) -> Action:
    # obs.storage          → current storage levels
    # obs.storage_pct      → fraction of capacity (0–1)
    # obs.forecast_demand  → expected demand today
    # obs.drought_active   → True during drought
    # config.capacities    → max capacity per reservoir
    # config.cost_matrix_array → transfer costs

    n = config.n_reservoirs
    transfers = [[0.0] * n for _ in range(n)]
    # ... your logic here ...
    return Action(transfers=transfers)

Then run it:

result = env.run_episode(my_agent)
score  = grade1(result)   # or grade2 / grade3
print(f"Score: {score:.4f}")

Grading a Custom Agent

from reservoir_env.environment import ReservoirEnv
from reservoir_env.tasks.task3_hard import get_task3_config, TASK_HORIZON
from graders.grader_task3 import grade

config = get_task3_config()
env    = ReservoirEnv(config, horizon=TASK_HORIZON, task_id="task3_hard")
result = env.run_episode(my_agent)
score  = grade(result)
print(f"Task 3 score: {score:.4f} / 1.0000")

---

🗂️ Project Structure

reservoir-water-env/
│
├── Dockerfile                        ← docker build + run
├── openenv.yaml                      ← OpenEnv spec metadata
├── requirements.txt
├── README.md
├── app.py                            ← Gradio HF Space UI
│
├── reservoir_env/
│   ├── config.py                     ← ReservoirConfig (fully configurable)
│   ├── models.py                     ← Pydantic: Observation, Action, Reward
│   ├── environment.py                ← reset / step / state
│   ├── stochastic.py                 ← seasonal truncated Normal sampler
│   └── tasks/
│       ├── task1_easy.py             ← 2 reservoirs, 30 days
│       ├── task2_medium.py           ← 3 reservoirs, 180 days, wet/dry
│       └── task3_hard.py             ← 5 reservoirs, 365 days, drought
│
├── graders/
│   ├── base_grader.py
│   ├── grader_task1.py               ← overflow + shortage rate
│   ├── grader_task2.py               ← day-fraction metric
│   └── grader_task3.py               ← weighted drought survival
│
├── baseline/
│   ├── agents.py                     ← Random, Greedy, GPT-4o agents
│   └── baseline_agent.py             ← runs all 3 tasks, prints score table
│
└── tests/
    └── test_validate.py              ← 13 OpenEnv validation tests

---

🧪 Running Tests

pytest tests/test_validate.py -v
# 13 passed in ~5s

---

📄 License

MIT — free to use, modify, and build upon.

---

🙏 Citation

If you use this environment in your research:

@misc{reservoir-openenv-2024,
  title  = {Reservoir Water Management Environment for AI Agent Benchmarking},
  year   = {2024},
  url    = {https://huggingface.co/spaces/YOUR_USERNAME/reservoir-water-env},
  note   = {OpenEnv-compliant real-world scheduling environment}
}