Closed-Loop CO2 Storage Control With History-Based Reinforcement Learning and Latent Model-Based Adaptation

Technical Summary: Closed-Loop CO2 Storage Control With History-Based Reinforcement Learning and Latent Model-Based Adaptation

Problem Formulation
The paper addresses the challenge of closed-loop management for geological CO2 storage (CCS), a process requiring adaptive control policies to handle uncertain reservoir behavior while relying on realistically available operational data. The authors formulate the control of CO2 injection and brine production as a finite-horizon, partially observable sequential decision problem (POMDP). The objective is to maximize storage-related performance (retained CO2, trapping quality, economic metrics) while adhering to pressure and containment constraints.

A central difficulty in CCS is the "partial observability" of the subsurface; operators typically access only sparse well-level measurements (rates, pressures) rather than full-field spatial states (pressure/saturation distributions). Furthermore, traditional open-loop optimization is limited by the high cost of repeated history matching and the uncertainty of geological realizations. The study seeks to develop controllers that are deployable in real-world settings (using only well data) yet robust enough to adapt to abnormal operating conditions such as injector failures, leakage, or changes in reservoir connectivity.

Methodology
The research employs a two-stage methodology involving model-free reinforcement learning (RL) and latent model-based adaptation.

1. Model-Free Controller Comparison:
   The authors train Soft Actor-Critic (SAC) agents under different information regimes to quantify the value of temporal context and privileged training signals:

   • Privileged-State Benchmark: Uses dense simulator fields and well data (an oracle upper bound).
   • Well-Only Baseline: Uses only current well observations (the deployable baseline).
   • History-Conditioned: Conditions the policy on a rolling history of well observations using a gated transformer-style encoder to capture temporal dependencies.
   • Masked-Critic Curriculum: An asymmetric setup where the actor sees only well data, but critics are trained with progressively masked spatial states to test dependence on privileged information.
   • Asymmetric Teacher-Student: Uses privileged spatial states to train "teacher" critics and distill value estimates into "student" critics and an actor that only sees public well history.

2. Latent Model-Based Adaptation:
   For abnormal scenarios, the authors propose a pipeline that reuses a nominal latent structure to reduce the need for expensive real-simulator interactions.

   • Latent Representation: Public well observations are mapped to a latent control state using a frozen history encoder.
   • World Models: Four latent world-model backbones are evaluated (Probabilistic GRU, Geometry-aware GRU, RSSM, and Koopman-inspired linear-residual) to predict latent transitions and rewards.
   • Adaptation Protocol: In abnormal scenarios (injector failure, leakage, connectivity shifts), a "residual adapter" corrects the nominal world model's predictions using a small budget of real abnormal transitions. The controller is then updated via "imagined rollouts" (Dreamer-style) within the latent space, minimizing the need for high-fidelity simulator calls.

Key Contributions

• Demonstration of History-Conditioning: The study shows that history-conditioned policies, using only deployable well-level information, can recover nearly all the performance of privileged-state benchmarks. This validates that temporal context in observation streams can compensate for the lack of full spatial state information.
• Asymmetric Training Efficacy: The teacher-student framework with privileged value supervision is shown to be an effective method for training deployable controllers that remain competitive with history-based models.
• Efficient Adaptation via Latent Models: The paper demonstrates that latent model-based retuning outperforms direct model-free retuning under the same limited real-simulator budget when facing abnormal conditions. This is particularly true when the mismatch involves shifts in both dynamics and reward (e.g., leakage).
• Latent Space Sensitivity: The work highlights that the choice of latent representation (e.g., pre-finetuning vs. post-finetuning) significantly impacts downstream model-based control performance, suggesting that latent design is integral to the adaptation problem.

Results

• Model-Free Performance: The well-only baseline was the weakest and most unstable. The history-conditioned and teacher-student models achieved high returns comparable to the privileged-state benchmark, significantly outperforming the masked-critic curriculum which showed instability.
• Nominal Retention (Scenario 0): Latent model-based controllers successfully retained the performance of the deployable model-free reference under nominal conditions, validating the viability of the latent pipeline.
• Abnormal Scenarios:
  • Scenario 1 (Injector Failure): Latent model-based retuning rapidly recovered performance, significantly outperforming direct model-free retuning.
  • Scenario 2 (Leakage): This was the most challenging case (dynamics + reward shift). The best model-based variants recovered to positive high-return regimes, while direct model-free retuning remained poor.
  • Scenario 3 (Connectivity Shift): Model-based adaptation again outperformed direct retuning, though the margin was smaller than in the leakage scenario.
• Budget Efficiency: The model-based approach achieved these gains using the same scenario-specific real-simulator interaction budget as the direct model-free approach, leveraging the pre-trained latent structure and imagined rollouts.

Significance and Claims
The authors claim that the proposed framework offers a "simulator-budget-aware" alternative to repeated online history matching and re-optimization for closed-loop CCS control. The work suggests that strong, deployable policies can be learned from public well-level information when temporal context is explicitly modeled. Furthermore, these policies can be effectively adapted to deviations in control interfaces or reservoir physics through latent model-based pipelines. The paper positions this approach as a practical solution for settings where online adaptation must remain computationally tractable, avoiding the prohibitive costs of maintaining large ensembles of geological models for repeated optimization. The authors modestly note that results are simulator-budget-controlled comparisons rather than exhaustive robustness studies, acknowledging the need for future work on multi-seed variability and richer monitoring streams.