Backpropagation through Time and Space: Learning Numerical Methods with Multi-Agent Reinforcement Learning
This addresses the challenge of developing efficient and generalizable numerical solvers for partial differential equations, which is incremental as it applies existing MARL concepts to a new domain.
The paper tackles the problem of learning numerical methods for hyperbolic conservation laws by treating numerical schemes as a Partially Observable Markov Game in multi-agent reinforcement learning, resulting in learned policies that are comparable to state-of-the-art numerics in Burgers' Equation and Euler Equations and generalize well to other setups.
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.