SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 12: Random training PDEs and ASMR refinements for the domain size generalization experiments. The solutions are on different scales as indicated by the colorbars on the left of the meshes.
We experiment with the abilities of ASMR to generalize to larger domains during inference for Poisson's equation. Such generalization is non-trivial due to the varying boundary conditions and complexities arising from domain scaling, yet extremely useful in practical scenarios where a policy is trained on small and relatively cheap training domains, and then applied to a much larger setups during inference. To generalize to larger domains, we modify the training PDEs to mimic larger mesh segments by altering boundary conditions and load functions. The means of the load function are sampled from a centered unit Gaussian, allowing components outside the mesh. We use domains with random holes for varied initial meshes and apply random Gaussian loads to selected boundary parts as 'inlets'. Exemplary training PDEs and ASMR refinements can be seen in Figure 12. We further add an L2 norm of  $3\mathrm{e} - 4$  to combat overfitting and omit the per-domain normalization in favor of a constant normalization factor of 100, i.e., use  $100 \cdot \hat{\mathrm{err}}(\Omega)$  instead of  $\mathrm{err}(\Omega)$  in Equation 3. These modifications can be seen as data augmentation and only affect the training environments without changing the ASMR algorithm. We evaluate the resulting policy on larger, spiral-shaped domains with initial elements of the same size as the evaluation domains. Figure 13 shows that ASMR consistently provides high-quality refinements as the domain increases size (left) while leading to more and more significant speedups when compared to the reference uniform refinement (right). Here, we use a spiral-shaped domain and load functions with 16 randomly placed components. A slice of these figures for a normalized squared error of 0.001 is provided in the main paper in Table 1.
Figure 14 shows how the same procedure scales to inference on a spiral mesh of size  $20 \times 20$  with a load function with 81 components. The refined mesh has more than 50 000 elements, which is several times larger than any refinement shown by previous work. Creating and solving this mesh using ASMR is roughly 100 times faster than solving the fine-grained reference  $\Omega^{*}$ . A close-up for the marked region is shown on the right side of Figure 15. The left side of Figure 15 compares ASMR trained on the generalization environments with the setup used throughout the paper, showing that the additional generalization capabilities only lead to a marginal decrease in performance on the original evaluation PDEs.
E.2 Same-scale Generalization Capabilities
We additionally visualize ASMR on Poisson's equation on domains of size  $1 \times 1$ , i.e., of the same size that is seen during training. Here, we utilize the regular training environments without the above augmentation and perform inference on 3 different domain types used in throughout the main experiments, plus a simple rectangular domain  $\Omega = (0,1)^2$  and randomly generated trapezoids. We sample 3 random domains per class, and use Gaussian Mixture Model load functions with 1, 3 and 5 components respectively. Figure 16 shows refinements of an ASMR policy with  $\alpha = 0.0075$  for the resulting  $3 \times 5$  problems. We find that ASMR generalizes across domains and load functions, which is likely a result of the Swarm RL setting, where each mesh element is governed by its own agent.
E.3 Runtime Comparison.
Finally, we compare the wallclock-time of our approach with that of directly computing the fine-grained uniform mesh  $\Omega^{*}$  on the evaluation PDEs. For ASMR, we measure the cumulative time of
Poisson - Domain Size Generalization
Figure 13: (Left) Pareto plot of the normalized squared error for ASMR evaluated on domains of size  $N \times N$ . While the initial number of elements increases with the size of the domain, fewer relative elements are needed to achieve a given error threshold, suggesting that there are fewer elements with a significant error in larger domains. (Right) Wallclock-time in seconds of ASMR (points, solid lines) for different numbers of elements compared to the uniform reference  $\Omega^*$  (dashed lines). ASMR provides high-quality refinements across different domain sizes, achieving larger and larger speedups when compared to the uniform mesh as the size of the domain increases.
creating an initial coarse mesh, iteratively solving the problem on this mesh, computing the resulting observation graphs after every step, feeding each observation graph to the policy to obtain a set of actions, and using each set of actions to refine the mesh a total of  $T = 6$  times. For the uniform mesh, we simply measure the time it takes to refine the coarse mesh 6 times and to subsequently solve the problem on the resulting mesh. We use a single 8-Core AMD Ryzen 7 3700X Processor for all measurements. Figure 17 shows the results for all tasks. We find that our approach is always significantly faster than computing the fine-grained mesh despite the comparatively large computational overhead. Further, the final resolution of the refined mesh produced by our method trades off the wallclock-time of the method, meaning that ASMR can be trained to generate coarser or finer meshes depending on task-specific computational budgets. Notably, for the Stokes flow equations, which use  $P_{1} / P_{2}$  Taylor-Hood-elements, our method is more than 30 times faster than  $\Omega^{*}$  even for highly refined final meshes. Since the Local Oracle baseline requires the calculation of  $\Omega^{*}$  and otherwise follows a similar iterative refinement procedure, its runtime is dominated by  $\Omega^{*}$ .
Figure 14: Visualization of a refinement produced by the policy of Figure 12 on a  $20 \times 20$  spiral domain. The presented mesh has 53 189 elements and ASMR produces it in about 10 seconds on a regular CPU, whereas a uniform refinement  $\Omega^{*}$  takes more than 20 minutes. A close-up of the marked region is shown on the right side of Figure 15. ASMR provides highly accurate refinements for domains that are significantly larger than those seen during training.
F.1 General Hyperparameters
We use the same hyperparameters across all methods and environments unless mentioned otherwise.
PPO. We largely follow the suggestions of [98] for our PPO parameters. We train each PPO policy for a total of 400 iterations. In each iteration, the algorithm samples 256 environment transitions and then trains on them for 5 epochs with a batch size of 32. The value function loss is multiplied with a factor of 0.5 and we clip the gradient norm to 0.5. The policy and value function clip ranges are chosen to be 0.2. We normalize the observations with a running mean and standard deviation. The discount factor is  $\gamma = 0.99$  and advantages are estimated via Generalized Advantage Estimate [99] with  $\lambda = 0.95$ . We compute an agent's advantage by subtracting the agent-wise value estimates from the combination of local and global returns in Equation 4.
DQN. For DQN-based approaches, we instead train for  $24 * 400 = 9600$  steps, where each step consists of executing an environment transition and then drawing a batch of 32 samples from the replay buffer for a single gradient update. We additionally draw 500 initial random replay buffer samples before the first training step. We keep 10000 transitions in the replay buffer, since each transition represents a full mesh and an action on each graph element. We experimented with both larger replay
Figure 15: (Left) Pareto plot of the normalized squared error for ASMR trained on regular (blue) and generalizing (green) environments evaluated on  $1 \times 1$  evaluation PDEs. (Right) A local region of the mesh in Figure 14. Adding data augmentation in the training environments allows generalization to significantly larger meshes during inference at the cost of slightly decreased refinement quality.
Buffers and more training steps in preliminary experiments, finding that neither significantly improve performance, but may lead to very long runtimes and large memory requirements. During training, we draw actions using a Boltzmann distribution over the predicted Q-values per agent, where we linearly decrease the temperature of the distribution from 1 to 0.01 in the first 4800 steps. We find that this action selection strategy leads to more correlated actions when compared to an epsilon greedy action sampling, which stabilizes the training for our iterative mesh refinement problems. We update the target networks using Polyak averaging at a rate of 0.99 per step. Further, we follow previous work [100] and include a number of common improvements for DQNs in our implementation. These are double Q-learning [101],ueling Q-networks [102] and prioritized experience replay [103].
Neural Networks. All networks are implemented in PyTorch [104] and trained using the ADAM optimizer [105] with a learning rate of 3.0e-4 unless mentioned otherwise. All MLPs use 2 hidden layers and a latent dimension of 64. We use separate MPNs for the policy and the value function. Each MPN consists of 2 message passing steps, where each update function is represented as an MLP with LeakyReLU activation functions. The policy and value function heads are additional MLPs with tanh activation functions acting on the final latent node features of the MPN. All message aggregations  $\oplus$  are mean aggregations. Additionally, we apply Layer Normalization [106] and Residual Connections [107] independently for the node and edge features after each message passing step.
F.2 Baseline-Specific Parameters
For Single Agent, we use a maximum refinement depth of 10 refinements per element to avoid numerical instabilities during simulation, skipping actions that try to refine elements that have been refined too often. We consider environment sequences of up to  $T = 400$  steps since the method marks only one element at a time. For Sweep, the agent is placed on a random mesh element for each training step and may decide not to refine this element, resulting in no change in the mesh. Here, we follow the proposed hyperparameters for this approach and train each rollout for 200 steps. As this approach is based on purely local agents, we adapt our input features per element to consist of our regular node features, the global resource budget proposed by the authors, the mean solution and area of the element's neighbors and the average distance to them. The global budget is controlled via a maximum number of elements  $N_{\mathrm{max}}$ , allowing to get refinements of different granularity. To accommodate for less overall changes in the mesh, we increase the number of environment transitions
Figure 16: ASMR refinements for different domains and Poisson's equation with a Gaussian Mixture Model load with 1, 3 and 5 components. Even though the ASMR policy is only trained on L-shaped domains with 3 components in the load function and no data augmentation, it generalizes to different domains and loads.
of PPO to 512, and the number of DQN steps to  $96*400 = 38400$ . Finally, we use a learning rate of  $1.0\mathrm{e - 5}$  instead of  $3.0\mathrm{e - 4}$  for the DQN variant of  $VDGN$  to stabilize its training.
F.3 Refinement Hyperparameters
The AMR methods considered in this work use different parameters to control the granularity of the final refined mesh. ASMR and  $VDGN$  use an element penalty  $\alpha$ , while Sweep considers an element budget  $N_{\mathrm{max}}$ . Single Agent varies the number of rollout steps  $T$ . For each learned method and task, we choose 10 different values for the refinement parameter that showcase a wide range of final mesh resolutions. For the Oracle, Maximum Oracle, and ZZ Error Heuristics we instead cover a range of up to 100 parameter thresholds  $\theta$ , yielding one aggregated evaluation result per threshold.
Table 2 lists the different ranges for these parameters for the different tasks. For stability purposes, we set a maximum number of 20 000 elements during training for all experiments except for the Sweep baseline, as this baseline uses its own element budget instead. If this number is surpassed, a constant penalty of 1000 is subtracted from the reward and the episode terminates early.
Figure 17: Wallclock-time in seconds of ASMR (points, solid lines) for different numbers of elements compared to the uniform reference  $\Omega^{*}$  (dashed lines). Each task is denoted by its color, and the full lines represent a quadratic regression of the wallclock-time of our method for different numbers of final elements. On average,  $\Omega^{*}$  contains about  $10^{5}$  elements, with concrete numbers varying depending on the domain. Our approach is significantly faster than the reference for all tasks, achieving speedups of factor 2-30 for reasonable numbers of elements.
Table 2: Ranges for the different refinement hyperparameters for all tasks. ASMR and  $VDGN$  apply an element penalty  $\alpha$ , but only ASMR scales the area of each element with its area in Equation 3. Sweep uses an element budget  $N_{\mathrm{max}}$ . Single Agent varies the number of rollout steps  $T$ , and the Oracle Error, Maximum Oracle Error and ZZ Error Heuristics use of different error thresholds  $\theta$ .
Method	Task
Laplace	Poisson	Stokes Flow	Lin. Elast.	Heat Diff.
ASMR (α)	[0.01, 0.3]	[0.002, 0.1]	[0.006, 0.15]	[0.01, 0.15]	[0.003, 0.3]
VDGN(-like) (α)	[2e-5, 5e-2]	[2e-5, 5e-2]	[3e-4, 5e-3]	[1e-5, 1e-2]	[5e-6, 5e-3]
Sweep (Nmax)	[200, 3000]	[200, 3000]	[200, 3500]	[500, 6000]	[400, 5000]
Single Agent (T)	[25, 400]	[25, 400]	[25, 400]	[25, 400]	[25, 400]
Oracle Error (θ)	[0.25, 1.00]	[0.1, 1.0]	[0.16, 1.0]	[0.02, 1.0]	[0.03, 1.0]
Max. Oracle Err. (θ)	[0.20, 1.0]	[0.2, 1.0]	[0.1, 1.0]	[0.01, 1.0]	[0.02, 1.0]
ZZ Error (θ)	[0.001, 1]	[0.002, 1]	[0.001, 1]	[0.001, 1]	[0.001, 1]
We provide additional visualizations for our method on all tasks, and for all methods on the Poisson task. All visualizations show the final refined mesh of the respective method for 5 different refinement levels on 3 randomly selected PDEs. For the RL methods, all policies are taken from the first repetition of the 10 random seeds conducted for the respective experiment.
G.1 ASMR Refinements
We visualize exemplary refinements of ASMR policies for all considered tasks in Figures 18 (Laplace's equation), 19 (Poisson's equation), 20 (Stokes equation), 21 (Linear Elasticity), and 22 (Heat Diffusion). Across all tasks, ASMR is able to provide highly accurate refinements for different numbers of total elements.
G.2 Baseline Comparisons
Figure 23 shows refinements for Single Agent for different total timesteps  $T$ , Figure 24 presents  $VDGN$  with different  $\alpha$  values. Figure 25 visualizes refinements of Sweep for a varying number of maximum elements  $N_{\mathrm{max}}$ . Figures 26, 27 and 28 show refinements of the Oracle, Maximum Oracle and ZZ Error Heuristics for different values of the threshold  $\theta$ .
The visualizations show that the RL baselines struggle to provide consistent high-quality refinements for different mesh resolutions. The Single Agent baseline sometimes focuses on uninteresting regions of the mesh or refines the same area too often. The  $VDGN$ -like baseline performs well in some cases, but collapses to no refinements or fully uniform refinements for some evaluation PDE. Sweep provides almost uniform refinements for most PDEs and element budgets, likely as a result of the misalignment in the environment transitions between training and inference. The Heuristics greedily refine the elements with the largest error estimates in their respective metric, regardless of the resulting decrease in error. This behavior generally leads to locally accurate refinements, but fails to effectively decrease the mesh error in some cases. Additionally, the heuristics act locally, which causes potential issues for PDEs with global dependencies [92] and conforming refinements. The ZZ Error Heuristic sometimes misses regions of interest, but provides a much smoother refinement than the Oracle Error Heuristic, which can be beneficial for the error reduction in some cases.
G.3 Element Markings
Figure 29 visualizes a full rollout of our method, including the markings of the elements after every step.
Figure 18: Final refined meshes of ASMR for randomly sampled PDEs for the Laplace equation for different element penalties  $\alpha$ .
$$
\alpha = 0. 2
$$
$$
\alpha = 0. 1
$$
$$
\alpha = 0. 0 7 5
$$
$$
\alpha = 0. 0 3
$$
$$
\alpha = 0. 0 1 5
$$
$\alpha = 0.03$
$\alpha = 0.01$
Figure 19: Final refined meshes of ASMR for randomly sampled PDEs for the Poisson equation for different element penalties  $\alpha$ .
$\alpha = 0.2$
Figure 20: Final refined meshes of ASMR for randomly sampled PDEs for the Stokes flow task for different element penalties  $\alpha$ .
$\alpha = 0.02$
Figure 21: Final refined meshes of ASMR for randomly sampled PDEs for the linear elasticity task for different element penalties  $\alpha$ . The visualizations show the deformed meshes, which are originally L-shaped.
Figure 22: Final refined meshes of ASMR for randomly sampled PDEs for the non-stationary heat diffusion task for different element penalties  $\alpha$ .
$T = 50$
$T = 100$
$T = 200$
$T = 300$
Figure 23: Final refined meshes of the Single Agent baseline for the Poisson equation on randomly sampled PDEs for different environment rollout lengths  $T$ .
$T = 400$
$\alpha = 0.05$
Figure 24: Final refined meshes of the  $VDGN$ -like baseline for the Poisson equation on randomly sampled PDEs for different element penalties  $\alpha$ .
$\alpha = 0.0005$
$\alpha = 5e - 05$
Figure 25: Final refined meshes of the Sweep baseline for the Poisson equation on randomly sampled PDEs for different maximum numbers of elements  $N_{\mathrm{max}}$ .
$$
N _ {\max } = 2 0 0
$$
$$
N _ {\max } = 5 0 0
$$
$$
N _ {\max } = 7 5 0
$$
$$
N _ {\max } = 1 5 0 0
$$
$$
N _ {\max } = 2 0 0 0
$$
Figure 26: Final refined meshes of the Local Oracle baseline for the Poisson equation on randomly sampled PDEs for different error thresholds  $\theta$ .
$$
\theta = 0. 9
$$
$$
\theta = 0. 6
$$
$$
\theta = 0. 5
$$
$$
\theta = 0. 4
$$
$$
\theta = 0. 2
$$
$\theta = 0.9$
$\theta = 0.6$
$\theta = 0.5$
$\theta = 0.4$
Figure 27: Final refined meshes of the Local Maximum Oracle baseline for the Poisson equation on randomly sampled PDEs for different error thresholds  $\theta$ .
$\theta = 0.2$
$\theta = 0.5$
Figure 28: Final refined meshes of the ZZ Error for the Poisson equation on randomly sampled PDEs for different error thresholds  $\theta$ .
$\theta = 0.25$
$\theta = 0.125$
$\theta = 0.0625$
$\theta = 0.03125$
Figure 29: Visualization of a full rollout of our method on a Poisson task, including the markings of the elements after every step. The figures in the first and third row show the markings, and the second and fourth row show the resulting refined meshes.