Learning Controllable Adaptive Simulation for Multi-scale Physics
Abstract
Simulating the time evolution of physical systems is pivotal in many scientific andengineering problems.
An open challenge in simulating such systems is their multiscale dynamics: a small fraction of the system is extremely dynamic, and requiresvery fine-grained resolution, while a majority of the system is changing slowly andcan be modeled by coarser spatial scales.
Typical learning-based surrogate modelsuse a uniform spatial scale, which needs to resolve to the finest required scale andcan waste a huge compute to achieve required accuracy.
In this work, we introduce
Learning controllable Adaptive simulation for Multi-scale Physics (LAMP) asthe first full deep learning-based surrogate model that jointly learns the evolutionmodel and optimizes appropriate spatial resolutions that devote more compute tothe highly dynamic regions.
LAMP consists of a Graph Neural Network (GNN)for learning the forward evolution, and a GNN-based actor-critic for learning thepolicy of spatial refinement and coarsening.
We introduce learning techniquesthat optimizes LAMP with weighted sum of error and computational cost asobjective, which allows LAMP to adapt to varying relative importance of error vs.computation tradeoff at inference time.
We test our method in a 1D benchmark ofnonlinear PDEs and a challenging 2D mesh-based simulation.
We demonstrate thatour LAMP outperforms state-of-the-art deep learning surrogate models with up to39.3% error reduction, and is able to adaptively trade-off computation to improvelong-term prediction error.
Introduction
Simulating the time evolution of a physical system is of vital importance in science and engineering(Lynch, 2008; Carpanese, 2021;
Sircombe et al., 2006; Courant et al., 1967; Lelievre & Stoltz, 2016).
Usually, the physical system has a multi-scale nature: a small fraction of the system is highly dynamic
and requires very fine-grained resolution to simulate accurately, while a majority of the system ischanging slowly.
Examples include hazard prediction in weather forecasting (Majumdar et al.
, 2021),disruptive instabilities in the plasma fluid in nuclear fusion (Kates-Harbeck et al., 2019), air dynamicsnear the boundary for jet engine design (Athanasopoulos et al., 2009), and more familiar examplessuch as wrinkles in a cloth (Pfaff et al., 2021) and fluid near the boundary for flow through thecylinder (Vlachas et al., 2022).
Due to the typical huge size of such systems, it is pivotal that thosesystems are simulated not only accurately , but also with as small of a computational cost as possible.
A uniform spatial resolution that pays similar attention to regions with vastly different dynamics, willwaste significant compute on slow-changing regions while insufficient for highly dynamic regions.
[Figure]
ˆ M t +1 at the next time step.
We use the reduction of both Error and Computation (mesh size), compared to a multi-step rollout without remeshing, as reward to learn the policy.
To accelerate physical simulations, deep learning (DL)-based surrogate models have recently emergedas a promising alternative to complement (Um et al., 2020) or replace (Li et al., 2021) classical solvers.
They reduce computation and accelerate the simulation with larger spatial (Um et al., 2020; Kochkovet al., 2021) or temporal resolution (Li et al., 2021), or via latent representations (Sanchez-Gonzalezet al., 2020; Wu et al., 2022).
However, current deep learning-based surrogate models typicallyassume a uniform or fixed spatial resolution, without learning how to best assign computation to themost needed spatial region.
Thus, they may be insufficient to address the aforementioned multi-scalechallenge.
Although adaptive methods, such as Adaptive Mesh Refinement (AMR) (Soner et al.
,2003; Cerveny et al., 2019) exist for classical solvers, they share similar challenge (e.g. slow) asclassical solvers.
A deep learning-based surrogate models, that is able to learn both the evolution andlearn to assign computation to the needed region, is needed.
In this work, we introduce Learning controllable Adaptive simulation for Multi-scale Physics (LAMP)as the first fully DL-based surrogate model that jointly learns the evolution model and optimizesappropriate spatial resolutions that devote more compute to the highly dynamic regions.
Our keyinsight is that by explicitly setting the error and computation as the combined objective to optimize,the model can learn to adaptively decide the best local spatial resolution to evolve the system.
Toachieve this goal, LAMP consists of a Graph Neural Network (GNN)-based evolution model forlearning the forward evolution, and a GNN-based actor-critic for learning the policy of discreteactions of local refinement and coarsening of the spatial mesh, conditioned on the local state and acoefficient β that weights the relative importance of error vs. computation.
The policy (actor) outputsboth the number of refinement and coarsening actions, and which edges to refine or coarsen, whilethe critic evaluates the expected reward of the current policy.
The full system is trained with analternating fashion, iterating between training the evolution model with supervised loss, and trainingthe actor-critic via reinforcement learning (RL).
Taken together, a single instance of evolution modeland actor-critic jointly optimizes reduction of error and computation for the physical simulation, andcan operate across the relative importance of the two metrics at inference time.
We evaluate our model on a 1D benchmark of nonlinear PDEs (which tests generalization across
PDEs of the same family), and a challenging 2D mesh-based simulation of paper folding.
In 1D,we show that our model outperforms state-of-the-art deep learning-based surrogate models in termsof long-term evolution error by up to 39.3%, and can adaptively tradeoff computation to improvelong-term prediction error.
On a 2D mesh-based simulation, our model can strategically chooseappropriate edges to refine or coarsen, and outperforms a baseline without remeshing.
Problem Setting and Related Work
We consider the numerical simulation of a physical system, following the notation introduced in (Pfaffet al., 2021).
The system’s state at time t is discretized into the mesh-based state M t =
( V t , E t ) , t =0 , 1 , 2 , ... , where E t is the mesh edges and V t is the states at the nodes at time t .
Each node i ∈
Vcontains the mesh-space coordinate u i and dynamic features q i .
Note that this representation of thephysical system is very general.
It includes Eulerian systems (Wu et al., 2022) where the mesh isfixed and the field q i on the nodes are changing, and Lagrangian systems (Sanchez-Gonzalez et al.
,2020; Pfaff et al., 2021) where the mesh coordinate in physical space is also dynamically moving(for this case, an additional world coordinate x
i is accompanying the mesh coordinate u i ).
Duringprediction, a simulator f (classical or learned) autoregressively predicts system’s state ˆ M t +1 at thenext time step based on its previous prediction ˆ M t at previous time step:
[Equation]
where ˆ M
=
M 0 is the initial state.
During the prediction, both the dynamic features V t at the meshnodes, and the mesh topology E t can be changing.
The error is typically computed by comparing theprediction and ground-truth after long-term prediction: error :=
ℓ ( ˆ M t , M t ) for a metric ℓ (e.g. MSE,
RMSE), and the computation cost (in terms of floating point operations, or FLOPs) typically scaleswith the size of the mesh (e.g. the number of nodes).
The task is to evolve the system long-term intothe future, with a low error and a constraint on the computational cost.
We review related work in
Appendix A.
Method
In this section, we detail our LAMP method.
We first introduce its architecture in Sec. 3.1.
Then weintroduce its learning method (Sec. 3.2), including learning objective and training, and technique tolet it learn to adapt to varying importance of error and computation.
3.1 Model architecture
The model architecture of LAMP consists of two components: an actor-critic which updates the meshtopology, and an evolution model which evolves the states defined on the mesh.
We will detail themone by one.
Actor-critic .
The actor-critic consists of a policy network f policy φ (with parameter φ ) which predictsthe probability of performing the spatial coarsening or refining actions, and a value network f value φwhich evaluates the long-term expected reward of the policy network:
[Equation]
where a t is the (refining and coarsening) action performed on the edges E t so that it will becomeˆ E t +1 .
The policy network f policy φ outputs the probability of performing such action and can samplefrom this probability.
v t estimates the “value” of the current policy starting from current state M t(for more information, see Sec.
3.2 below).
The explicit dependence on β (as the β in Eq. 7) allowsthe policy and value network to condition on the varying importance of error and computation.
Giventhe predicted mesh ˆ E t +
and the current node features V t on the current mesh E t , an interpolation isperformed which obtains the node features on the new mesh:
[Equation]
[Equation]
Now the new intermediate state ˆ M ′ t
= ( ˆ V ′ t , ˆ E t +1 ) is defined on the new mesh ˆ E t +1 .
Evolution model .
The second component is an evolution model f evo θ which takes as input theintermediate state M ′ t defined on ˆ E t +1 , and outputs the prediction of node features ˆ V t +1 for timet + 1 :
[Equation]
Note that in this stage, the mesh topology ˆ E t +1 is kept constant, and the evolution model f evo θ (withparameter θ ) learns to predict the state based on the current mesh.
Taken together, Eqs.
(2)(4)(5) evolve the system state from M t at time t to state ˆ M t +1 =(
ˆ V t +1 , ˆ E t +1 ) at t + 1 .
During inference, they are executed autoregressively following Eq.
(1),to predict the system’s future states ˆ M t , t = 1 , 2 , 3 , ... , given an initial state M 0 .
GNN architecture .
One requirement for the evolution model f evo θ , policy network f policy φ and valuenetwork f value φ is that they can operate on changing mesh topology E t .
Graph Neural Networks(GNNs) are an ideal choice that meets this requirement.
Since we represent the system’s state as mesh,we adopt MeshGraphNets (Pfaff et al., 2021) as the base architecture for the above three models.
Specifically, we encode V t as node features for the graph, and encode the mesh topology edges E tand world edges as two types of edges, and the edge features depend on the relative positions in themesh coordinates and world coordinates.
Based on the graph, a processor network that consists of Nlayers of message passing is performed to locally exchange and aggregate the information:
[Equation]
where Z ( v ) n i is the latent node vector on node i at layer n , and Z ( e ) n ij is the latent edge vectors on then th layer on the edge from node i to node j .
We have Z ( v )
i =
ˆ V
′ t
i and Z ( v ) N i
=
ˆ V
[Equation]
1 i are input andpredicted node features at time t and t + 1 , respectively, in Eq.
Figure 1 provides an illustrationof the architecture.
We use an independent processor for the evolution model, and share the processorfor the policy and value networks.
After the processor, the latent vectors are concatenated with βto feed into downstream decoders.
For the evolution model f evo θ , a decoder is operated on the latentstate and outputs the prediction ˆ V ′ t +1 on the nodes.
For the value network, an value MLP is operatedon all nodes, and a global pooling is performed to compute the overall estimated value.
For the policynetwork, we design the action decoder as follows.
Action representation .
To predict the action for policy network and its probability, we first need todesign the action space.
Note that compared to typical reinforcement learning problems, here theaction space is extremely high-dimensional and complicated: (1) each edge of the mesh may have theoption of choosing refinement or coarsening.
If there are thousands of edges N edge , then the possibleactions will be on the order of 2 N edge .
(2) Not all actions are valid, and many combinations of actionsare invalid.
For example, two edges on the same face of the mesh cannot be refined at the same time,nor can they be both coarsened.
To address this high-dimensionality action problem, we introducethe following design of action space, where for both refinement and coarsening, the policy networkfirst samples integers K re , K co ∈ { 0 , 1 , 2 , ...
K max } , and then independently samples K re edges toperform refinement and K co edges to perform coarsening, with proper filtering.
The combinedaction is a t =
( K re , e re 1 , e re 2 , ...
e re K refine , K co , e co 1 , e co 2 , ...
e co K co ) , where K re , K co ∈ { 0 , 1 , ...K max } , ande re k , e co k ∈ E t , k = 1 , 2 , .. .
The combined log-probability of choosing the action a t is given by:
[Equation]
This action representation offers several benefits compared to independent sampling of action oneach edge, detailed in Appendix B.
3.2 Learning
The ultimate goal of the learning for LAMP is to optimize the objective Eq.
(7) as follows:
[Equation]
for a wide range of β .
To achieve this, we break down the above objective into an alternative learningof two phases: learning the evolution model with objective L evo that minimizes long-term evolution error, and learning the policy with objective L policy that optimizes both the long-term evolution errorand computational cost.
Learning evolution .
In this phase, the evolution model f evo θ is optimized to reduce the multi-stepevolution error for the evolution model.
As before, we denote M t + s , t = 0 , 1 , 2 , ..., s = 0 , 1 , ...
S asthe state of the system at time t + s simulated by the ground-truth solver with very fine-grained mesh,and denote ˆ M t + s , t = 0 , 1 , 2 , ..., s = 0 , 1 , 2 , ..
S as the prediction by the current LAMP following thecurrent policy, up to a horizon of S steps into the future.
We further denote ˆ M ′′ t + s , t = 0 , 1 , 2 , ..., s =0 , 1 , 2 , ..S as the prediction by the current evolution model on the fine-grained mesh, where its meshis provided as ground-truth mesh E t
+ s at each time step.
Then the loss is given by:
[Equation]
Essentially, we optimize two parts of the evolution loss: (1) L evo s [ f policy φ , f evo θ ; ˆ M t ] (first term) which isthe evolution loss by following policy network f policy φ and evolution model f evo θ , starting at initial stateof ˆ M t for S steps.
(here α policy s is the coefficient for the s -step loss with loss function ℓ ).
This makessure that evolution model f evo θ adapts to the current policy f policy φ that designates proper computation.
The second part
of the loss, L evo s
[ I , f evo θ ; ˆ M ′′ t ] , is the evolution loss by using the ground-truthfine-grained mesh and evolved by the evolution model f evo θ , starting at initial state of fine-grainedmesh ˆ M ′′ t and evolve for s steps.
This encourages the evolution model to learn to utilize morecomputation to achieve a better prediction error.
Learning the policy .
In this phase, the policy network f policy φ learns to update the spatial resolution(refinement or coarsening of mesh) at each location, to improve both the computation and theprediction error .
Since the spatial refinement and coarsening are both discrete actions, and the metricof computation is typically non-differentiable, we use Reinforcement Learning (RL) to learn thepolicy.
Specifically, we model it as a Markov Decision Process (MDP), where the environment stateis the system’s state M t , the actions are the local refinement or coarsening at each edge of E t , andwe design the reward as the improvement on both the error and computation, between following thecurrent policy’s action, and an counterfactual scenario where the agent follows an identity policy Ithat does not update the mesh topology, starting on the initial state ˆ M t .
Concretely, the reward is:
[Equation]
Here C S [ · ] is a surrogate metric that quantifies the computation based on the predicted mesh topologyˆ E t +1 , ˆ E t +2 , ...
ˆ E t + S up to S steps into the future.
In this paper we use the number of nodes as thesurrogate metric for measuring the computation, since typically for the GNNs, the computation (interms of FLOPs) scales linearly with the number of nodes (since each node has a bounded numberof edges on the mesh, the number of edges thus also scales linearly with number of nodes, so willmessage passing and node updates).
To optimize the reward r t , we employ the standard REINFORCE algorithm as used in (Sutton et al.
,1999; Hafner et al., 2021) to update the policy network f policy φ , with the following objective:
[Equation]
Here H [ · ] is the entropy, which encourages the action to have higher entropy to increase exploration,where η here is a hyperparameter.
The sg ( · ) is stop-gradient.
Essentially, the first term in loss L policyencourages to increase the log-probability of actions that have a higher “advantage”, where theadvantage is defined as the difference between the current reward r t that follows the current action a ttaken, and the expected reward (value)
f value φ ( M t , β ) that follows the current policy starting from thecurrent state M t .
We can also think of it as an actor-critic where the critic tries to evaluate accurately the expected reward of the current policy, and an actor (policy) is trying to exceed that expectation.
To train the value network, we use MSE loss:
[Equation]
[Equation]
Learning to adapt to varying β .
In the overall objective (Eq. 7), the β stipulates the relativeimportance between Error and Computation.
β
= 0 means we only focus on minimizing Error,without constraint on Computation.
β
= 1 means we only focus on minimizing computation, withoutconsidering the evolution error.
In practice, we typically wish to improve both, with a β ∈ (0 , 1) thatputs more emphasis on one metric but still considers the other metric.
To allow LAMP to be able tooperate at varying β at inference time, during the learning of policy, we sample β uniformly within arange B ⊆ [0 , 1] (e.g. B can be [0 , 1] or [0 , 0 . 5] ), for different examples within a minibatch, and alsotrain the policy and value network jointly, where the total loss L policy is the weighted sum of the twopolicy and value losses:
[Equation]
where α value is a hyperparameter, which we set as 0.5.
In this way, the policy can learn a generic wayof spatial coarsening and refinement, conditioned on β .
For example, for smaller β that focuses moreon improving error, the policy network may learn to refine more on dynamic regions and coarsen less,sacrificing computation to improve prediction error.
Experiments
In the experiments, we set out to answer the following questions on our proposed LAMP: (1)
Can
LAMP learn to coarsen and refine the mesh, focusing more computation on the more dynamic regionsto improve prediction accuracy?
(2) Can LAMP improve the Pareto frontier of Error vs. Computation,compared to state-of-the-art deep learning surrogate models?
(3) Can LAMP learn to condition on theβ to change is behavior, and perform varying amount of refinement and coarsening depending on theβ ?
We evaluate our LAMP on two challenging datasets: (1) a 1D benchmark nonlinear PDEs , whichtests generalization of PDEs in the same family (Brandstetter et al., 2022); (2) a mesh-based papersimulation generated by the ArcSim solver (Narain et al., 2012).
Both datasets possess multi-scalecharacteristics where some parts of the system is highly dynamic, while other parts are changingmore slowly.
4.1 1D nonlinear family of PDEs
[Figure]
Data and Experiments .
In this section, we test LAMP’s ability to balance error vs. computationtested on unseen equations with different parameters in a given family.
We use the 1D benchmark in
Brandstetter et al.
(2022), whose PDEs are given by
[Equation]
The parameter for the PDE is given by p
= ( α, β, γ ) .
The term δ is a forcing term (Bar-Sinai et al.
with J = 5 , L = 16 and coefficients A j and ω j sampled uniformly from A j ∼ U [ − 0 .
5 , 0 . 5] ,ω j ∼ U
[ − 0 .
4 , 0 . 4] , ℓ j ∈ { 1 , 2 , 3 } , ϕ j ∼ U [0 , 2 π ) .
We uniformly discretize the space to n x = 200in [0 , 16) and uniformly discretize time to n t = 250 points in [0 , 4] .
Space and time are furtherdownsampled to resolutions of ( n t , n x ) ∈ { (250 , 100) , (250 , 50) , (250 , 33) } as initial resolution.
Weuse the E2 scenario in the benchmark, which tests the model’s ability to generalize to novel parametersof PDE with the same family.
Specifically, we have that the parameter p
=
(1 , η, 0) , η ∼ U [0 , 0 . 2] .
As our LAMP autoregressively simulate the system, it can refine or coarsen the mesh at appropriatelocations by the policy network f policy φ , before evolving to the next state with the evolution modelf evo θ .
We evaluate the models with the metric of Computation and long-term evolution Error.
Forthe computation, we use the average number of vertices throughout the full trajectory as a surrogatemetric, since the number of floating point operations typically scales linearly with the number ofvertices in the mesh.
For long-term evolution error, we use the cumulative MSE over rollout timesteps, starting with initial state with a bundle of S = 25 time steps t = [25 , 26 , ... 49] and predict therest 175 steps.
We compare LAMP with strong baselines of deep learning-based surrogate models,including CNNs, Fourier Neural Operators (FNO) (Li et al., 2021), and MP-PDE (Brandstetter et al.,2022) which is a state-of-the-art deep learning-based surrogate models for this task.
Our base neuralarchitecture is based on MeshGraphNets (Pfaff et al., 2021) which is a state-of-the-art GNN-basedmodel for mesh-based simulations.
We compare an ablation of our model that does not performremeshing (LAMP no remeshing), and a full version of our model (LAMP).
We perform three groupsof experiments, starting at initial number vertices of 33, 50 and 100 that is downsampled from the100-vertex mesh, respectively.
All the baselines do not perform remeshing, and our full model hasthe ability to perform remeshing that coarsen or refine the edges at each time step.
Note that for allmodels, our metric of cumulative MSE is computed on the full ground-truth mesh with 100 vertices,where the values of prediction are linearly interpolated onto the location of the ground-truth.
Thisprevents the model to “cheat” by reducing the number of vertices and only predict well on thosevertices.
Additional details are given in Appendix C.1.
[Table]
[Figure]
Results .
Table 1 shows the results.
We see that our LAMP outperforms all baselines by a largemargin, achieving an error reduction of 39.3%, 23.5% and 6.2% compared with the best performingbaseline of CNN, FNO and MP-PDE.
Importantly, we see that compared with an ablation with noremeshing, our full LAMP is able to significantly reduce error (by 6.3% error reduction for 33 initial vertices scenario, and 30.6% for 50 vertices scenario), with only slight increase of average number ofvertices (by 13.9% and 6.4% increase, respectively).
This shows the ability of LAMP to adaptivelytrade computation to improve long-term prediction error.
To investigate whether LAMP is able to focus computationon the most dynamic region, we visualize example trajectories of LAMP, as shown in Fig.
2 and more examples in
Fig.
7 in Appendix D.1.
Starting with 50 vertices, we testour model on different β , where smaller β focuses more onimproving error.
We see that with smaller β (e.g. β = 0 . 1in Fig. 2), LAMP is able to add more nodes (33 nodes) onthe most dynamic region, and only coarsen few (removing 3nodes in total).
With a larger β that focuses more on reducing computation, we see that LAMP refines less and coarsenmore, and only coarsen on the more slowly changing region.
Additionally, we visualize the error vs. number of verticesfor varying β for two different models over all test trajectories in Fig.
We see that with increasing β , LAMP isable to reduce the vertices more, with only slight increaseof error.
In summary, the above results show that LAMP isable to focus computation on dynamic regions, and able toadapt to different β at inference.
4.2 2D mesh-based simulation
Model
[Table]
Reported number is the MSE over 10 learned-simulator-steps roll-out averaged over 50 test trajectories.
Here we evaluate our LAMP’s ability on a more challenging setting with paper folding simulation.
The paper in this simulation is square-shaped and its boundary edge length is 1 .
During the simulation,4 corners of the paper receive different magnitude of force.
When generating ground-truth trajectoriesby the ArcSim solver (Narain et al., 2012), we set the minimum and maximal length of edges to be0.01 and 0.2.
We evaluate the models with the metric of computation and long-term evolution error.
Similar to the Section 4.1, we use the average number of nodes throughout the full trajectory as asurrogate metric for complexity.
As shown in Table 2, our model is able to add more resolution tothe high curvature region, and achieve better roll-out accuracy than the ablation without remeshing.
Additional details (including visualization) of the experiments and 2D triangular mesh remeshingactions are given in Appendix C.2.
In Fig.
6 in Appendix C.2 and Fig.
8 in Appendix D.2, we seethat our LAMP learns to add more mesh onto the more dynamic regions near to the folding part (withhigh curvature), showing LAMP’s ability to assign computation to the most needed region.
Conclusion
In this work, we have introduced LAMP, the first fully deep learning-based surrogate model thatjointly learns the evolution of physical system and optimizes assigning the computation to the mostdynamic regions.
In 1D and 2D datasets, we show that our method is able to adaptively performrefinement or coarsening actions, which improves long-term prediction error than strong baselines ofdeep learning-based surrogate models.
We hope our method provides a useful tool for more efficientand accurate simulation of physical systems.
A Related Works
Most classical solvers use a fixed mesh E t whose topology does not vary with time.
For example,the mesh E t ≡ E 0 can be a 2D or 3D regular grid, or an irregular mesh that is pre-generated at thebeginning of the simulation (Geuzaine & Remacle, 2009).
Classical Adaptive Mesh Refinement(AMR)
(Narain et al., 2012) addresses the multi-scale challenge by adaptively refine or coarsen themesh resolution, with heuristics based on local state variation.
Since they are based on classicalsolvers, they may not benefit from the many advantages that deep learning brings (GPU acceleration,less stringent spatial and temporal resolution requirement, explicit forward, etc.).
In contrast, our
LAMP is a deep-learning based surrogate model, and can benefit from the many advantages (e.g.speedup) offered by the deep learning framework.
Furthermore, since it directly optimize for a linearcombination of error and computation, it has the potential to directly optimize to a better error vs.computation tradeoff, nearer to the true Pareto frontier.
Deep-learning based surrogate models, although having achieved speedup compared to classicalsolvers, still typically operate on a fixed grid or mesh (Li et al., 2021; Sanchez et al., 2020;
Wu et al.
Zhao et al., 2022; Han et al., 2022), and have yet to exploit the multi-scale nature typical inphysical simulations.
One important exception is MeshGraphNets (Pfaff et al., 2021), which bothlearns how to evolve the state V t , and uses supervised learning to learn the spatial adaptation thatchanges E t .
However, since it uses supervised learning where the ground-truth mesh is providedfrom the classical solver with AMR, it cannot exceed the performance of AMR in terms of error vs.computation tradeoff, and has to interact with the classical solver in inference time to perform theadaptive mesh refinement.
In contrast, our LAMP directly optimizes for the objective, which usesreinforcement learning for learning the policy of refinement and coarsening, and has the potential tosurpass classical AMR and achieve a better error vs. computation tradeoff.
Moreover, a single trained
LAMP has the ability to adapt to the full range of relative importance β of error vs. computation atinference time, thus can be more versatile than MeshGraphNets with a fixed strategy.
B Model Architecture
Here we detail the architecture of LAMP, complementary to Sec. 3.1.
This architecture is usedthroughout all experiment, with just a few hyperparameter (e.g. latent dimension, message passingsteps) depending on the dimension (1D, 2D) of the problem.
B.1 Architecture
In this subsection, we first detail the base architecture that is common among the evolution modelf evo θ , the policy network f policy φ , and value network f value φ .
Then describe their respective aspects, andexplain the benefits of our action representation.
Base architecture .
We use MeshGraphNets Pfaff et al.
(2021) as our base architecture, for theevolution model f evo θ , the policy network f policy φ , and value network f value φ .
All three models have anencoder and a processor.
For the encoder, we encode both node features and edge features, uplift tolatent dimension using MLPs, to Z eij , Z vi .
For the processor, it consists of N message passing layers.
For each message-passing layer, we first update the edge features using current edge features, andconnected node features, Z ( e ) n +1 ij =
MLP ( v ) θ ( E nij , Z ( v ) n i , Z ( v ) n j ) , then we update the node featuresusing current node feature and connected edge features, Z ( v ) n +1 i =
MLP ( v ) θ ( Z ( v ) n i , (cid:80) j Z ( e ) n +1 ij ) .
Depending on the models and their intended functions, different models have different modules forpredicting the output, described as follows.
Evolution model .
For the evolution model, the output is at each node, and uses the last latent vectoron the node, feeds into a decoder, to predict the output ˆ V t +1 i
= Z ( v ) N i
(See Eq. 6) at the next timestep t + 1 .
Policy network and value network .
Both f policy φ and f value φ shares the same processor f processor φ ,which has N message-passing layers.
The final node feature is appended with β for the controllability.
The f policy φ takes the encoded graph as input and consist of two parts - a global mean pooling is used toobtain the global latent representation of the whole graph, which is then fed into a MLP ( k ) followedby a action specific linear layer to predict the probability distribution for different number of actionsshould be performed; an action specific MLP ( a ) is applied on each edge to predict the probability ofa certain action (i.e. split, coarse) to be applied on this edge.
The f value φ takes the same encoded graphas input, and perform a global mean pooling to obtain the global latent features for the graph, whichis then feed through a linear layer to predict the accuracy reward, and another linear layer to predictthe computation reward, the final predicted value is, v = loss + β · compute.
Benefits of our action representation .
Here we first elaborate the benefits of our action representation, compared to an independent sampling on each edge.
If offers the following benefits:
• The action space reduces from 2 N edge to N edge K max , where N edge is the number of edges.
In the case of N edge ∼ 1000 and K max ∼ 10 ≪ N edge , the difference in action spacedimensionality is significant, e.g. 2 1000 = 10 300 vs. 1000 10 = 10 30 .
Therefore, it is easier to credit assign the reward to appropriate action, with a smaller space of action.
• Compared with each edge performing action independently, now only K ≤ K max actions ofrefinement or coarsening can be performed.
Therefore, it will need to focus on a few actions
that can best improve the objective.
• The sampling of K will also make the policy more “controllable”, since now the K isexplicitly dependent on the β , and learning how many refinement or coarsening action to
take depending on β is much easier than figuring out which concrete independent actions to
take.
B.2 Architectural hyperparameters used in 1D and 2D experiments
Here we detail the architectures used in the 1D and 2D experiment.
A summary of the hyperparametersis also provided in Table 3.
B.2.1 1D nonlinear PDEs
For 1D experiment, our evolution model f evo θ has N = 3 message-passing layers in the processor andlatent dimension of 64.
It uses the SiLU activation (Elfwing et al., 2018).
The shared processor forf policy φ and f value φ has N = 3 message-passing layers and latent dimension of 64.
B.2.2 2D mesh-based simulation
For f evo θ , we set message passing layers for MeshGraphNets (Pfaff et al., 2021) to be 8 .
We alsomodel our message passing function with MLP having 56 units followed by SiLU activation (Elfwinget al., 2018).
Here, we share the same message passing function across all the 8 layers because wefound in our experiments that sharing the same message passing function achieved better performancethan having independent MLP for each of the layers.
The shared processor for f policy φ and f value φ has
N = 3 message-passing layers and latent dimension of 56.
C Experiment Details
In this section, we provide experiment details for 1D and 2D datasets. Firstly, we explain the use ofreward r t (in Eq. 13) as the value target instead of the typical recursive definition. Then we detail thetraining details for 1D and 2D in the following subsections. Use of r t as value target .
Different from typical RL learning scenario (computer games or robotics)where the episode always has an end, and the reward is bounded, here in physical simulations, thereare two distinct characteristics: (1) the rollout can be performed infinite time steps into the future,(2) as the rollout continues, at some point the error between predicted state and ground-truth statewill diverge, so the error is not bounded. Based on these two characteristics, a value target basedon the error for infinite horizon (e.g. the one used in Dreamer v2 (Hafner et al., 2021)) does notmake sense, since it will not be bounded. In our experiment, we also observe similar phenomena, inwhich both the value target and value prediction continue to increase indefinitely. Thus, we use theaverage reward in the S step rollout as the value target, which measures the error and computationimprovement within the rollout window we care about and proves to be much more stable.
C.1 1D nonlinear PDEs
The action space for 1D problem is composed with split and coarsen actions.
The coarse action isinitially defined on all edges, which is then sampled base on the predicted number of coarse actions,and probability of each edge to be coarsen.
Among those sampled edges, if two edges share acommon vertex, the right most one will be removed.
During training, to let the evolution model adapt to the varying size of the mesh, we performrandom node dropout, where we randomly sample 10% of the minibatch to perform dropout, and if aminibatch is selected for node dropout, for each example, randomly drop 0-30% of the nodes in themesh.
C.2 2D mesh-based simulation
In this subsection, we give details on remeshing actions and how to generate 2d mesh data based on
Narain et al.
We also details pre-training of our evolution function.
Action Space .
There are 3 kinds of actions defined in the 2D mesh simulation: split, flip, and coarsen.
Split action in 2D case can be performed on any edges that are shared with two triangular faces in amesh and also on the boundary edges of the mesh.
The split action results in 4 triangles (Fig. 4a).
Flip action is also performed on edges shared with two faces, but it also requires some additionalcondition, that is when sum of angles at vertices located at opposite side of edges is greater than π :see also Fig 4b.
In our remeshing function, all edges satisfying the condition are flipped.
Finally,coarsen action has relatively strong conditions.
One condition is that one of the source or target nodesof a coarsened edge needs to be of degree 4 and another condition is that all the faces connected tothe node of degree 4 needs to have acute angles except angles around the node.
See also Fig. 4c.
Wefilter the sampled actions from f policy φ base on aforementioned conditions to get sets of valid edges tosplit and coarse.
[Figure]
Generating data .
To generate ground-truth data, we use a 2D triangular mesh-based AMR simulator
Narain et al.
This simulator adopts adaptive anisotropic kernel method to automaticallyconform to the geometric and dynamic detail of the simulated cloth: see also Fig. 5.
In our experiment,each trajectory consists of 325 meshes (excluding initial mesh.)
Each mesh consists of vertices andfaces and every vertex is equipped with its 2D coordinates as well as 3D world coordinates.
Timeframe between two consecutive meshes are set to be 0 . 04 .
The edge length was set to range between0 .
01 and 0 .
Note that decision on respective remeshing actions is made based on the 2D coordinates. T= T=
T= T= T= For training and test dataset, we generate 1050 configurations specifying direction and magnitudeof forces applied to 4 corners of meshes, and generate 1000 trajectories for training data and use 50trajectories as test data based on the generated configurations.
When training our evolution model,we downsample data by half of the original data; we use every other meshes starting from an initialmesh in each trajectory.
The reason of downsampling is that we observed that cumulative RMSEover rollout with model trained with full trajectories blew up after iteration exceeded 160 steps.
Pre-training of Evolution function .
We model our evolution function f evo θ with MeshGraphNets(Pfaff et al., 2021).
The input of the model adopts a graph representation where a vertex is equippedwith 3D velocity computed from mesh information at both current and past time steps, and featurefor an edge consists of relative distance and its magnitude of boundary vertices of the edge.
Sincemesh topology varies during the forward iteration, we perform barycentric interpolation on the meshat past time step to get interpolated vertex coordinates corresponding to vertices at current time step.
When evaluating our model, we use rollout RMSE, taking the mean for all spatial coordinates, allmesh nodes, all steps in each trajectory, and all 50 trajectories in the test dataset.
The rollout RMSEachieved 4 . 45 × 10 − 2 with our best parameter setting.
We found that adding noise helped to improveits performance; we added the noise with scale 0 .
01 to vertex coordinates.
D Additional visualization
[Figure]
D.1 1D nonlinear PDE experiment
In this section, we provide additional randomly sampled results for the 1D nonlinear PDE experiment,as shown in Fig.
We see that similar to Fig.
2, our LAMP is able to add more vertices to thelocations with more dynamicity, while remove more vertices at more static regions.
Moreover, withincreasing β that emphasizes more on reducing computational cost, LAMP splits less and coarsensmore. [Table]
D.2 Additional visualization for
2D mesh- based simulation
In this section, we provide additional randomly sampled results for