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, 14].
In this approach, a neuralnetwork is trained to learn the L (or H ) of a system based on its configuration ( q, ˙ q ) .
The L is thenused along with the Euler-Lagrange (EL) equation to obtain the time evolution of the system.
Notethat the training of L NN s is performed 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 trajectoryof a multi-particle system [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, 15].
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 ) [16, 17,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 [18, 19, 20].
G NN s havealso been used to model rigid bodies mainly following two approaches, namely, particle-based [17]and lumped mass [20, 21] 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 [17].
Note that this approach is philosophically similar to mess-less methods such assmoothed-particle hydrodynamics (SPH) [22] or peridynamics (PD) [23], where the time-evolutionof a continuum body is simulated by discretizing the domain using particles.
This approach [17],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 [20, 24], 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 articulated rigidbodies such as 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 arbitrary system 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 ofarticulated rigid 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,in turn, computed based on this scalar quantity.
Among the energy-based approaches, Lagrangianformulation has been widely used to predict the dynamics of particles and rigid bodies by computingthe Lagrangian L of the system.
The standard form of Lagrange’s equation for a system withholonomic constraints is given by ddt (cid:16) ∂ L ∂ ˙ q (cid:17) − (cid:16)
∂
L ∂q (cid:17) = 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 andthe potential function from which generalized forces can be derived.
Accordingly, the dynamics ofthe system can be represented using EL equations as ¨
[Equation]
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 [15].
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 λ (see A.3) and substituting in Eq. 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 bearbitrarily large such as chains, ropes or trusses, that can bedivided into smaller constituent members.
This is because, it isgenerally assumed that extending L NN s to large structures is achallenging problem [15].
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 than 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 with orientation of each link represented by ϕ i =tan − 1 (cid:16) y i − y
i − 1 x i − x i − 1 (cid:17) can be considered.
Placing the origin at the beginning of first segment (see Figure1), the center of mass of i th segment ( x cmi , y cmi ) can be written in terms of generalized coordinatesas
[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, we employ the approach shown in Figure 2.
Specifically, wemodel a physical system as a graph.
Further, the Lagrangian of system is learned by decoupling thepotential and kinetic energy, each of which are learned by 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 based by minimizingthe loss on the acceleration predicted by the LG NN using EL equation with respect to the groundtruth.
In this section, we describe the LG NN architecture for rigid bodies in detail (See Figure 2 foran overview).
We empirically show that the dynamics of a rigid body can be learned by LG NN .
Inaddition, due to the inductive nature of the graph architecture, once trained on a small system, LG NNcan 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 nodes represent the connection between these rigid bodies.
This is in contrast to earlier approaches used for particle-based systems, where node represented the particle 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 2, 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 2, 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] [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 and || represents concatenation operation.
Thenode embeddings in a given layer l are learned as follows:
[Equation]
Here, N i = { u j | ( 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 representationsin the L th layer are denoted by z i
= h Li and z ij = h Lij 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 [Equation] ¨ 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.
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 JAX environment [19].
The graph architecture is implemented using the jraph package [25].
All the codes
related to dataset generation and training are made available in the Supplementary.
Software packages: numpy-1.20.
[-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-][Table]
• 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.
To compare theperformance of LG NN , we employ three baselines, namely, (i) a graph network simulator G NS , (ii) aLagrangian graph network (L GN ), and (iii) constrained Lagrangian neural network (CL NN ). G NSemploys a full graph network architecture [5, 12, 17] to predict the update in the position and velocityof node based on the present position and velocity.
G NS has been shown to be a versatile modelwith the capability to simulate a wide range of physical systems [17].
L GN and CL NN employs theexact same equations as LG NN for computing the acceleration and trajectory and hence has the sameinductive biases as LG NN in terms of the training and inference.
However, while L GN employs a fullgraph network, CL NN employs a feed-forward multilayer perceptron.• 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,which are then evaluated on other system sizes.
Further, to evaluate the zero-shot generalizabilityof LG NN to large-scale unseen systems, we simulate 8-, and 16-segment chain systems.
Further,to push the limits of LG NN , we evaluate the model trained on 4-segment chain to complex shapedtopologies involving truss members (long rigid members) and chains (short rigid members), whichhave more than 40 segments (see Figure 3).
The mass m i and moment of inertia I
i of all the membersare maintained to be the same for all the segments irrespective of their length.
The details of theexperimental systems are 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 ) ||
+ || 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.
To compare the performance of LG NN with baselines , G NS , L GN [12, 6] and
CL NN [13], we evaluate the evolution of energy violation and rollout error. It worth noting that G NSand L GN have been demonstrated only particle-based systems and not on rigid bodies. Hence, tomake a fair comparison, we give the same node and edge input features as provided for the LG NNfor both G NS and L GN , while training. All the models are trained on a 4-link system and evaluatedon all other systems.
In the case of CL NN , due to the fully connected architecture, the model is noinductive in nature.
Hence, the model is trained and tested on the same system only, that is, the 4-linksystem. Detailed architecture of each of these systems are provided in Section A.6. Figure 4 showsthe error in energy and rollout for LG NN , G NS , L GN , and CL NN . We observe that G NS , L GN and
CL NN have a larger error in comparison to LG NN as shown in Figure 4 for both energy and rollout error, establishing the superiority of LG NN . Two similar 4-link systems one with varying masses andmoment of inertia, and one subjected drag are evaluated in the Appendix A.8. Figures 7 and 12 showthat LG NN is able to infer the dynamics in both these systems.
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,on systems with 8-, and 16-segments that are unseen by the model.
In contrast, G NS exhibitsrelatively increased error in energy violation error and rollout error, although the error in L GNremains comparable for all systems.
This suggests that the inductive bias in terms of the EL equationsprevent the accumulation of error and allow improved generalization. However, the error in L GN isstill orders magnitude higher than LG NN . This suggests that the architecture employed in LG NN isleading improved learning of the dynamics of the system.
This confirms that LG NN can generalize tolarger unseen system sizes when trained on a significantly smaller system size. Note that the plots for
CL NN are not shown for 8 and 16-links as the architecture cannot exhibit generalizability to largersystem sizes.
5.3 Zero-shot generalizability
In the conventional L NN s employing feed forward MLPs, the training and test system have thesame number of particles and degrees of freedom.
In other words, an L NN trained for an n -particlesystem cannot be used to perform inference on an m -particle system.
In contrast, we show here that LG NN trained on a small 4 -link system can be used to perform forward simulations on other unseencomplex systems such as 100 -link system, and tensegrity structures. This ability to infer on differentunseen system sizes and topology is referred to as zero-shot generalizability. In order to analyzethe zero-shot generalizability of the trained LG NN to simulate complex real-world geometries andstructures, we evaluate the ability of LG NN to model the dynamics of tensegrity and lattice-likestructures (see Fig. 3).
Note that tensegrity structures are truss-like structures comprising of bothtension and compression-members.
The topology of a tensegrity structure is designed so that thecompression members are always bars and the tension members are always ropes.
Here, we analysethe ability LG NN to model the equilibrium dynamics of two complex tensegrity structures and thelattice-like shown in Figure 3.
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 5 shows the energy error androllout for both the complex structures and the lattice-like structure shown in Figure 3.
We note that the LG NN is able to generalize to a complex structure with varying bar lengths and topology withhigh accuracy.
Specifically, the energy violation and rollout error exhibits very low values for LG NN( ∼ 10 − 4 ).
Further, it saturates after a few initial timestep suggesting an equilibrium dynamics.
Incontrast, we observe that the error in G NS is very high and continues to increase until it reaches1, which is the maximum it can take.
This confirms the superior nature of LG NN to generalize toarbitrary topology, boundary conditions, and bar lengths, after training on a simple 4-segment chainwith constant length segments.
Note that since the initial configuration for the forward simulation isfixed, it is not possible to generate error bars 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 6.
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 6).
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.
Limitations and Future Works: From the mechanics perspective, the LG NN assumes the knowledgeof constraints.
Learning constraints directly from the trajectory is useful.
Similarly, extending LG NNto model contacts, collisions, and deformations allows more comprehesive learning of realisticsystems.
From the modeling perspective, in our message passing LG NN , all messages are providedequal important.
Attention heads in message-passing neural networks have been shown to improveperformance remarkably in several domains [26].
We plan to study the impact of attention in LG NNin 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 werechosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] As shaded region in plots
(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]
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
[Table]
A AppendixA.1 Experimental systems
Here, we describe the details of the experimental systems considered in the work, namely, n -linksystems that represent a chain when the bars are relatively long and represent a rope when the barsare very small in comparison to the total length of the chain/rope.
In addition, we also describe thecomplex systems, which can be considered as the tensegrity structures.
We describe the constraintsand the physics-based equations for predicting the Lagrangian of these systems.
A.1.1 n -link
In an n -link system, links are connected by hinges where they are free to rotate.
These links are rigid,thus, impose a distance constraint between both the end points as
[Equation]
(13) where, l i represents the length of the link connecting the ( i − 1)
th and i th node.
This constraint canbe differentiated to write in the form of a Pfaffian constraint as
[Equation]
Note that such constraint can be obtained for each of the n links considered to obtain the A ( q ) .
The Lagrangian of this system has contributions from potential energy due to gravity and kineticenergy.
Thus, the Lagrangian can be written as
[Equation]
15) where g represents the acceleration due to gravity in the y direction, I represents the moment ofinertia of link, ω represents the angular velocity and m represents the mass of the link.
A.2 Complex system
The complex system (with tensegrity structures) is a combination of the rope and rod system.
Here,components with segment size significantly lower than the size of the system are termed as rope andparts with single segment, where the length of the segment is comparable to that of the system, istermed as a rod.
These rods and ropes are connected such that ropes are in tension and rods are incompression holding the whole structure (see Fig. 3).
The constraints similar to the n -link system arepresent in this system as well.
The Lagrangian of this system is given as
[Equation]
16) where g represents the acceleration due to gravity in the y direction, I represents the moment ofinertia of link, ω represents the angular velocity and m represents the mass of the link.
A.3 Derivation of λ
{+From the EL equation, the acceleration is given by 
+}
[Equation]
17) and the constraints equation is given by A ( q )¨ q + ˙ A ( q ) ˙
q = 0 .
Substituting for ¨ q from Eq 17, theconstraint equation is modified as
[Equation]
(18) where A ( q ) and ˙
A ( q ) is written as A and ˙ A , respectively for simplicity.
The equation can be simplifiedas
[Equation]
A.4 Trajectory visualization
For visualization of trajectories of actual and trained models, videos are provided as supplementarymaterial.
The supplementary materials contains:
(a) 8-link system: Here, an 8-link chain is connected to supports at both ends is stretched downwardin initial condition.(b) 10-link system: Here, a 10-link chain is connected to support as one end and free from other end.(c) Tensegrity structure: A tensegrity structure is created using rope and rodes as described insection A.2 .
A.5 Implementation details
A.5.1 Dataset generation
{+Software packages: numpy-1.20.3, jax-0.2.24, jax-md-0.1.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
All the datasets are generated using the known Lagrangian for the system along with the constraints,as described in Section A.1.
For 4-link system, we create the training data by performing forwardsimulations.
A timestep of 10 − 5 s is used to integrate the equations of motion.
The velocity-verletalgorithm is used to integrate the equations of motion due to its ability to conserve the energy in longtrajectory integration.
The datapoints were collected every 10 − 3 time-step and for 10000 datapoints.
It should be noted thatin contrast to the earlier approach, here, we do not train from the trajectory.
Rather, we randomlysample different states from the training set to predict the acceleration.
A.6 Architecture
For LG NN , we use two different size of fully connected feed forward neural network for initialembeddings and message passing.
For initial embeddings, we use a linear layer to project the featuresto latent space with size 5.
For message passing, we use a MLPs with a single hidden layer with 10neurons.
We use squareplus as activation functions for all non-linear MLPs.
For G NS and L GN , we follow a similar architecture suggested in the literature [17, 20].
Specifically,
For initial embedding of node and edge features we use linear transformation with latent dimensionwith size 8.
We use message passing MLPs with a single layer with 128 neuron.
We use squareplusas activation functions for all non-linear MLPs.
For G NS the output is acceleration values at nodeand for L GN the output is lagrangian of the whole system.
For CL NN , we follow a similar architecture suggested in the literature [6, 13].
Specifically, we use afully connected feed forward neural network with 2 hidden layers each having 128 hidden units witha square-plus activation function.
For LG NN , all the MLPs consist of two layers with 5 hidden units,respectively.
Thus, LG NN has significantly lesser parameters than L NN .
A.7 Training details
The training dataset is divided in 75:25 ratio randomly, where the 75% is used for training and 25%is used as the validation set.
Further, the trained models are tested on its ability to predict the correcttrajectory, a task it was not trained on.
All models are trained for 10000 epochs with early stopping.
A learning rate of 10 − 3 was used with the Adam optimizer for the training.
• Lagrangian Graph Neural Network (LG NN )
[Table]
• Graph Neural Simulator (G NS ) and Lagrangian Graph Network (L GN )
[Table]
• Constraint Lagrangian Neural Network (CL NN )
[Table]
A.8 Rollout and energy error
Details of system with varying link properties
{+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.
+}