Learning to Discretize: Solving 1D Scalar Conservation Laws via Deep Reinforcement Learning

Conservation laws are considered to be fundamental laws of nature . It has broad application in many fields including physics , chemistry , biology , geology , and engineering . Solving the differential equations associated with conservation laws is a major branch in computational mathematics . Recent success of machine learning , especially deep learning , in areas such as computer vision and natural language processing , has attracted a lot of attention from the community of computational mathematics and inspired many intriguing works in combining machine learning with traditional methods . In this paper , we are the first to view numerical PDE solvers as a MDP and to use ( deep ) RL to learn new solvers . As a proof of concept , we focus on 1-dimensional scalar conservation laws . We deploy the machinery of deep reinforcement learning to train a policy network that can decide on how the numerical solutions should be approximated in a sequential and spatial-temporal adaptive manner . We will show that the problem of solving conservation laws can be naturally viewed as a sequential decision making process and the numerical schemes learned in such a way can easily enforce long-term accuracy . Furthermore , the learned policy network is carefully designed to determine a good local discrete approximation based on the current state of the solution , which essentially makes the proposed method a meta-learning approach . In other words , the proposed method is capable of learning how to discretize for a given situation mimicking human experts . Finally , we will provide details on how the policy network is trained , how well it performs compared with some state-of-the-art numerical solvers such as WENO schemes , and how well it generalizes . Our code is released anomynously at https : //github.com/qwerlanksdf/L2D . 1 INTRODUCTION . Conservation laws are considered to be one of the fundamental laws of nature , and has broad applications in multiple fields such as physics , chemistry , biology , geology , and engineering . For example , Burger ’ s equation , a very classic partial differential equation ( PDE ) in conservation laws , has important applications in fluid mechanics , nonlinear acoustics , gas dynamics , and traffic flow . Solving the differential equations associated with conservation laws has been a major branch of computational mathematics ( LeVeque , 1992 ; 2002 ) , and a lot of effective methods have been proposed , from classic methods such as the upwind scheme , the Lax-Friedrichs scheme , to the advanced ones such as the ENO/WENO schemes ( Liu et al. , 1994 ; Shu , 1998 ) , the flux-limiter methods ( Jerez Galiano & Uh Zapata , 2010 ) , and etc . In the past few decades , these traditional methods have been proven successful in solving conservation laws . Nonetheless , the design of some of the high-end methods heavily relies on expert knowledge and the coding of these methods can be a laborious process . To ease the usage and potentially improve these traditional algorithms , machine learning , especially deep learning , has been recently incorporated into this field . For example , the ENO scheme requires lots of ‘ if/else ’ logical judgments when used to solve complicated system of equations or high-dimensional equations . This very much resembles the old-fashioned expert systems . The recent trend in artificial intelligence ( AI ) is to replace the expert systems by the so-called ‘ connectionism ’ , e.g. , deep neural networks , which leads to the recent bloom of AI . Therefore , it is natural and potentially beneficial to introduce deep learning in traditional numerical solvers of conservation laws . 1.1 RELATED WORKS . In the last few years , neural networks ( NNs ) have been applied to solving ODEs/PDEs or the associated inverse problems . These works can be roughly classified into three categories according to the way that the NN is used . The first type of works propose to harness the representation power of NNs , and are irrelevant to the numerical discretization based methods . For example , Raissi et al . ( 2017a ; b ) ; Yohai Bar-Sinai ( 2018 ) treated the NNs as new ansatz to approximate solutions of PDEs . It was later generalized by Wei et al . ( 2019 ) to allow randomness in the solution which is trained using policy gradient . More recent works along this line include ( Magiera et al. , 2019 ; Michoski et al. , 2019 ; Both et al. , 2019 ) . Besides , several works have focused on using NNs to establish direct mappings between the parameters of the PDEs ( e.g . the coefficient field or the ground state energy ) and their associated solutions ( Khoo et al. , 2017 ; Khoo & Ying , 2018 ; Li et al. , 2019 ; Fan et al. , 2018b ) . Furthermore , Han et al . ( 2018 ) ; Beck et al . ( 2017 ) proposed a method to solve very high-dimensional PDEs by converting the PDE to a stochastic control problem and use NNs to approximate the gradient of the solution . The second type of works focus on the connection between deep neural networks ( DNNs ) and dynamic systems ( Weinan , 2017 ; Chang et al. , 2017 ; Lu et al. , 2018 ; Long et al. , 2018b ; Chen et al. , 2018 ) . These works observed that there are connections between DNNs and dynamic systems ( e.g . differential equations or unrolled optimization algorithms ) so that we can combine deep learning with traditional tools from applied and computational mathematics to handle challenging tasks in inverse problems ( Long et al. , 2018b ; a ; Qin et al. , 2018 ) .The main focus of these works , however , is to solve inverse problems , instead of learning numerical discretizations of differential equations . Nonetheless , these methods are closely related to numerical differential equations since learning a proper discretization is often an important auxiliary task for these methods to accurately recover the form of the differential equations . The third type of works , which target at using NNs to learn new numerical schemes , are closely related to our work . However , we note that these works mainly fall in the setting of supervised learning ( SL ) . For example , Discacciati et al . ( 2019 ) proposed to integrate NNs into high-order numerical solvers to predict artificial viscosity ; Ray & Hesthaven ( 2018 ) trained a multilayer perceptron to replace traditional indicators for identifying troubled-cells in high-resolution schemes for conservation laws . These works greatly advanced the development in machine learning based design of numerical schemes for conservation laws . Note that in Discacciati et al . ( 2019 ) , the authors only utilized the one-step error to train the artificial viscosity networks without taking into account the longterm accuracy of the learned numerical scheme . Ray & Hesthaven ( 2018 ) first constructed several functions with known regularities and then used them to train a neural network to predict the location of discontinuity , which was later used to choose a proper slope limiter . Therefore , the training of the NNs is separated from the numerical scheme . Then , a natural question is whether we can learn discretization of differential equations in an end-to-end fashion and the learned discrete scheme also takes long-term accuracy into account . This motivates us to employ reinforcement learning to learn good solvers for conservation laws . 1.2 OUR APPROACH . The main objective of this paper is to design new numerical schemes in an autonomous way . We propose to use reinforcement learning ( RL ) to aid the process of solving the conservation laws . To our best knowledge , we are the first to regard numerical PDE solvers as a MDP and to use ( deep ) RL to learn new solvers . We carefully design the proposed RL-based method so that the learned policy can generate high accuracy numerical schemes and can well generalize in varied situations . Details will be given in section 3 . Here , we first provide a brief discussion on the benefits of using RL to solve conservation laws ( the arguments apply to general evolution PDEs as well ) : • Most of the numerical solvers of conservation law can be interpreted naturally as a sequential decision making process ( e.g. , the approximated grid values at the current time instance definitely affects all the future approximations ) . Thus , it can be easily formulated as a Markov Decision Process ( MDP ) and solved by RL . • In almost all the RL algorithms , the policy π ( which is the AI agent who decides on how the solution should be approximated locally ) is optimized with regards to the values Qπ ( s0 , a0 ) = r ( s0 , a0 ) + ∑∞ t=1 γ tr ( st , at ) , which by definition considers the long-term accumulated reward ( or , error of the learned numerical scheme ) , thus could naturally guarantee the long-term accuracy of the learned schemes , instead of greedily deciding the local approximation which is the case for most numerical PDEs solvers . Furthermore , it can gracefully handle the cases when the action space is discrete , which is in fact one of the major strength of RL . • By optimizing towards long-term accuracy and effective exploration , we believe that RL has a good potential in improving traditional numerical schemes , especially in parts where no clear design principles exist . For example , although the WENO-5 scheme achieves optimal order of accuracy at smooth regions of the solution ( Shu , 1998 ) , the best way of choosing templates near singularities remains unknown . Our belief that RL could shed lights on such parts is later verified in the experiments : the trained RL policy demonstrated new behaviours and is able to select better templates than WENO and hence approximate the solution better than WENO near singularities . • Non-smooth norms such as the infinity norm of the error is often used to evaluate the performance of the learned numerical schemes . As the norm of the error serves as the loss function for the learning algorithms , computing the gradient of the infinity norm can be problematic for supervised learning , while RL does not have such problem since it does not explicitly take gradients of the loss function ( i.e . the reward function for RL ) . • Learning the policy π within the RL framework makes the algorithm meta-learning-like ( Schmidhuber , 1987 ; Bengio et al. , 1992 ; Andrychowicz et al. , 2016 ; Li & Malik , 2016 ; Finn et al. , 2017 ) . The learned policy π can decide on which local numerical approximation to use by judging from the current state of the solution ( e.g . local smoothness , oscillatory patterns , dissipation , etc ) . This is vastly different from regular ( non-meta- ) learning where the algorithms directly make inference on the numerical schemes without the aid of an additional network such as π . As subtle the difference as it may seem , meta-learning-like methods have been proven effective in various applications such as in image restoration ( Jin et al. , 2017 ; Fan et al. , 2018a ; Zhang et al. , 2019 ) . See ( Vanschoren , 2018 ) for a comprehensive survey on meta-learning . • Another purpose of this paper is to raise an awareness of the connection between MDP and numerical PDE solvers , and the general idea of how to use RL to improve PDE solvers or even finding brand new ones . Furthermore , in computational mathematics , a lot of numerical algorithms are sequential , and the computation at each step is expert-designed and usually greedy , e.g. , the conjugate gradient method , the fast sweeping method ( Zhao , 2005 ) , matching pursuit ( Mallat & Zhang , 1993 ) , etc . We hope our work could motivate more researches in combining RL and computational mathematics , and stimulate more exploration on using RL as a tool to tackle the bottleneck problems in computational mathematics . Our paper is organized as follows . In section 2 we briefly review 1-dimensional conservation laws and the WENO schemes . In section 3 , we discuss how to formulate the process of numerically solving conservation laws into a Markov Decision Process . Then , we present details on how to train a policy network to mimic human expert in choosing discrete schemes in a spatial-temporary adaptive manner by learning upon WENO . In section 4 , we conduct numerical experiments on 1-D conservation laws to demonstrate the performance of our trained policy network . Our experimental results show that the trained policy network indeed learned to adaptively choose good discrete schemes that offer better results than the state-of-the-art WENO scheme which is 5th order accurate in space and 4th order accurate in time . This serves as an evidence that the proposed RL framework has the potential to design high-performance numerical schemes for conservation laws in a data-driven fashion . Furthermore , the learned policy network generalizes well to other situations such as different initial conditions , mesh sizes , temporal discrete schemes , etc . The paper ends with a conclusion in section 5 , where possible future research directions are also discussed .

Summary: This paper proposes to use reinforcement learning for constructing discretziation stencils of numerical schemes. More specifically, the method focuses on the widely used WENO schemes, which are an established class of finite difference schemes. Within this context, the method aims for training models to infer the weighting for a specific stencil with eight flux terms.