Source-linked AI summary

Efficient Geothermal Well-Control Optimization via Diffusion-Surrogate Reinforcement Learning

Ruimin Dai, Guodong Chen, Randy Harsuko, Kunpeng Liu, Nori Nakata

arXiv:2608.28791v1cs.AI

TL;DR

Long-horizon EGS well-control optimization is costly because it combines high-dimensional decisions with repeated high-fidelity hydrothermal simulations. The paper constructs a diffusion-surrogate environment for PPO using reservoir temperature and pressure states, injection-rate actions, and surrogate economic rewards. On a fractured EGS benchmark, the surrogate-assisted policy achieves competitive control performance while reducing dependence on expensive simulations.

  • Problem

    Long-horizon EGS optimization involves high-dimensional control spaces and repeated expensive hydrothermal simulations, while fixed control sequences do not adapt explicitly to evolving reservoir states.

  • Method

    Conditional diffusion models predict reservoir-state transitions and a separate reward surrogate estimates economic returns for PPO policy training.

  • Results

    The PPO-Surrogate method achieves competitive well-control performance compared with direct COMSOL-trained PPO and existing optimization methods while reducing dependence on high-fidelity simulations.

  • Takeaways & Limitations

    Diffusion-based surrogate environments can support long-horizon reinforcement learning for geothermal well-control optimization with less reliance on expensive high-fidelity simulations.

Abstract

from arXiv · show

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.

I. INTRODUCTION

EGS well-control optimization must manage long-horizon economic trade-offs while adapting to evolving reservoir conditions. The proposed framework uses diffusion surrogates with PPO to reduce costly high-fidelity simulation during state-dependent policy learning.

  • Motivation: EGS extracts heat by circulating injected fluid through fractured hot-rock formations, with well operation governing pressure, flow, and heat transfer.Cold fluid enters through injection wells, absorbs heat, and returns through production wells for surface use.
  • Motivation: Long-horizon well-control decisions must balance cumulative thermal returns against injection, production, and operational costs.Net present value captures this trade-off across the full production horizon.
  • Challenges: Existing optimization faces high-dimensional control spaces and repeated hydrothermal simulations, making computation a major bottleneck.The difficulty increases when multiple wells and control stages create many nonlinear decision variables.
  • Proposed framework: The formulation learns a policy mapping evolving reservoir temperature and pressure fields to well-control actions instead of optimizing a fixed open-loop sequence.This explicitly incorporates changing reservoir states into later decisions.
  • Proposed framework: Conditional diffusion models approximate state transitions, while a separate reward surrogate estimates economic return for PPO training.The surrogate environment reduces the number of high-fidelity simulations required during policy learning.
  • Evaluation: On a fractured EGS well-control problem, the learned policy achieves competitive performance while requiring substantially fewer high-fidelity simulation evaluations.The evaluation examines surrogate accuracy, control performance, and computational cost.

II. PROBLEM STATEMENT

The problem statement models fractured EGS production through coupled fluid-flow and heat-transfer dynamics. These dynamics link injection operations to evolving reservoir pressure and temperature fields that can be used for control optimization.

  • Reservoir system: The benchmark EGS contains a low-permeability rock matrix intersected by highly permeable fractures that provide preferential circulation pathways.The surrounding hot rock matrix supplies thermal energy to the circulating fluid.
  • Framework overview: The proposed workflow combines high-fidelity data collection, surrogate modeling, reinforcement-learning training, and policy evaluation.Figure 2 presents these four stages as the framework overview.
  • Hydrothermal dynamics: Fluid transport is governed by mass conservation and Darcy flow, with well operation represented through mass sources or sinks.The formulation includes fluid density, porosity, Darcy velocity, matrix-fracture exchange, permeability, viscosity, pressure, gravity, and elevation.
  • Hydrothermal dynamics: Reservoir temperature evolves through an energy balance that accounts for advective and conductive heat transfer and matrix-fracture heat exchange.The temperature equation includes effective heat capacity, fluid heat capacity, thermal conductivity, and thermal sources or sinks.
  • Control dependence: Changing injection rates modifies pressure gradients and circulation pathways, thereby altering heat transport and the spatial temperature distribution.The coupled hydrothermal equations are solved with the high-fidelity COMSOL simulator to characterize reservoir states.

B. Reinforcement Learning Formulation

EGS well control is formulated as a finite-horizon Markov decision process whose state is the reservoir’s evolving thermal and hydraulic condition. A feedback policy selects injection rates to maximize discounted economic return over production.

  • State and action: At each control step, the state consists of reservoir temperature and pressure fields representing the coupled thermal and hydraulic conditions.These fields summarize the reservoir response from fluid-flow and heat-transfer processes.
  • State and action: A parameterized policy maps the current reservoir state to injection-rate actions for the injection wells.Applying the action changes pressure distribution and fluid circulation, producing the next reservoir state.
  • Objective: The reward uses net present value to balance thermal energy production against fluid injection and production costs.The objective is cumulative economic return over all control stages with discount factor γ.
  • Feedback control: The formulation learns a feedback policy that continuously responds to evolving reservoir temperature and pressure instead of following a predefined full-horizon sequence.This represents well control as a controlled dynamical process.

III. METHOD

The framework builds a surrogate environment that predicts reservoir temperature and pressure transitions from the current state and well-control action. Conditional diffusion models generate next-state fields, enabling multi-step rollouts without repeatedly invoking the high-fidelity simulator.

  • High-fidelity COMSOL simulations provide state transitions and economic returns used to construct the surrogate environment.
  • Diffusion-Based Surrogate Model: Conditional diffusion models predict the next reservoir temperature and pressure fields from the current state and well-control action.Two independent diffusion models are trained for temperature and pressure prediction.
  • Conditional DDPM: The DDPM forward process progressively adds Gaussian noise, while the learned reverse process reconstructs the target reservoir field under state-action conditioning.
  • 2D Conditional U-Net: A conditional 2D U-Net denoises spatial reservoir fields while preserving localized temperature and pressure structures through multi-scale features and skip connections.
  • Surrogate Rollouts: During inference, reverse diffusion generates the next state, which is returned to the RL agent for multi-step rollouts without repeatedly invoking COMSOL.

C. Reward Surrogate Model

The reward surrogate complements diffusion-based state prediction by estimating the stage-wise economic return required for reinforcement-learning training. It uses current state, action, and next state to provide rewards during surrogate-environment rollouts.

  • The reward surrogate estimates stage-wise NPV because diffusion-based state prediction does not directly provide the economic reward required by RL.
  • The reward model predicts the reward from the current state, applied action, and next reservoir state.
  • Training uses rewards computed from COMSOL simulations and minimizes mean squared error.
  • Together, the diffusion and reward surrogates provide the next state and reward required for PPO training.

D. PPO-Based Policy Optimization

PPO optimizes a state-dependent well-control policy through repeated rollouts in the learned surrogate environment. The policy receives predicted reservoir states and surrogate rewards, then is evaluated with the original COMSOL simulator.

  • The policy outputs continuous well-control actions from evolving reservoir states, while diffusion and reward surrogates provide predicted next states and rewards.
  • PPO Policy Update: PPO updates the control policy using the probability ratio between rollout and updated policies.
  • Clipped Objective: PPO restricts excessive policy updates with a clipped surrogate objective controlled by the clipping range ϵ.
  • Value Estimation: The value network estimates expected returns for reservoir states, providing the value-loss component of PPO training.
  • Training Procedure: Algorithm 1 trains temperature, pressure, and reward surrogates before generating surrogate trajectories and updating the policy and value network.
  • Evaluation: After training, the learned policy is evaluated using the original COMSOL simulator rather than only the surrogate environment.

A. Experimental Setup

The study uses a fractured EGS benchmark to compare surrogate-assisted PPO with direct simulator-based PPO and other well-control methods. Evaluation covers cumulative NPV, production responses, and computational performance.

  • The benchmark contains four injection wells and five production wells for evaluating fractured EGS well-control strategies.
  • The 12,000-day production period is divided into 20 control stages, with four injection rates as actions and 64 × 64 temperature and pressure fields as states.
  • 512 COMSOL trajectories provide 10,240 transitions, split into 8,700 training and 1,540 validation transitions for surrogate construction.
  • PPO-Surrogate is compared with DE, GPEME, SACOSO, SHPSO, MFSKT, and PPO-COMSOL, with final policies evaluated using the original COMSOL simulator.
  • PPO-COMSOL achieves the best cumulative NPV, while PPO-Surrogate ranks second and consistently outperforms evolutionary and surrogate-assisted baselines.
  • PPO-based strategies produce higher initial thermal energy rates but faster average-temperature declines than baseline methods, with similar responses for PPO-COMSOL and PPO-Surrogate.

C. Performance of Surrogate Model

The diffusion surrogate reproduces the main temperature and pressure structures generated by COMSOL across multiple control stages. PPO exploration also yields structured, state-dependent injection strategies rather than fixed or uniformly distributed actions.

  • Temperature and pressure fields generated recursively by the surrogate closely reproduce COMSOL’s main spatial patterns over the full production horizon.
  • Temperature errors concentrate near sharp thermal gradients, while pressure predictions preserve major distributions and localized well responses despite accumulated rollout errors.
  • High-return PPO trajectories concentrate within specific injection-rate regions, indicating structured control patterns across the feasible action range.
  • Optimized injection rates vary across wells and stages, allowing actions to adapt to evolving temperature and pressure fields rather than follow a fixed sequence.

E. Model Capacity and Data Size Analysis

Model capacity improves surrogate prediction more consistently than training-set size within the tested range, while combined temperature-pressure states improve control and reward prediction. Larger models, however, increase computational cost.

  • U-Net widths of 8, 16, and 32 and training sets of 32–480 trajectories are evaluated using mean autoregressive test RMSE for temperature and pressure.
  • Width-32 yields the lowest and most stable temperature RMSE, while widths 16 and 32 perform comparably for pressure prediction.
  • Increasing training trajectories does not consistently reduce error, suggesting model capacity has a stronger effect on accuracy within the tested range.
  • The combined temperature-pressure state achieves normalized return 2.17 for PPO-COMSOL, compared with 2.03 for temperature-only and 1.99 for pressure-only states.
  • Surrogate PPO requires substantially less computation than direct PPO while maintaining competitive well-control performance.
  • Increasing U-Net width from 8 to 32 raises model size from 0.52M to 7.17M parameters and 20-step DDPM cost from 1.23 to 19.50 TFLOPs.
  • The T+P reward model achieves MAE 0.0026 and RMSE 0.0049, versus 0.0029 and 0.0055 for temperature-only and 0.0031 and 0.0058 for pressure-only inputs.

V. RELATED WORK

Prior work uses expensive optimization or surrogate objective models for geothermal well control, while reinforcement learning learns state-dependent policies. This paper instead constructs a diffusion-based surrogate environment for PPO and reports competitive performance on a fractured geothermal benchmark.

  • Conventional and surrogate optimization: Conventional well-control optimization often requires many expensive simulator evaluations, motivating surrogate-based methods.Derivative-free optimization avoids simulator gradients but can demand large numbers of costly simulations; data-driven methods approximate objectives to improve search efficiency.
  • Limitations of prior methods: Existing surrogate-assisted methods primarily optimize predefined control sequences rather than learning state-dependent control policies.This distinction separates objective approximation for fixed schedules from reinforcement learning based on current reservoir states.
  • Reinforcement learning: Reinforcement learning formulates well control as sequential decision-making in which current actions affect future reservoir states and cumulative returns.Prior work applied RL to reservoir production and waterflooding control, including continuous policies based on reservoir pressure and saturation states.
  • Diffusion-surrogate framework: The proposed approach constructs a complete surrogate environment: conditional diffusion models predict temperature and pressure evolution, while a separate surrogate estimates economic rewards for PPO.This differs from methods that only approximate the optimization objective.
  • Evaluation: On a fractured geothermal benchmark, PPO-Surrogate achieves competitive well-control performance against existing optimization methods and remains close to PPO trained directly with COMSOL.The reported results support diffusion-based surrogate environments for long-horizon geothermal well-control optimization with reduced dependence on expensive high-fidelity simulations.
Loading 2608.28791v1…