Learning Rigid Body Dynamics with Lagrangian Graph Neural Network
Anonymous Author(s
AffiliationAddress
Abstract
Lagrangian and Hamiltonian neural networks (L NN and H NN respectively) encodestrong inductive biases that allow them to outperform other models of physicalsystems significantly.
However, these models have, thus far, mostly been limitedto simple systems such as pendulums and springs or a single rigid body suchas a gyroscope or a rigid rotor.
Here, we present a Lagrangian graph neuralnetwork (L GNN ) that can learn the dynamics of rigid bodies by exploiting theirtopology.
We demonstrate the performance of L GNN by learning the dynamicsof ropes, chains, and trusses with the bars modeled as rigid bodies.
L GNN alsoexhibits generalizability—L GNN trained on chains with a few segments exhibitsgeneralizability to simulate a chain with large number of links and arbitrary linklength.
We also show that the L GNN can simulate unseen hybrid systems includingbars and chains, on which they have not been trained on.
Specifically, we show thatthe L GNN can be used to model the dynamics of complex real-world structuressuch as the stability of tensegrity structures.
Finally, we discuss the non-diagonalnature of the mass matrix and it’s ability to generalize in complex systems.
Introduction and Related Works
Movements of a robotic arm, rolling ball, or falling chain can be characterized by rigid bodymotion [1, 2].
Understanding the dynamics of the motion is crucial in several applications includingrobotics, human-robot interaction, planning and computer graphics [3, 1].
Traditionally, the rigidbody mechanics is studied in the framework of classical mechanics, which relies on either forcebased or energy-based approaches [4].
Force-based approaches involve the computation of all theunknown forces based on the equations of equilibrium and hence is cumbersome for large structures.
Energy-based approaches present an elegant formalism which involve the computation of a scalarquantity representing the state of a system, namely, Lagrangian ( L = T − V ), which is the differencebetween the kinetic ( T ( q, ˙ q )) and potential ( V ( q )) energies, or Hamiltonian ( H = T + V ) , whichrepresents the total energy of the system.
This scalar quantity can, in turn, be used to predict thedynamics of the system.
However, the functional form governing this scalar quantity may not beknown a priori in many cases [5].
Thus, learning the dynamics of rigid bodies directly from thetrajectory can simplify and accelerate the modeling of these systems [5, 6, 7, 8].
Learning the dynamics of particles has received much attention recently using physics-informedapproaches [9].
Among these, Lagrangian neural networks (L NN s) and Hamiltonian neural networks(H NN s) are two physics-informed neural networks with strong inductive biases that outperform otherlearning paradigms of dynamical systems [10, 11, 12, 8, 6, 13, 7].
In this approach, a neural networkis trained to learn the L (or H ) of a system based on its configuration ( q, ˙ q ) .
The L is then used alongwith the EL equation to obtain the time evolution of the system.
Note that the training of L NN s isperformed by minimizing the error on the predicted trajectory with respect to the actual trajectory.
Thus, L NN s can effectively learn the Lagrangian directly from the trajectory of a multi-particlesystem [6, 13].
Most of the works on L NN has focused on relatively simpler particle-based systems such as springsand pendulums [6, 13, 7, 10, 14].
This approach models a rigid body, for instance a ball, as a particleand predicts the dynamics.
This approach thus ignores the additional rotational degrees of freedomof the body due to its finite volume.
Specifically, while a particle in 3D has three degrees of freedom(translational), a rigid body in 3D has six degrees of freedom (translational and rotational).
Thus,the dynamics and energetics associated with these degrees of motions are lost by modeling a rigidbody as a particle.
To the best of authors’ knowledge, thus far, only one work has attempted to learnrigid body dynamics using L NN s and H NN s, where it was demonstrated the dynamics of simple rigidbodies such as a gyroscope or rotating rotor can be learned [13].
However, the L NN s used in thiswork, owing to their fully connected MLP architecture, are transductive in nature.
An L NN trainedon a double-pendulum system or 3-spring system can be used only for the same system and doesnot generalize to a different system size such as 3-pendulum or 5-spring, respectively.
In realisticsituations the number of particles in a system can vary arbitrarily, and accordingly, a large numberof trained models might be required to model these systems.
An alternate approach to model these systems would be to use a graph neural network (G NN ) [15, 16,5], which, once trained, can generalize to arbitrary system sizes.
G NN s have been widely used tomodel physical and atomic systems extensively due to their inductive bias [17, 18, 19].
G NN s havealso been used to model rigid bodies mainly following two approaches, namely, particle-based [16]and lumped mass [19, 20] methods.
In the first approach, a rigid body is discretized into finitenumber of particles and the motion of the individual particles are learned to predict the dynamicsof rigid body [16].
Note that this approach is philosophically similar to mess-less methods such assmoothed-particle hydrodynamics (SPH) [21] or peridynamics (PD) [22], where the time-evolutionof a continuum body is simulated by discretizing the domain using particles.
This approach [16],although useful, have several limitations, namely, it does not (i) conserve physical quantities such asenergy when simulated over a long duration, and (ii) generalize to a different timestep of forwardsimulation than the one on which it is trained.
In the second approach, a rigid body is modeled as alumped mass [19, 23], the dynamics of which is learned by assuming this lumped mass as a particle.
For instance, the dynamics of a chain is modeled by discretizing the chain to smaller segments andmodeling each segment as a lumped mass.
As mentioned earlier, this approach leads to the loss ofadditional degrees of freedom that are associated with a rigid body.
Here, we present a Lagrangian graph neural network (LG NN ) framework that can learn the dynamicsof rigid bodies.
Specifically, exploiting the topology of a physical system, we show that a rigid bodycan be modeled as a graph.
Further, the Lagrangian of the graph structure can be learned directlyby minimizing the loss on the predicted trajectory with respect to the actual trajectory of the system.
The major contributions of the work are as follows.
• Topology aware modeling of rigid body.
We present a graph-based model for rigid bodies suchas in-extensible ropes, chains, or trusses.
Further, we demonstrate using LG NN that the dynamics
of these systems can be learned in the Lagrangian framework.
• Generalizability to arbitrary system sizes.
We show that LG NN can generalize to arbitrarysystem sizes once trained.
• Generalizability to complex unseen topology.
We demonstrate that the LG NN can generalize tounseen topology, that is, links with varying lengths, a combination of truss and chain structures,and different boundary conditions.
Altogether, we demonstrate that LG NN can be a strong framework for simulating the dynamics ofrigid bodies.
Dynamics of Rigid Bodies
The dynamics of a physical system can be represented as ¨ q = F ( q, ˙ q, t ) , where q, ˙ q ∈ R D is afunction of time ( t ) for a system with D degrees of freedom.
The future states or trajectory of thesystem can be predicted by integrating these equations to obtain q ( t + 1) and so on.
While thereare several physics-based methods for generating the dynamics of the system such as d’Alembert’sprinciple, Newtonian, Lagrangian, or Hamiltonian approaches, all these approaches result in theequivalent sets of equations [3].
The two broad paradigms for modeling the dynamics involve force- and energy-based approaches.
Energy-based approaches is an elegant framework, which relies on the computation of a single scalarquantity, for instance energy, that represents the state of system.
The dynamics of the system is, inturn, computed based on this scalar quantity.
Among the energy-based approaches, Lagrangian formulation has been widely used to predict the dynamics of particles and rigid bodies by computing the Lagrangian L of the system.
The standard form of Lagrange’s equation for a system with holonomicconstraints is given by ddt ( ∇ ˙ q L ) − ( ∇ q L ) = 0 , and the Lagrangian is L ( q, ˙ q, t )
=
T ( q, ˙ q, t ) −V ( q, t )with T ( q, ˙ q, t ) and V ( q, t ) representing the total kinetic energy of the system and the potential function from which generalized forces can be derived.
Accordingly, the dynamics of the system can berepresented using Euler-Lagrange (EL) equations as ¨ q i = ( ∇ ˙ q i ˙ q i L ) − 1 [ ∇ q i L − ( ∇ ˙ q i q i L ) ˙
q i ] .
Modified Euler-Lagrange Equation.
A modified version of the EL can be used in cases where someof the terms involved in the equation can be decoupled.
This formulation allows explicit incorporationof constraints (holonomic and Pfaffian) and additional dissipative terms for friction or drag [3, 1].
Inrigid body motion, Pfaffian constraints can be crucial in applications such as multi-fingered graspingwhere, the velocity of two or more fingers are constrained so that the combined geometry formed isable to catch or hold an object.
A generic expression of constraints for these systems that accountsfor both holonomic and Pfaffian can be A ( q ) ˙
q = 0 , where, A ( q ) ∈ R k × D represents k velocityconstraints.
In addition, drag, friction or other dissipative terms of a system can be expressed as anadditional forcing term in the EL equation.
It is worth noting that EL equation, by nature, is energyconserving.
Hence, the additional dissipative terms are crucial for modeling realistic systems withfriction and drag.
If these terms are not included, the system will essentially try to simulate an energypreserving trajectory, thereby resulting in huge errors in the dynamics [14].
Considering the additional forces mentioned above, the modified EL equation can be written as:
[Equation]
where A T forms a non-normalized basis for the constraint forces, λ ∈ R k , known as the Lagrangemultipliers, gives the relative magnitudes of these force constraints, Υ represents the non-conservativeforces, such as friction or drag, which are not directly derivable from a potential, and F representsany external forces acting on the system.
This equation can be modified to obtain ¨ q as:
[Equation]
where M = ∂∂ ˙ q ∂ L ∂ ˙
q represents the mass matrix, C = ∂∂q
∂
L ∂ ˙ q represents Coriolis-like forces, andΠ =
∂
L ∂q represents the conservative forces derivable from a potential.
Differentiating the constraintequation gives A ( q )¨ q + ˙ A ( q ) ˙ q = 0 .
Solving for λ and substituting in Equation 2, we obtain ¨ q as
[Equation]
For a system subjected to these forces, the dynamics can be learned using L NN by minimizing theloss on the predicted and observed trajectory, where the predicted acceleration ˆ¨ q is obtained using the
Equation 3.
It is worth noting that in this equation, M, C, and Π can be directly derived from the
L .
Constraints on the systems are generally known as they generally form part of the topology.
It isworth noting that there are some recent works that focus on learning constraints as well [8].
Lagrangian Mechanics for Rigid Bodies
In the case of particle systems such as spring or pendulum systems, the approach mentioned in
Sec.2 can be directly used in conjunction with an L NN to learn the dynamics.
In this case, themass matrix M ( q ) remains constant with only diagonal entries m ii in Cartesian coordinates.
Inducing this as a prior knowledge, wherein the masses are parameterized as a diagonal matrix isshown to simplify the learning process [13].
However, in the case of a rigid body, the mass matrix is non-diagonal in the Cartesian coordinates.
Further, the kinetic energy term T becomes afunction of both position and velocity.
In other words, the kinetic energy also becomes a functionof the topology.
This makes learning the dynamics a complex problem especially in real-worldcomplex structures such as trusses or tensegrities, which are a combination of bars, ropes, and chains.
To this extent, we briefly review the mechanics of a falling ropeor chain as an example.
Note that simple rigid bodies such asa gyroscope or rotating rotor has already been studied using
L NN s [13].
Of our special interest are rigid bodies that can be arbitrarily large such as chains, ropes or trusses, that can be dividedinto smaller constituent members.
This is because, it is generallyassumed that extending L NN s to large structures has been posed achallenging problem [14].
Traditionally, the mechanics of chainsor ropes are modeled using discrete models [2].
Figure 1 showsa discrete model of a rope of mass M and length L .
The ropeis discretized into n cylindrical rods or segments each havinga mass m i = M/n and length l
i
=
L/n .
These segments areconsidered to be rigid, and with a finite uniform cross-sectionalarea and volume.
In order to replicate realistic dynamics of arope, the l i should be significantly smaller that L .
Note that inthe case of a chain or truss, such artificial discretization is notrequired and the bars associated with each segment can be directly considered as a rigid body.
To formulate the L , the generalized coordinates orientation of each link represented by ϕ i = ( y i − y
i − 1 ) x i − x i − 1can be considered.
Placing the origin at the beginning of first segment (see Figure 1), the center ofmass of i th segment ( x cmi , y cmi ) can be written in terms of generalized coordinates as
[Equation]
Accordingly, the kinetic energy of the system is given by [2]
[Equation]
where I i =
112 m i
l 2 i represents the moment of inertia of the rigid segment i .
Similarly, the potentialenergy of the system can be expressed as:
[Equation]
where g represents the acceleration due to gravity.
Finally, the Lagrangian of the system can beobtained as L = T − V , which can be substituted in the EL equation to obtain the dynamics of therigid body.
Lagrangian Graph Neural Network
To learn the dynamics of rigid body, weemploy the approach shown in Figure 2.
Specifically, we model a physical system asa graph.
Further, the Lagrangian of systemis learned by decoupling the potential andkinetic energy, each of which are learnedby two G NN s, namely, G V and G T .
Finally,the Lagrangian is computed as L = T − V .
This framework is trained end-to-end basedby minimizing the loss on the accelerationpredicted by the LG NN using EL equationwith respect to the ground truth.
In thissection, we describe the LG NN architecturefor rigid bodies in detail (See Figure 3 for an overview).
We empirically show that the dynamics of arigid body can be learned by LG NN .
In addition, due to the inductive nature of the graph architecture,once trained on a small system, LG NN can generalize to arbitrary system sizes and topology. 4.1 Graph architecture
Graph structure.
Figure 1 shows a chain.
The (undirected) graph of the physical system isconstructed by considering the bars/segments of the chain as the edges and the connections as nodes.
Here, edges represent the rigid bodies and node represent the connection between these rigid bodies.
This is in contrast to earlier approaches used for particle-based systems, where node represented theparticle position and edge represented the connections between them.
Hereon, we use the notation
G ( U , E ) to to represent the graph representation of a rigid body with U and E as its node and edge sets.
Overview of the architecture.
As shown in Figure 3, we use two G NN s; one to predict the potentialenergies and the other to predict kinetic energies.
From these predictions the Lagrangian is computed.
The error on the Lagrangian is minimized through an RMSE loss function to jointly train both the
G NN s. The architecture of both the G NN s, shown in Figure 3, are identical.
Input features.
Each node u i ∈
U is characterized by its position q i =
( x i , y i , z i ) , and velocity( ˙ q i ).
Each edge e ij is characterized by its type t ij , and the relative differences in the positions( ∆ q ij =
q
i − q j ) of its connecting nodes, and ω ij = ∆
q ij × ∆ ˙ q ij .
The type t ij is a discrete variableand is useful in distinguishing edges of different characteristics within a system (Ex. moment inertiaor area of cross section of the edge).
Note that the velocity of a rigid body represented by an edge is afunction of the velocities of its end points in two and three dimensional spaces.
Hence, we do notexplicitly track edge velocities.
Pre-Processing.
In the pre-processing layer, we construct a dense vector representation for each nodev i ∈ U and edge e ij ∈ E using MLP s (multi-layer perceptrons).
The exact operation for potentialenergy is provided below in Eqs.7-8.
For kinetic energy, we input ˙
q i in Eq 7 instead of q i and ω ij in
Eq. 8 instead of ∆ q ij .
[Equation]
squareplus is an activation function.
Message passing.
To infuse structural information in the edge and node embeddings, we perform Llayers of message passing, wherein the embedding in each layer l ∈ [1 , · , L ] is computed as follows:
[Equation]
Here, W l E is a layer-specific learnable weight vector.
The node embeddings in a given layer l arelearned as follows:
[Equation]
Here, N i = { ( u i , u j ) ∈ E} denotes the edges incident on node u i .
Similar to W l E , W l U is alayer-specific learnable weight vector, which performs a linear transformation on the embedding ofeach incident edge.
Following L layers of message passing, the final node and edge representations in the L th layer aredenoted as z i =
h L i and z ij = h L ij respectively.
Potential and kinetic energy prediction.
The predicted potential energy of each edge (rigid body) iscomputed by passing its final layer embedding through an MLP, i.e., (cid:99) v ij =
MLP ( z i,j ) .
The globalpredicted potential energy of the rigid body system is therefore the sum of the individual energies,i.e., (cid:98) V = (cid:80) ∀ e ij ∈E (cid:98) v ij .
For kinetic energy, the computation is identical except that it occurs in theother G NN with parameters optimized for kinetic energy.
Loss function:
The predicted Lagrangian is simply the difference between the predicted kineticenergy and the potential energy.
Using Euler-Lagrange equations, we obtain the predicted acceleration(cid:98) ¨ q i ( t ) for each node u i .
The ground truth acceleration is computed directly from the ground truthtrajectory using the Verlet algorithm as:
[Equation]
The parameters of the G NN s are trained to minimize the RMSE loss over the entire trajectory T :
[Equation]
Since the integration of the equations of motion for the predicted trajectory is also performed usingthe same algorithm as: q ( t +∆ t ) = 2 q ( t ) − q ( t − ∆ t )+ ¨ q (∆ t ) 2 , this method is equivalent to trainingfrom trajectory/positions.
4.2 Extension to Hamiltonian graph neural network
It is worth noting that the LG NN presented here can be extended directly to H GNN .
Specifically,due to the modular nature of the LG NN , which predicts the kinetic and potential energies usingtwo separate graphs, the final computation of the graph output could be modified to compute the
Hamiltonian instead of the Lagrangian as H = T + V .
Empirical Evaluation
In this section, we evaluate the ability of LG NN to learn rigid body dynamics.
In addition, we evaluatethe ability of LG NN to generalize to larger unseen system sizes, complex topology, and realisticstructures such as tensegrity.
5.1 Experimental setup
• Simulation environment: All the training and forward simulations are carried out in the JAXenvironment [18].
The graph architecture is implemented using the jraph package [24].
All the codes
related to dataset generation and training are made available in the Supplementary.
Software packages: numpy-1. 3, jax jax-md-0. .20, jaxlib-0.1.73, {+jraph-0.0.1.dev
+}
Hardware: Memory: 16GiB System memory, Processor: Intel(R) Core(TM) i7-10750H CPU @2.60GHz
• Baselines: As outlined earlier, there are very few works on rigid body simulations using graph-basedapproaches, where the graph is used to model the topology of the rigid body.
As such, to comparethe performance of LG NN we employ a graph network simulator G NS , which employs a full graphnetwork architecture [5, 12, 16] to predict the update in the position and velocity of node based onthe present position and velocity.
G NS has been shown to be a versatile model with the capability tosimulate a wide range of physical systems [16].• Datasets and systems:
To evaluate the performance LG NN , we selected n -chain/rope systems,where n = (4 , 8 , 16) .
All the graph based models are trained only on 4-segment chain system, whichare then evaluated on other system sizes.
Further, to evaluate the zero-shot generalizability of tolarge-scale unseen systems, we simulate 8-, and 16-segment chain systems.
Further, to push the limitsof LG NN , we evaluate the model trained on 4-segment chain to complex shaped topologies involvingtruss members (long rigid members) and chains (short rigid members), which have more than 40segments (see Figure 4).
The mass m i and moment of inertia I
i of all the members are maintained tobe the same for all the segments irrespective of their length.
The details of the experimental systemsare given in Appendix.• Evaluation Metric: Following the work of [13], we evaluate performance by computing the relativeerror in (1) the trajectory, known as the rollout error , given by RE ( t ) =
||
ˆ q ( t ) − q ( t ) ||
2 / ( || ˆ q t )
|| 2 +|| q ( t ) || 2 ) and (2) energy violation error given by ||
ˆ H − H|| 2 /
( || ˆ H||
+ ||H|| 2 ).
In addition, we alsocompute the geometric mean of rollout and energy error to compare the performance of differentmodels [13].
Note that all the variables with a hat, for example ˆ x , represent the predicted valuesbased on the trained model and the variables without hat, that is x , represent the ground truth.• Model architecture and training setup:
For the graph architectures, namely, LG NN and G NS , allthe neural networks are modeled as one hidden layer MLPs with varying number of hidden units.
Forall the MLPs, a square-plus activation function is used due to its double differentiability.
In contrastto the earlier approaches, here, the training is not performed on trajectories.
Rather, it is performed on10000 data points generated from 100 trajectories for all the models.
This dataset is divided randomlyin 75:25 ratio as training and validation set.
The model performance is evaluated on a forwardtrajectory, a task it was not explicitly trained for, of 1 s .
Note that this trajectory is ∼ 2-3 orders ofmagnitude larger than the training trajectories from which the training data has been sampled.
Thedynamics of n -body system is known to be chaotic for n ≥ 2 .
Hence, all the results are averaged overtrajectories generated from 100 different initial conditions.
Detailed model architecture associatedwith each of the models and the hyperparameters used in the training are provided in the Appendix.
5.2 Comparison with baselines
Model performance.
Figure 5 shows the performance of LG NN in comparison to G NS .
First, weanalyze the performance of LG NN on the 4-link system on which it was trained.
We observe that
LG NN exhibits superior performance in terms of both energy conservation error and rollout error forthe 4-link system in comparison to G NS .
Specifically, LG NN exhibits ∼ 3 orders of magnitude betterperformance than G NS in terms of both energy violation error and rollout error.
Generalizability to different system sizes.
Now, we analyze the performance of LG NN , trained on 4link segment, on 8- and 16-link segments.
We observe that LG NN exhibits comparable performanceswith respect to the 4-segment model, in terms of both energy violation error and rollout error, onsystems with 8-, and 16-segments that are unseen by the model.
In contrast, G NS exhibits relativelyincreased error in energy violation error and rollout error, with the superiority in performance of
LG NN maintained by orders of magnitude.
This confirms that LG NN can generalize to larger unseensystem sizes when trained on a significantly smaller system size.
5.3 Zero-shot generalizability
In order to analyze the zero-shot generalizability of the trained LG NN to simulate complex real-worldgeometries and structures, we evaluate the ability of LG NN to model the dynamics of tensegritystructures.
Note that tensegrity structures are truss-like structures comprising of both tension and compression-members.
The topology of a tensegrity structure is designed so that the compressionmembers are always bars and the tension members are always ropes.
Here, we analyse the ability LG NN to model the equilibrium dynamics of two complex tensegrity structures shown in Figure 4.
To this extent, we use the LG NN trained on the 4-segment structure.
We convert the rigid bodystructure to an equivalent graph and use the trained LG NN to predict the dynamics of the structurewhen released from the original configuration under gravity.
Figure 6 shows the energy error androllout for both the complex structures shown in Figure 4.
We note that the LG NN is able to generalizeto a complex structure with varying bar lengths and topology with high accuracy.
Specifically, theenergy violation and rollout error exhibits very low values for LG NN ( ∼ 10 − 4 ).
Further, it saturatesafter a few initial timestep suggesting an equilibrium dynamics.
In contrast, we observe that the errorin G NS is very high and continues to increase until it reaches 1, which is the maximum it can take.
This confirms the superior nature of LG NN to generalize to arbitrary topology, boundary conditions,and bar lengths, after training on a simple 4-segment chain with constant length segments.
Note thatsince the initial configuration for the forward simulation is fixed, it is not possible to generate errorbars for the trajectory.
5.4 Nature of the learned mass matrix
Finally, we investigate the nature of the mass matrix of LG NN for different systems.
Note that inearlier approaches either the mass matrix was learned directly for a given system based on the ELequations [6], or it was assumed to be diagonal in the Cartesian coordinates [13], or the functionalform of kinetic energy was assumed [7].
In the present approach, we do not make any assumptionson the nature of the mass matrix.
In fact, for a rigid body, the mass matrix need not be diagonal innature and depends on the actual topology of the structure.
This raises an interesting question aboutthe nature of the mass matrix learned by the LG NN and how it generalizes to arbitrary topologies.
In order to investigate the nature of the mass matrix, we plot the mass matrix of the LG NN in Figure 7.
Note that the mass matrix is computed directly from the Lagrangian as M = ∇ ˙ q ˙ q L , where L isobtained from the LG NN .
First, we analyze the mass matrix of the 16-segment structure.
We observethat the mass matrix is banded with a penta-diagonal band as expected for a chain structure.
Now, weanalyze the mass matrix for a complex structure T1.
Interestingly, we observe that the mass matrixlearned is non-diagonal in nature and is congruent with the complex topology of the structure (see
Figure 7).
This confirms that the mass matrix of LG NN is learned on-the-fly during the forwardsimulation that provides the versatility for LG NN to simulation complex structures.
Conclusions
Here, we present a LG NN -based framework that can be used to simulate the dynamics of rigid bodies.
Specifically, we present the graph architecture, which allows the decoupling of kinetic and potentialenergies, that can be used to compute the Lagrangian of the system.
We show that LG NN can learnthe dynamics from a small 4-segment chain and then generalize to larger system sizes.
We alsodemonstrate the zero-shot generalizability of LG NN to arbitrary topology including a tensegritystructure.
From the mechanics perspective, the LG NN assumes the knowledge of constraints.
Learning constraints directly from the trajectory is useful.
Similarly, extending LG NN to model contacts, collisions,and deformations allows more comprehesive learning of realistic systems.
From the modeling perspective, in our message passing , all messages are provided equal important.
Attention heads inmessage-passing neural networks have been shown to improve performance remarkably in severaldomains [25].
We plan to study the impact of attention in in our future works.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] §
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes](b) Did you include complete proofs of all theoretical results?
[Yes]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimentalresults (either in the supplemental material or as a URL)?
[Yes]
In supplementary(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were [Table] chosen)?
[Yes]
[Table] (c) Did you report error bars (e.g., with respect to the random seed after running experimentsmultiple times)?
[Yes] As shaded region in plots
[Table] (d) Did you include the total amount of compute and the type of resources used (e.g., type ofGPUs, internal cluster, or cloud provider)?
[Yes]4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes](c)
Did you include any new assets either in the supplemental material or as a URL?
[Yes] Insupplementary(d)
Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[N/A](e)
Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applica-
ble
[-From the EL equation, the acceleration is given by 
-]
[-Software packages: numpy-1.20.3, jax-0.2.24, jax-md-0.1.20, jaxlib-0.1.73, jraph-0.0.1.dev
-]
[-Length: 0.81331408, 1.15980562, 0.8647536 and 1.
-]
[-Mass: 1.83244264, 1.18182497, 1.30424224, 1.43194502 and 1.
-]
[-Moment of inertial: 1.21233911, 1.18340451, 1.52475643 and 1.
-]