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

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 33 of its 51 references that are in the map.
- First posted on arXiv 1 Feb 2024.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 30 of 33)

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Fully Convolutional Networks for Semantic Segmentation](https://mlsanity.com/paper/1411.4038/) (2014)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Uformer: A General U-Shaped Transformer for Image Restoration](https://mlsanity.com/paper/2106.03106/) (2021)
- [Hamiltonian Neural Networks](https://mlsanity.com/paper/1906.01563/) (2019)
- [Learning the solution operator of parametric partial differential equations with physics-informed DeepOnets](https://mlsanity.com/paper/2103.10974/) (2021)
- [Learning Deep CNN Denoiser Prior for Image Restoration](https://mlsanity.com/paper/1704.03264/) (2017)
- [Lagrangian Neural Networks](https://mlsanity.com/paper/2003.04630/) (2020)
- [Inductive Biases for Deep Learning of Higher-Level Cognition](https://mlsanity.com/paper/2011.15091/) (2020)
- [PDE-Refiner: Achieving Accurate Long Rollouts with Neural PDE Solvers](https://mlsanity.com/paper/2308.05732/) (2023)
- [Towards Multi-spatiotemporal-scale Generalized PDE Modeling](https://mlsanity.com/paper/2209.15616/) (2022)
- [Respecting causality is all you need for training physics-informed neural networks](https://mlsanity.com/paper/2203.07404/) (2022)
- [Multi-Agent Reinforcement Learning is a Sequence Modeling Problem](https://mlsanity.com/paper/2205.14953/) (2022)
- [Tianshou: a Highly Modularized Deep Reinforcement Learning Library](https://mlsanity.com/paper/2107.14171/) (2021)
- [Jump-Start Reinforcement Learning](https://mlsanity.com/paper/2204.02372/) (2022)
- [Augmenting Physical Models with Deep Networks for Complex Dynamics Forecasting](https://mlsanity.com/paper/2010.04456/) (2020)
- [Clifford Neural Layers for PDE Modeling](https://mlsanity.com/paper/2209.04934/) (2022)
- [NOMAD: Nonlinear Manifold Decoders for Operator Learning](https://mlsanity.com/paper/2206.03551/) (2022)
- [Scientific multi-agent reinforcement learning for wall-models of turbulent flows](https://mlsanity.com/paper/2106.11144/) (2021)
- [Don't Start From Scratch: Leveraging Prior Data to Automate Robotic Reinforcement Learning](https://mlsanity.com/paper/2207.04703/) (2022)
- [Incorporating physical constraints in a deep probabilistic machine learning framework for coarse-graining dynamical systems](https://mlsanity.com/paper/1912.12976/) (2019)
- [Semi-supervised Invertible Neural Operators for Bayesian Inverse Problems](https://mlsanity.com/paper/2209.02772/) (2022)
- [Reward-Respecting Subtasks for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2202.03466/) (2022)

## Built on by (most cited first, 2 of 2)

- [Optimal Lattice Boltzmann Closures through Multi-Agent Reinforcement Learning](https://mlsanity.com/paper/2504.14422/) (2025)
- [Reinforcement Learning Closures for Underresolved Partial Differential Equations using Synthetic Data](https://mlsanity.com/paper/2505.11308/) (2025)

## Abstract (arXiv)

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.