Optimizing Chlorination in Water Distribution Systems via Surrogate-assisted Neuroevolution

This paper proposes a surrogate-assisted neuroevolution framework that combines NEAT and NSGA-II to optimize multi-objective chlorine injection strategies in complex water distribution systems, demonstrating superior performance over standard reinforcement learning methods while leveraging a neural network surrogate to bypass the computational costs of traditional hydraulic simulators.

The Gist

Imagine a city's water system as a massive, living circulatory system. Just like your body needs blood flowing smoothly to every organ, a city needs clean water flowing through miles of pipes to every home and business. But there's a catch: the water needs a "medicine" called chlorine to kill harmful bacteria.

The problem is tricky. If you don't add enough chlorine, people get sick. If you add too much, it becomes toxic and creates dangerous chemical byproducts. Plus, the water moves differently depending on how many people are showering, washing dishes, or watering lawns at any given moment. It's a chaotic, moving target.

For a long time, trying to figure out the perfect amount of chlorine to pump at the right time has been like trying to steer a ship in a storm while blindfolded. Traditional computer programs are too slow or too rigid to handle the chaos.

This paper introduces a clever new solution: Teaching a computer to "dream" the solution.

Here is how they did it, broken down into simple steps:

1. The "Video Game" Simulator (The Surrogate)

Real water pipes are expensive and dangerous to experiment with. You can't just turn the knobs randomly in a real city to see what happens. So, the researchers built a super-accurate digital twin (a simulator) of the water system.

However, running this simulator is like trying to run a high-end video game on a calculator—it takes way too long. To fix this, they trained a "surrogate" model. Think of this as a student who watches a teacher (the slow, perfect simulator) play the game thousands of times. The student learns to predict the outcome almost instantly. Now, instead of waiting hours for a result, the computer gets an answer in a split second.

2. The "Evolutionary" Coach (Neuroevolution)

Instead of programming a robot with strict rules (like "if water is low, add chlorine"), they used evolution, the same process nature uses to evolve animals.

They created a population of digital "brains" (neural networks).

• Generation 1: These brains were terrible. They injected chlorine randomly, like a toddler throwing darts.
• The Test: The "surrogate student" quickly predicted what would happen if these brains controlled the pipes.
• Survival of the Fittest: The brains that did the worst (causing sickness or wasting money) were deleted. The brains that did the best (keeping water safe and cheap) were allowed to "reproduce."
• Mutation: When they reproduced, they made tiny, random mistakes (mutations). Sometimes these mistakes were lucky and made the brain smarter.

Over hundreds of generations, the "brains" evolved into expert managers who knew exactly when and how much chlorine to add.

3. The "School Curriculum" (Curriculum Learning)

Here is the secret sauce. If you try to teach a child calculus before they know how to add 1+1, they will fail. The researchers realized the same thing happened to their AI.

So, they used a curriculum:

1. First Grade: They taught the AI only one thing: Don't let the chlorine levels get too high or too low.
2. Second Grade: Once they mastered that, they added a second rule: Make sure the water tastes the same everywhere (fairness).
3. Third Grade: Then they added: Don't waste money (cost).
4. Fourth Grade: Finally, they added: Don't turn the valves on and off too jerkily (smoothness).

By teaching them step-by-step, the AI learned to balance all these competing goals much better than if they had been thrown into the deep end immediately.

4. The "Dream Team" (The Result)

The result wasn't just one "perfect" solution. Instead, the AI produced a menu of options (called a Pareto front).

• Option A: Super cheap, but slightly less safe.
• Option B: Extremely safe, but costs a bit more.
• Option C: The perfect middle ground.

City planners can now look at this menu and choose the policy that fits their budget and safety needs.

Why This Matters

This method beat the standard "Reinforcement Learning" methods (like the AI that plays video games) and random guessing. It found creative solutions that humans might never think of.

The Big Picture:
This paper shows that by combining evolution (nature's trial-and-error), surrogate models (fast AI approximations), and step-by-step learning, we can solve incredibly complex real-world problems. It's not just about water; this same "recipe" could help optimize traffic lights, power grids, or even how we build hospitals.

In short: They taught a computer to evolve its own brain, step-by-step, to keep our tap water safe, clean, and affordable.

Technical Summary

1. Problem Statement

Water Distribution Systems (WDS) are critical infrastructure requiring the maintenance of microbiological safety through disinfectant residuals, typically chlorine. However, managing chlorine levels is a complex control problem due to:

• Nonlinear Dynamics: Fluid interactions, reaction kinetics, and varying consumer demands create a noisy, non-stationary environment.
• Conflicting Objectives: Operators must minimize total chlorine usage (cost) while ensuring concentrations remain within safe bounds (avoiding toxicity) and are spatially/temporally homogeneous (fairness).
• Computational Infeasibility: High-fidelity simulators like EPANET are accurate but computationally expensive, making them unsuitable for direct training of reinforcement learning (RL) agents or evolutionary algorithms that require millions of evaluations.
• Contamination Risks: Systems must react to random contamination events (pathogens/organic matter) which are not directly observable but must be mitigated through proactive injection strategies.

2. Methodology

The authors propose an Evolutionary Surrogate-Assisted Prescription (ESP) framework combined with Neuroevolution and Curriculum Learning. The approach consists of three main components:

A. Surrogate Modeling (The "Meta-Simulator")

To bypass the computational cost of EPANET, the authors train a neural network surrogate to emulate the simulator's state transitions.

• Architecture: A unidirectional LSTM ($S_\theta$) predicts state increments ($\Delta_t$) rather than raw values to reduce long-horizon error.
• Knowledge Distillation: A "Teacher-Student" framework is employed. A bidirectional LSTM (Teacher, $T_\phi$) with lookahead context generates non-causal targets. The causal Student model is trained using a hybrid loss function combining:
  • Hard loss (Ground truth vs. prediction).
  • Soft loss (Teacher vs. prediction).
  • Rollout Loss: A critical term that minimizes error over multi-step horizons (5, 10, 20, 50 steps), ensuring the surrogate remains accurate over time.
• Iterative Refinement: The surrogate is not static. It is continuously fine-tuned using data generated by the best-performing evolved agents, creating a virtuous cycle where better agents generate better data, which in turn improves the surrogate.

B. Neuroevolution with NEAT

Instead of standard gradient-based RL, the control policies (prescriptors) are evolved using NEAT (Neuroevolution of Augmenting Topologies).

• Mechanism: NEAT evolves both the weights and the topology of neural networks simultaneously. It uses speciation to maintain diversity, preventing premature convergence.
• Input/Output: Agents observe chlorine concentrations at 17 monitoring nodes and flow rates at 2 pipes. They output injection decisions for 5 booster stations.
• Fitness Evaluation: Agents are evaluated against the surrogate model (not the real EPANET simulator) during evolution to save computational resources.

C. Multi-Objective Optimization & Curriculum Learning

The system optimizes four conflicting objectives using NSGA-II (Non-dominated Sorting Genetic Algorithm II):

1. Bound Violations: Keeping chlorine within safe limits ($0.2 - 0.4$ mg/L).
2. Fairness: Minimizing spatial variance in chlorine concentration.
3. Smoothness: Minimizing drastic changes in injection rates between timesteps.
4. Cost: Minimizing total chlorine injected.

Curriculum Learning Strategy:
To prevent the agents from getting stuck in local optima or finding unintuitive tradeoffs, objectives are introduced incrementally:

1. Start with Bound Violations and Fairness.
2. Add Smoothness.
3. Finally, add Cost.
   This step-by-step approach allows the population to find robust niches before tackling the full complexity of the problem.

3. Key Contributions

1. Surrogate-Assisted Framework for WDS: Demonstrated that a knowledge-distilled LSTM surrogate can effectively replace expensive hydraulic simulators for training control agents, significantly reducing computational cost.
2. Curricular Multi-Objective Neuroevolution: Showed that introducing objectives sequentially (Curriculum Learning) yields superior Pareto-optimal solutions compared to optimizing all objectives simultaneously or using standard RL (PPO).
3. Automatic Regularization via Co-Optimization: Discovered that the co-optimization of the predictor (surrogate) and prescriptor (agent) leads to "automatic regularization." As agents improve, the surrogate learns to produce smoother state predictions, which in turn encourages the agents to generate smoother control policies.
4. Innovation through Speciation: The iterative fine-tuning of the surrogate allows the evolutionary process to discover new "species" of agents that explore previously unreachable regions of the objective landscape.

4. Results

The framework was tested on a WDS topology from the AI for Drinking Water Chlorination Challenge (IJCAI 2025).

• Surrogate Accuracy: The "HSR" loss configuration (Hard + Soft + Rollout) achieved the lowest Horizon MSE, significantly outperforming models without rollout loss. The surrogate successfully captured complex, noisy dynamics.
• Performance Comparison (vs. Baselines & PPO):
  • PPO (Proximal Policy Optimization): Failed to learn effective policies, collapsing to near-zero injections due to the noisy composite reward signal.
  • Constant/Random Injection: Performed poorly on fairness, cost, or infection risk.
  • Curricular NSGA-II: Dominated all baselines. It achieved a ~40% reduction in bound violations compared to high-dose constant injection while maintaining similar infection risk.
  • Trade-offs: The curricular approach improved infection risk by nearly 30% and reduced fairness/smoothness metrics by 4-5x compared to non-curricular NSGA-II.
• Pareto Fronts: The method generated a diverse set of Pareto-optimal policies, allowing decision-makers to select strategies based on specific priorities (e.g., cost-saving vs. maximum safety).

5. Significance and Future Work

• Practical Impact: This approach offers a viable pathway to deploy AI-driven control systems for real-world water infrastructure, ensuring microbiological safety while minimizing chemical costs and byproduct formation.
• Generalizability: The ESP framework is applicable to other spatiotemporal optimization problems in transportation, communication, and construction where high-fidelity simulation is too costly for direct learning.
• Future Directions: The authors suggest extending simulations to longer timeframes (months/years) to better generalize against rare contamination events, expanding to diverse network topologies, and automating the curriculum ordering process.

In conclusion, the paper successfully demonstrates that surrogate-assisted neuroevolution with curriculum learning outperforms standard reinforcement learning in complex, multi-objective control tasks, providing a robust, diverse, and computationally efficient solution for optimizing water distribution systems.