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

arXiv: https://arxiv.org/abs/2603.28074 · HTML: https://mlsanity.com/paper/2603.28074/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in RL with diffusion and flow models, one of 15 directions in Deep reinforcement learning methods.
- Builds on 21 of its 30 references that are in the map.
- First posted on arXiv 30 Mar 2026.

## Placed in

- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Builds on (most cited first, 21 of 21)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Gaussian Error Linear Units (GELUs)](https://mlsanity.com/paper/1606.08415/) (2016)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Modern Koopman Theory for Dynamical Systems](https://mlsanity.com/paper/2102.12086/) (2021)
- [Linearly-Recurrent Autoencoder Networks for Learning Dynamics](https://mlsanity.com/paper/1712.01378/) (2017)
- [Control of chaotic systems by Deep Reinforcement Learning](https://mlsanity.com/paper/1906.07672/) (2019)
- [Effective control of two-dimensional Rayleigh--B'enard convection: invariant multi-agent reinforcement learning is all you need](https://mlsanity.com/paper/2304.02370/) (2023)
- [Distributed Control of Partial Differential Equations Using Convolutional Reinforcement Learning](https://mlsanity.com/paper/2301.10737/) (2023)
- [Koopman-Assisted Reinforcement Learning](https://mlsanity.com/paper/2403.02290/) (2024)
- [Direct data-driven forecast of local turbulent heat flux in Rayleigh-B'enard convection](https://mlsanity.com/paper/2202.13129/) (2022)
- [Efficient Dynamics Modeling in Interactive Environments with Koopman Theory](https://mlsanity.com/paper/2306.11941/) (2023)
- [Multi-agent reinforcement learning for the control of three-dimensional Rayleigh-B'enard convection](https://mlsanity.com/paper/2407.21565/) (2024)
- [Koopman-Based Surrogate Modelling of Turbulent Rayleigh-B'enard Convection](https://mlsanity.com/paper/2405.06425/) (2024)
- [Surrogate Modeling of 3D Rayleigh-Benard Convection with Equivariant Autoencoders](https://mlsanity.com/paper/2505.13569/) (2025)
- [Oceananigans.jl: A Julia library that achieves breakthrough resolution, memory and energy efficiency in global ocean simulations](https://mlsanity.com/paper/2309.06662/) (2023)
- [Control of Rayleigh-B'enard Convection: Effectiveness of Reinforcement Learning in the Turbulent Regime](https://mlsanity.com/paper/2504.12000/) (2025)
- [Physics-informed Neural-operator Predictive Control for Drag Reduction in Turbulent Flows](https://mlsanity.com/paper/2510.03360/) (2025)
- [Plug-and-Play Benchmarking of Reinforcement Learning Algorithms for Large-Scale Flow Control](https://mlsanity.com/paper/2601.15015/) (2026)

## Abstract (arXiv)

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'enard 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.