Efficient Geothermal Well-Control Optimization via Diffusion-Surrogate Reinforcement Learning
Authors: Ruimin Dai, Guodong Chen, Randy Harsuko, Kunpeng Liu, Nori Nakata
Organizations: School of Computing Clemson University Clemson, SC, USA · Energy Geoscience Division Lawrence Berkeley National Laboratory University of California Berkeley Berkeley, CA, USA
Real-time decision-making for enhanced geothermal systems (EGS) is challenging because long-term production periods involve high-dimensional control spaces and a large number of time-consuming high-fidelity hydrothermal simulations. Reinforcement learning provides a natural framework for state-dependent sequential control, but direct policy training with numerical simulators is computationally expensive. To address this issue, we propose a diffusion-surrogate guided reinforcement learning framework for long-horizon EGS well-control optimization. The reservoir temperature and pressure fields are used as system states, while injection rates are selected as control actions. A learned surrogate environment is constructed using conditional diffusion models to predict the evolution of reservoir temperature and pressure fields and a separate reward model to estimate the corresponding economic return. The surrogate environment is then integrated with Proximal Policy Optimization (PPO) for efficient policy training. Experiments on a fractured EGS benchmark show that the diffusion surrogate can accurately reproduce reservoir-state evolution over multiple control stages. The resulting surrogate-assisted PPO policy achieves competitive well-control performance compared with direct simulator-based PPO and existing optimization methods, while substantially reducing the dependence on expensive high-fidelity simulations. These results demonstrate the potential of diffusion-based surrogate environments for efficient reinforcement learning in geothermal well-control optimization.
Figures & tables
Fig. 1 : Schematic illustration of an enhanced geothermal system (EGS). Cold fluid is injected into the geothermal reservoir through an injection well, heated as it circulates through the subsurface formation, and recovered through a production well for surface energy utilization.
Fig. 2 : Overview of the proposed diffusion-surrogate reinforcement learning framework for EGS well-control optimization, including data collection, surrogate environment construction, PPO policy training, and high-fidelity policy evaluation.
Fig. 3 : Conditional diffusion process for reservoir-state prediction
Parameter
Value
Initial temperature ( ∘ C)
200
Initial pressure (MPa)
30
Temperature of injected fluid ( ∘ C)
20
Bottom-hole pressure of producers (MPa)
30
Injection rate range of injectors ( m3/s )
[0,200]×10−3
Depth (m)
2,500
TABLE I : Parameter configurations of the fractured geothermal energy system.
Component
Parameter
Value
Diffusion
Spatial resolution
64×64
U-Net channels
16 , 32 – 64 – 128 – 256
Training transitions
8,700
Validation transitions
1,540
Diffusion steps
200
Optimizer
AdamW
TABLE II : Main hyperparameter settings of the diffusion surrogate and PPO.
Method
3000 d
6000 d
12000 d
PPO-COMSOL
1.8201
1.9906
2.1749
PPO-Surrogate
1.7295
1.8776
2.0100
MFSKT
0.7343
1.2947
1.9407
SHPSO
0.6901
1.2676
1.8969
SACOSO
0.6992
1.2170
1.8588
GPEME
0.5757
1.0444
1.6725
TABLE III : Cumulative NPV achieved by different well-control methods at different production stages.
Fig. 4 : Production performance under different well-control strategies: (a) average production temperature and (b) thermal energy production rate.
Fig. 5 : Multi-step temperature-field prediction by the diffusion surrogate. The rows show the COMSOL reference, surrogate prediction, and absolute prediction error at different control stages.
Fig. 6 : Multi-step pressure-field prediction by the diffusion surrogate. The rows show the COMSOL reference, surrogate prediction, and absolute prediction error at different control stages.
Fig. 7 : Injection-rate trajectories explored during PPO optimization for the four injection wells. The color indicates the normalized return of each trajectory.
Fig. 8 : Effects of training-set size and model capacity on autoregressive prediction accuracy for reservoir temperature and pressure.
Method
Return (Normalized)
PPO-temperature (T)
2.03
PPO-pressure (P)
1.99
PPO-T-P
2.17
PPO-surrogate-T
1.96
PPO-surrogate-P
1.91
PPO-surrogate-T-P
2.01
TABLE IV : Ablation study on different reservoir state representations.
Fig. 9 : Computational cost of PPO-COMSOL and PPO-Surrogate during policy optimization.
Fig. 10 : Computational scaling of the diffusion surrogate under different model capacity and training-set sizes. (a) Number of model parameters for different U-Net widths. (b) Measured training time and training throughput for 20k optimization steps. (c) Theoretical computational cost of DDPM rollout. (d) Peak training memory usage under different U-Net widths and training-set sizes.
Fig. 11 : Comparison between true and predicted scaled rewards using different reservoir inputs: (a) temperature (T), (b) pressure (P), and (c) combined temperature and pressure (T+P).
Fig. 12 : Reward prediction performance under different reservoir inputs. (a) Overall MAE and RMSE for models using temperature (T), pressure (P), and combined temperature-pressure (T+P) inputs. (b) Reward MAE at different control stages over the 20-stage production horizon.
Reinforcement learning control uses reward-driven learning to make decisions in dynamic systems, from robots and vehicles to energy and infrastructure.