Swarm Reinforcement Learning For Adaptive Mesh Refinement

Summary

This paper formulates h-adaptive mesh refinement (AMR) as a decentralized partially-observable Markov decision process (Dec-POMDP), and proposes a methods ASMR with parameter-sharing among agents and individual rewards to find refinement strategies. Refinement policies are parameterized by message-passing graph neural networks (MPN) and trained by Q-learning (specifically, DQN) and proximal policy optimization (PPO). Evaluation was conducted on various 2D elliptic partial differential equations (PDEs) with different domains, showing that the proposed approach is competitive with traditional error-based AMR heuristics and outperforms modified versions of existing learning-based approaches.

Strengths

Originality: One main novelty of this paper is the use of independent multi-agent reinforcement learning with individual rewards/returns for AMR. The proposed individual return is a novel solution to the posthumous credit assignment problem that arises in the AMR context, whereby an agent vanishes upon taking a refinement action. It also has the advantage of avoiding the difficulties of the regular multi-agent credit assignment problem that arises from using a single team reward. Quality: Overall, this paper covers most of the aspects of a high-quality applied paper, as it provides sufficient motivation and background for the application and provides a rigorous formulation and empirical evaluation. Clarity: The description of the problem, the proposed method, the notation, and experimental setup and results are all clearly written. See below for more comments about clarity. Significance: The results in this paper are significant as they provide a path to improving AMR for the finite element method, which is a critical tool in engineering and applied sciences for solving PDEs.

Weaknesses

The characterization of the scale of experiments in this paper, in contrast to previous work, should be made more precise. The abstract of this paper says that previous work "scale only to simple toy examples", which leads readers to expect that this paper deals with real-world problems on large domains. However, the experiments in this paper were conducted only on simple elliptic PDEs on meshes with up to thousands of elements, which appears to be on the same level of complexity and scale as previous work (e.g., AMR on simple hyperbolic equations with up to thousands of elements). The paper's Limitation section also acknowledges this fact, so the abstract and main paper should be made consistent. The proposed approach seems to be extendable to include coarsening actions, but this is not demonstrated in this paper and hence it is hard to judge whether the proposed approach holds promise if extended. There is existing work that support coarsening (refs 28,29), so one would expect this to be a requirement for this application area. One would expect that an optimal strategy for the moving heat source problem requires coarsening actions, since regions that were refined due to proximity to the source may need to be coarsened once the source has moved away from them. Experimental evaluation includes randomization over domains, boundary conditions and initial conditions. However, for trained policies to be useful in applications to real-world problems, one should show that trained policies can work on larger meshes and longer step counts than those seen in training. The proposed method uses global features such as the number of mesh elements and environment time step, which this leads to concerns about what would happen if the trained policies are run on test cases where the number of elements and step count go out of the training distribution. It was confusing to read the phrase "PPO version of VDGN" (line 245) and "VDGN...produce better results with PPO" (line 265), since PPO is a policy optimization method that is fundamentally different from value-based methods, so it is unclear how exactly PPO can be used in combination with VDGN. Upon looking closer, it appears that some of the learning-based baselines described in Section 4 do not match the methods in the previous work (refs 27,28), so it is misleading to label them as such. For the "Argmax" method, (ref 27) trains a stochastic policy on a discrete action space consisting of all the available elements, but the implementation in this paper "predicts a continuous action for each mesh element." It is not clear what is that continuous action and what RL algorithm is used to train this baseline. For the baseline labeled as "VDGN", it appears that the implementation in this paper is an actually an on-policy policy optimization method based on PPO, using a decomposed value function. But that is fundamentally different from VDGN in (ref 28), which looks like an off-policy method based on Q-learning. Additional suggestions: - line 37: "shared observations" imply that agents have the same observation, but that is not the case since each agent is a node in the graph and have different observations. Perhaps the authors mean to say agents have the same observation space. - It is hard to see what is gained by calling the formulation an "Adaptive Swarm Markov Decision Process", as opposed to the more standard terminology of a Dec-POMDP with parameter-sharing among agents. - line 257: not clear what is meant by "holistic" in this context. Perhaps the authors mean that RL can find optimal strategies for long-term objectives

Questions

Suggestions: - Demonstrate that trained policies generalize to larger meshes and more time steps than those in training. This is necessary for policies trained on tractable small problems to be deployed in real FEM simulations. - some of the learning-based baselines implemented in this paper do not accurately reflect the methods in the referenced papers, so they should not be described as though they represent those works. - Include empirical evidence that the approach extends to coarsening actions.

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The statement of limitations should include the absence of empirical results on coarsening, and the lack of other classes of PDEs such as hyperbolic equations.

Summary

This paper proposes a novel MDP formulation and policy architecture for adaptive mesh refinement. The MDP formulation (ASMDP) defines the components of the MDP in order to account for the changing number of agents across timesteps, and the reward is formulated in a manner to make credit assignment easier and to account for particular aspects of the AMR task (e.g., accounting for the area associated with an element). The paper applies RL algorithms (PPO, DQN) with a graph neural network policy (with a discrete action space of whether or not to refine a particular node). The experiments consider a set of 2d mesh refinement problems and evaluate primarily with top 0.1% error against other learned and oracle baselines, showing that ASMR outperforms the other learned baselines.

Strengths

### Significance - Problem setting: FEM plays an important role in many fields and (given the potential scale of such problems) improvements in the accuracy / computational efficiency could have a significant impact on the rate of progress or computation required - Contribution: The proposed problem formulation / algorithm significantly improves performance and reliability on a sample of AMR tasks as demonstrated in the experiments and appendix, which constitutes a significant contribution. ### Originality - Aspects of the MDP formulation (reward function, aspects of the handling of a time-varying number of agents) appear to be novel contributions. - The use of the particular policy architecture appears to be novel (though the use of GNNs for the policy isn’t novel) - Extensive experiments investigating the effectiveness of the proposed formulation/algorithm in a variety of of AMR setting (demonstrating sota performance of the algorithm) ### Quality Approach - The application of RL in this setting is reasonable. The problem formulation / algorithm are formulated reasonably. Evaluation - The main experimental claim is that ASMR outperforms existing learned methods for AMR, and this is evaluated extensively through experiments across different tasks with appropriate metrics, and the results support the claim. - A number of important additional questions are addressed experimentally by the paper (e.g., ablations, use of top 0.1% error directly as reward, runtime comparison w/ just computing $\Omega^{*}$) ### Clarity - The paper is exceptionally well-written and clear. Figures and tables are explained clearly and contribute to the understanding of the method / experiments. - Sufficient details are given in the appendix to reproduce the paper, and code is provided in the supplement, which appears to be reasonably well-designed / implemented.

Weaknesses

### Originality Algorithm - Ultimately many aspects of the problem formulation / algorithm are based on prior work, and the novel contributions are (1) specific aspects of the problem formulation (e.g., reward function), (2) the application of existing methods to the problem (some of which are novel) (e.g., policy architecture, RL algorithms), and (3) extensive experimental evaluation. ### Quality Algorithm - See questions section Evaluation - Further evaluation of the generalization abilities of the algorithm would be beneficial (e.g., building on D.4 with quantitative metrics) because (I assume) that an RL algorithm for AMR would in practice be applied outside its training distribution. - Relatedly, it’s not clear to me based on the experimental results that practitioners using AMR for FEM would prefer ASMR to the local oracle or local maximum oracle approaches in practice (perhaps ASMR is faster excluding training, but is the reduction in simulation time worth the less consistent performance?). Does increasing the amount of training data (or other parameters) improve the performance of ASMR beyond these oracle baselines across (more of) the tasks? ### Clarity - The setting / setup of section D.4 on the generalization capabilities is unclear to me

Questions

Problem formulation / algorithm 1. Does it make sense to formulate the problem as a SwarMDP / special case of a DecPOMDP? It seems that in practice (both at training and inference time) the problem is fully observable and decision making is performed in a centralized fashion (which would make it an MDP). Is that not the case? Perhaps this is sort of a semantics question about what constitutes the problem formulation vs solution. 2. Relatedly, assuming the SwarMDP formulation, in what sense is the problem partially observable to individual agents/elements? It seems that each agent makes decisions based on an observation space capturing information about the full state of the system (i.e., the observation graph taken as input seems to fully capture the state of the system). Is that not the case? 3. Given that $\Omega^{\*}$ and $\Omega^{0}$ are available, it seems like the problem could be formulated as a supervised learning problem (SL for AMR is discussed in the related work, but I don’t think the references answer the following question). Given this, why is it necessary to formulate the problem as an RL problem? Is it simply that mapping from $\Omega^{0}$ to $\Omega^{*}$ is best done in an iterative manner (due to the variable number of nodes) (thereby making it a sequential decision making problems best addressed by RL)? Evaluation 1. In the runtime comparison of section D.3, it seems like either of the processes being compared could be parallelized effectively on the GPU. Is that not the case? Why is the comparison performed on CPUs? 2. Related to that runtime comparison, what is the additional runtime cost of executing the local oracle approaches? Does it make sense to consider that in the comparison? 3. In figure 14 (right), would you elaborate on why ASMR underperforms the baselines?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

See weaknesses and questions sections

Summary

This paper builds on recent advances in learned adaptive mesh refinement method to scale to complex physical simulations. Instead of formulating AMR as a reinforcement learning problem with a single agent, the authors formulate AMR as a swarm reinforcement learning problem, in which multiple agents collaborate and share observations each other. By allowing the agents to be split into new agents in AMR process, a dense reward signal is efficiently feed to policy function to decide cells to refine. The experiments demonstrate that the method achieves competitive error against traditional error-based refinement strategies and generalizes well across different problem domains.

Strengths

The paper is placed in the current growing literature on ML-based Adaptive Mesh Refinement (AMR), to enable scaling to large and complex systems. The authors propose formulating AMR as swarm reinforcement learning problem. The paper conducts various range of experiments and compare to RL-based baselines. The proposed method outperforms existing RL-based methods in most of the cases. The experiments are thorough, covering all the important ablations I could think of while reading the paper.

Weaknesses

While the paper is generally well written, notations introduced in Section 3 are ambiguous. Below I list a few which made the paper harder to understand * $M^{t,k}$ seems not to meet $\forall$ $j, $ $\sum M^{t, k}_{i, j} = 1$ since it seems to be defined by Hadmard product. * Although definitions for state and observation are provided, they do not look well-defined; a local view of observation graph can have infinitely many choices. * The definition of $J_{I}^{t}$ is also ambiguous; the domain of the policy comprises the observation and action spaces while reward is defined on state space and action space. * What do $V$ and $E$ in the observation graph represent? As someone not familiar with this field, I found the subsection on "Systems of Equations" in Section 4 described very clearly PDE systems defined on various types of (maybe complex) shapes. However it is still unclear to me how challenging solving AMR for those PDEs is, although they are claimed as being challenging in the introduction section. Can you expand on the details of PDEs and give some reasons that make solving AMR for those types of PDEs difficult? For the evaluation of the wall-clock time of ASMR, I believe I don’t see plot of the proposed method in the figure in Appendix D.3. If so, it is difficult to compare the proposed method against other baselines.

Questions

It would be interesting to see a study of how the hyperparameters of MPN (for policy function) impact the wall-clock time and error estimate. For example, * Does large message passing steps decrease the error estimate and improves the scalability to much larger number of mesh elements? * How much does the wall-clock time increase as increasing the message passing steps?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

Yes, the paper discusses technical limitations in the main text.

Summary

This paper presents a novel framework for adaptive mesh refinement for solving PDEs describing physical systems. The refinement is done as a Markov Decision Process in a Swarm RL setting, with each element of the mesh being an agent in the swarm. The agents' action space is binary and a learned policy decides whether an element should be refined or not. The system is supervised by a task-agnostic reward.

Strengths

1. The problem of neural-AMR for efficient simulation is an important one, and the proposed solution is novel and promising. 2. The ability to have variable number of agents as the process progresses allows for finer refinements with more degrees of freedom. 3. The proposed refinement strategy doesn't require task-specific heuristics/rewards. 4. The improvement in speed as observed in the supplemental is impressive, especially for Fluid Flow. 5. Code is provided.

Weaknesses

1. While the method handles complex PDEs, it does so only for 2D domains (as observed by the authors). it'd be interesting to see the results in 3D domains. 2. While Message Passing Networks do the job, it might be that Message Passing Attention Networks/ Graph Attention Networks as recently proposed might improve the results. I think an experiment in that direction is important to justify MPNs over more recent techniques.

Questions

Can the method be extended/adapted for Mesh Movement instead of AMR? It'd extend the applicability of the proposed framework.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Same as those written in the weakness section - primarily, the results are only shown on 2D domains.

Summary

The article presents an RL-based framework for adaptive mesh refinement in FEM. In this framework, each element in the mesh is perceived as an agent from RL's perspective. The observation for each agent is then constructed using a variety of local and global features relevant to the PDE of interest. The reward function is setup to account for the impact of agent's action in terms of error reduction and increase in compute cost, and also for future agents that are a result of a decision at current time (due to mesh refinement). The method is demonstrated to work for a variety of PDEs and be somewhat competitive with traditional error-based strategies.

Strengths

Use of graph neural networks helps tackle unstructured meshes. The agent mapping $M$ helps account for future agents that result from an agent due to mesh refinement. Reward formulation helps account for error reduction and increase in compute cost due to mesh refinement.

Weaknesses

While there is some novelty in terms of reward formulation and agent mapping, I find that the biggest weakness of the method is that it is only demonstrated to work for simple problems on meshes with only thousands of elements.

Questions

None

Rating

4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

The authors only talk about how the formulation can be used for de-refinement, but do not actually include it in their method. The authors do not compare their method to existing estimators such as the one by Zienkiewicz-Zhu, which is fairly inexpensive and commonly used. The demonstrated generalizations are for simple PDEs for relatively small meshes. Thus, I find the use of the phrase "complex simulations" in the abstract arguable.

Acknowledgement of rebuttal, main concerns on complexity claim and lack of coarsening remain; question remains about whether baselines match prior work.

Regarding claim of complexity: I appreciate that the author ran an additional experiment on more than 50k elements shown in Figure 2 of the rebuttal PDF. However, this single result by itself is not sufficient to justify the way the abstract was written, which makes it sound as though the overall complexity of experiments in this paper is another level higher than previous work. The mostly stationary problems and the smooth nature of heat diffusion considered in this paper are arguably simpler than time-dependent problems where the solution feature cross over multiple elements between re-mesh steps. Regarding lack of coarsening and the heat diffusion problem: The lack of coarsening is a significant weakness, since it is critical for more complex time-dependent PDEs with rapidly moving features beyond the pedagogical problems used in this paper. The smoothness of the heat diffusion problem used in this paper is not a good representative example of the level of difficulty of time-dependency that needs to be handled in other problems, e.g. Euler equations. Regarding generalization to larger meshes: I appreciate the new result, and I believe this needs to included in the main paper because this kind of generalization is a necessary condition for using RL-trained AMR policies in practice. Regarding whether or not the baselines in this paper match exactly with prior work: I appreciate the clarification by the authors. For the baseline labeled "Argmax", is the use of PPO in this paper the same as the use of PPO in the ref [27]? In ref [27], the method has a global action space over all elements, and PPO is used as the single-agent RL for this action space. Is this the case for the "PPO implementation of Argmax" in this paper? Or is the PPO applied independently for each individual element (akin to the decentralized PPO used for the proposed method in this paper)? If applied independently, then this quite different from ref [27]. For the baseline labeled "VDGN", if the implementation shown in the main paper is the "PPO version" created by the authors, then it is necessary to revise the label names in Figure 4 and 5 to make the distinction clear. Upon closer reading, I see the authors use an MPN for the baseline labeled "VDGN", whereas ref [28] has a graph attention network. This architectural difference alone makes it hard to label the baseline as "VDGN".

Summary of rebuttal period

We again thank all reviewers for their time and effort in reviewing our paper, and appreciate that the majority vote for acceptance. Unfortunately, some of our rebuttals where left unanswered, so we would like to summarize the main concerns and lines of discussion during the review. **********************************************************************Generalization experiments (Tf6s, LCc4, qjoP)********************************************************************** We show in the rebuttal PDF that our method generalizes to significantly larger meshes on more complex domains during inference when slightly adapting the training setup. Here, the method provides accurate refinements and is more than $100$x faster than a uniform solution. **Mesh coarsening and task complexity (Tf6s, fwaD, ymhS, qjoP)** Our submission focuses on unstructured conforming meshes and predominantly stationary problems. Here, while coarsening is difficult to realize due to the non-local nature of refinements, it's typically redundant given the problem is stationary. Instead, the challenge for the policy is to find multiple precise refinement actions and manage the resulting large number of mesh elements. Our experiments show that while our method works well on this setting, the baselines break down for large meshes. This result strongly indicates that the experiments are at a higher level of complexity in terms of finding good refinement strategies than those considered in existing work. In the revision, we will clarify our emphasis on stationary problems and increased refinement steps to distinguish our paper from related work that considers non-stationary problems but simpler domains and shallow refinements. **Zienkiewicz-Zhu error estimator (Tf6s)** Results for the Zienkiewicz-Zhu error estimator for the Poisson problem are provided in the rebuttal PDF. We find that it performs poorly when applied out of the box. With additional tuning of the initial mesh, the method performs competitive with ours in terms of maximum remaining error, but is considerably worse on the mean error across the mesh. **Baseline correctness (qjoP)** We followed the implementations of the learned baselines as outlined in their original papers. Additionally, we implemented PPO and DQN variants of these baselines, keeping the core methodology and adapting it to a policy gradient and Q-learning setting respectively. We provide ablations over both algorithms and use the better of the two per method for all experiments. We adapted our “Argmax” baseline post-submission to more closely match the original paper, finding that its performance aligns with that in our initial submission. We use the same Message Passing Network architecture for all methods for comparability, and will include additional experiments on Graph Attention Networks in the revision. We believe that the new results provided in the rebuttal PDF in combination with the clarification of the scope and complexity of our experiments opposed to previous work addresses the main concerns of the reviewers.