Reinforcement Learning for Adaptive Mesh Refinement

1 INTRODUCTION . Numerical simulation of PDEs via the finite element method ( FEM ) ( Brenner & Scott , 2007 ) plays an integral role in computational science and engineering ( Reddy & Gartling , 2010 ; Monk et al. , 2003 ) . Given a fixed set of basis functions , the resolution of the finite element mesh determines the trade-off between solution accuracy and computational cost . For complex systems with large variations in local solution characteristics , uniform meshes can be computationally inefficient due to their suboptimal distribution of mesh density , under-resolving regions with complex features such as discontinuities or large gradients and over-resolving regions with smoothly varying solutions . For systems with multi-scale properties in particular , attempting to resolve these features with uniform meshes can be challenging even on the largest supercomputers . To achieve more efficient numerical simulations , adaptive mesh refinement ( AMR ) , a class of methods that dynamically adjust the mesh resolution during a simulation to maintain equidistribution of error , is used to significantly increase accuracy relative to computational cost . Existing methods for AMR share the same iterative process of computing a solution on the current mesh , estimating refinement indicators , marking element ( s ) to refine , and generating a new mesh by refining marked elements ( Bangerth & Rannacher , 2013 ; Červený et al. , 2019 ) . The optimal algorithms for error estimation and marking in many problems , especially evolutionary PDEs , are not known ( Bohn & Feischl , 2021 ) , and deriving them is difficult for complex refinement schemes such as hp-refinement ( Zienkiewicz et al. , 1989 ) . As such , the current state-of-the-art is guided largely by heuristic principles that are derived by intuition and expert knowledge ( Zienkiewicz & Zhu , 1992 ) , but choosing the best combination of heuristics is complex and not well understood . We advance the novel notion that adaptive mesh refinement is fundamentally a sequential decisionmaking problem in which a sequence of greedy decisions based on instantaneous error indicators does not constitute an optimal sequence of decisions for the actual goal of achieving high cumulative or terminal accuracy . In time-dependent problems for example , an error estimator by itself can not preemptively refine elements which would encounter complex features in the next time step . This means that the optimality of a refinement decision depends on the accuracy of the future solution and that selecting an element which yields the largest reduction in error at the current time step may not be the optimal decision over the entire simulation . Whether and how optimal AMR strategies can be found by directly optimizing a long-term performance objective are open questions . Given this perspective , we formulate AMR as a Markov decision process ( MDP ) ( Puterman , 2014 ) ( Figure 1 ) and propose a reinforcement learning ( RL ) ( Sutton & Barto , 2018 ) approach that explicitly trains a mesh refinement policy to optimize a performance metric , such as final solution error . In contrast to most , if not all , benchmark problems and complex applications of RL ( Mnih et al. , 2015 ; Brockman et al. , 2016 ; Osband et al. , 2019 ; Berner et al. , 2019 ; Vinyals et al. , 2019 ) , AMR poses a new challenge as the sizes of both the state and the set of available actions depend on the current number of mesh elements , which changes with each refinement action at every MDP time step . While one may define a fixed and bounded state and action space given a finite refinement budget , doing so is very inefficient as the policy ’ s input-output dimensions would have to accommodate the full exponentially large space but only subspaces ( with increasing size ) are encountered during simulation . In many practical applications , one would routinely encounter input dimensions on the order of millions or billions of degrees of freedom . This motivates the design of efficient policy architectures that leverage the correspondence between the current mesh state and valid action set . In this paper , we make the following conceptual , methodological , and experimental contributions : 1 ) We formally define an MDP with effective variable-size state and action spaces for AMR ( Section 3.2 ) ; 2 ) We propose three policy architectures—with differing generality , inductive bias , and capacity for modeling interaction—that operate on such variable-size spaces ( Section 4 ) ; 3 ) As a path toward potentially solving large and complex problems on which RL can not tractably be trained , we investigate the generalizability of policies trained on small representative features with known analytic solutions and the effectiveness of policies trained using a novel reward formulation that can be applied to problems without known analytic solutions ( Section 5 ) ; 4 ) Our experiments demonstrate for the first time that RL can be competitive with , and sometimes outperform , a greedy refinement strategy based on the widely-used Zienkiewicz-Zhu-type error estimator ; moreover , we show that an RL refinement policy can generalize to higher refinement budgets and larger meshes , transfer effectively from static to time-dependent problems , and can be effectively trained on more complex problems without readily-available ground truth solutions . ( Section 6 ) . 2 RELATED WORK . The formulation of problems in numerical analysis as statistical learning problems can be traced at least as far in time as to Poincaré ( Poincaré , 1912 ; Diaconis , 1988 ) . Contemporary works have employed neural networks as powerful function approximators in existing numerical PDE and linear system solvers to achieve faster convergence rates , generalize to different boundary conditions or larger problems , and approximate underresolved features in coarse-grained simulations ( Hsieh et al. , 2018 ; Luz et al. , 2020 ; Bar-Sinai et al. , 2019 ) . Our work focuses on optimizing a finite element space rather than components of a numerical solver . To the best of our knowledge , no prior work has formulated adaptive mesh refinement as a sequential decision-making problem and proposed a reinforcement learning approach ( Sutton & Barto , 2018 ) . Previous work at the intersection of neural networks and mesh-based simulation trained neural networks to predict mesh densities , sizes , or error fields for use by downstream mesh generators ( Dyck et al. , 1992 ; Chedid & Najjar , 1996 ; Zhang et al. , 2020 ; Pfaff et al. , 2020 ; Chen & Fidkowski , 2020 ) . Brevis et al . ( 2020 ) apply supervised learning to find an optimal parameterized test space without modifying the degrees of freedom . Bohn & Feischl ( 2021 ) show theoretically that the estimation and marking steps of AMR for an elliptic PDE can be represented optimally by a recurrent neural network , but model optimization was left as an open question . Recent studies have leveraged the effectiveness of graph neural networks ( GNN ) ( Sperduti & Starita , 1997 ; Gori et al. , 2005 ; Scarselli et al. , 2008 ) at representing relational structure to predict PDE dynamics on general unstructured and non-uniform meshes ( Alet et al. , 2019 ; Belbute-Peres et al. , 2020 ; Pfaff et al. , 2020 ) . Previous work on graph generation and formation have employed GNNs as the policy model in an RL context with applications to biological and social network datasets ( You et al. , 2018 ; Trivedi et al. , 2020 ) . Learning a policy for unbounded variable-size state and action spaces is a rare—if not new— problem for RL , which has been typically applied to environments with fixed-size observation and small bounded action spaces in almost all benchmark problems ( Mnih et al. , 2015 ; Brockman et al. , 2016 ; Osband et al. , 2019 ) . While there are notable applications where the available action set varies with state ( Berner et al. , 2019 ; Vinyals et al. , 2019 ) , they do not face the challenge of potentially millions of possible actions that arises in large-scale AMR . The technique of growing action spaces ( Farquhar et al. , 2020 ) maintains a fixed action space size within each episode , whereas both state and action space sizes change at every time step within an episode in AMR . 3 BACKGROUND AND FORMULATION . 3.1 FINITE ELEMENT METHOD . Our mesh adaptation strategy is implemented in a FEM-based framework ( Brenner & Scott , 2007 ) . In FEM , the domain Ω ⊂ RD is modeled with a mesh that is a union of E nonoverlapping subsets ( elements ) such that Ω : = ⋃ Ωk where k ∈ N : k 6 E. The solution on these elements is represented using polynomials ( basis functions ) which are used to transform the governing equations into a system of algebraic equations via the weak formulation . AMR is a commonly used approach to improve the trade-off between the solution accuracy , which depends on the shape and sizes of elements , and the computational cost , which depends on the number of elements . The most ubiquitous method for AMR is h-refinement , whereby elements are split into smaller elements ( refinement ) or multiple elements coalesce to form a single element ( derefinement ) . In practical applications with unknown true solutions , the conventional AMR approach is to take greedy refinement decisions based on a posteriori error estimators , which rely on the numerical solution and its derived quantities on the current mesh , without regard to long-term optimality . 3.2 AMR AS A MARKOV DECISION PROCESS . We formulate AMR with spatial h-refinement1 as a Markov decision process M : = ( O , Nmax , A , R , P , γ ) with each component defined as follows . Each episode consists of T RL time steps : for time-dependent PDEs , T spans the entire simulation and there may be multiple underlying PDE evolution steps per RL step ; for static problems , T is an arbitrary number of steps at which RL can act . Consider a time step t when the current mesh has Nt ≤ Nmax ∈ N elements . Each element i is associated with an observation oit ∈ O and the global state is st : = [ o 1 t , . . . , o N t ] ∈ ONt . We define O : = Rd such that each element ’ s observation is a tensor of shape d : = l × w × c that includes the values and refinement depths of a local window centered on itself . For brevity , let St denote the current global state spaceONt . We denote an action by at ∈ At : = { 0 , 1 , . . . , Nt } ⊂ A : = { 0 , 1 , . . . , Nmax } , where 0 means “ do-nothing ” and i 6= 0 means refine element i . Given the current state and action , the MDP transition P consists of : 1 ) refining the selected element into multiple finer elements ( which increases Nt ) if a refinement budget B is not exceeded and the selected element is not at the maximum refinement depth dmax ; 2 ) stepping the finite element simulation forward in time ( for time-dependent PDEs only ) ; 3 ) computing a solution on the new finite element space . When a true solution is available at training time , the reward at step t is defined as the change in error from the previous step , normalized by the initial error to reduce variation across function classes : rt : = ( ‖et−1‖2 − ‖et‖2 ) /‖e0‖2 , ( 1 ) where error e is computed relative to the true solution . With abuse of notation , we shall use e to indicate the error norm . The ground truth is not needed to deploy a trained policy on test problems . When the true solution is not readily available , as is the case for most non-trivial PDEs , one may run a reference simulation on a highly-resolved mesh to compute equation 1 , but this approach can be prohibitively expensive for training on large-scale simulations . Instead , we propose the use of a surrogate reward rt : = ‖ut , refine−ut , no-refine‖2 , the normed difference between the estimated solution u with and without executing the chosen refinement action . This surrogate , which is an upper bound 1Polynomial p-refinement can be formulated in a similar way . r-refinement ( Huang & Russell , 2010 ; Dobrev et al. , 2019 ) can be formulated as an RL problem but is not treated in this work . on the true reward and effectively acts as an estimate of the error reduction , is only used at training time to minimize computational effort , whereas at test time , the effectiveness of trained policies is evaluated using the error computed with respect to a highly-resolved reference simulation . Our objective to find a stochastic policy π : St → ∆ ( At ) to maximize the objective J ( π ) : = Ea∼π ( ·|s ) , st+1∼P ( ·|a , st ) [ T∑ t=1 γtrt ] . ( 2 ) Aside from γ ∈ ( 0 , 1 ) , this objective is equivalent to maximizing total error reduction : e0 − efinal . Although the size of the state vector and set of valid actions changes with each time step due to the varying Nt , this MDP is well-defined since one can define the global state space as the set of all possible ON , N < Nmax , and likewise for the action space . Hence , the policy is navigating through subspaces of increasing size during an episode . Moreover , the exact 1:1 correspondence between the number of observation components and the number of valid actions calls for designing a dedicated policy architecture for AMR , which we present below in Section 4 . We work with the class of policy optimization methods as they naturally admit stochastic policies that could benefit AMR at test time : a stochastic refinement action could reveal the need for further refinement in a region that appears flat on a coarse mesh . We build on the policy gradient algorithm ( Sutton et al. , 2000 ; Schulman et al. , 2017 ) to train a policy πθ ( parameterized by θ ) using batches of trajectories { τb : = { ( st , at , rt ) k } Tt=1 } Kk=1 generated by the current policy .

Summary: This paper proposes an application of reinforcement learning for adaptive mesh refinement in large-scale finite element simulations of complex physical systems. The authors suggest to formulate the mesh refinement problem as a MDP and propose different policy architectures for scalable application of reinforcement learning. Experiment results demonstrate that the proposed RL approaches outperform existing baselines, and can generalize well to situations of different finement budgets and larger meshes.