Improving the Accuracy of Coarse-grained Partial Differential Equations with Grid-based Reinforcement Learning

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 involve heuristics and empirical closure terms to account for the missing information. We propose Closure-RL, 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 the framework through numerical solutions of the advection equation and the Burgers' equation. The results demonstrate improved accuracy for in- and out-of-distribution test cases as well as a significant speedup compared to fine grained simulations.