Year: 2021

Abstract:
Large-scale finite element simulations of complex physical systems governed by partial differential equations (PDE) crucially depend on adaptive mesh refinement (AMR) to allocate computational budget to regions where higher resolution is required. Existing scalable AMR methods make heuristic refinement decisions based on instantaneous error estimation and thus do not aim for long-term optimality over an entire simulation. We propose a novel formulation of AMR as a Markov decision process and apply deep reinforcement learning (RL) to train refinement policies directly from simulation. AMR poses a new problem for RL as both the state dimension and available action set changes at every step, which we solve by proposing new policy architectures with differing generality and inductive bias. The model sizes of these policy architectures are independent of the mesh size and hence can be deployed on larger simulations than those used at train time. We demonstrate in comprehensive experiments on static function estimation and time-dependent equations that RL policies can be trained on problems without using ground truth solutions, are competitive with a widely-used error estimator, and generalize to larger, more complex, and unseen test problems.

heading: 1 INTRODUCTION
text: 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 cannot 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 cannot 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).

heading: 2 RELATED WORK
text: 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.

heading: 3 BACKGROUND AND FORMULATION

heading: 3.1 FINITE ELEMENT METHOD
text: 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.

heading: 3.2 AMR AS A MARKOV DECISION PROCESS
text: 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.

heading: 4 POLICY ARCHITECTURES FOR VARIABLE STATE-ACTION SPACES
text: We propose three policy architectures, each with different inductive biases, that address the challenge of variable size state vector s ∈ RNt×d and action set {0, 1, . . . , Nt}, both of whose sizes changes with number of elements Nt within an episode. These architectures are compatible with any stochastic policy gradient algorithm. We focus on the special case of 1:1 correspondence between the number of observations that compose each global state and the number of available actions at that state. Although not treated in this work, these policy architectures can be easily extended to the general case of 1:k correspondence—e.g., k = 2 to include derefinement actions.

heading: 4.1 INDEPENDENT POLICY NETWORK
text: The Independent Policy Network (IPN) handles the 1:1 correspondence by mapping each observation to a probability for the corresponding action. Let fθ : Rd 7→ R be a function parameterized by θ. Given a matrix of observations s := [o1, . . . , oN ] ∈ RN×d, we define the policy as
π(·|s) = softmax ( fθ(o 1), . . . , fθ(o N ) ) . (3)
For example, using a neural network with hidden layer W ∈ Rd×h with h nodes, output layer H ∈ Rh×1, and activation function σ, the discrete probability distribution over N actions conditioned on s is defined by softmax (σ(sW )H).
IPN applies to meshes of any size since the set of trainable parameters θ is independent of N , but it has two main limitations. Firstly, it makes a strong assumption of locality as the action probability at an element does not depend on the observations at other elements. This assumption also appears in existing AMR methods that estimate error independently at each element; in fact, the output probabilities of IPN may be viewed as normalized error estimates. Secondly, the permutation equivariance of this architecture—i.e., π(aµ(i)|(oµ(1), . . . , oµ(N))) = π(ai|s) for any permutation operator µ : [N ] 7→ [N ]—means that one cannot use the ordering of inputs to represent spatial relations among elements, which would be necessary for refining an element based on neighboring conditions. We mitigate this problem by defining each element’s observation as an image tensor that includes neighborhood information and using a convolutional network layer, but this may face difficulties on unstructured meshes with non-quadrilateral elements (Červený et al., 2019).

heading: 4.2 HYPERNETWORK POLICY
text: The hypernetwork policy captures higher-order interaction among inputs via the function form π(·|s) = softmax ( fgφ(s)(o 1), ..., fgφ(s)(o N ) ) . (4)
The main policy network weights θ are now the output of a hypernetwork (Ha et al., 2017) gφ : RN×d 7→ Rdim(θ), parameterized by φ, which produces mixing among the inputs s ∈ RN×d. Continuing with the example in IPN, where the policy network’s first layer is W ∈ Rd×h, a hypernetwork with two layers can be instantiated as [∑N i=1 (sU)i,: ] V = W where U ∈ Rd×h1 and V ∈ Rh1×(d×h) are the trainable parameters φ, and Mi,: denotes the i-th row of matrix M . The output W can then be used as part of θ in equation 3.
This increased generality comes with more difficulty in the choice of architecture (specific form of gφ), which affects the extent to which it captures interaction among inputs. It does not contain an inductive bias for the local nature of interactions seen in classical applications of AMR. In fact, the use of a summation from i = 1 to N in the example above means that complete global information affects each local refinement decision, which is an extremely strong inductive bias.

heading: 4.3 GRAPH NETWORK POLICY
text: We build on graph networks (Scarselli et al., 2008; Battaglia et al., 2018) to address both the issue of interaction terms and spatial relation among elements. Specifically, we construct a policy based on Interaction Networks (Battaglia et al., 2016), which is a special case without global attributes2. At each step, the mesh is represented as a graph G = (V,E). Each vertex vi in V = {vi}i=1:N corresponds to element i and is initialized to be the observation oi. E = {(ek, rk, sk)}k=1:Ne is a set of edges with attributes ek between sender vertex sk and receiver vertex rk. An edge exists between two vertices if and only if they are spatially adjacent. We define the initial edge attribute ek as a one-hot vector indicator of the difference in refinement depth between rk and sk.
Graph networks capture the relations between nodes and edges via the inductive bias of its internal update rules. A single forward pass through the graph policy involves one or many rounds of message passing (Algorithm 1 in the Appendix). Each round is defined by the following sequence of computations: 1) Each edge attribute ek is updated by learned function ϕe using local node information via êk ← ϕe(ek, vrk , vsk); 2) For each node i, we denote by Êi := {(êk, rk, sk)}rk=i the set of all edges with node i as the receiver, and all updated edge attributes are aggregated into a single feature ēi ← ρe→v(Êi) by aggregation function ρe→v (e.g., element-wise sum); 3) Then, each node attribute is updated by v̂i ← ϕv(ēi, vi) using learned function ϕv . Each additional round increases the size of the local neighborhood that determines node attributes. Finally, we map each node attribute to a scalar using learned function ψ, apply a global softmax over all nodes, and interpret the value at each node i as the probability of choosing element i for refinement.
These update rules allow the graphnet policy to address both limitations of the IPN and the hypernetwork policy. Cross terms arise in the forward pass due to mutual updates of edge and node attributes using local information. The order of cross terms increases with each message-passing round. Local spatial relations between mesh elements are included by construction in the initial edge attributes, so there is no need to include numerical spatial information in each element’s observation vector.

heading: 5 EXPERIMENTAL SETUP
text: Our experiments assess the ability of RL, using the proposed policy architectures, to find AMR strategies that generalize to test function classes that differ from the training class, generalize to variable mesh sizes and refinement budgets, and extend to more complex problems without readilyavailable ground truth solutions. We define the FEM environment in Section 5.1, the train-test procedure in Section 5.2, and the implementation of our method and baselines in Section 5.3.

heading: 5.1 AMR ENVIRONMENT
text: MFEM. We use MFEM (Anderson et al., 2021; MFEM), a modular open-source C++ library for FEM, to implement the MDP for AMR. We ran experiments on two classes of AMR problems: static and time-dependent. In the static case, the objective of mesh refinement is to minimize the L2 error norm of projecting a variety of test functions onto a two-dimensional finite element space. In the time-dependent case, the functions are projected onto the finite element space, and a PDE of the form ∂u∂t +∇·F(u) = 0 is solved on a periodic domain using the finite element framework.
2While not demonstrated in this work, it is conceivable to use function coefficients and initial/boundary conditions as global attributes to improve generalization in a graphnet policy for AMR.
In contrast to the static case, the numerical error accumulated at each time step propagates with the physical dynamics and determines future error. Two types of PDEs were used: the linear advection equation, where F(u) = cu with c = [1, 0], and the nonlinear Burgers equation, where F(u) = cu2 with c = [1, 0.3]. The advection equation is used as there is an analytic solution to provide a ground truth, whereas the Burgers equation is used as the resulting solutions are representative of more complex phys-
ical systems that include shock and rarefaction waves, albeit without a readily-available analytic solution. The solution is represented using continuous (or discontinuous) second-order Bernstein polynomials for the static (or time-dependent) case, and the initial mesh is partitioned into nx × ny quadrilateral elements.
True solutions. We aim to show that RL policies have the potential to be deployed on larger and more complex problems without a priori known solutions by either training these policies on small representative features with known solutions or utilizing the proposed surrogate reward at training time. As such, we defined a collection of parameterized function classes, each exhibiting features such as sharp discontinuities and smooth variations, from which we randomly sample ground truth functions f : [0, 1]2 7→ R to initialize each episode. The collection, shown in Figure 2 and defined precisely in Appendix A.1.2, includes: bumps, circles, steps, and steps2 (a combination of two steps). These functions with closed form allow us to compute the error and reward at train time for static and advection problems. In the case of Burgers equation where the exact solution is not readily-available, we either use reference simulations on a highly-resolved mesh to act as a “ground truth” or employ the surrogate reward (defined in Section 3.2) to compute the reward for training.
In the static case, the true solution is fixed and each simulation time step is an RL step. For the time-dependent PDE cases, the initial solution is transported through the periodic domain and the ratio of simulation time steps to RL steps is set such that a feature advecting at unit velocity returns to its original position after 10 RL steps. We set refinement budget B = 10 for static problems, B = 20 for advection, and B = 50 for Burgers; episode length at train time equals B. Due to the Gibbs phenomena in FEM, using smooth polynomial approximations to solve hyperbolic systems containing discontinuities can introduce spurious oscillations which, in turn, can cause the simulation to become unstable. Therefore, we limit the true solutions to smooth functions (e.g., bumps, circles) for the advection and Burgers cases. Due to the nonlinearity of Burgers equation, initially smooth solutions can develop discontinuities in finite time. This behavior is resolved using the flux-corrected transport (FCT) approach (Boris & Book, 1997).

heading: 5.2 EXPERIMENTS AND PERFORMANCE METRIC
text: We conducted the following experiments to compare RL policies with baselines:
In-distribution: Train and test on true solutions sampled from the same function class. Out-of-distribution: For problems such as Burgers equation where training on multiple initial conditions (ICs) may be expensive, we show the effectiveness of policies trained on a single initial condition (IC) when tested on multiple random ICs, either with or without fine-tuning.
Generalization: 1) Static→advection: Policies trained on static functions are tested on advection. 2) Budget↑: Policies trained with a small refinement budget B (20 on static and 10 on advection) are test with B = 50, 100. 3) Size↑: Policies trained on an 8× 8 initial mesh are tested on 16× 16 and 64× 64 initial meshes with and without preserving the relative solution and mesh length scales. We define the performance of a given refinement policy in an episode in the static case as (einitial−efinal)/einitial, where einitial (or efinal) is the error norm at the beginning (or end) of an episode, to remove the variation in the error due to different true solution classes and random function initialization within each class. In the time-dependent case, without any refinement, the error may increase over the course of the simulation due to the accumulation of discretization error. Hence, given a refinement policy that achieves efinal at episode termination, we define its performance as (eno-refine, final − efinal)/einitial, where eno-refine, final is the final error without any refinement.
We tuned hyperparameters by training on a multitask scenario for static and advection, and individually for Burgers; the procedure and chosen values are given in Appendix A.3. For every experiment and every policy architecture, we trained four independent policies with different random seeds. For each test case, we report the mean and standard error—over the four independent policies and with different simulator seeds —of the mean performance metric over 100 test episodes. At each test episode, we ensured that all methods faced the same initial condition (which differs across episodes).

heading: 5.3 IMPLEMENTATION AND BASELINES
text: We describe the high-level implementation here and provide complete details in Appendix A.2. All policy architectures use a convolutional neural network with the same architecture as the input layer. The IPN has two fully-connected hidden layers with h1 and h2 nodes and ReLU activation, followed by a softmax output layer. Its action on input states is described in Section 4.1. The Graphnet policy is implemented with the Graph Nets library (Battaglia et al., 2018). Each input state consists of node observation tensors, all edge vectors, and the adjacency matrix. Node tensors are first passed through an Independent block, after which multiple Interaction networks (Battaglia et al., 2016) act on both node and edge embeddings to produce a probability at each node (see Section 4.3). The Hypernet policy is parameterized by matrices U ∈ Rd×h1 , V ∈ Rh1×(d×h), and Y ∈ Rd×h, where h1 and h are design choices. U and V act on input state s to produce the main policy weights W ∈ Rd×h (see Section 4.2) while Y acts on s to produce a bias b ∈ Rh, so that the main policy’s first hidden layer is ReLU(sW + b). Output probabilities are computed in the same way as IPN.
Baselines. The ZZ policy uses a Zienkiewicz-Zhu-type recovery-based error estimator (Zienkiewicz & Zhu, 1992) and refines the element with the largest estimated error. The TrueError policy refines the element where the error of the numerical solution with respect to the true solution is largest. It is not the theoretical upper bound on performance because refining the element with largest error does not necessarily result in largest reduction of error but it does effectively pose an upper bound on the performance of an instantaneous error estimator. The GreedyOptimal policy performs one-step lookahead by checking all possible outcomes of refining each element individually and chooses the element whose refinement would result in the lowest error at the next step. In comparison to RL policies, only the ZZ policy can feasibly be deployed at test time as TrueError cannot be deployed on systems without known solutions and GreedyOptimal is intractable for real applications.

heading: 6 RESULTS
text: We find that the proposed methods achieve performance that is competitive with baselines and, more importantly, generalize well to larger refinement budgets and mesh sizes, and transfer effectively from a static problem to a time-dependent problem. Videos of policies on advection and Burgers can be viewed at https://sites.google.com/view/iclr2022-amr.

heading: 6.1 IN-DISTRIBUTION
text: Static functions (Figure 3a). RL policies either meet or significantly exceed the performance of ZZ on all function classes. Notably, both IPN and Graphnet outperform ZZ significantly on steps by spending the limited refinement budget only on regions with discontinuities (Figure 2c). On the smoother function classes such as bumps where ZZ is known to perform well, all three policy architectures have comparable performance to ZZ. Overall, IPN outperforms both Graphnet and Hypernet, while Hypernet performed poorly (albeit still better than random) on all classes except
bumps. This suggests that capturing higher-order interaction among observations, each of which already contains local neighborhood information, is unnecessary for estimation of static functions as they only have a local domain of influence. Hypernetwork policies converged to the behavior of making no refinements on at least one out of four independent runs on all classes except bumps. This could be attributed to the inherent difficulty of choosing and training a highly nonlinear model.
Advection (Figure 3b). As explained above, we limit the true solutions to smooth functions (bumps and circles) in the advection case. Graphnet significantly outperformed ZZ on circles and is comparable to TrueError on bumps, while IPN is comparable to ZZ on both functions. Hypernet is comparable to ZZ on circles but has high variance across independent runs. Graphnet’s higher performance than other methods indicates that its inductive bias can better represent the local geometric relations between neighboring mesh elements along the circle.
Burgers equation. In experiments with a single bump function (visualized in Figure 4c) as the fixed initial condition (IC) in both train and test, IPN trained with both the exact and surrogate rewards outperformed all baselines (Figure 4a). The policy trained using the surrogate reward resulted in similar performance to the policy trained with the exact reward, indicating that the surrogate reward can be effectively used to train policies without the need for a ground truth solution, as is necessary for general random ICs.

heading: 6.2 OUT-OF-DISTRIBUTION (OOD)
text: Figure 4b shows the performance of IPN RL policies and baselines on Burgers equation with random ICs. We observe that policies trained on a fixed IC using either the exact or surrogate reward (labeled “Fixed IC, ER” and “Fixed IC, SR”) generalize well to random unseen ICs and still outperform baselines. Moreover, policies that were pretrained with the surrogate reward on the fixed IC for 2k episodes and fine-tuned with the surrogate reward on random ICs for another 2k episodes performed the best (“pretrained, SR”). We truncated training around 4-5k episodes due to computational limits, but we see that performance can improve with more training episodes. Policies trained only on random ICs with the surrogate reward were not as performant, indicating that the training time was not sufficient and that pretraining on a fixed IC is a more efficient approach. We do not envision that one would train and test on highly different function classes in applications involving static and advection problems, but we include experimental results of all combinations of train and test function classes in Appendix B.

heading: 6.3 GENERALIZATION
text: Static→advection (Figure 3c). All static-trained policies demonstrated comparable performance to ZZ and TrueError when tested on advection-bumps, while both IPN and Graphnet significantly outperformed ZZ on advection-circles. Surprisingly, static-trained IPN significantly outperforms advection-trained IPN when tested on advection-circles, and the static-trained Hypernet does so as well on advection-bumps, while static-trained Graphnet maintains comparable performance to its advection-trained counterpart (Figure 3b vs. Figure 3c). Figure 12 shows that a static-trained policy on bumps withB = 10 correctly refines the region of propagation on advection-bumps withB = 50.
Budget↑ (Figure 5). RL policies trained with low refinement budget generalize to test cases with higher budget. In the static case, comparing Figure 3a (B = 10) with Figure 5a (B = 50) shows that
the performance of RL policies relative to ZZ is generally preserved by the increase in refinement budget. Figures 5 and 8 show that an IPN trained withB = 10 makes qualitatively correct refinement decisions when allowed B = 100 during test. In the advection case (Figure 5b), Graphnet trained with B = 20 significantly outperforms both ZZ and TrueError when tested with B = 50 on bumps and comes within the margin of error of TrueError on circles. Figure 7 shows that an IPN trained with B = 20 correctly allocates a higher budget B = 100 to the limited region of propagation.
Size↑ (Figure 6). In the static case, the relative performance of RL policies that were trained with an 8 × 8 initial mesh (Figure 3a) is generally preserved when deployed on a 16 × 16 initial mesh (Figure 6a). All policy architectures outperform ZZ on bumps, while IPN and Graphnet still outperform ZZ on steps. IPN and Graphnet were comparable to ZZ on 8 × 8 but underperformed on 16 × 16 on circles. Nonetheless, Figure 10 shows that IPN makes qualitatively correct refinements. On advection, relative performance is preserved on circles while IPN and Graphnet deproved slightly on bumps (Figure 3b vs. Figure 6b). Without preserving the solution-to-mesh length scales, the prior tests emulate deploying a policy trained on coarser versions of the simulations. When tested on 64 × 64 and 200 × 200 initial meshes that preserves the relative solution-to-mesh length scales, which emulates deploying a policy trained on a small subset of a highly-resolved simulation, Figures 6c and 6d shows that IPN is competitive with baselines on bumps and even outperforms TrueError on circles.

heading: 7 CONCLUSION
text: We present a novel formulation of adaptive mesh refinement as a Markov decision process and propose new policy architectures for scalable application of reinforcement learning. Our experiments on static and time-dependent problems demonstrate that RL policies can outperform a policy based on the widely-used ZZ-type error estimator, and in some cases even outperform a policy based on the exact true error, suggesting that RL can potentially provide efficiency gains beyond the reach of existing AMR approaches. We demonstrated that these RL policies generalize to different refinement budgets and larger meshes, transfer from static to time-dependent settings, and generalize to more complex problems even when trained without ground truth rewards. Moreover, because these RL policies do not use problem-specific knowledge or domain expertise as input, our results provide a path for learning novel AMR strategies for cases such as hp-refinement that currently lack effective solutions. Future work can extend our RL methods to include derefinement actions, sampling multiple elements at each time step, and taking a multi-agent perspective that views each element as an agent who acts concurrently with all other agents.
REPRODUCIBILITY STATEMENT
We have described the experimental protocol involving tuning, training, and testing in Section 5.2; the high-level implementation of the proposed method in Section 5.3; the specific implementation details of the proposed method in Appendix A.2; the experimental test cases in Appendix A.1.2 and Table 1; the detailed tuning procedure in Appendix A.3; and the complete specification of all hyperparameters in Appendix A.3, Table 2, and Table 3.

heading: A EXPERIMENTAL SETUP
text: A.1 ENVIRONMENT DETAILS
A.1.1 MFEMCTRL
To interface between the MFEM framework and the RL environment, we developed MFEMCtrl, a C++/Python wrapper for the AMR and FEM capabilities in MFEM. MFEMCtrl is used to convert solutions to observations, apply refinement decisions, and calculate errors.
The initial mesh is partitioned into nx × ny = 8× 8 elements for static and advection experiments and 10×10 for Burgers equation. Generalization experiments on larger initial mesh used nx×ny = 16 × 16 or 64 × 64. The true solution is projected onto the finite element space by interpolation to the nodes of the Bernstein basis functions. After each refinement action, the solution is projected again onto the refined mesh (for the static case) or integrated in time until the next refinement action (for the time-dependent case). The maximum refinement depth is fixed by the parameter dmax such that the maximally-refined mesh consists of 2dmaxnx × 2dmaxny elements. dmax was set to 3 for static experiments, whereas dmax = 2 for advection and dmax = 1 for Burgers equation due to the time step restrictions imposed by the Courant-Friedrichs-Lewy (CFL) condition of the finest elements.
Observation. The observation consisted of the solution and the depth of each element. Since the gradients of the solution are, by definition, a function of the solution, the observation does not include the gradients as they can be implicitly learned. The solution/depth of each element was observed by interpolating the functions to a local equispaced mesh (image) centered around each element, shown by the white box in Figure 1. Each element’s observation is a l×w×c tensor where l = w = lelement + 2lcontext is the spatial observation window with lelement = 16 sampled points inside the element and lcontext = 4 sampled points in a coordinate direction outside the element. We chose c = 2 channels so that estimated function values and element depths are observed, while gradients are omitted since the policy network can in principle estimate gradients from the value channel. To impose a 1:1 map between each observation and possible action, we append a dummy o0 to state s corresponding to action 0. At most one refinement is allowed per MDP step.
Table 1: Parameterized true solutions
Parameter [min, max]
Bumps (static)
cx [0.2, 0.9] cy [0.2, 0.9] w [0.05, 0.2] n {1, . . . , 6}
Bumps (advection)
cx [0.3, 0.7] cy [0.3, 0.7] w [0.005, 0.05] n {1, . . . , 4} cx 0.5
Bumps (Burgers) cy 0.5 single IC w 0.05
n 1
cx [0.3, 0.7] Bumps (Burgers) cy [0.3, 0.7]
random IC w [0.005, 0.05] n {1, . . . , 4}
Circles (static)
cx [0.2, 0.8] cy [0.2, 0.8] r [0.05, 0.2] w [0.1, 1.0] n {1, . . . , 6}
Circles (advection)
cx [0.3, 0.7] cy [0.3, 0.7] r [0.05, 0.2] w [0.03, 0.05] n {1, . . . , 4}
Steps and Steps2 o [0, 1.0] θ [0, π/2] n {1, . . . , 6}
A.1.2 GROUND TRUTH FUNCTIONS
Bumps
n ∼ Uniform[nmin, nmax] cx,i ∼ Uniform[cx,min, cx,max], i = 1, . . . , n cy,i ∼ Uniform[cy,min, cy,max], i = 1, . . . , n wi ∼ Uniform[wmin, wmax], i = 1, . . . , n
f(x, y) = n∑ i=1 exp ( − (x− cx,i) 2 + (y − cy,i)2 wi )
Circles
n ∼ Uniform[nmin, nmax] cx,i, cy,i ∼ Uniform[cmin, cmax], i = 1, . . . , n
ri ∼ Uniform[rmin, rmax], i = 1, . . . , n wi ∼ Uniform[wmin, wmax], i = 1, . . . , n
f(x, y) = n∑ i=1 exp
( − ( √ (x− cx,i)2 + (y − cy,i)2 − ri)2
wi
)
Steps
n ∼ Uniform[nmin, nmax] θ ∼ Uniform[θmin, θmax] oi ∼ Uniform[omin, omax], i = 1, . . . , n
f(x, y) = n∑ i=1 1 + tanh [100(oi − (x+ y tan θ))]
Steps2
n ∼ Uniform[nmin, nmax] θi ∼ Uniform[θmin, θmax], i = 1, . . . , n oi ∼ Uniform[omin, omax], i = 1, . . . , n si := (x− 0.5) cos θi − (y − 0.5) cos θi
f(x, y) = 1
2 n∑ i=1 1 + tanh [100(si − oi)]
A.1.3 REWARD
Let ut denote the true solution at time t, let ût denote the estimated solution on the mesh at time t. For a given mesh at time t − 1, a given time evolution of the true solution from t − 1 to t, and a refinement action at (which may be “do-nothing”), let ût,refine denote the estimated solution on the mesh at time t that has undergone refinement action at, and let ût,no-refine denote the estimated solution on the mesh at time t without that refinement. We have two reward definitions:
1. Delta norm reward with true solution
rt := (et−1 − et)/e0 (5) et := ‖ut − ût‖2 (6)
2. Surrogate reward
rt := ‖ût,refine − ût,no-refine‖2 (7)
We used the first reward definition for static and advection experiments where analytic true solutions are available, and for Burgers experiments involving a single initial condition and pre-computed reference data that acts as the true solution. We used the second reward definition for all other Burgers experiments.
A.2 IMPLEMENTATION
We used standard policy gradient (Sutton et al., 2000) for all experiments except for experiments on Burgers equation and generalization of 8x8-trained advection policies to 64x64 test meshes. We used PPO (Schulman et al., 2017) for the latter two cases. We trained for 20k episodes on static problems, 10k episodes on advection problems, 2k episodes on Burgers equation with a single IC, and 4k episodes on Burgers equation with random ICs. The Burgers experiment with pretraining used 2k episodes on a single IC and a further 2k on random ICs. Each episode is initialized with refinement budget B, where B = 10 for static problems, B = 20 for advection, and B = 50 for Burgers.
IPN. For efficient computation on a batch of B trajectories, where each trajectory b consists of T environment steps and each step tb consists of a variable-sized global state s ∈ RNtb×d, we merge the variable dimension with the batch and time dimension to form an input matrix whose dimensions are [ ∑B b=1 ∑T t=1Ntb , d]. The output is reshaped into a “ragged” matrix of logits with dimensions [B × T,Ntb ], where the row lengths vary for each batch and time step. A softmax operation over each row produces the final action probabilities at each step.
Graphnet policy. The first graph layer is an Independent recurrent block that passes the input node tensors through a convolutional layer followed by a fully-connected layer, to arrive at node embeddings. This is followed by two recurrent passes through an InterationNetwork (Battaglia et al., 2016) where fully-connected layers are used for edge and node update functions. A final InteractionNetwork output layer followed by a global softmax over the graph produces a scalar at each node, which is interpreted as the probability of selecting the corresponding element for refinement. Except for the input node feature vi ∈ Rd and output node scalar, all internal node (edge) embeddings have the same size, denoted as dim(v) (dim(e)). We fixed dim(e) = 16 for both static and advection and tuned dim(v) (Table 2).
Hypernet policy. We fixed the main network’s hidden layer dimension at h = 64 and tuned the hypernetwork’s hidden layer dimension h1 (Table 2).
A.3 HYPERPARAMETERS
For both static and advection problems, we tuned a subset of all hyperparameters for all methods by the following procedure to handle the large set of policy architectures and ground truth functions. Chosen values of tuned hyperparameters are given in Table 2; all other hyperparameters have the same values for all methods and are listed below. We conducted tuning in a multi-task setup, where we train a single policy on functions randomly sampled from all ground truth classes, with randomly sampled parameters according to Appendix A.1.2. This is done separately on static and advection
problems. The tuning process is coordinate descent where the best parameter from one sweep is used for the next sweep. We start with exploration decay div ∈ {100, 500, 1000, 5000} (a lower bound on exploration was enforced by using behavioral policy π̃(at|st) = (1 − )π(at|st) + /Nt with decaying linearly from start to end by div episodes). Next we tune the size of hidden layers in the policy network (over (h1, h2) ∈ {(128, 64), (256, 64), (128, 128), (256, 256)} for IPN, node representation dimension dim(v) ∈ {32, 64, 128, 256} for Graphnet, and h1 ∈ {16, 32, 64, 128} for Hypernet). Lastly, we tune the learning rate α ∈ {5 · 10−5, 10−4, 5 · 10−4, 10−3, 5 · 10−3}. For Graphnet and Hypernet, we inherit the best div from IPN because optimal exploration depends in large part on the complexity of the environment, which is the same across all policy architectures.
Separately for the static and advection cases, all three policy architectures have the same values for all other hyperparameters. These are: policy gradient batch size 8, initial exploration lower bound start = 0.5, final exploration lower bound end = 0.05, discount factor γ = 0.99, convolutional neural network layer with 6 filters of size (5, 5) and stride (2, 2).
For Burgers experiments and advection experiments on generalization from 8× 8 to 64× 64 initial mesh sizes, we used a more comprehensive population-based hyperparameter search with successive elimination for all methods. We start with a batch of nbatch tuples, where each tuple is a combination of hyperparameter values, with each value sampled either log-uniformly from a continuous range or uniformly from a discrete set. We train independently with each tuple for nepisode episodes, eliminate the lower half of the batch based on their final performance, then initialize the next set of nepisode episodes with the current models for the remaining tuples. We use the hyperparameters of the last surviving model. Chosen values are shown in Table 3.
The hyperparameter ranges are: discount factor in {0.1, 0.5, 0.99}, policy entropy coefficient in (10−3, 1.0), GAE λ in {0.85, 0.90, 0.95}, learning rate in (10−5, 5 · 10−3, PPO in (0.01, 0.5), value loss coefficient in {0.1, 0.5, 1.0}, IPN h1 in {128, 256}, and IPN h2 in {64, 128, 256}.
A.4 COMPUTING INFRASTRUCTURE AND RUNTIME
Experiments were run on Intel 8-core Xeon E5-2670 CPUs, using one core for each independent policy training session. Average training time with 20k episodes in the static case was approximately 6 hours for IPN and Hypernet, and 9 hours for Graphnet. Average training time with 10k episodes
in the advection case was approximately 14 hours for IPN and Hypernet, and 18 hours for Graphnet. Policy decision times are shown in Table 4.

heading: B ADDITIONAL RESULTS
text: OOD. In the static case (Figures 13a to 13c), IPN policies trained on circles transfer well to bumps (and vice versa). Hypernet policies performed poorly overall even in the case of in-distribution, and consequently does not show comparable performance when transferring across function classes. In the advection case (Figures 13d to 13f), both IPN and Graphnet policies trained on bumps significantly outperformed ZZ when tested on circles (compare to ZZ in Figure 3b).

Review summaries:
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.

For various complicated problems governed by PDEs (e.g. solid/fluid interactions, aerodynamics, elasticity, backscattering, etc.) the computational cost can become prohibitive even for one inquiry, let alone parametric study. In the same time, mesh refinement is crucial to achieve acceptable accuracy. To mitigate such challenges, one solution is to use adaptive mesh refinement (AMR), for which the mesh refined only in the regions that numerically are sensitive to error propagation. For example for boundary layer models or shock-boundary layer interactions, one much capture the dynamics in high-gradients region of the solution, which are typically in vicinity of the walls or are part of the solution, while in very far regions of the domain, a coarse mesh is sufficient.  Authors recognize that the process of AMR, refinement at each step, can be formulated a Markov decision process (MDP) and hence utilize reinforcement learning (RL) to train refinement policies directly from simulation. But this in turn poses a new challenge, at each step, the dimension of state (number of elements) and action space may (and should) alter. They propose suitable policy updates to overcome this challenge and come up with three different architectures for the implementation. Three test cases are used for experiments (static, advection and Burger) and authors compare all three architectures with each other as well as some traditional AMR methods to demonstrate the performance of the proposed method. They also carry out extra tests on the same set of PDEs to show generalization and out of distribution capabilities of the method for both static and transient PDEs. The paper is well-written and two tricks for RL are impressive (using Nmax to take care of varying dimension of state/action space and use of surrogate reward for training). However, I have some major reservations that I'll explain below in the Main Review.