SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Niklas Freymuth $^{1*}$  Philipp Dahlinger $^{1}$  Tobias Würth $^{2}$  Simon Reisch $^{1}$   
Luis Kärger $^{2}$  Gerhard Neumann $^{1}$
<sup>1</sup> Autonomous Learning Robots, Karlsruhe Institute of Technology, Karlsruhe  
<sup>2</sup>Institute of Vehicle Systems Technology, Karlsruhe Institute of Technology, Karlsruhe
Adaptive Mesh Refinement (AMR) enhances the Finite Element Method, an important technique for simulating complex problems in engineering, by dynamically refining mesh regions, enabling a favorable trade-off between computational speed and simulation accuracy. Classical methods for AMR depend on heuristics or expensive error estimators, hindering their use for complex simulations. Recent learning-based AMR methods tackle these issues, but so far scale only to simple toy examples. We formulate AMR as a novel Adaptive Swarm Markov Decision Process in which a mesh is modeled as a system of simple collaborating agents that may split into multiple new agents. This framework allows for a spatial reward formulation that simplifies the credit assignment problem, which we combine with Message Passing Networks to propagate information between neighboring mesh elements. We experimentally validate our approach, Adaptive Swarm Mesh Refinement (ASMR), on challenging refinement tasks. Our approach learns reliable and efficient refinement strategies that can robustly generalize to different domains during inference. Additionally, it achieves a speedup of up to 2 orders of magnitude compared to uniform refinements in more demanding simulations. We outperform learned baselines and heuristics, achieving a refinement quality that is on par with costly error-based oracle AMR strategies.
The Finite Element Method (FEM) is a widely used numerical technique in engineering and applied sciences for solving complex partial differential equations [1, 2, 3, 4]. The method discretizes the continuous problem domain into smaller, finite elements, allowing for an efficient numerical solution. A key aspect of the FEM for complex systems is Adaptive Mesh Refinement (AMR), which dynamically refines regions of high solution variability, allowing for a favorable trade-off between computational speed and simulation accuracy [5, 6, 7]. As problems in engineering grow more complex, the FEM and especially Adaptive Mesh Refinement (AMR) techniques have become increasingly important tools in providing computationally tractable yet precise solutions. Applications of AMR include fluid dynamics [8, 9, 10, 11, 12, 13], structural mechanics [14, 15, 16, 17], and astrophysics [18, 19, 20]. Yet, classical approaches for AMR usually rely on either problem-dependent or more general but potentially suboptimal error indicators, or require expensive error estimates [21, 22, 23, 24, 25, 26, 12]. In either case, they can be cumbersome to use in practice.
To address this issue, we formalize AMR as a Reinforcement Learning (RL) [27] problem. Following previous work [28, 29, 30], each refinement step encodes the state of the current simulation as local observations that we feed to RL agents, who then determine which elements of a mesh to refine. However, previous work has issues with scalability due to an expensive inference process [28],
\*correspondence to niklas.freymuth@kit.edu
Figure 1: Given a mesh  $\Omega^t$ , an observation graph  $\mathcal{G}_{\Omega^t}$  encodes the elements as graph nodes and the neighborhood between elements as edges. The graph is given to a learned policy  $\pi$ , which marks mesh elements for refinement. A remesher refines the mesh, and spatial rewards  $\mathbf{r}(\Omega_i^t)$  are calculated for all agents  $i$  based on the quality of the refinement. This process is iterated for several steps until the mesh is fully refined.
misaligned objectives and high variance in the state transitions [30], and noisy reward signals [29]. To mitigate these shortcomings and scale to more complex problems, we formulate AMR as a Swarm RL [31, 32] problem. We extend the Swarm RL framework to per-agent rewards, a shared observation space, and the option of splitting agents into new agents as required in the AMR process. Additionally, we introduce a novel spatial reward formulation that provides a dense reward signal for the refinement of each mesh element, simplifying the credit assignment problem for swarm systems. Our policy is based on Message Passing Networks (MPNs) [33], a class of Graph Neural Networks (GNNs) [34, 35, 36, 37] that has proven to be effective for physical simulations [38, 39]. The resulting method, Adaptive Swarm Mesh Refinement (ASMR), consistently produces highly refined meshes with thousands of elements while applying to arbitrary Partial Differential Equations (PDEs). A high-level overview is given in Figure 1.
Experimentally, we show the effectiveness of our approach on a suite of PDEs that require complex and challenging refinement strategies, including a non-stationary heat diffusion problem and a linear elasticity task. We implement our tasks as OpenAI gym [40] environments. The experiments use static meshes with conforming triangular elements and corresponding h-adaptive refinements [41, 42] due to their importance in engineering applications [43, 44, 45, 46]. Here, accurate meshes require multiple precise refinement steps and thousands of elements. We implement and compare to current state-of-the-art RL methods for AMR [28, 29, 30] that have been shown to work well on dynamic tasks where shallow mesh refinement and coarsening is sufficient<sup>1</sup>. We observe that these methods struggle with finer static meshes, whereas ASMR yields stable and consistent refinements across tasks. To further evaluate our method's effectiveness, we compare it to the popular Zienkiewicz-Zhu Error Estimator (ZZ Error) estimate and a traditional AMR heuristic requiring oracle error estimates. ASMR not only outperforms both learned and traditional methods that lack oracle data but also achieves performance comparable to oracle-based heuristics. Furthermore, ASMR is 2 to 100 times faster than computing the uniform mesh on which the oracle information is based, and demonstrates robust generalization capabilities across different domains and initial conditions. We conduct a series of ablations to show which parts of the approach make it uniquely effective, finding that spatial rewards per agent are preferable to a shared global reward signal for all agents.
To summarize our contributions, we (1) propose a novel Markov Decision Process (MDP) formulation that naturally integrates local rewards for swarms where agents may split into new agents over time; (2) combine this formulation with MPNs and a novel spatial reward formulation to reliable and efficiently scale learned AMR to static meshes with thousands of elements on multiple levels of refinement; (3) showcase our approach's effectiveness on a suite of PDEs with challenging refinement
<sup>1</sup>We publish the first codebase on RL for AMR, including all methods and tasks presented in this paper, to facilitate research in this direction. The code is available at https://github.com/NiklasFreymuth/ASMR.
problems. Our method surpasses state-of-the-art RL methods and the popular ZZ Error heuristic, and achieves refinement quality comparable to oracle-based AMR strategies, all without requiring costly error estimates during inference.
Learned Physics Simulation. A considerable body of work deals with directly learning to simulate physical systems with neural networks. These approaches typically learn from data generated by some underlying ground-truth simulator and train the network to predict the (change in) quantities of interest during a simulation. Such learned physics simulators are fully differentiable and often orders of magnitude faster than their classical counterparts [35, 38], lending them to use cases such as Inverse Design [47, 48, 49]. Researchers have developed simulators based on simple feed-forward networks [50, 51] and Convolutional Neural Networks [52, 53, 54, 55, 56, 57, 58, 59, 60]. Closely related to our method are Graph Network Simulators (GNSs) [33, 38, 61, 62, 63, 64, 39, 65], which utilize GNNs to encode physical problems as a graph on which to compute quantities of interest per node. Here, a recent method [66] jointly learns a GNS and an AMR strategy on mesh edges to allow for simulation on different resolutions.
Physics-Informed Neural Networks [67, 68, 69] are mesh-free methods designed to directly train neural networks to satisfy the governing equations of a physical system. They share the goal of using deep neural networks to solve PDEs, yet differ in their approach in that they directly approximate the equations rather than providing a mesh for a classical solver. Thus, AMR strategies provide a more robust, flexible, and risk-averse approach to solving complex physics problems that demand high precision and accuracy [11, 70, 13]. As Physics-Informed Neural Networks also operate on geometric domains, they have been extended to GNN architectures [71, 72, 73]. In this work, we do not learn to solve a system of equations directly, but rather propose an efficient mesh refinement for a classical solver.
Supervised Learning for AMR. Applications of supervised learning for AMR include directly calculating an error per mesh element with a Multilayer Perceptron (MLP) [11] and predicting mesh densities from domain images [70]. Additionally, recurrent networks have been used to find optimal marking strategies for second-order elliptical PDEs [74]. Another body of work speeds up the computation of Dual Weighted Residual [75, 76] error estimators by substituting expensive parts of the procedure with neural networks. Here, recent methods [77, 78] consider learning a metric tensor from solution information that can then be used in existing refinement procedures [24]. Other approaches [79, 13] employ neural networks to solve the strong form of the adjoint problem and use hand-crafted features to compute error estimates directly. We instead leverage the fact that RL can optimize non-differentiable rewards, enabling us to directly learn a refinement strategy in an iterative manner instead of learning error estimators or other specific facets of AMR.
Reinforcement Learning for AMR. Current research in Reinforcement Learning for AMR involves various approaches, such as optimizing the positions of mesh elements [80], predicting a global threshold for heuristic-based refinement using existing error estimates [81], and generating quadrilateral meshes by iteratively extracting elements from the problem domain [82].
We instead directly manipulate the mesh elements themselves. This approach presents a unique challenge in that the size of the observation and action spaces is constantly changing during refinement. Existing methods [28, 29, 30] derive their observation spaces from the mesh geometry and the solution computed on the mesh. These methods are generally designed for non-stationary PDEs and thus include mesh coarsening operations that are required for time-dependent problems. While our method can easily be extended to coarsening and time-dependent problems, this work instead considers static meshes and mostly stationary problems due to their prevalence in engineering [43, 44]. Here, the challenge lies in finding multiple levels of accurate refinements rather than shallow or local time-dependent refinement and coarsening. The earliest of these methods [28] treats the entire mesh as an action and observation space for a Single Agent, which uses a GNN-based policy to provide a categorical action that selects a single element for refinement. Single Agent demands solving the system of equations after each refinement step, significantly increasing inference time while reducing the amount of information per environment sample. Another approach [30] iteratively selects a random element during training and uses an MLP policy to determine its marking based on local and global features. During inference, the method performs a Sweep for all mesh elements in parallel.
Stokes Flow
Figure 2: Exemplary ASMR refinements. The heatmaps represent the normalized quantities of interest. ASMR provides complex and accurate refinements for different tasks. From left to right: Laplace's equation requires a homogeneous refinement near the source at the inner boundary. For Poisson's equation, a multi-modal load function causes multiple distinct regions of interest. The Stokes flow uses more complex shape functions and requires high precision near the inlet on the left and the rhomboid holes. For the linear elasticity task, we consider both the deformation and the resulting stress as quantities of interest. The heat diffusion task requires accurate refinements on the path of the moving heat source to predict an accurate solution of the final step.
Linear Elasticity
Heat Diffusion
This procedure speeds up inference but causes a misalignment in the environment transition between training and inference [29]. Additionally, during training, the agent is randomly assigned a new element after each action, leading to high variance in the state transitions.
Most similar to our method are Value Decomposition Graph Networks (VDGN) [29] which frame AMR as a cooperative multi-agent problem by setting a maximum refinement depth and number of agents.  $VDGN$  employs a Value Decomposition Network [83] to circumvent the posthumous credit assignment problem [84] of vanishing agents. Though theoretically efficient in training and inference, the method's performance depends on the quality of the value decomposition, which becomes more difficult for larger meshes. In summary, existing RL methods do not utilize the spatial nature of AMR and thus only scale to either simple or comparatively shallow refinements. We instead formulate AMR as a Swarm Reinforcement Learning problem with spatial rewards, naturally integrating changing observation and action spaces while also providing a strong feedback signal to all agents.
3 Adaptive Swarm Mesh Refinement
In the following, we introduce the individual components of ASMR, including our novel Adaptive Swarm Markov Decision Process (ASMDP) and spatial reward function. We consider each element  $\Omega_i^t\in \Omega^t$  of a mesh  $\Omega^t$  to be an agent in a swarm system. The agent's state is its position in the mesh, as well as boundary conditions and other PDE-dependent quantities. The agent's observation consists of a local view of a graph  $\mathcal{G}_{\Omega^t}$  where each node represents a mesh element and each edge the neighborhood of two elements. We train a simple MPN-based policy  $\pi (\mathbf{a}|\mathcal{G}_{\Omega^t})$  that computes a joint action vector  $\mathbf{a}\in \mathcal{A}^N$  for each mesh element by passing messages along the observation graph, as detailed in Appendix A. The action vector is used for refinement, and the process is repeated with the refined mesh for a given number of steps. Since our policy uses a GNN, it is equivariant to permutation and can handle varying numbers of agents by construction. Figure 1 provides a schematic overview of our method.
Adaptive Swarm Markov Decision Process. We adapt the SwarMDP framework [31, 32] to incorporate action and observation spaces of changing size as necessary for AMR, agent-wise rewards and mappings between agents over time. The resulting framework is conceptually simpler than, e.g., a decentralized partially observable MDP with dummy states [29], makes the permutation-equivalence of the agents explicit and naturally integrates both the agent-dependent reward and the mapping between agents. Formally, we define an Adaptive Swarm Markov Decision Process (ASMDP) as a tuple  $\langle \mathbb{S},\mathbb{O},\mathbb{A},T,\mathbf{r},\xi ,\mathbf{M}\rangle$ . Here,  $\mathbb{S}$  is the state space of the system,  $\mathbb{O}$  is the space of observations for this state space, and  $\mathbb{A}$  is the action space for the system of agents. Let  $S^N\subset \mathbb{S}$ ,  $\mathcal{O}^N\subset \mathbb{O}$ ,  $\mathcal{A}^N\subset \mathbb{A}$  denote the subsets of the state, observation, and action spaces with exactly  $N$  agents. The transition function  $T:S^{N}\times \mathcal{A}^{N}\to S^{K}$  maps to a new system state with a potentially different number of agents, and  $\mathbf{r}:S^N\times \mathcal{A}^N\to \mathbb{R}^N$  is a per-agent reward function. The observation graph of the agents is calculated from their states via the observation function  $\xi :S^{N}\to \mathcal{O}^{N}$ . To
accommodate changing numbers of agents throughout an episode, we define an agent mapping  $\mathbf{M}^t\in [0,1]^{N\times K}$  with  $\sum_{i}\mathbf{M}_{ij}^{t} = 1$  for all  $j = 1,\dots ,K$  that specifies how agents evolve at time step  $t$ . Each entry  $\mathbf{M}_{ij}^{t}$  describes whether agent  $i$  at step  $t$  progresses into agent  $j$  at step  $t + 1$ . The influence of each agent at step  $t$  on the reward, in terms of all successor agents is responsible for up to step  $t + k$ , can then be computed via the matrix multiplication  $\mathbf{M}^{t,k}\coloneqq \mathbf{M}^t\mathbf{M}^{t + 1}\ldots \mathbf{M}^{t + k - 1}$ .
The usual objective in RL is to find a policy  $\pi : \mathcal{O}^N \times \mathcal{A}^N \to [0,1]$  that maximizes the return, i.e., the expected discounted cumulative future reward  $J^t \coloneqq \mathbb{E}_{\pi(\mathbf{a}|\xi(\mathbf{s}))}\left[\sum_{k=0}^{\infty} \gamma^k r(\mathbf{s}^{t+k}, \mathbf{a}^{t+k})\right]$  for a discount factor  $\gamma$  and scalar reward  $r(\mathbf{s}^{t+k}, \mathbf{a}^{t+k})$  at step  $t + k$ . Adapting this to varying numbers of agents within a single episode, as necessary for e.g., the refinement of mesh elements, yields
$$
J _ {i} ^ {t} := \mathbb {E} _ {\pi (\mathbf {a} | \xi (\mathbf {s}))} \left[ \sum_ {k = 0} ^ {\infty} \gamma^ {k} \left(\mathbf {M} ^ {t, k} \mathbf {r} \left(\mathbf {s} ^ {t + k}, \mathbf {a} ^ {t + k}\right)\right) _ {i} \right], \tag {1}
$$
for agent  $i$  at step  $t$ . Intuitively, this return represents the discounted sum of rewards of all agents that agent  $i$  is responsible for. We set  $V_{i}(\mathbf{s}^{t}) = \mathbf{r}(\mathbf{s}^{t},\mathbf{a}^{t})_{i} + \gamma \sum_{j}\mathbf{M}_{ij}^{t}V_{j}(T(\mathbf{s}^{t},\mathbf{a}^{t}))$  with  $\mathbf{a}\sim \pi (\xi (\mathbf{s}))$  for training the value function and derive the targets for  $Q$ -functions analogously.
Agents and Observations. Given a domain  $\Omega$  and a mesh  $\Omega^t := \{\Omega_i^t \subseteq \Omega | \bigcup_i \Omega_i^t = \Omega\}$ , we view each mesh element  $\Omega_i^t$  as an agent. Each element's action space comprises a binary decision to mark it for refinement. These markings are provided to a remesher, which refines all marked elements, yielding a finer mesh  $\Omega^{t+1} = \{\Omega_j^{t+1}\}_j$ . Here,  $\Omega_j^{t+1} := \Omega_i^t$  for no refinement and  $\Omega_j^{t+1} \subsetneq \Omega_i^t$  with  $\bigcup_j \Omega_j^{t+1} = \Omega_i^t$  if  $\Omega_i^t$  is refined. The remesher may also refine unmarked elements to assert a conforming solution [41], i.e., to make sure that elements of the mesh align with each other at the boundaries and interfaces to ensure continuity of solution variables between adjacent elements. We define the mapping for an agent to its successor agents as the indicator function  $\mathbf{M}_{ij}^t := \mathbb{I}(\Omega_j^{t+1} \subseteq \Omega_i^t)$ . In other words, an agent maps to all future agents that it spawns, or equivalently, an element is responsible for all sub-elements that it refines into over time. While we focus on mesh refinement in this work, this mapping can be extended to coarsening by setting, e.g.,  $\mathbf{M}_{ij}^t := \mathbb{I}(\Omega_j^{t+1} \subseteq \Omega_i^t) + \left( \mathbb{I}(\Omega_i^t \subsetneq \Omega_j^{t+1}) / \left( \sum_k \mathbb{I}(\Omega_k^t \subsetneq \Omega_j^{t+1}) \right) \right)$ .
For encoding the observations, we use an observation graph  $\mathcal{G}_{\Omega^t} = \mathcal{G} = (\mathcal{V},\mathcal{E},\mathbf{X}_{\mathcal{V}},\mathbf{X}_{\mathcal{E}})$ , which is a bidirectional directed graph with mesh elements as nodes  $\mathcal{V}$  and their neighborhood relation as edges  $\mathcal{E}\subseteq \mathcal{V}\times \mathcal{V}$ . Node and edge features of dimensions  $d_{\mathcal{V}}$  and  $d_{\mathcal{E}}$  are given as  $\mathbf{X}_{\mathcal{V}}:\mathcal{V}\to \mathbb{R}^{d_{\mathcal{V}}}$  and  $\mathbf{X}_{\mathcal{E}}:\mathcal{E}\rightarrow \mathbb{R}^{d_{\mathcal{E}}}$ . Further details can be found in Appendix B.
Reward. A good refinement strategy trades off the accuracy of the solution of the mesh  $\Omega^t$  with its total number of elements  $\Omega_i^t\in \Omega^t$ . We define an error per element as the difference in the solution of this element compared to a solution using a fine-grained reference mesh  $\Omega^{*}$  [28]. We consider  $\Omega^{*}$  to be optimal, but very slow to compute due to a large number of elements. However, we only require the reference mesh  $\Omega^{*}$  for the reward calculation, not during inference. For each element  $\Omega_i^t$  we then integrate over the evaluated differences of all midpoints  $p_{\Omega_m^*}\in \Omega_m^*$  of reference elements  $\Omega_m^*$  that fall into it, scaling each by the area  $\mathrm{Area}(\Omega_m^*)$  of its respective element. This procedure results in an error estimate
$$
\operatorname {e r r} \left(\Omega_ {i} ^ {t}\right) \approx \sum_ {\Omega_ {m} ^ {*} \subseteq \Omega_ {i} ^ {t}} \operatorname {A r e a} \left(\Omega_ {m} ^ {*}\right) \left| u _ {\Omega^ {*}} \left(p _ {\Omega_ {m} ^ {*}}\right) - u _ {\Omega^ {t}} \left(p _ {\Omega_ {m} ^ {*}}\right) \right|, \tag {2}
$$
where  $u_{\Omega^*}$  denotes the solution on the fine mesh and  $u_{\Omega^t}$  the solution on the current mesh. We note that this error estimate can be efficiently calculated using a  $k$ -d tree [85] and that it is generally applicable for a large range of PDEs. Problem-specific error estimates may be used instead to include domain knowledge. To get an error estimate that is consistent across different geometries, we normalize the error with the total error of the elements of the initial mesh  $\Omega^0$ , i.e.,  $\mathrm{err}(\Omega_i^t) = \hat{\mathrm{er}} (\Omega_i^t) / \sum_{\Omega_j^0\in \Omega^0}\hat{\mathrm{er}} (\Omega_j^0)$ . We then formulate a local reward per element as
$$
\mathbf {r} \left(\Omega_ {i} ^ {t}\right) := \frac {1}{\operatorname {A r e a} \left(\Omega_ {i} ^ {t}\right)} \left(\operatorname {e r r} \left(\Omega_ {i} ^ {t}\right) - \sum_ {j} \mathbf {M} _ {i j} ^ {t} \operatorname {e r r} \left(\Omega_ {j} ^ {t + 1}\right)\right) - \alpha \left(\sum_ {j} \mathbf {M} _ {i j} ^ {t} - 1\right), \tag {3}
$$
where  $\alpha$  is a hyperparameter that penalizes adding new elements.
This reward function evaluates whether a refinement decreases the overall error by enough to justify the extra resources required, with a reward of 0 for unrefined elements. Thus, the reward maximizes
$\alpha = 0.1$
$\alpha = 0.05$
$\alpha = 0.03$
Figure 3: Final refinements of ASMR on a Heat Diffusion problem for different values of the element penalty  $\alpha$ . All refinements focus on the relevant parts of the problem, and lower element penalties lead to more fine-grained meshes.
$\alpha = 0.01$
$\alpha = 0.005$