Code Structure Overview

MORL4Water models a multi-objective water resource management environment by simulating interconnected components like reservoirs, irrigation districts, power plants, inflows, and catchments. It is organized around facilities that exchange water via flows, all within an environment managed by a custom (MO)Gym-compatible class.

Top-level, general structure of the codebase:

morl4water
├── core
│   ├── envs
│   │   └── water_management_system.py
│   ├── models
│   │   ├── catchment.py
│   │   ├── facility.py
│   │   ├── flow.py
│   │   ├── irrigation_district.py
│   │   ├── objective.py
│   │   ├── power_plant.py
│   │   ├── reservoir.py
│   │   ├── reservoir_with_pump.py
│   │   └── weir.py

Core Environment: WaterManagementSystem

The WaterManagementSystem class, defined in core/envs/water_management_system.py, implements a custom environment that is fully compatible with mo-gymnasium.Env. It simulates a complex water infrastructure system over time.

Key Features

  • Inherits from gym.Env

  • Supports multi-objective settings

  • Fully vectorized observation and reward spaces

  • Modular design via a list of facilities and flow objects

Components of the Gym Environment API

__init__()

  • Initializes the environment with:

    • A list of components (facilities and flows)

    • A dictionary of reward signals

    • Start date and timestep duration

    • Optional timestamp and custom reward objectives

  • Determines:

    • observation_space and action_space

    • reward_space

    • max_capacities (used to normalize observations)

reset(seed, options)

  • Resets the environment state:

    • Sets date and timestep to initial values

    • Resets each facility’s internal state

    • Returns normalized observation and info dictionary

step(action)

  • Core simulation logic:

    • Propagates actions to ControlledFacility objects

    • Steps through all components (facilities, flows)

    • Aggregates rewards based on objectives

    • Advances the environment time

    • Returns:

      • observation: normalized state of ControlledFacility, which represents volume of water stored at the specific timestep in the Resevoir

      • reward: vector of objective values

      • terminated: True if the capacity of any Facility is exceeded

      • truncated: False unless simulation ends

      • info: dictionary with simulation metadata

Model Components

Facility and ControlledFacility

  • Base classes in core/models/facility.py

  • Facility is for passive components (e.g. irrigation districts)

  • ControlledFacility is for components with actions (reservoirs, reservoirs with a pump and weir)

  • Provide common methods:

    • reset(), step(), determine_reward(), determine_observation()

Flow Routing: flow.py

  • Flow, Inflow, and Outflow classes define directed edges between components

  • Handle delay, capacity, and direction of water flow

  • Connect upstream and downstream Facility instances

Objective Functions: objective.py

  • Provides a library of scalarizable reward functions:

    • deficit_minimised, scalar_identity, etc.

  • Used by each facility to compute its local reward contribution

Specialized Components

  • Reservoir, ReservoirWithPump (storage and release logic)

  • PowerPlant (hydropower generation)

  • IrrigationDistrict (demand satisfaction)

  • Catchment (rainfall accumulation)

  • Weir (overflow mechanisms)

Each inherits from either Facility or ControlledFacility and overrides reward and flow handling logic.

Interaction Flow

  1. User constructs an environment (e.g., via create_nile_river_env())

  2. Gym agent calls reset() and step(action)

  3. Each component updates state, computes reward

  4. WaterManagementSystem aggregates outputs and steps simulation

Extensibility

  • New components can be added by subclassing Facility

  • Objectives can be extended in objective.py

  • Works seamlessly with MO-RL algorithms via gymnasium

Data Flow in the Environment

The environment simulates time-dependent water movement across interconnected infrastructure components. Each timestep, the environment performs the following sequence:

  1. External Inputs
    Components such as Inflow and Catchment introduce water into the system using external time-series data (e.g., rainfall or river inflows).

  2. Processing Nodes

    • Reservoir, PowerPlant, and Weir components store, transform, or release water based on physical constraints and action inputs.

    • IrrigationDistrict components withdraw water based on predefined or data-driven demand.

  3. Routing

    • Flow, Inflow, and Outflow classes connect components, enforcing directionality, capacity limits, and transport delays.

  4. State Updates

    • Each facility and flow updates its internal state.

    • The environment collects and normalizes observations across all ControlledFacility components.

  5. Reward Calculation

    • Rewards are computed at the facility level using objective functions defined in objective.py.

    • The environment aggregates and normalizes these values into a vectorized reward output.

Code Structure