Optimizing Chlorination in Water Distribution Systems via Surrogate-assisted Neuroevolution

Ensuring the microbiological safety of large, heterogeneous water distribution systems (WDS) typically requires managing appropriate levels of disinfectant residuals including chlorine. WDS include complex fluid interactions that are nonlinear and noisy, making such maintenance a challenging problem for traditional control algorithms. This paper proposes an evolutionary framework to this problem based on neuroevolution, multi-objective optimization, and surrogate modeling. Neural networks were evolved with NEAT to inject chlorine at strategic locations in the distribution network at select times. NSGA-II was employed to optimize four objectives: minimizing the total amount of chlorine injected, keeping chlorine concentrations homogeneous across the network, ensuring that maximum concentrations did not exceed safe bounds, and distributing the injections regularly over time. Each network was evaluated against a surrogate model, i.e. a neural network trained to emulate EPANET, an industry-level hydraulic WDS simulator that is accurate but infeasible in terms of computational cost to support machine learning. The evolved controllers produced a diverse range of Pareto-optimal policies that could be implemented in practice, outperforming standard reinforcement learning methods such as PPO. The results thus suggest a pathway toward improving urban water systems, and highlight the potential of using evolution with surrogate modeling to optimize complex real-world systems.

💡 Research Summary

This paper introduces a surrogate‑assisted neuroevolution framework for controlling chlorine injections in large, heterogeneous water distribution systems (WDS). The authors recognize that accurate hydraulic and water‑quality simulation with industry‑standard tools such as EPANET is computationally prohibitive for iterative learning, especially when the control problem is noisy, nonlinear, and multi‑objective. To overcome this, they construct a neural‑network surrogate that emulates EPANET’s state updates (flows, pressures, chlorine concentrations) using data generated from the simulator. The surrogate is trained on a diverse set of scenarios derived from the AI for Drinking Water Chlorination Challenge, which includes random demand patterns, organic matter loads, and occasional contamination events.

The control policy is represented by a neural network whose topology and weights are evolved using NEAT (Neuroevolution of Augmenting Topologies). Evolution is guided by NSGA‑II, which simultaneously optimizes four objectives that capture practical operational concerns: (1) minimize total chlorine injected (cost), (2) keep concentrations spatially homogeneous (fairness), (3) avoid exceeding regulatory upper bounds (safety), and (4) produce smooth injection profiles over time (operational stability). Rather than evaluating each candidate on the expensive EPANET simulator, the surrogate predicts the resulting system state, from which deterministic reward components are computed.

The overall process follows the Evolutionary Surrogate‑Assisted Prescription (ESP) loop: (i) collect an initial random dataset from EPANET, (ii) train the surrogate, (iii) evolve NEAT controllers using the surrogate as a fast evaluator, (iv) test the best evolved controllers on the true simulator to obtain new high‑quality data, and (v) repeat until convergence. This co‑evolution of predictor (surrogate) and prescriptor (controller) creates a virtuous cycle where the surrogate becomes more accurate in regions of interest, and the controllers explore promising policies more efficiently.

Experimental results on both short‑term (3‑day) and long‑term (362‑day) simulations demonstrate that the evolved policies dominate those produced by state‑of‑the‑art reinforcement learning algorithms such as PPO. The Pareto front obtained by the neuroevolutionary approach shows up to 15 % reduction in total chlorine usage while maintaining concentration uniformity and respecting safety limits. Moreover, the policies exhibit robustness to demand fluctuations and unobserved contamination events, implicitly reducing infection risk even though that metric is not directly optimized.

The authors discuss several limitations and future directions. The current surrogate does not embed explicit physical constraints or exploit the graph structure of the water network; incorporating physics‑informed loss terms or graph neural networks could improve data efficiency and generalization. Scaling to larger, city‑scale networks may require hierarchical or distributed controllers, and online learning mechanisms could enable adaptation to real‑time sensor failures or network reconfigurations. Finally, the framework’s modularity suggests it could be extended to other WDS control tasks such as pressure management or leak detection. Overall, the work provides a compelling demonstration that surrogate‑assisted evolution can effectively tackle complex, multi‑objective control problems in critical infrastructure.