Latent Task-Specific Graph Network Simulators

## Abstract

Abstract Simulating dynamic physical interactions is a critical challenge across multiple scientific domains, with applications ranging from robotics to material science.
For mesh-based simulations, Graph Network Simulators (GNSs) pose an efficient alternative to traditional physics-based simulators.
Their inherent differentiability and speed make them particularly well-suited for inverse design problems.
Yet, adapting to new tasks from limited available data is an important aspect for real-world applications that current methods struggle with.
We frame mesh-based simulation as a meta-learning problem and use a recent Bayesian meta-learning method to improve GNSs adaptability to new scenarios by leveraging context data and handling uncertainties.
Our approach, latent task-specific graph network simulator, uses non-amortized task posterior approximations to sample latent descriptions of unknown system properties.
Additionally, we leverage movement primitives for efficient full trajectory prediction, effectively addressing the issue of accumulating errors encountered by previous auto-regressive methods.
We validate the effectiveness of our approach through various experiments, performing on par with or better than established baseline methods.
Movement primitives further allow us to accommodate various types of context data, as demonstrated through the utilization of point clouds during inference.
By combining GNSs with meta-learning, we bring them closer to real-world applicability, particularly in scenarios with smaller datasets.

## 1 Introduction

Simulating physical systems is a fundamental challenge across a variety of scientific fields, with applications ranging from structural mechanics (Yazid et al., ; Zienkiewicz & Taylor, ; Stanova et al., ) over fluid dynamic (Chung, ; Zienkiewicz et al., ; Connor & Brebbia, ) to electromagnetism (Jin, ; Polycarpou, ; Reddy, ).
Mesh-based simulations are often chosen for these tasks due to the computational efficiency and accuracy of the underlying finite element method (Brenner & Scott, ; Reddy, ).
Yet, the diversity of modeled problems usually necessitates the development of task-specific simulators to accurately capture the relevant physical quantities (Reddy & Gartling, ).
Despite these efforts, these specialized simulators can be slow and cumbersome, especially for larger simulations (Paszynski, ; Hughes et al., ).

As a result, data-driven models have gained traction as an appealing alternative (Guo et al., ; Da Wang et al., ; Li et al., ). Among these, general-purpose Graph Network Simulators (GNSs) have become increasingly popular (Battaglia et al., ; Pfaff et al., ; Allen et al., , ; Linkerhägner et al., ). Building on Graph Neural Networks (GNNs) (Scarselli et al., ; Wu et al., ; Bronstein et al., ), GNSs encode the simulated system as an interaction graph between nodes, predicting their dynamics.
These models offer a significant speed advantage over classical simulators (Pfaff et al., ) while being fully differentiable, making them highly effective for applications like inverse design (Allen et al., ; Xu et al., ).

GNSs are commonly trained through simple next-step supervision (Battaglia et al., ; Pfaff et al., ; Allen et al., ).
During inference, entire trajectories are unrolled by iteratively predicting per-node dynamics from an initial system state.
This process is susceptible to accumulating errors over time, especially as the input distribution diverges from the training set (Brandstetter et al., ; Han et al., ).
While data augmentation strategies exist to offset this issue (Pfaff et al., ; Brandstetter et al., ), they neither correct mistakes once they have been made nor effectively address the challenges posed by partially known initial system states (Linkerhägner et al., ).
Moreover, these models usually require large amounts of data to train, which is an issue in real-world scenarios where data is often sparse and the need for efficient adaptation to new tasks is crucial (Linkerhägner et al., ).
In this work, we reformulate learned mesh-based simulation as a trajectory-level meta-learning problem that uses mesh states as a context set to address these limitations.
We employ a Bayesian meta-learning approach based on non-amortized task posterior approximations (Volpp et al., , ) for rapid adaptation to new task properties and uncertainties.
We further mitigate the issue of error accumulation through the use of Probabilistic Dynamic Movement Primitivess (ProDMPs) (Schaal, ; Paraschos et al., ; Li et al., ) to represent higher-order dynamics of mesh nodes on a trajectory level.
Combined, these methods allow us to model a distribution over unknown material properties and induced simulation trajectories from a given context set.
When this context set contains sufficient information, the method is able to accurately determine specific system properties and adapt the simulation accordingly.
We visualize an example in Figure [1](#S1.F1).
A more detailed overview of our approach, called Latent Task-Specific Graph Network Simulator (LTSGNS) (^1^11Code available at [https://github.com/PhilippDahlinger/ltsgns_ai4science](https://github.com/PhilippDahlinger/ltsgns_ai4science)) is shown in Figure [2](#S3.F2).

We validate the effectiveness of LTSGNS on challenging deformable object simulations, showing superior prediction quality compared to MeshGraphNet (MGN), a state-of-the-art GNS.
Furthermore, we showcase that our method can incorporate real-world observations such as point-clouds.
This feature is particularly useful for applications dealing with sparse and incomplete data sets.
The ability to accommodate such data types makes the model more applicable in practical scenarios and opens new avenues for research that require robust and generalizable graph network simulators.

Figure: Figure 1: Illustration of the latent space of Latent Task-Specific Graph Network Simulator (LTSGNS) describing material properties of the mesh. (Left) Given an initial mesh as a context, the task posterior depicted as a black and white contour plot is spread out to include different possible mesh deformations. Thus, each sample of the posterior approximation (colored Gaussian components) results in a different yet plausible simulation outcome. (Right) When the material property can be inferred from additional context information, the posterior collapses to a unimodal distribution that represents deformations that are compatible with this context.
Refer to caption: /html/2311.05256/assets/x1.png

## 2 Related Work

Graph Network Simulators.
Recent research has increasingly focused on training deep neural networks for physical simulations as such models can yield significant speedups over traditional solvers while being fully differentiable (Pfaff et al., ; Allen et al., ), making them a natural choice for e.g., model-based Reinforcement Learning (Mora et al., ) and Inverse Design problems (Baqué et al., ; Durasov et al., ; Allen et al., ).
Examples use Convoluational Neural Networks (CNNs) for fluid and aerodynamic flow simulations  (Tompson et al., ; Guo et al., ; Chu & Thuerey, ; Zhang et al., ; Bhatnagar et al., ; Ummenhofer et al., ) and train either on a supervised loss or with the help of Generative Adversarial Networks (Goodfellow et al., ) in an adversarial fashion (Kim et al., ; Xie et al., ).
A popular class of learned neural simulators are GNSs (Battaglia et al., ; Sanchez-Gonzalez et al., ), a special type of GNN (Scarselli et al., ; Bronstein et al., ) that is designed to handle graph-structured physical data by modeling relations between arbitrary entities.
Here, applications include particle-based simulations (Li et al., ; Sanchez-Gonzalez et al., ), atomic force prediction (Hu et al., ) and fluid dynamic problems (Brandstetter et al., ).
Most notably, they have been applied to the mesh-based prediction of deformable objects (Pfaff et al., ; Weng et al., ; Han et al., ; Fortunato et al., ; Linkerhägner et al., ).
Here, extensions include handling rigid objects (Allen et al., , ) and integrating learned adaptive mesh refinement strategies (Plewa et al., ; Yang et al., ; Freymuth et al., ) into the simulator (Wu et al., ).
Closely related to our work, another extension utilizes additional sensory information to ground simulators to improve long-term predictions (Linkerhägner et al., ).

#### Meta Learning.

Meta-learning (Schmidhuber, ; Thrun & Pratt, ; Vilalta & Drissi, ; Hospedales et al., ) extracts inductive biases from a training set of related tasks in order to increase data efficiency on unseen tasks drawn from the same task distribution.
In contrast to other multi-task learning methods, such as transfer learning (Krizhevsky et al., ; Golovin et al., ; Zhuang et al., ), which typically merely fine-tune or combine standard single-task models, meta-learning makes the multi-task setting explicit in the model architecture (Bengio et al., ; Ravi & Larochelle, ; Andrychowicz et al., ; Volpp et al., ; Santoro et al., ; Snell et al., ).
This allows the resulting meta-models to learn how to learn new tasks with only a few context examples.
A popular variant is the model-agnostic meta-learning (MAML) family (Finn et al., ; Grant et al., ; Finn et al., ; Kim et al., ), which employs standard single-task models and formulates a multi-task optimization procedure.
The neural process (NP) model family (Garnelo et al., , ; Kim et al., ; Gordon et al., ; Louizos et al., ; Volpp et al., ) is a complementary approach in the sense that it builds on a multi-task model architecture (Heskes, ; Bakker & Heskes, ), but employs standard gradient based optimization algorithms (Kingma & Ba, ; Kingma & Welling, ; Rezende et al., ; Zaheer et al., ).
Recently, Volpp et al. () demonstrated the importance of accurate task posterior inference for efficient meta-learning by combining an NP-like architecture with more powerful inference and optimization schemes (Arenz et al., ; Lin et al., ; Arenz et al., ).

#### Motion Primitives.

Movement Primitives (MPs) (Schaal, ; Paraschos et al., ) allow for compact and smooth trajectory representation via a set of basis functions.
Recent methods combine MPs with neural networks to increase their expressiveness (Seker et al., ; Bahl et al., ; Li et al., ).
Among these, ProDMPs (Li et al., ) introduce a novel set of basis functions that sidestep an otherwise expensive numerical integration procedure in the training pipeline while being fully differentiable.
Additionally ProDMPs can be queried at arbitrary points in time, making them particularly suitable for our approach.

## 3 Latent Task-Specific Graph Network Simulators

Figure: Figure 2: Schematic of the LTSGNS architecture. We compute relative node positions stored in the edges of initial mesh and obtain latent node features through a GNN. Combining the node features with a latent task-specific variable, we compute a trajectory per node using ProDMPs. To get the latent variable $z$ for the current context, we approximate the task posterior with a Gaussian Mixture Model (GMM). This requires the gradient of the likelihood with respect to $z$.
Refer to caption: /html/2311.05256/assets/x3.png

#### Graph Network Simulators

A Message Passing Network (MPN) (Sanchez-Gonzalez et al., ; Pfaff et al., ) consists of a series of message passing steps which iteratively update latent node and edge features based on the graph topology.
Given a graph $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathbf{X}_{\mathcal{V}},\mathbf{X}_{\mathcal{E}})$ with nodes $\mathcal{V}$, edges $\mathcal{E}$ and associated vector-valued node and edge features $\mathbf{X}_{\mathcal{V}}$ and $\mathbf{X}_{\mathcal{E}}$, each step is given as

$$ $\displaystyle\textbf{h}^{k+1}_{e}$ $\displaystyle=f^{k}_{\mathcal{E}}(\textbf{h}^{k}_{v},\textbf{h}^{k}_{u},\textbf{h}^{k}_{e})\text{,}\quad\textbf{h}^{k+1}_{v}=f^{k}_{\mathcal{V}}(\textbf{h}^{k}_{v},\bigoplus_{e=(v,u)}\textbf{h}^{k+1}_{e})\text{,}\quad\textrm{with~{}}e=(u,v)\in\mathcal{E}\text{.}$ $$

Here, $\textbf{h}_{v}^{0}$ and $\textbf{h}_{e}^{0}$ are embeddings of the system state per node and edge, and $\oplus$ is a permutation-invariant aggregation such as a sum, max, or mean operator. Each $f^{l}_{\cdot}$ is a learned function such as a small Multilayer Perceptron (MLP).
The network’s final output is a node-wise learned representation $\mathbf{h}$ that encodes local information about the graph topology and the predicted dynamics of the respective parts of the simulated system.

GNSs encode the system state as a graph, feed it through a Message Passing Network (MPN) and interpret the outputs per node as dynamics that can be used to forward the simulation in time using e.g., a forward-Euler integrator.
The state encoding usually uses relative distances and velocities rather than absolute ones, as the resulting translation-invariance allows for better generalization over local areas (Sanchez-Gonzalez et al., ).
When parts of the simulation are known, such as e.g., the positions of a robot’s end-effector for a planned trajectory, only the remaining nodes are predicted.
Existing GNSs usually minimize a next-step Mean Squared Error (MSE) per node during training and produce longer trajectories by iteratively applying the resulting forward dynamics (Pfaff et al., ).
As this iterative dependence on previous predictions causes errors to accumulate over time, carefully tuned implicit de-noising strategies are often added during training (Sanchez-Gonzalez et al., ; Pfaff et al., ; Brandstetter et al., ).
Here, we instead use ProDMPs to directly predict a compact trajectory representation per system node in a single step, reducing the effect of error accumulation similar to, e.g., temporal bundling (Brandstetter et al., ).

#### Meta-Learning and Graph Network Simulators.

We view GNS as a meta-learning problem, where each task consists of simulating a deformable object with unknown material properties over time.
Our goal is to learn a simulator which is adaptable to a specific scenario by observing context data.
Following the notation of Volpp et al. (), the meta-dataset $\mathcal{D}=\mathcal{D}_{1:L}$ consists of simulation tasks $\mathcal{D}_{l}=\{\mathcal{G}_{l},x_{l,1:T},\bm{y}_{l,1:T}\}$ with an initial mesh $\mathcal{G}_{l}$ and time steps $x_{l,t}$ of the node positions $\bm{y}_{l,t}\in\mathbb{R}^{N\times D}$. We use $N$ as the number of nodes and $D$ for the world dimension of the simulation (usually $D=3$ or $D=2$).
In contrast to standard meta-learning, we additionally use the initial graph of the system as a task-level context.
Given the initial positions of the deformable mesh and rigid collider, the task is to predict the node positions $\bm{y}_{l,t}$ at time steps $x_{l,t}$.
We note that the initial graph $\mathcal{G}_{l}$ does not contain the full system state, as we consider the material properties of the deformable object to be unknown.

#### Model architecture.

Our model likelihood $L=p_{\bm{\theta}}(\bm{y}_{l,t}\mid x_{l,t},\bm{z}_{l},\mathcal{G}_{l})$ is parametrized by a global parameter $\bm{\theta}\in\mathbb{R}^{d_{\bm{\theta}}}$ and defines the probability distribution over targets $\bm{y}_{l,t}$ at corresponding timesteps $x_{l,t}$, conditioned on a latent task descriptor $\bm{z}_{l}\in\mathbb{R}^{Z}$ and the initial graph $\mathcal{G}_{l}$, cf. Fig.  [2](#S3.F2). Given the graph $\mathcal{G}_{l}$, the initial mesh of a deformable object and a tractable rigid collider, we connect both objects based on proximity, compute relative distances between nodes and store this information in the edges.
We then use a MPN to encode latent features $\bm{h}\in\mathbb{R}^{H}$ per node. We combine this encoding with the global latent variable $\bm{z}_{l}\in\mathbb{R}^{Z}\sim q_{\bm{\phi}_{l}}(\bm{z}_{l})$ by concatenating $\bm{z}_{l}$ to every node feature $\bm{h}$.
Intuitively, the latent variable $\bm{z}_{l}$ encodes the material properties and high-level deformations of the respective simulation, but does not focus on individual nodes.
We use a simple MLP as the node-level decoder to yield final predictions per node.

Instead of iteratively predicting dynamics for the current simulation step like existing GNS (Pfaff et al., ; Allen et al., ), we use ProDMP (Li et al., ) to predict a representation of the full trajectory.
The node-wise ProDMP weights $w\in\mathbb{R}^{W}$ define a trajectory over the full time horizon, and the ProDMP framework allows an efficient backpropagation from the trajectories to the weights $w$.
As such, we can predict the node positions $\bm{y}_{l,t}$ at the desired time steps $x_{l,t}$ and use a node-wise Gaussian log-likelihood between the given and the predicted node positions to fit our latent space as detailed below.

#### Model predictions.

Under our model, the predictive distribution for a target task $\mathcal{D}_{*}$, given a set of context examples $\mathcal{D}^{c}_{*}\subset\mathcal{D}_{*}$,
is given by

$$ $p_{\bm{\theta}}(\bm{y}_{*,1:T}\mid x_{*,1:T},\mathcal{D}^{c}_{*})=\int\prod_{t=1}^{T}p_{\bm{\theta}}(\bm{y}_{*,t}\mid x_{*,t},\mathcal{G}_{*},\bm{z}_{*})p_{\bm{\theta}}(\bm{z}_{*}\mid\mathcal{D}^{c}_{*})\odif{\bm{z}_{*}},$ (1) $$

where the task posterior distribution is given in terms of the model and a prior distribution $p(\bm{z}_{*})$ over task descriptors by means of Bayes’ theorem as

$$ $p_{\theta}(\bm{z}_{*}\mid\mathcal{D}_{*}^{c})=\prod_{t=1}^{T^{c}}p_{\bm{\theta}}(\bm{y}_{*,t}^{c}\mid x_{*,t}^{c},\bm{z}_{*},\mathcal{G}_{*})p(\bm{z}_{*})\,/\,Z_{*}^{c}(\bm{\theta})\equiv\tilde{p}_{\bm{\theta}}(\bm{z}_{*})\,/\,Z_{*}^{c}(\bm{\theta}).$ (2) $$

Here, $Z_{*}^{c}(\bm{\theta})$ denotes the marginal likelihood of the context data, i.e.,

$$ $Z_{*}^{c}(\bm{\theta})\equiv p_{\bm{\theta}}(\bm{y}_{*,1:T^{c}}^{c}\mid x_{*,1:T^{c}}^{c},\mathcal{G}_{*})=\int\prod_{t=1}^{T^{c}}p_{\bm{\theta}}(\bm{y}_{*,t}^{c}\mid x_{*,t}^{c},\bm{z}_{*},\mathcal{G}_{*})p(\bm{z}_{*})\odif{\bm{z}_{*}}.$ (3) $$

For reasonably complex models, the marginal likelihood and, thus, the posterior distribution is intractable and requires approximation.
As shown by Volpp et al. (), the predictive accuracy is highly dependent on the accuracy of the task posterior approximation, causing us to mimic their approach and employ an expressive full-covariance Gaussian Mixture Model (GMM) of the form

$$ $p_{\bm{\theta}}(\bm{z}_{*}\mid\mathcal{D}_{*})\approx q_{\bm{\phi}_{*}}(\bm{z}_{*})=\sum_{k=1}^{K}w_{*,k}\mathcal{N}(\bm{z}_{*}\mid\bm{\mu}_{*,k},\bm{\Sigma}_{*,k}).$ (4) $$

We use $K$ mixture components with corresponding weights $w_{*,k}$, means $\bm{\mu}_{*,k}$, and covariance matrices $\bm{\Sigma}_{*,k}$, which we collective denote by $\bm{\phi}_{*}$.
To fit the variational distribution $q_{\bm{\phi}_{*}}(\bm{z}_{*})$, we also follow Volpp et al. () and use the Trust Region Natural Gradient Variational Inference (TRNG-VI) method, specifically SEMTRUX (Arenz et al., ).
This requires samples of $\nabla_{z}\tilde{p}_{\bm{\theta}}(\bm{z}_{*})$, which can be readily obtained using standard automatic differentiation tools (Paszke et al., ).

#### Meta-training.

The aim of meta-learning is to automatically encode inductive biases towards the task distribution extracted from the meta-dataset $\mathcal{D}$ in the task-global parameter $\bm{\theta}$.
To this end, we maximize w.r.t. $\bm{\theta}$ the log marginal likelihood, which is given as the sum of the per-task log marginal likelihoods

$$ $\log Z_{l}(\bm{\theta})\equiv\log p_{\bm{\theta}}(\bm{y}_{l,1:T}\mid x_{l,1:T},\mathcal{G}_{l})=\log\int\prod_{t=1}^{T}p_{\bm{\theta}}(\bm{y}_{l,t}\mid x_{l,t},\bm{z}_{l},\mathcal{G}_{l})p(\bm{z}_{l})\odif{\bm{z}_{l}}.$ (5) $$

As discussed above, the marginal likelihood is intractable, which is why we employ an evidence lower bound (ELBO) of the form

$$ $\text{ELBO}_{l}(\bm{\theta})=\mathbb{E}_{q_{\bm{\phi}_{l}}(\bm{z}_{l})}\left(\sum_{t=1}^{T}\log p_{\bm{\theta}}(\bm{y}_{l,t}\mid x_{l,t},\bm{z}_{l},\mathcal{G}_{l})+\log\frac{p(\bm{z}_{l})}{q_{\phi_{l}}(\bm{z}_{l})}\right)\leq\log Z_{\ell}(\bm{\theta})$ $$

as a surrogate objective for maximization of the log marginal likelihood.
The efficiency of this optimization scheme increases with the tightness of the lower bound, which is in turn controlled by the task posterior approximation quality of the variational distribution $q_{\bm{\phi}_{l}}$ (Bishop, ; Volpp et al., ).
Therefore, we also employ the expressive GMM-TRNG-VI approximation procedure to fit $q_{\bm{\phi}_{l}}$ during meta-training.
The resulting ELBO can then be efficiently approximated using a Monte-Carlo estimation of the expectation, and be optimized using standard gradient-based optimization (Kingma & Welling, ; Rezende et al., ; Volpp et al., ; Kingma & Ba, ).
After meta-training, we fix the parameters $\bm{\theta}$ and use them for predictions on unseen tasks.

## 4 Experiments

Figure: Figure 3: (Left) Rollout and (Right) Last Step MSE of the different methods on the Deformable Plate task. Both metrics are plotted on a logarithmic scale. LTSGNS outperforms MGN and its MP variant for both metrics from a single context point, further improving its performance when provided with additional context information. During inference, point cloud information can be used as a context set, making for easier context acquisition at the cost of slightly worse predictions. Overall, the performance metrics for both Rollout and Last Step MSE are fairly similar. However, the Last Step MSE exhibits increased variance across all methods, making it less consistent in comparison.

#### Setup.

We base our experimental setup on that of Linkerhägner et al. (), adapting it to meta-learning and the prediction of full trajectories using movement primitives.
To represent the system as a graph, we employ one-hot encoding to differentiate between deformable objects and colliders, and encode relative distances between neighboring nodes in their edges.
We do not use explicit world edges, but include the distances between nodes in mesh space.
We utilize the context set $\mathcal{D}^{c}_{*}$, comprising the initial mesh and $C$ randomly sampled simulation states, to predict node-wise ProDMP parameters $w$. These parameters encode the full mesh trajectory, and can be queried to obtain the system state at arbitrary timesteps.

All models are optimized using Adam (Kingma & Ba, ) with a learning rate of $5\times 10^{-4}$.
For the MPNs, we use $5$ separate message passing steps for all methods.
For LTSGNS, we repeat each step $2$ times in an inner loop to increase the receptive field of the individual nodes, as there is no iterative prediction that could otherwise pass implicit information.
Each block uses a latent dimension of $128$, LeakyReLU activations and a $1$-layer MLPs for its node and edge updates.
We additionally apply Layer Normalization (Ba et al., ) and Residual Connections (He et al., ) independently for both node and edge updates.
We repeat each experiment for $5$ random seeds and report the mean and standard deviation over these seeds.
Each run’s results are averaged over all trajectory steps and test set trajectories.
We evaluate the models using the Rollout MSE, which is the average MSEs of all simulation steps, and the Last Step MSE, which is the error of the final simulation step.
Additional details on our experimental setup are provided in Appendix [A](#A1).

#### Tasks.

We consider a simpler $2$-dimensional Deformable Plate and a more challenging $3$-dimensional Tissue Manipulation task (Linkerhägner et al., ).
Both tasks use Simulation Open Framework Architecture (SOFA) (Faure et al., ) to generate the underlying ground truth data, and use triangular surface meshes for the simulation.
While the initial meshes are known, both tasks use materials with a randomized and unknown Poisson’s ratio (Lim, ) that governs whether the material contracts or expands under deformation.
The Deformable Plate task simulates different trapezoids that are deformed by a circular collider with constant velocity and varying size and starting position.
Each trajectory consists of a mesh with $81$ nodes that is deformed over $50$ timesteps, and we use $675/135/135$ trajectories for training, evaluation and testing respectively.
The Tissue Manipulation task simulates a common scenario in surgical robotics where a piece of tissue is deformed by a gripper.
Here, the gripper starts attached to a random position of the object and moves in a random direction with constant velocity. The simulated mesh has $361$ nodes, and we use $600/120/120$ training, evaluation and testing trajectories with $100$ steps each.
All tasks are normalized to be in $[-1,1]$.

#### Baselines and Ablations.

We use MGN (Pfaff et al., ) as our main baseline.
MGN iteratively predicts the velocities of the current simulation step to generate the next mesh state.
It is trained to minimize the $1$-step MSE over node velocities and crucially employs Gaussian input noise (Brandstetter et al., ) to prevent error accumulation over time and thus generalize from $1$-step predictions to larger rollouts during inference.
We follow previous work (Linkerhägner et al., ) and set the standard deviation of the input noise to $0.01$.
We experiment with both MGN without information about the material properties, and with a variant that includes this additional information, called MGN(M).
The latter renders the simulation deterministic with respect to the initial system state and sets an upper performance limit for the standard MGN model.
Additionally, we compare to an MGN(MP) variant that incorporate ProDMPs to directly predict a trajectory for each node feature instead of iteratively predicting the next state.
Since this method directly predicts the full trajectory, it does not use Gaussian input noise.
Building on our model’s adaptability to different context sets, we conduct an ablation experiment focused on practical applications.
Specifically, we test the use of point clouds as the context set during inference.
This is particularly relevant as point clouds can be readily generated from depth cameras in real-world settings, while mesh states cannot.
Importantly, this adjustment requires no modifications to the existing training process.

Figure: Figure 4: (Left) Rollout and (Right) Last Step MSE of the different methods on the Tissue Manipulation task. Both metrics are plotted on a logarithmic scale. LTSGNS yields more accurate simulations with increasing context size, clearly outperforming the MGN baselines for $5$ or more context points.

Results.
Figure [3](#S4.F3) shows results for the Deformable Plate task.
We find that LTSGNS outperforms MGN even when provided with just a single context point.
The model’s performance continues to improve as the size of the context set increases.
Specifically, for a context set with $10$ points, LTSGNS outperforms MGN (M) even though the latter has direct access to the ground truth material information.
Although the performance drops when using point cloud data as context instead of system states, LTSGNS still outperforms MGN and benefits from the addition of more contextual information.
Similarly, Figure [4](#S4.F4) evaluates the Tissue Manipulation task.
Here, LTSGNS again improves with an increasing context size, outperforming all MGN baselines for $5$ or more context points.
Providing material properties still improves MGN, but the difference is less significant than for the Deformable Plate task, presumably because the dynamics are overall harder to predict even with this additional information.
We provide additional results for larger context sets in Appendix [B.1](#A2.SS1).

Figure: Figure 5: Final simulation step of an exemplary test trajectory for the Tissue Manipulation task across different methods. Blue denotes the ground position of the deformable object, while the wireframe and yellow shading outline the predicted mesh. Only LTSGNS provides accurate predictions, which further improve as the size of the context set increases.
Refer to caption: /html/2311.05256/assets/x4.png

Figure [5](#S4.F5) visualizes exemplary final simulation steps, supporting the findings above.
LTSGNS accurately simulates the object’s deformation from a single context point, and further improves when provided with additional context information.
Opposed to this, MGN, even when provided with explicit material properties or when combined with ProDMPs, fails to produce consistent meshes.
Appendix [B.2](#A2.SS2) shows visualizations for full rollouts.

## 5 Conclusion

We introduce Latent Task-Specific Graph Network Simulator (LTSGNS), a novel Graph Network Simulator that employs meta-learning and movement primitives for accurate probabilistic predictions in physical simulations.
Our model uses meta-learning and movement primitives to effectively addresses the issue of error accumulation and dynamically adapt to context information during inference.
LTSGNS is also able to accommodate sensory inputs like point clouds during inference, broadening its applicability in real-world scenarios.
We experimentally validate the effectiveness of our approach compared to existing baselines, particularly in tasks involving uncertain material properties.
The model’s ability to produce distributions over trajectories adds an extra layer of robustness, making it a valuable tool for both academic research and practical applications where data is sparse.

#### Limitations and Future Work

We currently consider each trajectory as a task, and require data from either states or point clouds of this trajectory to fit our model during inference.
However, as generating such data is often impractical in real-world scenarios, we plan to instead group different material properties into tasks.
This shift will enable the model to encode abstract material behavior, without being tied to a particular mesh topology or simulation.
We further aim to extend our approach to accommodate longer simulations to fully capitalize on the benefits of incorporating movement primitives.
Combining both, we plan to apply our model to real-world deformations.
Here, an additional focus will be on integrating online re-planning of trajectories, thereby enhancing prediction accuracy when live sensory information is available.

## 6 Acknowledgments

This work was funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) – Project 5339.
This work was supported by funding from the pilot program Core Informatics of the Helmholtz
Association (HGF). The authors acknowledge support by the state of Baden-Württemberg through
bwHPC, as well as the HoreKa supercomputer funded by the Ministry of Science, Research and the
Arts Baden-Württemberg and by the German Federal Ministry of Education and Research.

## Appendix A Experimental Protocol

In order to promote reproducibility, we offer comprehensive information regarding our experimental methodology. In Table[1](#A1.T1), we give our used hyperparameters for the experiments. To obtain a more detailed description of the tasks and the dataset, please refer to Linkerhägner et al. () Appendix B.

In this paper, we use a training approach that adapts its batch size dynamically based on the context set size. To achieve this, we define a batch cost as an upper bound, calculated as 0.8 plus 0.2 times the batch size. The batch size is then meticulously chosen to align with the specified batch cost, allowing us to optimize our training process effectively. Specifically, for the deformable plate task, we set a target total batch cost of 60, whereas for the tissue manipulation task, we aim for a batch cost of 300. The training for the deformable plate task was conducted on an NVIDIA GeForce RTX 3080 GPU, while the tissue manipulation task utilized an NVIDIA A100 GPU for efficient processing and optimization. The availability of more powerful hardware enabled us to opt for a larger batch cost in the context of the tissue manipulation task.

For each trajectory, we generate 30 auxiliary training tasks, as inspired by the approach outlined in Volpp et al. (), Appendix A3.2. In this process, we randomly select pairs of $(x_{l,t},\bm{y}_{l,t})$ for each auxiliary training task. Additionally, we designate one of the first 30 time steps as the time steps for constructing the graphs denoted as $\mathcal{G}_{l}$.
During the testing phase, the tasks consistently make use of the initial mesh for the corresponding $\mathcal{G}_{l}$, and we adapt the context size $C$ differently for each evaluation scenario.

**Table 1: Table listing the hyperparameters and configurations of the experiments**
| Parameter | Value |
| --- | --- |
| Batches per epoch | $500$ |
| Epochs | $1000$ |
| Node feature dimension | $128$ |
| Latent variable dimension | $8$ |
| Decoder dimension (deformable plate) | $128$ |
| Decoder dimension (tissue manipulation) | $256$ |
| Message passing blocks | 5 |
| Message passing repeats (MGN) | 1 |
| Message passing repeats (LTSGNS) | 2 |
| MPN Aggregation function | Mean |
| Learning rate | $5.0e-4$ |
| Gaussian likelihood standard deviation | $0.01$ |
| Number of $z$ samples for ELBO estimation | $32$ |
| Number of GMM components | 3 |
| Auxiliary train tasks per trajectory | $30$ |
| Activation function | Leaky ReLU |
| SEMTRUX component KL bound | 0.01 |
| Number of ProDMP basis functions | $10$ |

## Appendix B Additional Results

### B.1 Evaluations.

We additionally show how LTSGNS performs for larger context sizes in Figures [6](#A2.F6) (Deformable Plate) and [7](#A2.F7) (Tissue Manipulation).
While the inclusion of more context information generally enhances performance, it sometimes diminishes performance for the Last Step MSE when point clouds are used as the context.
We hypothesize that this effect occurs because a larger context size shifts the balance between the likelihood $L$ and the prior $p(z_{l})$, thereby magnifying any existing inaccuracies in the model that may exist in more complex tasks.

Figure: Figure 6: (Left) Rollout and (Right) Last Step MSE for larger context sizes ($C=20$, $C=30$) compared to the methods presented in Figure [3](#S4.F3) for the Tissue Manipulation task.

Figure: Figure 7: (Left) Rollout and (Right) Last Step MSE for larger context sizes ($C=20$, $C=30$) compared to the methods presented in Figure [4](#S4.F4) for the Tissue Manipulation task.

### B.2 Visualizations.

We provide additional visualizations for all methods over different timesteps.
Figure [6](#A2.F6) shows a test trajectory for the Deformable Plate task, and Figure [9](#A2.F9) visualizes the same for the Tissue Manipulation task.
Both figures show that LTSGNS performs similar to MGN for a single context point, but significantly improves performance when provided with additional information.
Especially for $C=10$ data points in the context set, the predictions are visually consistent with the ground truth.

Figure: Figure 8: Simulation over time of an exemplary test trajectory of the Deformable Plate task across different methods. Blue denotes the ground position of the deformable object, while the wireframe and yellow shading outline the predicted mesh.
Refer to caption: /html/2311.05256/assets/x5.png

Figure: Figure 9: Simulation over time of an exemplary test trajectory of the Tissue Manipulation task across different methods. Blue denotes the ground position of the deformable object, while the wireframe and yellow shading outline the predicted mesh.
Refer to caption: /html/2311.05256/assets/x6.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. How does ProDMP generate smooth trajectories based on the predefined conditions of the initial state? Please give detailed justification and explanation.

2. Could the author provide a detailed explanation of how a meta-learning problem can contribute to simulating new scenarios?
soundness: 3
strengths: 1. This paper shows a clear motivation for initial state uncertainty and data limitation, which are all critical problems in related research fields.

2. Consider the "node-level latent features," which is, to the best of my knowledge, a novel method for solving such a problem.

3. The results of the new simulation task in the paper are convincing for the proposed method.
summary: This paper proposes a graph network simulator for mesh-based simulation on material study. The framework is constructed on a meta-learning problem and applies conditional Neural Processes to address data limitations. This paper shows both qualitative and quantitative experiments.
weaknesses: 1. Some methodology details are unclear, especially in the "Probabilistic Dynamic Movement Primitives" section and "Meta-Learning and Graph Network Simulators."

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: 1. The model's architecture is not clearly explained, and it is unclear why certain modules are necessary. For example, from the results, it seems that MGN, even without history information, can surpass M3GN in performance. This raises questions about the value of incorporating historical information in M3GN. Moreover, the experimental results do not clearly demonstrate the necessity or advantages of using a meta-learning scheme. A thorough analysis on how meta-learning benefits model performance would be valuable, including ablation studies comparing model performance with and without meta-learning.
2. The authors claim that the baseline MGN does not incorporate historical information, which appears inaccurate. In certain datasets, MGN does include history. For a fair comparison, the MGN baseline should also be evaluated with historical data to assess its impact on performance.
3. The results section only reports the average MSE across all time steps. It would be helpful to provide a comparison of MSE over the number of prediction steps, as this would give insight into the model's performance stability over time as claimed in the paper.
4. Based on Figure 3, the proposed M3GN method does not appear to use ground truth collider information. If this is the case, does the collider state being predicted by the mode? How accurate is the collider state prediction, especially when history steps are limited? Additionally, including collider ground truth (as in MGN) is actually intuitive and makes sense, as the primary goal of developing a simulation model is to understand how a solid deforms under varying contact forces and obstacle displacements. Predicting these external forces may not be necessary for achieving this objective.
5. It would be informative to visualize the node-level latent task descriptions learned by the model. Such visualizations could help in understanding how task-specific information is represented.
6. The datasets used in this paper have relatively small node counts compared to those in previous MGN studies or those used in other related papers. When the number of nodes increases significantly, it is concerned that M3GN may struggle due to the large number of historical steps required. Comparing M3GN’s memory usage with MGN’s would provide a more comprehensive evaluation. 
7. The authors consider each trajectory as a separate task with varying context sizes. However, this approach may not align with the broader goals of meta-learning, as tasks are typically defined by consistent properties such as the same material setting. Currently, the meta-learning setup seems more focused on adapting to different context sizes rather than generalizing across diverse tasks.
8. As the input context size changes, will the number of predicted steps vary as well? If so, the model’s ability to generalize to different context sizes is unclear, and it may not be as flexible as MGN in this respect. Any experiments or evaluation on this aspect? Additionally, splitting single data points into multiple input-output sets seem to increase the effective amount of training data for M3GN, potentially creating an unfair comparison with MGN which use less training data.
9. The authors do not specify how material properties are incorporated. Also, it is unclear whether the test data involve material properties that are in-distribution or out-of-distribution relative to the training data. Providing this information is crucial for evaluating the model's generalization capabilities.
10. The authors mention that material node features are not added to M3GN. Given that these features enhance MGN's performance, it would be useful to understand the rationale for this exclusion and perform related ablation study.
11. Although the authors mention other methods in related work besides MGN, these methods are not included in the baselines. Some of these methods have better accuracy and efficiency. Including these additional baselines would provide a clearer view of M3GN’s comparative performance.
12. Will the data used in this study be publicly available? Making the dataset accessible would facilitate further research and replication studies.
soundness: 2
strengths: The paper takes a novel approach to enhancing rollout stability by predicting entire future mesh states, and it incorporates a meta-learning scheme to improve adaptability within the simulation framework.
summary: This paper introduces Movement-primitive Meta-MeshGraphNet (M3GN), a model for simulating object deformations in data-limited scenarios. M3GN combines meta-learning and movement primitives to improve the adaptability and accuracy of Graph Network Simulators (GNSs) by framing mesh-based simulation as a meta-learning task.
weaknesses: While the approach appears novel, the rationale behind certain modules in the model is unclear, and the results do not provide sufficient evidence to justify their inclusion. Also, the paper is not clearly written and sometimes hard to follow. The detailed comments and suggestions are listed below.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What is the timestep for simulation?

2. A figure illustraing all the relation and symbols of input, output can be added. Fig.3 Right is not information for undertanding the task setting.
soundness: 3
strengths: This work aims to address two important problems in learning-based simulation:

1. It treats the simulation as a trajectory-level meta-learing problem and use trajectory history as the context to predict future trajectories.

2. It mitigates the problem of error accumulation by using ProDMP to directly predict the full simulation trajectories.

The paper is well structured and written.
summary: In this paper, the authors propose a graph network simulator that combines movement primitives and trajectory-level meta-learning. The network uses the simulation history as the context information to predict the deformation for objects with unknown properties. They also use probabilistic dynamic movement primitives to represent the future trajecteries and directly predicts the full simulation trajectories instead of iteratively predicting the next-step. Experiments show that it outperforms STOA in different simulation tasks. Abalation studies validate the effectivenss of the design choice.
weaknesses: 1. Some descriptions are unclear and some important details are missing.
(1) in line 242, "graph edges between the deformable object and collider are added based on physical proximity to model interactions
 between objects." what is the physical proximity exactly? Since the deformation mesh node position for the end timestep is unknown, I suppose we cannot use that to compute the distance. Whether this edge creation is done only for known timesteps or if it's updated during prediction?

(2) in line 231, why is the term c_1y_1(t) + c_2y_2(t) only depending on the inital conditions? What is the representation of the pre-computed basis fuction \phi?

2. More detailed description of the training/val/test split should be added. Specify how trajectories are divided between training, validation, and test sets. What are different between training and test? Clarify if test trajectories involve different objects, material properties, or initial conditions than training trajectories. In the limitation part, it is claimed ''We currently consider each trajectory as a task, and require initial states of this trajectory as a context set during inference."

3. Since the method needs a trajectory with simulated states as context, the author better include a runtime comparison between your method (including context computation) and traditional simulators for predicting the same number of future timesteps and discuss the trade-offs between computation time and accuracy compared to traditional simulators.

---

confidence: 5
contribution: 3
initial_rating: 5
presentation: 4
questions: See above.
soundness: 3
strengths: Strength:
1. Adopting meta-learning to deal with dynamic prediction tasks is novel, especially the concept of regarding each tajectory as a new task is interesting. 

2. The authors consider past information and the eventual state of the collider as the condition to predict the subsequent movement trajectory, which make the network infer the future from the past rather than remember the dynamic behaviour of a certain material. In addition, predicting the whole rest path by a single forward pass could significantly improve the efficiency, compared with previous Graph-based single timestep prediction.
summary: This paper first propose a meta-learning framework to efficiently learn generalizable mesh-based dynamic prediction tasks. Different from previous graph neural simulators which predict the state updates in a step-by-step manner, the proposed M3GN targets to predict the whole trajectories by a conditional neural process to effectively diminish the error accumulation issue.
weaknesses: Weakness:
1. This paper is highly related to the Graph-based Neural Simulators. However, in the related work section, the latest advancements in this field are not included, and most of the work discussed is from 2023 or earlier. This could make the paper appear somewhat outdated. I believe this section could benefit from a more comprehensive overview of the field, especially more works from 2024. Below are two of the latest advancements about Graph Network Simulators that I recommend the authors to discuss them in Section 2.1 ,or better, use them as baselines for comparison. However, given the tight rebuttal timeline, it is also tolerated that concurrent works were not included for comparison.

    (1) "DEL: Discrete Element Learner for Learning 3D Particle Dynamics with Neural Rendering" 2024   ..  This work integrate traditional Newton mechanics into the graph network design to benefit from mechanics priors for longer term prediction.

    (2) "Equivariant graph neural operator for modeling 3d dynamics" 2024  ..  This paper deal with dynamic prediction tasks as trajectory-level rather than next-step level by operator learning, which is somewhat relavent with this reviewing work. Also, it handle the equivariant issues. 

2. For Equation 3, does it use past trajectory collider states when encoding z because I saw that you seem to only use the latest state, or does it rely solely on the historical information of the deformed object? I believe it would be more reasonable to use all the historical information of the collider here as well, since the deformation of the mesh is passive. 

3. If this method is trained on an elastic dataset, can it generalize directly to elastoplastic materials? I believe it would be worthwhile to discuss the generalization across different materials in the experiments, rather than limiting it to variations in mechanical parameters within the same material. 

4. Line 276 mentions that the context information z is concatenated with the node features. Is the same z concatenated to each node?

5. Finally, the neural network predicts a set of weights, and the shape of the weight matrix is 𝑇, 𝐷,3. Which basis functions are these weights applied to in order to obtain the predicted trajectory? Are they precomputed from the historical trajectory? If yes, how?

In appendix A.2 "Initially, we integrate a relative goal position as part of the node weights w" What's the exact mean of the relative goal position? 

I will raise the score if most of concerns are well addressed by the authors.

Meta-review: **Summary** The paper propose a Meta-MeshGraphNet model for simulate object deformations. The approach is to meta-learn a graph-network simulator across different types of deformations with varying object materials. For each trajectory, the model is conditioned on the context of initial steps of the simulation that allows to implicitly infer deformation parameters.

**Strengths** The authors demonstrate that graph-network simulators are able to infer deformation parameters from the few initial steps and predict the rest of the simulation for the first time in the literature. They show a variety of deformation problems such as Planar Bending, Deformable Plate, Tissue Manipulation and Teddy Bear Falling.

**Weaknesses** The main weakness and concern is that the paper presents using the history of first few steps of the simulation as the novel approach to meta-learn the deformation parameters, however the previous approaches MeshGraphNet (MGN) and GNS used conditioning on history of several simulation time steps. The reviewers also had questions about rationale behind certain modules, namely ProDMPs. Reviewers point out lack of justification for using ProDMPs and the ablations with and without ProDMPs.

**Decision**  Rejection. Although making the graph network across different deformation parameters is novel, the paper will require substantial reframing to the main claim that MGN and GNS baselines also include the history (Figure 3) and updating the main results (Figures 5 and 6) to include MGN and GNS with history. Therefore, the paper cannot be accepted in its current form.

Decision: Reject