Closure Discovery for Coarse-Grained Partial Differential Equations Using Grid-based Reinforcement Learning

TL;DR

The paper tackles the challenge of accurate yet efficient simulations of multi-scale PDEs by proposing Closure-RL, a grid-based reinforcement learning framework that learns a centralized FCN policy to discover closures for coarse-grid simulations. By treating each coarse-grid cell as a local-agent with a shared policy and local rewards, the method learns forcing terms that compensate discretization errors, bridging CGS and FGS dynamics without differentiable solvers. Empirical results on 2D advection and Burgers’ equations show substantial reductions in mean absolute error compared to CGS and competitiveness with, or superiority to, higher-order coarse schemes, including strong generalization to in- and out-of-distribution initial conditions and velocity fields. The approach is computationally favorable relative to full fine-grid solves and is adaptable to various discretizations, suggesting broad applicability for rapid, reliable PDE simulations with learnable closures.

Abstract

Reliable predictions of critical phenomena, such as weather, wildfires and epidemics often rely on models described by Partial Differential Equations (PDEs). However, simulations that capture the full range of spatio-temporal scales described by such PDEs are often prohibitively expensive. Consequently, coarse-grained simulations are usually deployed that adopt various heuristics and empirical closure terms to account for the missing information. We propose a novel and systematic approach for identifying closures in under-resolved PDEs using grid-based Reinforcement Learning. This formulation incorporates inductive bias and exploits locality by deploying a central policy represented efficiently by a Fully Convolutional Network (FCN). We demonstrate the capabilities and limitations of our framework through numerical solutions of the advection equation and the Burgers' equation. Our results show accurate predictions for in- and out-of-distribution test cases as well as a significant speedup compared to resolving all scales.

Figures (12)

• Figure 1: Illustration of the training environment with the agents embedded in the CGS. The reward measures how much the action taken improves the CGS.
• Figure 2: Results of Closure-RL applied to the advection equation. The relative MAEs are computed over 100 simulations with respect to (w.r.t.) the FGS. The shaded regions correspond to the respective standard deviations. The violin plot on the right shows the number of simulation steps until the relative MAE w.r.t. the FGS reaches the threshold of $1\%$.
• Figure 3: Example run comparing CGS, Closure-RL and FGS with the same ICs. The concentration is sampled from the test set and the velocity components are randomly sampled from $\mathcal{D}^{Vortex}_{Train}$. Note that the agents of the Closure-RLmethod have only been trained up to $n=100$. However, they qualitatively improve the CGS simulation past that point.
• Figure 4: Results of Closure-RL applied to the Burgers' equation. The relative MAEs are computed over 100 simulations with respect to the FGS. The shaded regions correspond to the respective standard deviations. The violin plot on the right shows the number of simulation steps until the relative MAE w.r.t. the FGS reaches the threshold of $10\%$.
• Figure 5: The IRCNN backbone takes the current global observation and maps it to a feature tensor. This feature tensor is passed into two different convolutional layers that predict the per-discretization point action-distribution-parameters and state-values. In the case of a PDE with three spatial dimensions, the architecture would need to be based on three-dimensional convolutional layers instead.