SlowGuess/NeurIPS_2023_OCR huggingface dataset:
[92] Walter A Strauss. *Partial differential equations: An introduction*. John Wiley & Sons, 2007.
[93] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304-29320, 2021.
[94] Petar Velickovic, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pietro Lio, and Yoshua Bengio. Graph attention networks. stat, 1050:4, 2018.
[95] Alfio Quarteroni and Silvia Quarteroni. Numerical models for differential problems, volume 2. Springer, 2009.
[96] Volker John et al. Finite element methods for incompressible flow problems, volume 51. Springer, 2016.
[97] Olek C Zienkiewicz, Robert Leroy Taylor, and Jian Z Zhu. The finite element method: its basis and fundamentals. Elsevier, 2005.
[98] Marcin Andrychowicz, Anton Raichuk, Piotr Stanczyk, Manu Orsini, Sertan Girgin, Raphaël Marinier, Leonard Hussenot, Matthieu Geist, Olivier Pietquin, Marcin Michalski, Sylvain Gelly, and Olivier Bachem. What matters for on-policy deep actor-critic methods? a large-scale study. In International Conference on Learning Representations, 2021.
[99] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438, 2015.
[100] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
[101] Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q-learning. In Proceedings of the AAAI conference on artificial intelligence, volume 30, 2016.
[102] Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. *Dueling network architectures for deep reinforcement learning*. In International conference on machine learning, pages 1995–2003. PMLR, 2016.
[103] Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay. In *Yoshua Bengio and Yann LeCun*, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016.
[104] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32, 2019.
[105] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
[106] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. stat, 1050:21, 2016.
[107] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
A Message Passing Network Architecture
Given a graph  $\mathcal{G} = (\mathcal{V},\mathcal{E},\mathbf{X}_{\mathcal{V}},\mathbf{X}_{\mathcal{E}})$ , Message Passing Networks (MPN) [33, 38, 39] are GNNs consisting of  $L$  Message Passing Steps. Each step  $l$  receives the output of the previous step and updates the features  $\mathbf{X}_{\mathcal{V}}$ ,  $\mathbf{X}_{\mathcal{E}}$  for all nodes  $v\in \mathcal{V}$  and edges  $e\in \mathcal{E}$ . Using linear embeddings  $\mathbf{x}_v^0$  and  $\mathbf{x}_e^0$  of the initial node and edge features, the  $l$ -th step is given as
$$
\mathbf {x} _ {e} ^ {l + 1} = f _ {\mathcal {E}} ^ {l} (\mathbf {x} _ {v} ^ {l}, \mathbf {x} _ {u} ^ {l}, \mathbf {x} _ {e} ^ {l}), \mathrm {w i t h} e = (u, v),
$$
$$
\mathbf {x} _ {v} ^ {l + 1} = f _ {\mathcal {V}} ^ {l} (\mathbf {x} _ {v} ^ {l}, \bigoplus_ {e = (v, u) \in \mathcal {E}} \mathbf {x} _ {e} ^ {l + 1}).
$$
The operator  $\oplus$  is a permutation-invariant aggregation such as a sum, max, or mean operator. Each  $f^l$  is a learned function that we generally parameterize as a simple MLP. The network's final output is a learned representation  $\mathbf{x}_v^L$  for each node  $v \in \mathcal{V}$ .
B Systems of Equations
In its most general form, the FEM is used to approximate the solution  $\pmb{u}(\pmb{x})$  that satisfies the weak formulation  $\forall \pmb{x} : \forall \pmb{v}(\pmb{x}) : a(\pmb{u}(\pmb{x}), \pmb{v}(\pmb{x})) = l(\pmb{v}(\pmb{x}))$  of the underlying system of equations for the set of test functions  $\pmb{v}$ . In the following, we describe the specific equations and boundary conditions used for our experiments.
B.1 Laplace's Equation
Let  $\Omega$  be a domain with an inner boundary  $\partial \Omega_0$  and an outer boundary  $\partial \Omega_1$ . We seek a solution  $u(\pmb{x})$  that satisfies the weak formulation of the Laplace Equation
$$
\int_ {\Omega} \nabla u (\boldsymbol {x}) \cdot \nabla v (\boldsymbol {x}) d \boldsymbol {x} = 0
$$
for all test functions  $v(\pmb{x})$ . Additionally, the solution has to satisfy the Dirichlet boundary conditions
$$
u (\boldsymbol {x}) = 0, \boldsymbol {x} \in \partial \Omega_ {0} \quad \text {a n d} \quad u (\boldsymbol {x}) = 1, \boldsymbol {x} \in \partial \Omega_ {1}.
$$
We use a unit square  $(0,1)^2$  for the outer boundary  $\partial \Omega_0$  of the domain and add a randomly sampled square hole, whose borders are considered to be the inner boundary  $\partial \Omega_1$ . The size of the hole is sampled from the uniform distribution  $U(0.05, 0.25)^2$ , and its mean position is sampled from  $U(0.2, 0.8)^2$ . We add the closest distance to the inner boundary as an additional node feature.
B.2 Poisson's Equation
The weak formulation of the considered Poisson problem is given as
$$
\int_ {\Omega} \nabla u (\boldsymbol {x}) \cdot \nabla v (\boldsymbol {x})   d \boldsymbol {x} = \int_ {\Omega} f (\boldsymbol {x}) v (\boldsymbol {x})   d \boldsymbol {x} \quad \forall v.
$$
Here,  $f(\pmb{x}) : \Omega \to \mathbb{R}$  denotes the load function and  $v(\pmb{x})$  the test function. In addition to the weak formulation, the solution must be zero on the boundary  $\partial \Omega$  of the domain  $\Omega$ . We model Poisson's Equation on L-shaped domains  $\Omega$ , using a rectangular cutoff whose lower left corner is sampled from  $p_0 \sim U(0.2, 0.95)^2$ , resulting in a domain  $\Omega = (0,1)^2 \setminus (p_0 \times (1,1))$ . On this domain, we sample a Gaussian Mixture Model with 3 components. Each component's mean is sampled from  $U(0.1, 0.9)^2$ , and we use rejection sampling to ensure that all means lie within the domain. The components' covariances are determined by first drawing diagonal covariances, where each dimension is drawn independently from a log-uniform distribution  $\exp(U(\log(0.0003, 0.003)))$ . The diagonal covariances are then rotated by a random angle in  $U(0, 180)$  to produce Gaussians with a full covariance matrix. The component weights are drawn from the distribution  $\exp(N(0,1)) + 1$  and subsequently normalized, where the 1 in the end is used to ensure that all components have relevant weight. The evaluation of the load function  $f$  at the respective face midpoint is added as a node feature.
B.3 Stokes flow
Let  $\pmb{u}(\pmb{x})$  be the velocity field and  $p(\pmb{x})$  the pressure field. We consider a Stokes flow of a fluid through a channel. Therefore, we seek a solution  $u$  and  $p$ , which satisfy the weak formulation of the Stokes flow without a forcing term
$$
\nu \int_ {\Omega} \nabla \boldsymbol {v} \cdot \nabla \boldsymbol {u} d \boldsymbol {x} - \int_ {\Omega} (\nabla \cdot \boldsymbol {v}) p d \boldsymbol {x} = 0 \quad \forall \boldsymbol {v}
$$
$$
\int_ {\Omega} (\nabla \cdot \boldsymbol {u}) q   d \boldsymbol {x} = 0 \quad \forall q,
$$
wherein  $\pmb{v}(\pmb{x})$  and  $q(\pmb{x})$  denote the test functions [95]. We define the inlet-profile as
$$
\boldsymbol {u} (x = 0, y) = u _ {\mathrm {P}} y (1 - y) + \sin (\varphi + 2 \pi y).
$$
At the outlet, the gradient of velocity  $\nabla u(x = 1,y) = 0$  is set to zero. Additionally, we assume a no-slip condition  $\pmb{u} = \mathbf{0}$  at all boundaries except for the inlet and the outlet. For stability purposes, we use  $P_{1} / P_{2}$  Taylor-Hood-elements, i.e., quadratic shape functions for the velocity and linear shape functions for the pressure [96]. We sample the quadratic part  $u_{\mathrm{P}}$  of the velocity inlet from a log-uniform distribution  $\exp (U(\log (0.5,2)))$ . The class of domains uses a unit square for the outer boundary and 3 rhomboid holes with length 0.4 and height 0.2 whose centers are set to  $y\in \{0.2,0.5,0.8\}$  in y-direction and randomly sampled from  $U(0.3,0.7)$  in x-direction. We optimize the meshes for the prediction accuracy of the velocity in  $x$  and  $y$  direction and calculate the overall error as the norm of these errors.
B.4 Linear Elasticity
We are looking for the steady-state deformation of a solid under stress, due to displacements at the boundary of the part  $\partial \Omega$ . Here, we are interested in both the norm of the deformation and the norm of the stress. The weak formulation of the considered problem on the domain  $\Omega$  without body forces is given as [97]
$$
\int_ {\Omega} \boldsymbol {\sigma} (\varepsilon (\boldsymbol {u})): \varepsilon (\boldsymbol {v})   \mathrm {d} \boldsymbol {x} = 0.
$$
Here,  $\pmb{u}(\pmb{x})$  is the displacement field,  $\pmb{v}(\pmb{x})$  is the test function, and  $\varepsilon(\pmb{u}) = \frac{1}{2} (\nabla \pmb{u} + (\nabla \pmb{u})^\top)$  is the strain tensor.  $\sigma(\varepsilon)$  is the stress tensor, which is given as  $\sigma(\varepsilon) = 2\mu \varepsilon + \lambda \mathrm{tr}(\varepsilon)\pmb{I}$  in a linear-elastic and isotropic case. The Lamé parameters  $\lambda = \frac{E\nu}{(1 + \nu)(1 - 2\nu)}$  and  $\mu = \frac{E}{2(1 + \nu)}$  can be calculated with the problem specific Young's modulus  $E = 1$  and the Poisson ratio  $\nu = 0.3$ . The displacement  $\pmb{u}(x = 0,y) = \pmb{u}_0$  on the left side of the boundary is specified by a task-dependent parameter  $\pmb{u}_0$ , whereas the displacement  $\pmb{u}(x = L,y) = 0$  is set to zero on the right boundary. The stress  $\pmb{\sigma} \cdot \pmb{n} = 0$  is zero normal to the boundary at both the top and bottom of the part. We use the same class of L-shaped domains as in the Poisson problem in Section B.2 and set  $u_P$  by drawing a random angle from  $U[0,\pi]$  to pull on the domain from different angles, and add random magnitude from  $U(0.2,0.8)$ . We add the task-dependent displacement  $u_P$  as a feature to all nodes. We are interested in the norm of the displacement field  $u$  and the resulting Von-Mises stress, giving us a 2-dimensional objective. We weight both dimensions equally in the reward.
B.5 Non-stationary Heat Diffusion
We consider a non-stationary thermal diffusion problem defined by the weak formulation
$$
\int_ {\Omega} \frac {\partial u}{\partial t} d x + \int_ {\Omega} a \nabla u \cdot \nabla v d x = \int_ {\Omega} f v d x \quad \forall v,
$$
wherein  $u$  denotes the temperature,  $v$  the test function,  $a$  the thermal diffusivity and  $f$  a heat distribution, given as
$$
f = q \exp \left(- 1 0 0 \left(\left(x - x _ {p} (\tau)\right) + \left(y - y _ {p} (\tau)\right)\right)\right).
$$
The position of the maximum heat entry  $\pmb{p}_{\tau}(\tau) = (x_{p}(\tau), y_{p}(\tau))$  is changing over time, while its magnitude is scaled by a factor  $q$ . The temperature  $u \in \partial \Omega$  is set to zero on all boundaries. For the
time-integration, the implicit Euler method is applied. We use a total of  $\tau_{\mathrm{max}} = 20$  time steps in  $\{0.5,\dots ,10\}$ , a scaling factor of  $q = 1000$  and a diffusivity  $a = 0.001$ . The position of the heat source at step  $\tau$  is linearly interpolated as  $\pmb{p}_{\tau} = \pmb{p}_0 + \frac{\tau}{\tau_{\mathrm{max}}} (\pmb{p}_{\tau_{\mathrm{max}}} - \pmb{p}_0)$ , where the start and goal positions  $\pmb{p}_0$  and  $\pmb{p}_{\tau_{\mathrm{max}}}$  are randomly drawn from the domain. To create our domains, we start with 10 points that are equidistantly placed on a circle with center  $(0.5,0.5)$  and radius 0.4. Each point is distorted by a random value drawn from  $U(-0.2,0.2)^2$ . We then normalize the resulting points to be in  $(0,1)^2$  and calculate the convex hull. The result is a family of convex polygons with up to 10 vertices. We measure the error and solution of the final simulation step, and provide the distance to the start and end position of the heat source as additional node features for each element.
All experiments are repeated for 10 random seeds with randomized PDEs and network parameters. All domains are normalized to be in  $(0,1)^2$  unless mentioned otherwise. The initial meshes are created using meshpy<sup>2</sup>. For practical purposes, we add an element threshold  $\beta_{\mathrm{max}}$  in our environments, and terminate an episode with a large negative reward when this threshold is exceeded. We train all policies on 100 training PDEs and evaluate the resulting final policies on 100 different evaluation PDEs that we keep consistent across random seeds for better comparability. All experiments are run for up to 2 days on 8 cores of an Intel Xeon Platinum 8358 CPU. In terms of total compute, we train 4 different learned methods, namely the 3 RL baselines and our method, on 5 separate tasks. Each experiment is repeated for 10 different target mesh resolutions and 10 repetitions, resulting in  $5 \cdot 4 \cdot 10 \cdot 10 = 2000$  main experiments. Additionally, we use a similar amount of compute for the combined ablations, preliminary experiments and heuristics.
C.2 Maximum Reward
Equation 3 scales the reduction in error of each element by its area. This modification encourages the policy to focus on smaller elements, effectively shifting the objective from an reduction in average error across the mesh to a minimization of error densities. An alternate way to phrase this objective is to make the reward depend on the reduction in maximum error per element. For this, we modify the error estimate per element of Equation 2 to read
$$
\mathbf {\hat {e r}} (\Omega_ {i} ^ {t}) \approx \max _ {\Omega_ {m} ^ {*} \subseteq \Omega_ {i} ^ {t}} \left| u _ {\Omega^ {*}} (p _ {\Omega_ {m} ^ {*}}) - u _ {\Omega^ {t}} (p _ {\Omega_ {m} ^ {*}}) \right|,
$$
and subsequently drop the area scaling and replace the sum in Equation 3 with a maximum, i.e.,
$$
\mathbf {r} ^ {\prime} \left(\Omega_ {i} ^ {t}\right) := \left(\operatorname {e r r} \left(\Omega_ {i} ^ {t}\right) - \max  _ {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).
$$
While conceptually simpler than our reward formulation, evaluating the decrease in maximum error only optimizes this objective, which may result in worse meshes when looking at, e.g., the mean error. The left side of Figure 6 compares this alternate reward formulation to that of ASMR. We find that both reward schemes perform similarly. Therefore, we opt for the reward function defined in Equation 3 for simplicity and easier comparison with the baseline methods.
D.1 Proximal Policy Optimization and Deep Q-Networks.
The left side of Figure 7 shows results on Poisson's Equation for PPO and DQN as the RL backbone for all learned methods. We find that PPO outperforms DQN, suggesting that an on-policy objective is favorable for the changing observation and action spaces of AMR. We use a mean instead of a sum for the agent mapping of the targets of the  $Q$ -values for the DQN experiments with ASMR as this
$^{2}$ https://github.com/inducer/meshpy
Figure 7: Pareto plot of normalized squared errors and number of final mesh elements on Poisson's equation for (Left) PPO and DQN for all RL baselines and (Right) GAT and MPN for ASMR and  $VDGN(-\text{like})$ . PPO generally results in better performance. ASMR works well for both MPNs and GATs, while the  $VDGN-$ like baseline is more stable when using MPNs.
seems to experimentally increase training stability. For the  $VDGN$ -like variant that uses PPO, we factorize the value function instead of the Q-function, i.e., we define the value function of the full mesh as the sum of value functions of the individual mesh elements. We choose PPO for all other experiments as it leads to better performance for all methods.
D.2 Message Passing and Graph Attention Networks.
The right side of Figure 7 compares MPNs and GATs for ASMR and  $VDGN$ . For ASMR, the performance between MPNs and GATs is comparable, while  $VDGN$  seems to be more stable and produce better refinements when using MPNs. We use MPNs for the other experiments as it seems to benefit  $VDGN$  while decreasing the performance of our method.
D.3 Initial Meshes for the ZZ Error.
Figure 8 compares the ZZ Error Heuristic when directly applied to the initial mesh to variants that instead start each refinement procedure by uniformly refining either once or twice. We find that the method greatly benefits from two initial uniform refinements, likely because the heuristic may not detect gradients for interesting parts of the domain if the corresponding elements are too coarse. Given these results, we use the twice refined version for all experiments, noting that the RL based methods avoid having to tune the initial mesh by design.
D.4 Target Mesh Resolutions
All RL methods use some parameter to control the number of target elements of the final refined mesh. ASMR and  $VDGN$  use an element penalty  $\alpha$ , Sweep uses a budget  $N_{\mathrm{max}}$ , and Single Agent different numbers of rollout steps  $T$ . We visualize evaluations for different target resolutions in Figure 9. The results indicate that ASMR provides meshes with consistent numbers of elements for a given target resolution, while the other RL methods produce meshes with inconsistent numbers of elements over target resolutions. The concrete target resolution parameters for all experiments are found in Table 2.
D.5 Ablations.
Node Features. ASMR utilizes both task-dependent information, such as the evaluation of the load function for Poisson's Equation, and the local solution  $u(x)$  per mesh element as part of its observation graph. Here, we experiment how the performance is affected if either of these features is left out. Additionally, we consider a variant where we include explicit  $(x,y)$  positions of each element midpoint as node features. The results are shown on the left of Figure 10. We find that
Poisson - ZZ Error Heuristic
Figure 8: Pareto plot of normalized squared errors and number of final mesh elements on Poisson's equation for the Zienkiewicz-Zhu Error Estimator (ZZ Error) Heuristic when using either no, 1, or 2 initial uniform refinements. The ZZ Error Heuristic produces better refinements when provided with a finer initial mesh at the cost of not being able to produce meshes with very few elements.
Elements  $(\times 10^{3})$
Figure 9: Pareto plot of normalized squared errors and number of final mesh elements for the linear elasticity task for all RL methods. Small blue dots indicate a policy trained on a coarse target mesh resolution, which corresponds to large element penalties  $\alpha$  for ASMR and  $VDGN$ -like, a small budget  $N_{\mathrm{max}}$  for Sweep and a low number of rollout steps  $T$  for Single Agent. Large red dots correspond to a finer target meshes, and the medium-sized purple dots interpolate between the two. Details on the target resolution parameters are found in Table 2. We find that ASMR provides high-quality refinements with consistent numbers of final mesh elements for any given target resolution, whereas the other methods yield poor-quality refinements (Sweep) or inconsistent results (VDGN-like, Single Agent) for similar target resolutions and different random seeds.
Figure 10: Pareto plot of normalized Top  $0.1\%$  of errors and number of final mesh elements for Poisson's Equation. (Left) Omitting either the solution or evaluation of the load function per element leads to a decrease in refinement quality. Explicit node positions slightly decreases performance, likely because they cause the observation to lose equivariance w.r.t., e.g., reflection. (Right) Performance is reduced for fewer training PDEs, but stabilizes around 100 PDEs. Interestingly, ASMR achieves better-than performance when using a single training PDE, meaning that it can generalize from a single training example. This ability is likely a result of our spatial problem formulation.
both the task-dependent features and the solution are important for the performance of our approach. Omitting positional features slightly improves performance, presumably because the features assign a fixed position to each mesh element, causing the observation graph to no longer be equivariant to rotation, translation and reflection. Interestingly, ASMR provides reasonable refinements even without solution information, suggesting that the RL algorithm is able to detect relevant regions of the PDE from just an encoding of the domain and the boundary conditions and forcing functions.
Number of Training PDEs Since calculating the fine-grained reference  $\Omega^{*}$  is slow for large meshes and complex tasks, we want to minimize the number of unique PDEs that we need during training. We use 100 PDEs in our other experiments, and additionally visualize results for 1, 10 and 1000 training PDEs on the right of Figure 10. We find that fewer than 100 PDEs lead to less stable and reliable results, and that there is only a minor advantage in using 1000 PDEs compared to our 100. Noticeably, a single training PDE results in suitable refinements, which hints at significant generalization capabilities that are likely granted by our spatial treatment of the underlying task.
D.6 Alternate Error Metrics
Section 5 evaluates all approaches on the normalized squared error of the mesh. This metric captures both the average error across the domain, leading to a low mean error, and outliers, thus punishing a high maximum error. Here, we additionally present normalized mean and maximum error metrics to provide a more thorough nuanced evaluation. The first directly quantifies the average absolute error of the mesh, which makes it easy to interpret and less sensitive to outliers. The maximum error metric measures the worst-case performance, which is crucial for applications where a single high-error prediction could be costly. Since the maximum remaining error is susceptible to outliers, we approximate it as the average of the Top  $0.1\%$  of errors of all integration points  $p_{\Omega_m^*}$ . For comparability across PDEs, we normalize both metrics by the respective error of the initial mesh  $\Omega^0$ .
Figure 11 displays the results for all tasks and both alternate metrics. The general trends for both metrics are consistent with that of the squared error in Section 5, with ASMR outperforming all learned baselines while being on par with or better than the Heuristics in most cases. The Oracle Error Heuristic performs particularly well on the mean error metric, as it selects elements with high integrated error for refinement. Conversely, the Maximum Oracle Error Heuristic excels on the top  $0.1\%$  error metric, as it specifically targets elements with a high maximum error. Notably, the mean error metric tends to favor more uniform meshes due to its lower sensitivity to outliers, enabling baselines like Sweep, which generally produce relatively uniform meshes, to yield better performance here when compared to the other metrics.
Figure 11: Pareto plot of (left) normalized mean errors and (right) normalized top  $0.1\%$  errors compared to number of final mesh elements across different tasks. Performance on both metrics is highly correlated and generally consistent with that of the squared error in Figures 4 and 5. ASMR outperforms all learned baselines on both metrics and all tasks.
E Generalization Capabilities and Runtime Experiments
E.1 Domain Size Generalization