Multi-Agent Reinforcement Learning for Adaptive Mesh Refinement

The finite element method (Brenner and Scott, 2007) models the domain of a PDE with a mesh that consists of nonoverlapping geometric elements. Using the weak formulation and choosing basis functions to represent the PDE solution on these elements, one obtains a system of linear equations that can be numerically solved. The shape and sizes of elements determine solution error, while the computational cost is primarily determined by the number of elements. Improving the trade-off between error and cost is the objective of AMR. This work focuses purely on $h$-refinement, in which refinement of a coarse element produces multiple new fine elements (e.g., isotropic refinement of a 2D quadrilateral element produces four new elements), whereas de-refinement of a group of fine elements, restores an original coarse element. In both cases, the original coarse element or fine elements are removed from the FEM discretization.

Each element is identified with an agent, denoted by $i\in\{1,\dotsc,n_{t}\}$, where $n_{t}$ is variable across time step $t=0,1,\dotsc,T_{\max}$ within each episode of length $T_{\max}$, and $n_{t}\in[N_{\max}]$, where $N_{\max}$ is constrained by $\text{depth}_{\max}$ to be $n_{x}\cdot n_{y}\cdot 4^{\text{depth}_{\max}}$ in the case of quadrilateral elements. Let each element $i$ have its own individual observation space ${\mathcal{S}}^{i}$ and action space $\mathcal{A}^{i}$. Let $s$ denote the global state, $o^{i}$ and $a^{i}$ denote agent $i$’s individual observation and action, respectively, and $a\vcentcolon=(a^{1},\dotsc,a^{n_{t}})$ denote the joint action by $n_{t}$ agents. In the case of AMR, all agents have the same observation and action spaces. Let $R$ denote a single global reward for all agents and let $P$ denote the environment transition function, both of which are well-defined for any number of agents $n_{t}\in[N_{\max}]$. Let $\gamma\in(0,1]$ be the discount factor. The FEM solver time at episode step $t$ is denoted by $\tau$ (we omit the dependency $\tau(t)$ for brevity), which advances by $\tau_{\text{step}}$ in increments of $d\tau=0.002$ during each discrete step $t\rightarrow t+1$ up to final simulation time $\tau_{f}$.

For element $i$ at time $t$, let $u^{i}_{t}$ and $\hat{u}^{i}_{t}$ denote the true and numerical solution, respectively, and let $c^{i}_{t}\vcentcolon=\lVert u^{i}_{t}-\hat{u}^{i}_{t}\rVert_{2}$ denote the $\mathcal{L}_{2}$ norm of the error. Let $c_{t}\vcentcolon=\sqrt{\sum_{i=1}^{n_{t}}(c^{i}_{t})^{2}}$ denote the global error of a mesh with $n_{t}$ elements. Let $\text{depth}(i)=0,1,\dotsc,\text{depth}_{\max}$ denote the refinement depth of element $i$. Let $d_{t}\vcentcolon=\sum_{s=0}^{t}\text{DoF}_{s}$ denote the cumulative degrees of freedom (DoF) of the mesh up to step $t$, which is a measure of computational cost, and let $d_{\text{thres}}$ be a threshold (i.e., constraint) on the cumulative DoF seen during training.

In the paradigm of centralized training with decentralized execution (CTDE) (Oliehoek and Amato, 2016), global state and reward information is used in centralized optimization of a team objective at training time, while decentralized execution allows each agent to take actions conditioned only on their own local observations, independently of other agents, both at training time and at test time. One simple yet effective way to implement this in value-based MARL is Value Decomposition Networks (VDN) (Sunehag et al., 2018). VDN learns within the class of global action-value functions $Q(s,a)$ that decompose additively:

where $Q^{i}$ is an individual utility function representing agent $i$’s contribution to the joint expected return.

This decomposition is amenable for use in Q-learning (Watkins and Dayan, 1992), as it satisfies the individual-global-max (IGM) condition:

This means the individual maxima of $Q^{i}$ provide the global maximum of the joint $Q$ function for the Q-learning update step (Watkins and Dayan, 1992; Mnih et al., 2015), which scales linearly rather than exponentially in the number of agents. Using function approximation for $Q^{i}_{\theta}$ with parameter $\theta$, the VDN update equations using replay buffer $\mathcal{B}$ are:

where $\bot$ is the stop-gradient operator.

In a graph attention layer, each node is updated by a weighted aggregation over its neighbors: weights are computed by self-attention using node features as queries and keys, then applied to values that are computed from node and edge features.

Self-attention weights $\hat{a}^{ij}$ for each node $i$ are computed as follows (see Figure 2): 1) we define queries, keys, and values as linear projections of node features, via weight matrices $W^{q}$, $W^{k}$, and $W^{v}$ (all $\in\mathbb{R}^{d\times\dim(v)}$) shared for all nodes; 2) for each edge $(i,j)$, we compute a scalar pairwise interaction term $a^{ij}$ using the dot-product of queries and keys; 3) for each receiver node $i$ with sender node $j\in\mathcal{N}_{i}$, we define the attention weight as the $j$-th component of a softmax over all neighbors $k\in\mathcal{N}_{i}$:

We use these attention weights to compute the new feature for each node $i$ as a linear combination over its neighbors $j\in\mathcal{N}_{i}$ of projected values $W^{v}v^{j}$. Edge features $e^{ij}$, with linear projection using $W^{e}\in\mathbb{R}^{d\times\dim(e)}$, capture the relational part of the observation:

Despite being a special case of the most general message-passing flavor of graph networks (Battaglia et al., 2018), graph attention networks separate the learning of $\hat{a}^{ij}$, the scalar importance of interaction between $i$ and $j$ relative to other neighbors, from the learning of $b^{ij}$, the vector determining how $j$ affects $i$. This additional inductive bias reduces the functional search space and can improve learning with large receptive fields around each node, just as attention is useful for long-range interactions in sequence data (Vaswani et al., 2017).

Multi-head graph attention layer. We extend graph attention by building on the mechanism of multi-head attention (Vaswani et al., 2017), which uses $H$ independent linear projections $(W^{q,h},W^{k,h},W^{v,h},W^{e,h})$ for queries, keys, values and edges (all projected to dimension $d/H$), and results in $H$ independent sets of attention weights $\hat{a}^{ij,h}$, with $h=1,2,\dotsc,H$ . This enables attention to different learned representation spaces and was found to stabilize learning (Veličković et al., 2018). The new node representation with multi-head attention is the concatenation of all output heads, with an output linear projection $W^{o}\in\mathbb{R}^{d\times d}$. Section A.1 shows the multi-head versions of eqs. 7, 8, 9 and 10.

We use the graph attention layer to define the Value Decomposition Graph Network (VDGN), shown in Figure 1. Firstly, an independent graph network block (Battaglia et al., 2018) linearly projects nodes and edges independently to $\mathbb{R}^{d}$.

Each layer $F$ of VDGN involves two sub-layers: a multi-head graph attention layer followed by a fully-connected independent graph network block with ReLU nonlinearity. For each of these sub-layers, we use residual connections (He et al., 2016) and layer normalization (Ba et al., 2016), so that the transformation of input graph $g$ is $\text{LayerNorm}(g+\text{SubLayer}(g))$. This is visualized in Figure 3. We define a VDGN stack as $S\vcentcolon=F_{L}(F_{L-1}(\dotsm F_{1}(g)\dotsm))$ with $L$ unique layers without weight sharing, then define a single forward pass by $r$ recurrent applications of $S$: $P(g)=S(S(\dotsm S(g)\dotsm))$ with $r$ instances of $S$. Finally, to produce $|\mathcal{A}|$ action-values for each node (i.e., element) $i$, we apply a final graph attention layer whose output linear projection is $W^{o}_{\text{out}}\in\mathbb{R}^{|\mathcal{A}|\times d}$, so that each final node representation is $\hat{v}^{i}\in\mathbb{R}^{|\mathcal{A}|}$, interpreted as the individual element utility $Q(o^{i},a^{i})$ for all possible actions $a^{i}\in\mathcal{A}$.

VDGN is trained using eq. 1 in a standard Q-learning algorithm (Mnih et al., 2015) (see the leftward process in Figure 1). Since VDGN fundamentally is a $Q$ function-based method, we can extend it with a number of independent algorithmic improvements (Hessel et al., 2018) to single-agent Deep Q Network (Mnih et al., 2015) that provide complementary gains in learning speed and performance. These include double Q-learning (Van Hasselt et al., 2016), dueling networks (Wang et al., 2016), and prioritized replay (Schaul et al., 2015). Details are provided in Section A.2. We did not employ the other improvements such as noisy networks, distributional Q-learning, and multi-step returns because the AMR environment dynamics are deterministic for the PDEs we consider and the episode horizon at train time is short.

Methods for FEM and AMR should respect symmetries of the physics being simulated. For the simulations of interest in this work, we require a refinement policy to satisfy two properties: a) spatial equivariance: given a spatial rotation or translation of the PDE, the mesh refinement decisions should also rotate or translate in the same way; b) time invariance: the same global PDE state at different absolute times should result in the same refinement decisions. By construction of node and edge features, and the fact that graph neural networks operate on individual nodes and edges independently, we have the following result, proved in Appendix B.

In applications where a user’s preference between minimizing error and reducing computational cost is not known until test time, one cannot a priori combine error and cost into a single scalar reward at training time. Instead, one must take a multi-objective optimization viewpoint (Hayes et al., 2022) and treat cost and error as separate components of a vector reward function $\mathbf{R}=(R^{c},R^{e})$. The components encourage lower DoFs and lower error, respectively, and are defined by

The objective is $\mathbb{R}^{2}\ni\mathbf{J}(\theta)\vcentcolon=\mathbb{E}_{s_{t+1}\sim P(\cdot|a,s_{t}),a\sim\pi_{\theta}(\cdot|s)}[\sum_{t=0}^{T}\gamma^{t}\mathbf{R}_{t}]$, which is vector-valued. We focus on the widely-applicable setting of linear preferences, whereby a user’s scalar utility based on preference vector $\omega$ is $\omega^{T}\mathbf{R}$ (e.g., $\omega=[0.5,0.5]$ implies the user cares equally about cost and error). At training time, we randomly sample $\omega\in\Omega$ in each episode and aim to find an optimal action-value function $\mathbf{Q}^{*}(s,a,\omega)\vcentcolon=\arg_{Q}\sup_{\pi\in\Pi}\omega^{T}\mathbb{E}_{\pi}[\sum_{t=0}^{T}\gamma^{T}\mathbf{R}_{t}]$, where $\arg_{Q}$ extracts the vector $E_{\pi}[\dotsc]$ corresponding to the supremum. We extend VDGN with Envelop Q-learning (Yang et al., 2019), a multi-objective RL method that efficiently finds the convex envelope of the Pareto front in multi-objective MDPs; see Section A.3 for details. Once trained, $\mathbf{Q}^{*}$ induces the optimal policy for any preference $\omega$ according to the greedy policy $a^{*}=\operatorname*{argmax}_{a}\omega^{T}\mathbf{Q}^{*}(s,a,\omega)$.

We use MFEM (Anderson et al., 2021; mfem, [n.d.]) and PyMFEM (mfem, [n.d.]), a modular open-source library for FEM, to implement the Markov game for AMR. We ran experiments on the linear advection equation $\frac{\partial u}{\partial t}+\nu\nabla{\cdot}u=0$ with random initial conditions (ICs) for velocity $\nu$ and solution $u(0)$, solving it using the FEM framework on a two-dimensional $L^{2}$ finite element space with periodic boundary conditions. Each discrete step of the Markov game is a mesh re-grid step, with $\tau_{\text{step}}$ FEM simulation time elapsing between each consecutive step. The solution is represented using discontinuous first-order Lagrange polynomials, and the initial mesh is partitioned into $n_{x}\times n_{y}$ quadrilateral elements. Appendix C contains further FEM details on the mesh partition and the construction of element observations.

Linear advection is a useful benchmark for AMR despite its seeming simplicity because the challenge of anticipatory refinement can be made arbitrarily hard by increasing the $\tau_{\text{step}}$ of simulation time that elapses between two consecutive steps in the Markov game (i.e., between each mesh update step). Intuitively, an optimal refinement strategy must refine the entire connected region that covers the path of propagation of solution features with large solution gradients (i.e., high error on a coarse mesh), and maintain coarse elements everywhere else. Hence, the larger the $\tau_{\text{step}}$, the harder it is for distant elements, which currently have low error but will experience large solution gradients later, to refine preemptively. But such long-distance preemptive refinement capability is exactly the key for future applications in which one prefers to have few re-meshing steps during a simulation due to its computational cost. Moreover, the existence of an analytic solution enables us to benchmark against error threshold-based baselines under the ideal condition of having access to perfect error estimator.

Metric. Besides analyzing performance via error $c$ and cumulative DoFs $d$ individually, we define an efficiency metric as $\eta=1-\sqrt{\tilde{c}^{2}+\tilde{d}^{2}}\in[0,1]$, where a higher value means higher efficiency. Here, $\tilde{c}\vcentcolon=\frac{c-c_{\text{fine}}}{c_{\text{coarse}}-c_{\text{fine}}}$ is a normalized solution error and $\tilde{d}\vcentcolon=\frac{d-d_{\text{coarse}}}{d_{\text{fine}}-d_{\text{coarse}}}$ is normalized cumulative degrees of freedom (a measure of computational cost). Here, the subscripts “fine“ and “coarse” indicate that the quantity is computed on a constant uniformly fine and coarse mesh, respectively, that are held static (not refined/de-refined) over time. The uniformly fine and coarse meshes themselves attain efficiency $\eta=0$. Efficiency $\eta=1$ is unattainable in principle, since non-trivial problems require $d>d_{\text{coarse}}$.

The graph attention layer of VDGN was constructed using the Graph Nets library (Battaglia et al., 2018). We used hidden dimension $64$ for all VDGN layers (except output layer of size $|\mathcal{A}|$), and $H=2$ attention heads. For $\text{depth}_{\max}=1$, with initial $n_{x}=n_{y}=16$, we chose $L=2$ internal layers, with $R=3$ recurrent passes. Each Markov game step has $\tau_{\text{step}}=0.25$, $\tau_{f}=0.75$ (hence $T_{\max}=3$). For $\text{depth}_{\max}=2$, with $n_{x}=n_{y}=8$, we used $L=R=2$. Each Markov game step has $\tau_{\text{step}}=0.2$, $\tau_{f}=0.8$ (hence $T_{\max}=4$). For each training episode, we uniformly sampled the starting position and velocity of a 2D isotropic Gaussian wave as the initial condition. The FEM solver time discretization was $d\tau=0.002$ throughout. See Appendix C for further architectural details and hyperparameters.

We compare with the class of local error-based Threshold policies, each member of which is defined by a tuple $(\theta_{r},\theta_{d},\tau_{\text{step}})$ as follows: every $\tau_{\text{step}}$ of simulation time, all elements whose true error exceed $\theta_{r}$ are refined, while those with true error below $\theta_{d}$ are de-refined. These policies represent the ideal behavior of widely-used AMR methods based on local error estimation, in the limit of perfectly accurate error estimation.

As discussed in Section 6.1, each mesh update incurs a computational cost, which means AMR practitioners prefer to have a long duration of simulation time between each mesh update step. Figure 4 shows the growth of global error versus simulation time of VDGN and Threshold policies with different $\tau_{\text{step}}$ between each mesh update step. In the case of $\text{depth}_{\max}=1$, VDGN was trained and tested using the longest duration $\tau_{\text{step}}=125d\tau$ (i.e., it has the fewest mesh updates), but it matches the error of the most expensive threshold policy that updates the mesh after each $\tau_{\text{step}}=1d\tau$ (see Figure 4(a)). This is possible only because VDGN preemptively refines the contiguous region that will be traversed by the wave within $125d\tau$ (e.g., see Figure 5). In contrast, Threshold must update the mesh every $1d\tau$ to achieve this performance, since coarse elements that currently have negligible error due to their distance from the incoming feature do not refine before the feature’s arrival.

Moreover, in the case of $\text{depth}_{\max}=2$, the agents learned to choose level-1 refinement at $t=1$ for a region much larger than the feature’s periphery, so that these level-1 elements can preemptively refine to level 2 at $t=2$ before the feature passes over them. This is clearly seen in Figure 6. This enabled VDGN with $\tau_{\text{step}}=100d\tau$ (fewest update steps) to have error growth rate close to that of Threshold with $\tau_{\text{step}}=1d\tau$ (see Figure 4(b)).

Symmetry. Comparing Figure 12 with Figure 5, we see that VDGN policies are equivariant to rotation of initial conditions. Reflection equivariance is also visible for the opposite moving waves in Figure 10. Translation equivariance can be seen in Figure 13. Note that perfect symmetry holds only for rotation by integer multiples of $\pi/2$ and translation by integer multiples of the width of a level-0 element. Symmetry violation from mesh discretization is unavoidable for other values.

Figure 7 shows that VDGN unlocks regions of the error-cost landscape that are inaccessible to the class of Threshold policies in all of the mesh configurations that were tested. We ran a sweep over refinement threshold $\theta_{r}\in[5\times 10^{-3},\dotsc,5\times 10^{-8},5\times 10^{-15}]$ with de-refinement threshold $\theta_{d}=4\times 10^{-15}$. In the case of $\text{depth}_{\max}=1,2$ with $500$ solver steps, and $\text{depth}_{\max}=1$ with $2500$ solver steps, Figures 7(a), 7(b) and 7(c) show that VDGN lies outside the empirical Pareto front formed by threshold-based policies, and that VDGN Pareto-dominates those policies for almost every value of $\theta_{r}$: given a desired error (cost), VDGN has much lower cost (error). The “In-distribution” group in Table 1 shows that VDGN has significantly higher efficiency than Threshold policies for all tested threshold values, for $\text{depth}_{\max}=1,2$.

To understand the optimality of VDGN policies, we further compared multi-objective VDGN to brute-force search for the best sequence of refinement actions in an anisotropic 1D advection problem with $n_{x}=64,n_{y}=1$ and two mesh update steps. To make brute-force search tractable, we imposed the constraint that a contiguous region of $n$ elements are refined at each step (while all elements outside the region are de-refined). We searched for the starting locations of the region that resulted in lowest final global error. By varying $n$, this procedure produces an empirical Pareto front of such brute-force policies in the error-cost landscape, which we plot in Figure 7(d). For multi-objective VDGN, we trained a single policy and evaluated it with $100$ randomly sampled preferences $\omega=[\alpha,1-\alpha]$ where $\alpha\sim\text{Unif}[0,1]$. Figure 7(d) shows that a single multi-objective VDGN policy produces a Pareto front (o) that approaches the Pareto front formed by brute force policies (o). Moreover, we see that Threshold policies with various refinement thresholds (o) are limited to a small section of the objective landscape, whereas VDGN unlocks previously-inaccessible regions.

Longer time. At training time for VDGN, each episode consisted of approximately 400-500 FEM solver steps. We tested these policies on episodes with 2500 solver steps, which presents the agents with features outside of its training distribution due to accumulation of numerical error over time. Table 1 shows that: a) VDGN maintain highest top performance; b) the performance ratio between VDGN and the best Threshold policy actually increases in comparison to the case with shorter time (e.g., 1.47 in the “Depth 1 2500 steps” column, as opposed to 1.10 in the “Depth 1” column). This is because the error of threshold-based policies accumulates quickly over time due to the lack of anticipatory refinement, whereas VDGN mitigates the effect. Figure 15 shows that VDGN sustains anticipatory refinement behavior in test episodes longer than training episodes.

Out-of-distribution test problems. Even though VDGN policies were trained on square meshes with 2D isotropic Gaussian waves, we find that they generalize well to initial conditions and mesh geometries that are completely out of the training distribution. On anisotropic Gaussian waves (Figure 8), ring-shaped features (Figure 9), opposite-moving waves (Figure 10), star-shaped meshes (Figure 11), VDGN significantly outperforms Threshold policies without any additional fine-tuning or training (see the “Generalization” group in Table 1). Figure 14 shows qualitatively that a policy trained on quadrilateral elements shows rational refinement decisions when deployed on triangular elements.

Scaling. Since VDGN is defined by individual node and edge computations with parameter-sharing across nodes and edges, it is a local model that is agnostic to size and scale of the global mesh. Figures 16 and 17 show that a policy trained on $n_{x}=n_{y}=16$ can be run with rational refinement behavior on an $n_{x}=n_{y}=64$ mesh.