R EINFORCEMENT L EARNING FOR A DAPTIVE M ESH R EFINEMENT
Anonymous authors
Paper under double-blind review
A
BSTRACT
Large-scale ﬁnite element simulations of complex physical systems governed by partial differential equations (PDE) crucially depend on adaptive mesh reﬁnement (AMR) to allocate computational budget to regions where higher resolution is required.
Existing scalable AMR methods make heuristic reﬁnement 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 reﬁnement 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 in- ductive 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
I NTRODUCTION
Numerical simulation of PDEs via the ﬁnite element method (FEM) (Brenner & Scott, 2007) plays an integral role in computational science and engineering (Reddy & Gartling, 2010; Monk et al., 2003).
Given a ﬁxed set of basis functions, the resolution of the ﬁnite 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 inefﬁcient 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 efﬁcient numerical simulations, adaptive mesh reﬁnement (AMR), a class of methods that dynamically adjust the mesh resolution during a simulation to maintain equidistribution of error, is used to signiﬁcantly increase accuracy relative to computational cost.
Existing methods for AMR share the same iterative process of computing a solution on the current mesh, estimating reﬁnement indicators, marking element(s) to reﬁne, and generating a new mesh by reﬁning marked elements (Bangerth & Rannacher, 2013; ˇCervený 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 difﬁcult for complex reﬁnement schemes such as hp -reﬁnement
(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 reﬁnement 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 reﬁne elements which would encounter complex features in the next time step.
This means that the optimality of a reﬁnement 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 reﬁnement policy to optimize a performance metric, such as ﬁnal 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 reﬁnement action at every MDP time step.
While one may deﬁne a ﬁxed and bounded state and action space given a ﬁnite reﬁnement budget, doing so is very inefﬁcient 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 efﬁcient 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 deﬁne 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
; 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 ﬁrst time that RL can be competitive with, and sometimes outperform, a greedy reﬁnement strategy based on the widely-used Zienkiewicz-Zhu-type error estimator; moreover, we show that an RL reﬁnement policy can generalize to higher reﬁnement 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).
R ELATED 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 ﬁnite element space rather than components of a numerical solver.
To the best of our knowledge, no prior work has formulated adaptive mesh reﬁnement 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 ﬁelds 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 ﬁnd 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 ﬁxed-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 ﬁxed action space size within each episode, whereas both state and action space sizes change at every time step within an episode in AMR.
B ACKGROUND AND F ORMULATION 3. F INITE ELEMENT METHOD
Our mesh adaptation strategy is implemented in a FEM-based framework (Brenner & Scott, 2007).
In FEM, the domain Ω ⊂ R D is modeled with a mesh that is a union of E nonoverlapping subsets ( elements ) such that Ω :=
(cid:83)
Ω k where k ∈ N : k (cid:54) 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 -reﬁnement, whereby elements are split into smaller elements (reﬁnement) or multiple elements coalesce to form a single element (dereﬁnement).
In practical applications with unknown true solutions, the conventional AMR approach is to take greedy reﬁnement 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. AMR AS A M ARKOV DECISION PROCESS
We formulate AMR with spatial h -reﬁnement 1 as a Markov decision process M :=
( O , N max , A , R, P, γ ) with each component deﬁned 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 N t ≤ N max ∈ N elements.
Each element i is associated with an observation o
it ∈
O and the global state is s t :=
[ o 1 t , . .
, o Nt ] ∈ O N t .
We deﬁne O :=
R d such that each element’s observation is a tensor of shape d := l × w × c that includes the values and reﬁnement depths of a local window centered on itself.
For brevity, let S t denote the current global state space O N t .
We denote an action by a t ∈
A t := { 0 , 1 , . . . , N t } ⊂
A := { 0 , 1 , . . .
, N max } , where 0 means “do-nothing” and i (cid:54) =means reﬁne element i .
Given the current state and action, the MDP transition P consists of:
1) reﬁning the selected element into multiple ﬁner elements (which increases N t ) if
a
reﬁnement budget B is not exceeded and
the
selected element is not at the maximum reﬁnement depth d max ; 2) stepping the ﬁnite element simulation forward in time (for time-dependent PDEs only); 3) computing a solution on the new ﬁnite element space.
When a true solution is available at training time, the reward at step t is deﬁned as the change in error from the previous step, normalized by the initial error to reduce variation across function classes:
[Equation]
(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 r t := (cid:107) u t, reﬁne − u t, no-reﬁne (cid:107) 2 , the normed difference between the estimated solution u with and without executing the chosen reﬁnement action.
This surrogate, which is an upper bound
Polynomial can be formulated a r be an RL problem but is not treated 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 ﬁnd a stochastic policy π : S t → ∆( A t ) to maximize the objective
[Equation]
Aside from γ ∈ (0 , 1) , this objective is equivalent to maximizing total error reduction: e 0 − e ﬁnal .
Although the size of the state vector and set of valid actions changes with each time step due to the varying N t , this MDP is well-deﬁned since one can deﬁne the global state space as the set of all possible O N , N < N max , 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 beneﬁt AMR at test time: a stochastic reﬁnement action could reveal the need for further reﬁnement in a region that appears ﬂat 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 := { ( s t , a t , r t ) k } Tt =1 } Kk =1 generated by the current policy.
P OLICY ARCHITECTURES FOR VARIABLE STATE - ACTION SPACES
We propose three policy architectures, each with different inductive biases, that address the challenge of variable size state vector
s ∈ R N t × d and action set { 0 , 1 , . .
. , N t } , both of whose sizes changes with number of elements N t 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 dereﬁnement actions.
4. I NDEPENDENT POLICY NETWORK
The Independent Policy Network (IPN) handles the 1:1 correspondence by mapping each observation to a probability for the corresponding action.
Let f θ : R d (cid:55)→ R be a function parameterized by θ .
Given a matrix of observations s :=
[ o 1 , . . .
, o N ] ∈ R N × d , we deﬁne the policy as
[Equation]
For example, using a neural network with hidden layer W ∈ R d × h with h nodes, output layer H ∈ R h ×, and activation function σ , the discrete probability distribution over N actions conditioned on s is deﬁned by softmax ( σ ( s W ) 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 ) ))
=
π ( a i | s ) for any permutation operator µ : [ N ] (cid:55)→ [ N ] —means that one cannot use the ordering of inputs to represent spatial relations among elements, which would be necessary for reﬁning an element based on neighboring conditions.
We mitigate this problem by deﬁning each element’s observation as an image tensor that includes neighborhood information and using a convolutional network layer, but this may face difﬁculties on unstructured meshes with non-quadrilateral elements ( ˇ Cervený et al., 2019).
4. H YPERNETWORK POLICY
The hypernetwork policy captures higher-order interaction among inputs via the function form
[Equation]
The main policy network weights θ are now the output of a hypernetwork (Ha et al., 2017)
g φ : R N × d (cid:55)→ R dim ( θ ) , parameterized by φ , which produces mixing among the inputs s ∈ R N × d .
Continuing with the example in IPN, where the policy network’s ﬁrst layer is W ∈ R d × h , a hypernetwork with two layers can be instantiated as (cid:104)(cid:80) Ni =1 ( s U ) i, : (cid:105) V = W
where U ∈ R
d × h 1 and
V ∈ R h 1 × ( d × h ) are the trainable parameters φ , and M i, : 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 difﬁculty in the choice of architecture (speciﬁc 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 reﬁnement decision, which is an extremely strong inductive bias.
4. G RAPH NETWORK POLICY
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.
Speciﬁcally, we construct a policy based on Interaction Networks (Battaglia et al., 2016), which is a special case without global attributes 2 .
At each step, the mesh is represented as a graph G =
( V, E ) .
Each vertex v i in V = { v i } i =1: N corresponds to element i and is initialized to be the observation o i .
E
= { ( e k , r k , s k ) } k
=1: N e is a set of edges with attributes e k between sender vertex s k and receiver vertex r k .
An edge exists between two vertices if and only if they are spatially adjacent.
We deﬁne the initial edge attribute e k as a one-hot vector indicator of the difference in reﬁnement depth between r k and s k .
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 deﬁned by the following sequence of computations: 1) Each edge attribute e k is updated by learned function ϕ e using local node information via ˆ e k ← ϕ e ( e k , v r k , v s k ) ; 2) For each node i , we denote by ˆ E i :=
{ (ˆ e k , r k , s k ) }
r k = 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 , v i ) 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 reﬁnement.
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.
E XPERIMENTAL SETUP
Our experiments assess the ability of RL, using the proposed policy architectures, to ﬁnd AMR strategies that generalize to test function classes that differ from the training class, generalize to variable mesh sizes and reﬁnement budgets, and extend to more complex problems without readilyavailable ground truth solutions.
We deﬁne 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.
5. AMR ENVIRONMENT
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 reﬁnement is to minimize the
L 2 error norm of projecting a variety of test functions onto a two-dimensional ﬁnite element space.
In the time-dependent case, the functions are projected onto the ﬁnite element space, and a PDE of the form ∂u∂t + ∇· F ( u ) = 0 is solved on a periodic domain using the ﬁnite element framework.
[Figure]
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 ) =
c u with c =
[1 , 0] , and the nonlinear Burgers equation, where F ( u ) =
c
u 2 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 physical 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 n x × n y 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 deﬁned 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 (cid:55)→ R to initialize each episode.
The collection, shown in Figure 2 and deﬁned 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 (deﬁned in Section 3.2) to compute the reward for training.
In the static case, the true solution is ﬁxed 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 reﬁnement 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 ﬁnite time.
This behavior is resolved using the ﬂux-corrected transport (FCT) approach (Boris & Book, 1997).
5. E XPERIMENTS AND PERFORMANCE METRIC
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 ﬁne-tuning.
Generalization : 1) Static → advection : Policies trained on static functions are tested on advection.
2) Budget ↑ : Policies trained with a small reﬁnement 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 ×and 64 × 64 initial meshes with and without preserving the relative solution and mesh length scales.
We deﬁne the performance of a given reﬁnement policy in an episode in the static case as ( e initial − e ﬁnal ) /e
initial , where e initial (or e ﬁnal ) 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 reﬁnement, the error may increase over the course of the simulation due to the accumulation of discretization error.
Hence, given a reﬁnement policy that achieves e ﬁnal at episode termination, we deﬁne its performance as ( e no-reﬁne, ﬁnal − e ﬁnal ) /e
initial , where e no-reﬁne, ﬁnal is the ﬁnal error without any reﬁnement.
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).
5. I MPLEMENTATION AND BASELINES
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 h 1 and h 2 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 ﬁrst 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 ∈ R
d × h 1 , V ∈ R h 1 × ( d × h ) , and Y ∈ R d × h , where h 1 and h are design choices.
U and V act on input state s to produce the main policy weights
W ∈ R
d
× h (see Section 4.2) while Y acts on s to produce a bias b ∈ R h , so that the main policy’s ﬁrst hidden layer is ReLU ( s W + 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 reﬁnes the element with the largest estimated error.
The TrueError policy reﬁnes 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 reﬁning 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 reﬁning each element individually and chooses the element whose reﬁnement 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.
R ESULTS
We ﬁnd that the proposed methods achieve performance that is competitive with baselines and, more importantly, generalize well to larger reﬁnement 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 .
6. I N - DISTRIBUTION
Static functions (Figure 3a).
RL policies either meet or signiﬁcantly exceed the performance of
ZZ on all function classes.
Notably, both IPN and Graphnet outperform ZZ signiﬁcantly on steps by spending the limited reﬁnement 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
[Figure]
of the same function class.
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 inﬂuence.
Hypernetwork policies converged to the behavior of making no reﬁnements on at least one out of four independent runs on all classes except bumps .
This could be attributed to the inherent difﬁculty 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 signiﬁcantly 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 ﬁxed 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.
[Figure]
Visualization of resulting meshes for Burgers equation with a ﬁxed bump IC.
6. O UT - OF - DISTRIBUTION (OOD)
Figure 4b shows the performance of IPN RL policies and baselines on Burgers equation with random
ICs.
We observe that policies trained on a ﬁxed 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 ﬁxed IC for 2k episodes and ﬁne-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 sufﬁcient and that pretraining on a ﬁxed IC is a more efﬁcient 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.
6. G ENERALIZATION
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 signiﬁcantly outperformed ZZ on advection-circles .
Surprisingly, static-trained IPN signiﬁcantly 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 with B = 10 correctly reﬁnes the region of propagation on advection-bumps with B = 50 .
Budget ↑ (Figure 5).
RL policies trained with low reﬁnement budget generalize to test cases with higher budget.
In the static case, comparing Figure 3a ( B = 10 ) with Figure 5a ( B = 50 ) shows that
[Figure]
the performance of RL policies relative to ZZ is generally preserved by the increase in reﬁnement budget.
Figures 5 and 8 show that an IPN trained with B = 10 makes qualitatively correct reﬁnement decisions when allowed B = 100 during test.
In the advection case (Figure 5b), Graphnet trained with B = 20 signiﬁcantly 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.
[Figure]
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 reﬁnements.
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.
C ONCLUSION
We present a novel formulation of adaptive mesh reﬁnement 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 efﬁciency gains beyond the reach of existing AMR approaches.
We demonstrated that these RL policies generalize to different reﬁnement 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-speciﬁc knowledge or domain expertise as input, our results provide a path for learning novel AMR strategies for cases such as hp -reﬁnement that currently lack effective solutions.
Future work can extend our RL methods to include dereﬁnement 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.
R EPRODUCIBILITY 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 speciﬁc implementation details of the proposed method in Appendix A.2; the experimental test cases in Appendix A.1.
and Table 1; the detailed tuning procedure in Appendix A.3; and the complete speciﬁcation of all hyperparameters in Appendix A.3, Table 2, and Table
A E XPERIMENTAL SETUP A. E NVIRONMENT DETAILS A.1. MFEMC TRL
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 reﬁnement decisions, and calculate errors.
The initial mesh is partitioned into n x × n y = 8 × 8 elements for static and advection experiments and 10 × 10 for Burgers equation.
Generalization experiments on larger initial mesh used n x × n y =
× 16 or 64 × 64 .
The true solution is projected onto the ﬁnite element space by interpolation to the nodes of the Bernstein basis functions.
After each reﬁnement action, the solution is projected again onto the reﬁned mesh (for the static case) or integrated in time until the next reﬁnement action (for the time-dependent case).
The maximum reﬁnement depth is ﬁxed by the parameter d max such that the maximally-reﬁned mesh consists of 2 d max n x × 2 d max n y elements.
d max was set to 3 for static experiments, whereas d max = 2 for advection and d max = 1 for Burgers equation due to the time step restrictions imposed by the Courant-Friedrichs-Lewy (CFL) condition of the ﬁnest elements.
Observation.
The observation consisted of the solution and the depth of each element.
Since the gradients of the solution are, by deﬁnition, 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 = l element +2 l context is the spatial observation window with l element = 16 sampled points inside the element and l context = 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 o 0 to state s corresponding to action 0.
At most one reﬁnement is allowed per MDP step.
[Table]
A.
[Table]
G ROUND TRUTH FUNCTIONS Bumps
n ∼ Uniform [ n min , n max ] c x,i ∼ Uniform [ c x, min , c x, max ] , i = 1 , . .
. , n c y,i ∼ Uniform [ c y, min , c y, max ] , i = 1 , . . . , n w i ∼ Uniform [ w min , w max ] , i = 1 , . . . , n f ( x, y )
=
n (cid:88) i = exp (cid:18) − (
[Equation]
−
[Equation] [Equation]
w
[Equation]
(cid:19)
Circles n ∼ Uniform [ n min , n max ] c x,i , c y,i ∼ Uniform [ c min , c max ] , i = 1 , . . . , n r i ∼ Uniform [ r min , r max ] , i = 1 , . . . , n w i ∼ Uniform [ w min , w max ] , i = 1 , . . . , n f ( x, y )
=
n (cid:88) i = exp (cid:32) − ( (cid:112) ( x
[Equation]
c x,i ) 2 + ( y
[Equation]
i
[Equation]
)w i (cid:33)
Steps n ∼ Uniform [ n min , n max ] θ ∼ Uniform [ θ min , θ max ] o i ∼ Uniform [ o min , o max ] , i = 1 , . .
. , n f ( x, y ) =
n (cid:88) i = 1 + tanh
[100( o i − ( x + y tan θ ))]
Steps n ∼ Uniform [ n min , n max ] θ i ∼ Uniform [ θ min , θ max ] , i = 1 , . .
n o i ∼ Uniform [ o min , o max ] , i = 1 , . .
. , n s i :=
( x − 0 . 5) cos θ i − ( y − 0 .
cos θ i f ( x, y )
=
1n (cid:88) i = 1 + tanh
[100( s i − o i )]
A.1. R EWARD
Let u t 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 reﬁnement action a t (which may be “do-nothing”), let ˆ u t, reﬁne denote the estimated solution on the mesh at time t that has undergone reﬁnement action a t , and let ˆ u t, no-reﬁne denote the estimated solution on the mesh at time t without that reﬁnement.
We have two reward deﬁnitions:
Delta norm reward with true solution
[Equation]
. Surrogate reward
[Equation]
We used the ﬁrst reward deﬁnition 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 deﬁnition for all other
Burgers experiments.
A. I MPLEMENTATION
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 reﬁnement budget B , where B = 10 for static problems, B = 20 for advection, and B = 50 for
Burgers.
IPN.
For efﬁcient computation on a batch of B trajectories, where each trajectory b consists of T environment steps and each step t b consists of a variable-sized global state s ∈ R N tb × d , we merge the variable dimension with the batch and time dimension to form an input matrix whose dimensions are [ (cid:80) Bb =1 (cid:80)
Tt =
1 N t b , d ] .
The output is reshaped into a “ragged” matrix of logits with dimensions [ B × T, N t b ] , where the row lengths vary for each batch and time step.
A softmax operation over each row produces the ﬁnal action probabilities at each step.
Graphnet policy.
The ﬁrst 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 ﬁnal
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 reﬁnement.
Except for the input node feature v i ∈ R d and output node scalar, all internal node (edge) embeddings have the same size, denoted as dim( v ) (dim( e )).
We ﬁxed dim ( e ) = 16 for both static and advection and tuned dim ( v ) (Table 2).
Hypernet policy.
We ﬁxed the main network’s hidden layer dimension at h = 64 and tuned the hypernetwork’s hidden layer dimension h 1 (Table 2).
A. H YPERPARAMETERS
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 (cid:15) div ∈ { 100 , 500 , 1000 , 5000 } (a lower bound on exploration was enforced by using behavioral policy
˜ π ( a t | s t ) =
(1 − (cid:15) )
π ( a t | s t ) + (cid:15)/N t with (cid:15) decaying linearly from (cid:15) start to (cid:15) end by (cid:15) div episodes).
Next we tune the size of hidden layers in the policy network (over ( h 1 , h 2 ) ∈ { (128 , 64) , (256 , 64) , (128 , 128) , (256 , 256) } for IPN, node representation dimension dim ( v ) ∈ { 32 , 64 , 128 , 256 } for Graphnet, and h 1 ∈ { 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 (cid:15) 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 (cid:15) start = 0 .
5 , ﬁnal exploration lower bound (cid:15) end = 0 .
05 , discount factor γ = 0 .
99 , convolutional neural network layer with 6 ﬁlters of size (5 , 5) and stride (2 , 2) .
[Table]
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 n batch 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 n episode episodes, eliminate the lower half of the batch based on their ﬁnal performance, then initialize the next set of n episode 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 coefﬁcient in (10 − 3 , 1 . 0) , GAE λ in { 0 . 85 , 0 . 90 , 0 . 95 } , learning rate in (10 − 5 , 5 · 10 − 3 , PPO (cid:15) in (0 . 01 , 0 . 5) , value loss coefﬁcient in { 0 . 1 , 0 . 5 , 1 . 0 } , IPN h 1 in { 128 , 256 } , and IPN h 2 in { 64 , 128 , 256 } .
[Table]
A. C OMPUTING 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.
[Table]
[Figure]
B A DDITIONAL RESULTS
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).
[Figure]
[Figure] [Figure] [Figure]