Koopman-based surrogate modeling for reinforcement-learning-control of Rayleigh-Benard convection

Abstract: Training reinforcement learning (RL) agents to control fluid dynamics systems is computationally expensive due to the high cost of direct numerical simulations (DNS) of the governing equations. Surrogate models offer a promising alternative by approximating the dynamics at a fraction of the computational cost, but their feasibility as training environments for RL is limited by distribution shifts, as policies induce state distributions not covered by the surrogate training data. In this work, we investigate the use of Linear Recurrent Autoencoder Networks (LRANs) for accelerating RL-based control of 2D Rayleigh-Bénard convection. We evaluate two training strategies: a surrogate trained on precomputed data generated with random actions, and a policy-aware surrogate trained iteratively using data collected from an evolving policy. Our results show that while surrogate-only training leads to reduced control performance, combining surrogates with DNS in a pretraining scheme recovers state-of-the-art performance while reducing training time by more than 40%. We demonstrate that policy-aware training mitigates the effects of distribution shift, enabling more accurate predictions in policy-relevant regions of the state space.

• The paper demonstrates that policy-aware LRAN surrogates mitigate distribution shift in RL-driven control of 2D Rayleigh-Bénard convection.
• The methodology integrates a Koopman-inspired LRAN with two surrogate training paradigms to balance computational efficiency and control accuracy.
• Results reveal that combining surrogate pretraining with DNS finetuning cuts training time by over 40% while achieving near state-of-the-art performance.

Koopman-Based Surrogate Modeling for Reinforcement Learning Control of Rayleigh-Bénard Convection

Introduction and Motivation

The application of deep reinforcement learning (RL) to control nonlinear partial differential equation (PDE) systems, especially fluid dynamics exemplified by Rayleigh-Bénard convection (RBC), is frequently hindered by the prohibitive computational cost of direct numerical simulation (DNS). Surrogate models offer an attractive alternative by approximating dynamics at a substantially reduced computational burden. However, a critical challenge arises due to distribution shift: as the RL policy learns, it generates state-action distributions that diverge from those represented in static surrogate training datasets, potentially degrading surrogate reliability and the performance of policies trained therein.

This work investigates whether and how Koopman-inspired Linear Recurrent Autoencoder Networks (LRANs) can mitigate these challenges, enabling efficient and accurate RL-based control of 2D RBC. Two surrogate training paradigms are systematically evaluated: (1) static random-action dataset training, and (2) a policy-aware scheme with data continuously generated from an evolving policy. The study rigorously benchmarks control efficacy, model fidelity, and training efficiency, providing concrete recommendations for integrating surrogate modeling with RL in complex PDE control scenarios.

Methods and Experimental Design

Rayleigh-Bénard Convection and Control Setting

The target dynamical system, 2D RBC, consists of a fluid layer heated from below and cooled from above, driving convection governed by the incompressible Navier-Stokes equations. Control is effected via 12 thermal actuators at the lower boundary, each with bounded action mapped to local temperature modulation, with the overarching control objective being minimization of the time-averaged Nusselt number ($Nu$), a canonical measure of convective heat transfer.

The policy is parameterized as an MLP ingesting coarsened state fields and optimized with PPO, using reward defined as a normalized negative $Nu$.

LRAN-Based Surrogate Architecture

The LRAN, leveraging the Koopman operator framework, encodes state fields through deep convolutional neural nets into a latent space, which is evolved linearly in time. To handle control, actuator actions are incorporated into the latent dynamics via a linear affine transformation at each step.

Figure 1: Extension of the LRAN architecture to incorporate control actions as additional inputs into latent dynamics.

A normalized sequence reconstruction loss with temporal discount controls optimization. Data augmentation via translation and reflection, exploiting horizontal periodicity, is used to expand the expensive DNS-generated dataset.

Surrogate Training Paradigms

• Random-Action Surrogate: Trained solely on precomputed datasets arising from random actuator signals, covering the typical uncontrolled/impulsively driven state distribution.
• Policy-Aware Surrogate: Trained in an iterative, interleaved fashion, where new data is continually sampled from the states and actions generated by a co-evolving RL policy, thus better tracking the evolving training distribution and reducing distribution shift.

  Figure 2: Training loop of the policy-aware surrogate training scheme with alternating surrogate and policy optimization.

Evaluation Protocols

Two principal experimental protocols are employed:

1. Exclusive Surrogate Training: The RL agent is trained to convergence using surrogate rollouts only, and evaluated with ground-truth DNS.
2. Pretraining Scheme: The agent is first trained in the surrogate environment, then further finetuned via a capped number of interactions with DNS for policy refinement.

Results

Control Performance and Surrogate Characteristics

Both surrogates yield agents substantially outperforming uncontrolled and random-action baselines. However, surrogate-only training does not reach the control quality of a DNS-trained agent, primarily due to imperfections in surrogate predictive accuracy for policy-induced distributions not present in static data.

Figure 3: Control performance (test $Nu$) of policies trained on surrogates as a function of interaction steps.

The policy-aware surrogate demonstrates slower initial learning—attributable to its less accurate modeling of the "four-cell" initial state distribution—but enables further improvement and surpasses the random-action surrogate when policy-induced "two-cell" states become prevalent, ultimately attaining an $Nu$ close to DNS-trained performance. The random-action surrogate, by contrast, plateaus early and cannot model these less frequently visited states.

Pretraining with either surrogate, followed by DNS finetuning, restores state-of-the-art control, with a reduction in training time exceeding 40% relative to DNS-only training. This acceleration is due to surrogate rollouts being 25x faster than DNS integration and strategic scheduling of DNS for late-stage finetuning.

Qualitative Surrogate Behavior and Distribution Shift

Both surrogates fail to perfectly emulate all relevant state transitions. The random-action surrogate is unable to represent states arising from the "cell-merging" control strategy since they are underrepresented in random-action data, leading to poor generalization for high-performing policies. The policy-aware surrogate captures these states but at the cost of mild overfitting, causing degraded predictions in the initial distribution (see Figure 4).

Figure 4: Qualitative comparison: surrogate rollouts show how random-action surrogate cannot model cell merging, whereas policy-aware surrogate accurately predicts two-cell states.

Discussion and Implications

This study elucidates critical factors for surrogate modeling in RL for high-dimensional PDE systems. Policy-aware surrogate training is essential to mitigate distribution shift, ensuring accurate forward prediction as the agent's behavior departs from initial data distributions. Nevertheless, naive policy-aware training risks overfitting to current policy trajectories, underscoring the need for balanced data curation or regularization.

From a practical standpoint, leveraging LRAN surrogates as part of a pretraining pipeline delivers substantial reductions in wallclock and compute cost without degrading control performance. This enables scalable RL control for computationally intensive systems, unlocking faster prototyping and broader applicability of data-driven control methodologies in scientific and engineering domains.

Theoretically, the findings highlight limits of static-data surrogate methods and motivate further study of adaptive, uncertainty-aware, or hybrid surrogate constructions to handle nonstationary distributions and policy evolution in RL settings.

Conclusion

Koopman-inspired LRAN surrogates enable efficient RL control of Rayleigh-Bénard convection, particularly when trained in a policy-aware manner to address distribution shift. Their integration as pretraining environments, followed by focused DNS finetuning, matches state-of-the-art control at substantially decreased computational cost. Future directions include generalization to higher Rayleigh numbers, alternative fluid systems, and refinement of surrogate training strategies to further mitigate overfitting and distribution mismatch for robust policy optimization.