Backpropagation through Time and Space: Learning Numerical Methods with Multi-Agent Reinforcement Learning

TL;DR

This paper presents BPTTS, a novel method for training neural networks to learn numerical schemes for PDEs in a multi-agent RL setting, enabling efficient and generalizable solutions for hyperbolic conservation laws.

Contribution

Introduction of BPTTS, a method for training spatio-temporal neural networks in MARL to learn numerical methods for PDEs, addressing non-stationarity via gradient flow across space and time.

Findings

Learned numerical policies match state-of-the-art methods.

Policies generalize well to different simulation setups.

Applicable to hyperbolic conservation laws like Burgers' and Euler equations.

Abstract

We introduce Backpropagation Through Time and Space (BPTTS), a method for training a recurrent spatio-temporal neural network, that is used in a homogeneous multi-agent reinforcement learning (MARL) setting to learn numerical methods for hyperbolic conservation laws. We treat the numerical schemes underlying partial differential equations (PDEs) as a Partially Observable Markov Game (POMG) in Reinforcement Learning (RL). Similar to numerical solvers, our agent acts at each discrete location of a computational space for efficient and generalizable learning. To learn higher-order spatial methods by acting on local states, the agent must discern how its actions at a given spatiotemporal location affect the future evolution of the state. The manifestation of this non-stationarity is addressed by BPTTS, which allows for the flow of gradients across both space and time. The learned numerical policies are comparable to the SOTA numerics in two settings, the Burgers' Equation and the Euler Equations, and generalize well to other simulation set-ups.