Equivariant Graph Mechanics Networks with Constraints
Authors: Wenbing Huang
Jiaqi Han
Yu Rong
Tingyang Xu
Fuchun Sun
Junzhou Huang
Venue: ICLR 2022 Poster (2022)

Abstract: Learning to reason about relations and dynamics over multiple interacting objects is a challenging topic in machine learning. The challenges mainly stem from that the interacting systems are exponentially-compositional, symmetrical, and commonly geometrically-constrained.
Current methods, particularly the ones based on equivariant Graph Neural Networks (GNNs), have targeted on the first two challenges but remain immature for constrained systems. 
In this paper, we propose Graph Mechanics Network (GMN) which is combinatorially efficient, equivariant and constraint-aware. The core of GMN is that it represents, by generalized coordinates, the forward kinematics information (positions and velocities) of a structural object. In this manner, the geometrical constraints are implicitly and naturally encoded in the forward kinematics. Moreover, to allow equivariant message passing in GMN, we have developed a general form of orthogonality-equivariant functions, given that the dynamics of constrained systems are more complicated than the unconstrained counterparts. Theoretically, the proposed equivariant formulation is proved to be universally expressive under certain conditions. Extensive experiments  support the advantages of GMN compared to the state-of-the-art GNNs in terms of prediction accuracy, constraint satisfaction and data efficiency on the simulated systems consisting of particles, sticks and hinges, as well as two real-world datasets for molecular dynamics prediction and human motion capture.

Published as a conference paper at ICLR 2022
EQUIVARIANT GRAPH MECHANICS NETWORKS WITH CONSTRAINTS
Wenbing Huang"!, Jiaqi H
2 Yu Rong Tingyang Xu', Fuchun Sun Junzhou Huang* Institute for AI Industry Research (AIR), Tsinghua University
2 Beijing National Research Center for Information Science and Technology (BNRist), Department of Computer Science and Technology. Tsinghua University
3 Tencent Al Lab
4 Department of Computer Science and Engineering, University of Texas at Arlington
hwenbingg 126 com, sanjq21@maill twwngnui..n.n.... edu. en, yu. yu. rong@hotmmil con t iingynggxu@tencemn com, fcsunêmail wetimmaaa...snnniiii tsinghua edu. jzhuangguta. edu
ABSTRACT
Learning reason about relations and dynamies over multiple interacting objects is a challenging topic in machine learning. The challenges mainly stem from that symmetrical, and com- the interacting systems are Current methods, particularly the ones based on monly
equivariant Graph Neural Networks (GNNs), have targeted on the first two chal- lenges but remain immature for constrained systems. In this paper, we propose Graph Mechanics Network (GMN) which combinatorially efficient, equivariant and constraint-aware. The core of GMN is that it represents, by generalized coor- dinates, the forward kinematics information (positions and velocities) of a struc- tural object. In this manner, the geometrical constraints rally encoded the forward kinematics. Moreover, to allow equivariant message passing in GMN, we have developed general form of orthogonality-euuivrriaa functions, given that the dynamics of constrained systems are more complicated than the unconstrained counterparts. Theoretically, the proposed equivariant for- mulation is proved to be universally expressive under certain conditions. Exten- sive experiments support the advantages of GMN compared to the state-of-the-art GNNs in terms of prediction accuracy, constraint satisfaction and data efficiency on the simulated systems consisting particles, sticks and hinges, as well as two real-world datasets for molecular dynamics prediction and human motion capture.
implicitly and natu-
INTRODUCTION
Representing and reasoning about the rela- tions and dynamics of group of interacting objects among the core aspects of human intelligence (Tenenbaum et al. 2011; Ding et al. 2021] Gan et al. 2021). As a motivat- consider the N-body sys- tem (Kipf et al., 2018) where the movement of single charged particle is attracted or repelled by other charged particles. Physi- cists have revealed that this process can be modeled by Newton's laws along with the Coulomb force. One may wonder, however, if we can teach machine to rediscover the underlying physics by solely observing the particles' states. This thinking has inspired the study learning to model interacting systems, which now is
Hinge
ing example,
(a) N-body System
(b) Constrined N-body System
Figure 1: N-body vs. constrained N-body (red/blue balls denote positive/negative charges).
prevailing topic in machine learning (Battaglia et al. TToma Thomas et al. 2018 Köhler et al. 2019
Sanchez- et al. 2019 Fuchs et al. 2020) Martinkus et al. 2021 Satorras et al.) [02
Equal contributions: Wenbing Huang and Jiaqi Han; x9 Corresponding authors: Yu Rong and Fuchun Sun.
This work is done when Jiaqi Han works as intern in Tencent AI Lab.

---

Published as a conference paper at ICLR 2022
Learning to model interacting systems is challenging. First, the systems combinatorially complex, on account of that objects can be composed in combinatorially many possible arrange- ments (Battaglia et al. 2016). This challenge, to some extent, can be addressed by making use of Graph Neural Networks (GNNs) (Wu et al.] 2020). By regarding objects as nodes and interactions as edges, GNNs extract information via message passing, which is able to characterize arbitrarily ordered objects and combinatorial relations. The second challenge is related to an important sym- metry in physics: the model we use should be equivariant to any Euclidean transformation (transla- tion/refiection/rotatioo of the input. This complies with the fact that physics rules keep unchanged regardless of the reference coordinate system. Several works (Fuchs al.) 2020; Saaorras et al.) 2021 have investigated equivariance upon GNNs and exhibited remarkable benefits on N-body,
Another challenge, though less explored, is that the systems could be geometrically constrained. Geometric constraints arise in common practical systems, for example, in robotics where the joints of mechanical arms should be linked one by one, or in biochemistry where the atoms of molecules are connected by chemical bonds. When modeling these constrained systems, it is crucial to en- force the model to output legal predictions. For instance, the lengths/angles of chemical bonds in a molecular closely determine its chemical property which will be changed dramatically if the struc- tural constraints are broken. As mentioned above, equivariant GNNs (Fuchs et al.) 2020; [Satorras] al. 202T) have achieved desirable performance on the N-body system- this system, nevertheless, lacks of constraint. Considering constraints is not easy, as the dynamics of all elements within a constrained system evolve in joint and complex manner. Unfortunately, the best of our knowl- edge, there is no research (particularly among equivariant GNN methods) that learns to model the dynamical systems of multiple interacting objects under geometrical constraints.
In this paper, we propose Graph Mechanics Network (GMN) that can tackle the above three chal- lenges simultaneously: I. To cope with the combinatorial complexity, GMN takes advantage of graph models to encode the states of and interactions between objects, analogous to previous ap- proaches. I. For the constraint satisfaction, GMN resorts to generalized coordinates, a well known notion in conventional mechanics to model the dynamics of structural objects. Here, structural object (such as the stick and hinge in Fig. [ is defined as a set of multiple rigidly connected par- ticles. In GMN, the constraints are implicitly encoded in the forward kinematics that describes the Cartesian states as the function of generalized coordinates. This strategy can inherently maintain the constraints and requires external regulation. III. GMN is equivariant. In GMN, we need to compute the interaction forces in the Cartesian space (see Eq. (5), and infer the accelerations of the generalized coordinates based on the inverse dynamics (see Eq. (). both of which learned by general form of equivariant functions (see Eq. TT]. Notably, the proposed equivariant formulation is more efficient to compute compared to the previous versions based on spherical harmonics (Thomas et al. 2018 Fuchs et al., 2020), or more general than the one developed by EGNN (Satorras et al., 2021 More importantly, we have theoretically discussed when and how our formulation can uni- versally approximate any equivariant function.
To evaluate the advantages of GMN, we have constructed a simulated dataset composed of three types of objects: particles, sticks and hinges, which is complex case of the N-body system (Kipf jet al., 2018) and can be used as building blocks for common systems. Under various scenarios in terms of different ratios of object types and different numbers of training data, we empirically verify the superiority of GMN compared to statc-of-the, the-art models in prediction accuracy, constraint satisfaction and data efficiency. In addition, we also test the effectiveness of GMN on two real-world datasets: MD17 (Chmiela et al. 2017) and CMU Motion Capture (CMU 2003).
RELATED WORK
Learning to simulate complex physical systems has been shown to greatly benefit from using GNNs. Interaction Network (IN) proposed by |Battallia et al. (2016) is the first attempt for this purpose, and it can be deemed as a special kind of GNNs to learn how the system interacts and how the states of particles evolve. Later researches have extended IN in different aspects: HRN (Mrowca et al.) 2018) utilizes hierarchical graph convolution for tackling objects of various geometrical shapes and materials, NRI (Kipf et al. 2018) further explicitly infers the interactions with the help of varia- tional auto-encoder, and Hamiltonian graph networks Saancez-Goozzllee et al. 2019) equip GNNs with ordinary differential equations and Hamiltonian mechanics for energy conservation. However,
2

---

Published as a conference paper at ICLR 2022
all above approaches have ignored the symmetry in physics and the GNN models they use not Euclidean equivariant. There is a subset of models (Ummenhofer et al. 2019. Sanchez-Gonzalez et al. 2020 Pfaff et al.] 2020) that partially implement symmetries, but they only enforce translation equivariance but not rotation equivariance. On the other hand, it is nontrivial to enforce rotation equivariance. Vensor-Field networks (Thomas et al., (2088) uses filters built from spherical harmon- ics to allow 3D rotation cquivariance, and this idea has been developed by SE(3) Transformer (Fuchs ct al.
2020) that further takes the attention mechanism into account. Anther class of works (Finzi 2020a. Hutchinson et al. 2 resorts to the Lie convolution for the equivariance on any Lie group, based on lifting and sampling. Recently, EGNN (Satorras et al. (02T) has proposed simple yet effective form of equivariant message passing graphs, which does not require computationally
et al.
expensive higher-order representations while still achieving better performance on N-body.
As interpreted before, ensuring geometrical constraints is crucial for many practical systems, which, nevertheless, seldom investigated in aforementioned works. Although several attempts (Yang) et al. 2020, Finzi et al., 2020b) have been proposed for learning to enforce constraints, they explic- illy augment the training loss with soft Lagrangian regulation and thus are completely data-driven and have guarantee of generalization for limited training data. DeLaN (Lutter et al., 2019) also employs generalized coordinates to describe the kinematics of the rigid object. Nevertheless, it only target on the physical process of single object, other than the complex systems with multiple rigid and structural objects that are the focus of this paper. In DPI-Net (Li et al. 2018), the Box Bath task does share the similar strategy to us by first predicting the canomical coordinates of the box and then using the forward kinematic model to obtain the Cartesian positions. However, the passing messages in DPI-Net are scalars other than directional vectors (positions, velocities, and accelerations) used in our work. After all, both DeLaN and DPI-Net never study equivariant models.
3 GRAPH MECHANICS NEURAL NETWORK
We begin with introducing the N-body system Kipt et al. 2018). as illustrated in Fig. (a). This system consists of N interacting particles { of the same mass, and the kinematics states of each particle are defined as S, = (æ:, v;), where 2;, v. E R are the position and velocity vectors, respectively. There could be certain non-vector information of each particle (such as charge), which is represented by a e-channel feature h; E R [ Suppose the system we study is conservative and the dynamics is driven by the interaction force fú E R3 between any pair of particles and j. According to Newton's second law, the acceleration of particle i, a, E IR33 is proportional to the aggregated force from other particles Ej fij- All symbols will be specified with superscript t for temporal denotations, e.g., 2 indicating the position of particle at time 1. In this paper, we mainly concerned with the prediction task: we need to seek out function (f((S) (.)) given ) the initial states (S ===1 and features [6 to forecast the future states { J. Ji=I at time T.
As presented in Introduction, two kinds of inductive biases have been explored previously. The first one is to apply the graph structure to capture the distribution of particle states and their interac- tions (Battaglia et al. 2016 Kipf et al. 2018), where, particularly, the particle states S (along with h?) are as node features and the interaction forces f as edge messages. In this way, the transition function Φ boils down to GNN model. The second inductive bias is that ø should be equivariant to any transsation/rreeetttoo/aaoooo of the input states. By saying equivariance, we imply
s(llo S., ( 9 ( 1. Vg E O(3), V Vh..
(1)
Here, (3) defines the 3D orthogonal group (Fuchs al. 2020) that consists of translation, reflec- tion and rotation transformations.
S denotes to perform transformation on the states Sp, and it is instantiated as R(x + b) for the position and Rof for the velocity, where R E 3 is the orthogonal matrix and b E R is the translation vector.
Several works (Thomas al. 2018 Köhler et al. 2019 Fuchs al., 2020; Satorras et al., 2021) have investigated both inductive biases, among which EGNN (Satorras et al., 2021) has achieved promising performance on N-body. The typical process in EGNN iterates the following steps:
Although hss could be of multiple dimensions, we still call it as non-vector value since the dimension of h is distinct from that of the kinematics vector, for example æ;, the latter of which has intrinsic geometry and should obey translation/rotation equivariance. This is also why do not denote h; in bold.
3

---

Published as a conference paper at ICLR 2022
where the superscript denotes the I-th layer; the acceleration a returned by the message ag- gregation of Pegnn in Eq. 2 is adopted for the (3) update of the velocity v in Eq.
a hi = Σ Pozun(a ( 1.111,hh ,h
(2)
(multiplied by scalar " E R), followed by the ren- (4) ovation of the position æx. in Eq. 4 The formu- lation of Pognn is physically reasonable, since the interaction (actually the Coulomb force) between particles i and] truly depends on their relative position z 1-1
of = (h) 1 ))= o. _. 1 o
10 a
xt 1-1 a/ node features h/!-
and hn-1, and edge feature ejs. To enable equivariant message passing, EGNN has developed a specific form of Pegin let af be equivariant and K be invariant in terms of the input T which will be presented in Eq. [0] Notice that the computations for f; and hy are actually by two different functions, and we have abbreviated them into one in Eq. 2 (and also Eq. |]aater), since they share the same inputs; besides, their parameters are shared following EGNN.
3.1 OUR GENERAL ARCHITECTURE
Despite the desirable performance on N-body, existing methods are incapable of maintaining the geometric constraint. In this section, we design a general architecture that intrinsically meets the requirement of geometry constraints by making use of the generalized coordinates.
We define O {Pi) a structural object composed of ng rigidly connected particles. Fig. (b) illustrates two examples of the structural object, the stick with 2 connected particles and the hinge with 3 particles. To preserve the distance, the dynamics of the two particles updated a joint way, rather than fulfilling the independent process in EGNN (Eq, and Besides, the force each particle within O₂ will indirectly influence the dynamics of others through the physical connections. This requires us to analyze the dynamics of the particles in O₂ as a whole, which is implemented by generalized coordinates. There could be multiple generalized coordinates for each Oz, some located in the Cartesian space but some in the angle space. For instance, the states of stick can be decoupled by two independent sets of generalized coordinates: the state of particle as the Cartesian coordinates and the relative rotation angles of particle 2 to I as the angle coordinates. For conciseness, this section only focus on the Cartesian part which essentially determines the local coordinates in On, with providing full examples in $ 3.3 We denote the position, velocity and acceleration of the generalized Cartesian coordinates as q qu and qu. respectively.
stick should be
We now detail how to update the states of Ox- In Fig. 2 we first compute the in- teraction force between each particle and others, and aggregate information of all par- ticles within each structural object to infer the acceleration of the generalized coordi- nates (which is termed as the generalized acceleration henceforth) by inverse dynamics. Then, the dynamical updates are carried out in the space of the generalized coordinates. Finally, the up- dated generalized coordinates will be projected back to the particles' states via the forward kinematics.
Carteslan Coordlinates x/
Forward Kinematics
fl h
Σ
(1 R// . e.s))
(5)
Generlized Cocetinates 9 4L
Gereralizes/Acoceeaazioog
d
Σ p2 p2(f/ af 2
(6)
Inverse Dynamics
iEO,
/ Interaction Force
d +( Σ
11 W + t
(7)
Cartesian Coortinetes x/-1
iE0+
d de the
(8)
Forward Kinematics
Genoralzzdl Coordinalesa .4t-
2/. v. FK g.) Vi Ok:
Figure 2:
The flowchart of our GMN.
where, the elements of the system can be described in two views, {O as the object-level view and [ as the particle-level view; for distinc- tion. we index the structural object with the subscript k and particles with i or j. We explain each equation separately.
Interaction force (Eq. 5 The interaction force f is computed analogous to Eq. 2
EGNN
4

---

Published as a conference paper at ICLR 2022
straightly regards the interaction force to be the acceleration of each particle in Eq. 2 But here, given the constraints between particles, we record the force as an intermediate variable that will contribute to the inference of the generalized acceleration in the next step. Inverse dynamics (Eq. 6). This step is the core of our methodology. The generalized accel- eration i is dependent the forces f. all particles within Oz. and their relative positions Thi x d/11 and relative velocities o/= = 0/-1 q with regard to the generalized coordinates. The formulation of Eq. is physics-inspired. In Appendix (Eq, 26), we have analyt- ically derived the dynamics of hinges, where the acceleration of the hinge indeed related to the forces, relative positions and relative velocities of all particles in each hinge. This is reasonable in mechanics, since the cross product T3 x f yields the torque, and the relative velocity 1 is related to the centrifugal force of particle around qu. both of which influence the acceleration qt. Different from the analytical form which is always complex and hard to compute in practice, we will employ a learnable and equivariant function with universal expressivity. The details are in $3.2
Generalized update 788) The updates of the generalized coordinates are akin to Eq. and in EGNN, as the dimensions of the generalized coordinates have been made independent. Notice that in Eq. the scalar factor for of. takes as input the summation of all hidden features, which is generalized form of Eq. for multiple particles in O₁.
Forward kinematics (Eq. 9). Once the generalized coordinates have been refreshed, we can derive the states of all particles in O2 by proceeding the forward kinematics. Different system of O2 could have different type of the forward kinematics. Here we denote it as the function FK(() in general, while providing the specifications in $ 3.3
Our method will reduce to EGNN if setting the generalized coordinates as the states of each particle and utilize the identify map in Eq. G and
Alg. has summarized the updates for all objects.
3.2 EQUIVARIANT MESSAGE PASSING
As shown in Eq. [ the Euclidean equivariance on the estimation function 6 is necessary for ensuring the physical symmetry. When considering this property in our case, we demand the interaction force (5 and generalized acceleration (Eq. O to be equivariant with respect to orthogonal transforma-
tions, while other equations are already equivarian f EGNN (Satorras et al. 2021) has developed a particular orthogonality-eqiinn
form for the acceleration output in Eq. 2
Pegnn (x, h) == 20 .0 (1]==|1, h))
(10)
where (-) is an arbitrary Multi-Layer Perceptron (MIP) with parameter w. and we have abbrevi- ated other non-vector terms in Eq. 2 as h. This formulation does satisfy the rotation equivariance, but it is unknown it can be generalized to functions (such as B with multiple input vectors, and more importantly, its representation completeness is never explored rigorously. In this section, propose a general form of
functions with necessary theoretical guarantees.
Without loss of generality, the target function would like to enforce equivariance is denoted as kp h): Rdxmm x T >> Rdxm', We define the below formulation,
c == Zo w ( z.h).
(11)
It is easy to justify the function vp( ZZ h) in Eq. 11 is equivariant to any orthogonal matrix i.e.,
sp(OZ h) = O h), VO E gaxd o 0 = I. Apparently Eq. 11 reduces to Eq, by setting the number of vectors in Z 1, namely, m m' 1. We immediately have the following theory.
Theorem 1. If m 2 d and the row rank of Z is full, i.e. rank Z) d, then for any continuous orthogonality-eqnivvriin function ((Z, h). there must exist an MLP ow satisfying Il<p(Z, h) ((Z, < e arbitrarily small error e.
The proof employs the universality of MLP (Cybenko 1989 Hornik 1991). with the entire details deferred in Appendix. Theorem T nutritive, as it characterizes the rich expressivity of formulating the equivariant function via Eq. II The condition holds in general for the function like Eq 6q whose input = (fr, Dgi, Uss) E R is of full rank when the force, position and velocity expand the whole space. Indeed, this condition holds with probability 1, stated by the following corollary.
'The translation equivariance holds naturally as we use relative positions in Eq. 5 making the force, ac- celeration, velocity to be invariant with respect to translation. After the addition with the previous position in Eq. 8 the updated position de is translation cquivariant.
5

---

Published as a conference paper at ICLR 2022
Corollary 1. Assume 2 d, and also the entries of Z are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in R. Then, almost surely, the conclusion of Theorem holds.
When the condition m 2 d is invalid, the universal approximation still maintains if restricted in the linear subspace expanded by the columns of Z. Corollary 2. For any continuous
function (Z,, h) located in the linear subspace expanded by the columns of Z. the conclusion of Theorem holds universally.
Corollary tells that the message passing &Pegn in Eq. 10 is not universally expressive since m < d, and can only fit the vectors parallel to æ (i.e. the relative position). Yet, (Pognn is still physically complete for the Coulomb force that is oriented by the relative position between two particles.
In our experiments. we find that more stable performance is delivered by further adding the normal- ization term specifically when m m> 1:
(12) where || || computes the Frobenius norm. Notice that adding normalization does not change the conclusions in Theorem [ and its two corollaries, since the norm is also function of z Z that can be approximated by MLP. We implement yf in Eq. 5 and y22 Eq. by using the general formulation Eq. 12] where
<Z = Z /T///lz Z [F.
" --1-1 and Z = (f): o o1-1), respectively. Note that for the
update of hidden feature hy in Eq. we do not need equivariance but invariance, hence we set Z|| h
b! - Z Z Z Z
1-1] by just keeping the invariant part in Eq. 12
3.3 IMPLEMENTATION OF FORWARD KINEMATICS
Implementation of hinges. A hinge, displayed in Fig (a), consists of three particles O, 1, 2, and two sticks 01 and 02. The freedom degrees of this system can be explained in this way: particle 0 moves freely, and particles and 2 can only rotate round particle 0 owing to the length constraint by the two sticks. Hence, the generalized coordinates include the states of particle 0 denoted as qo E IR and the rotation Euler angles of stick 01 as Po1 E IR and 02 as 002 RR3 As qu are Cartesian, the dynamical update fol- lows directly Eq. 5 With the forces {3 estimated in Eq. 5 and acceleration d by Eq. the angle acceleration of 0o1 is calculated as
2'
a Ouz
d
0n2 (x/) ( (b=) 11 accord- ing to rigid mechanics, where the numerator is the relative torque, the denominator is the mo- ment of inertia under unit mass, and x means the cross product. Then the angle velocity is
Xye
Yo-
dn dn
2
Xax
x
(a) Hinge
(b) Stick
0/2 w (hi + +h1 + [of=")0fn + 0/1 simi- lar to Eq. We do not need to record Ooi. since the forward kinematics can be conducted without it. detail, the forward kinematics in Eq. iis
Figure 3: Illustrations of hinges and sticks.
#1 d rot(O', )=yT of 4 0/2 æ/1.
(13)
where rot(O%)) indicates the rotation matrix around the direction of 6/11 by absolute angle |
The dynamic updates for particle is similar. We put the whole details into Alg. Min Appendix.
Implementation of sticks. For a stick with particles I and 2 in Fig (b), we choose the center as the
3
generalized Cartesian coordinate qu and the rotation of particle 1 6or and particle 2 0u2 (001
022)
as the generalized angle coordinates. The dynamics propagation of qu is given by Eq. 5 5 There
are two choices to compute Th one using the general form in Fq. 6 and the other one leveraging a
simplified version as T Lieo, (p(f)) by explicitly omitting the relative position and velocity. The physical motivation of introducing the simplified version is that the acceleration of stick center is only affected by the forces according to the theorem of the motion of the center of mass. The angle accelerations are 0 (1/ x f + 1/22" x f2) 1-1 1 + |0/2 (12), and the velocity becomes 0, w (M1 + h ) + 0 The states of partieles and 2 renewed as the same Eq. 13
Learnable FK. Besides the above hand-crafted FK, we propose a learnable variant by replacing Eq. 7 with the update of = -ohh--1)ee / + p(q - ff). ) a! = x + v/, where is the
equivariant function via Eq. 12 Full details provided in Appendix

---

Published as a conference paper at ICLR 2022
Table 1:
Prediction error (x10-2) on various types of systems. The header of each column "p, 8, ht denotes the scenario with p isolated particles, sticks and h hinges. Results averaged across runs.
[Triin] = 500
[Train] 1500
1,2,0
2.0.1
3,2,1
0,10,0
5,3,3
1,2,0
2,0,1
3,2.1 0,10,0
5,3,3
Linear
8.23.10.0
7.55
9.76.10.00
11.36.100
11.62.000
8.22.1000
7.55.0.00
9.7610000
11.361000
11.62.1000
GNN
5.33.1007
5.00s
7.58 10.0
9.8310.0
1.77
1.61 3.23 1.73.aa 7.97_a04
7,91
TFN
11.54 9..7
11.66.0.08
13.43.001
12.23
5.86. 4.97 10.23 8,51
10.75 11.21
SF
5.54-2006
5.14
8,95
11.42:001
11.59=001
5.02:200 4.68 8,39
10.82
10.85
RF
3.50 3.07
5.25
7.59
7.7320.30
1.
2. 2. 3.80
5.31
5.662027
EGNN
2.81+
2.27++0a4
4.67
+0.7
4.75+ous
4.59+uum7
2.59+ 1.86+ua22 a.54
2.79
3.25+
EGNNReg
2.9
2.66+ta.
7.01
5.03+aous
6.31
2.74 1.58+00a 2.62
3.03=mu7
3.07=004
GMN
1.84
2.02....0
2.48
10.04
2.92.10.0
1..0
1.68
1.47.0..0 2.101.00 2.32
2.86 1.001
4 EXPERIMENTS
4.1
SIMULATION DATASET: CONSTRAINED N-BODY
Datasets. We inherit the 3D extension of Fuchs et al. 2020) based on the N-body simulation introduced in Kipf et al. (0088) For each trajectory, we provide the initial states of the system 1==1, connected by sticks or hinges. The task is to predict the final positions { of the particles when T = 1000. The validation and testing sets contain 2000 trajectories. We evaluate the prediction error by the MSE metric. Compared with the simulation conducted in Fuchs et al. 2020); Satorras et al.) (202T) our dataset is more challenging in three senses: 1. We consider systems with multiple scales, including 5, 10, and 20 particles in total, respectively. 2. We introduce to the system the dynamics of hinges and sticks (depicted in Appendix [B), and construct various combinations between these objects. 3. We investigate the performance of each model across different scales of training set, e.g. 500 and 1500, to see how the models perform with scarce or relatively abundant training data. The system consisting of p isolated particles, sticks and h hinges, is abbreviated as (p. h).
{ the particles' charges {c E 1)]=1 and configuration indicating which particles are
Implementation details. Following Satorras et al. (021), D. we use a linear mapping of the scale of initial velocity ||e|]] |2 the input node feature h.. The edge feature is provided by a concatenation of the product of charges c,ey and an edge type indicator Tıy. where I5 is valued as 0 if node and j are disconnected, if connected by stick, and 2 if connected by a hinge. Note that this edge type indicator is an augmentation over the original setting in EGNN, designed to enforce EGNN and other baselines the ability to distinguish different types of edges, namely, with or without constraints. Other settings including the apper-prrameters are introduced in Appendix
Comparison with SOTAS. Table reports the performance of GMN and various compared models: EGNN (Satorras et al. 202T and its regulated version EGNNReg, SE(3)-Transforme Transformer (Fuchs et al.) 2.20). Radial-Field (RF) (Köhler al., 2019), Fiisor-Field-Netoorr (TFN) (Thomas et al.] 2018). and other two baselines, GNN and the Linear prediction (Satorras et al. 2021). Regarding EGNN and EGNNReg, they share the same backbones (i.e. (P and 1 ) and training hyper-parameters (learning rates, layer number, etc) with our GMN for a fair comparison. For EGNNReg, we explic- itly involve a regularization term during training by enforcing the geometrical constraints, namely preserving the lengths between two particles sticks and hinges; the regularization factor is ranged from 0.01 to 0.1, where the value giving the best performance is selected. The default settings of -E(3)-Transformer and TFN perform poorly our experiments, hence have tried our best to tune their hyper-parameter by validation. From T we have these observations:
1. GMN achieves the best performance in all scenarios, suggesting its general superiority. 2. GMN is more robust when the complexity of the system increases. On |Train] = 500, for example, the performance of GMN degenerates slightly by increasing the number of particles and hinges (e.g. from (1,2,0) to (3,2,1)), while other methods (such as EGNN) drops significantly. 3. Reducing the training size will hinder the performance of all compared methods remarkably. On the contrary, GMN still performs promisingly in general. For instance, on (3,2,1), EGNNReg becomes much worse from 2.62 to 7.01 when the volume of training data is decreased from 1500 to 500, whereas the change of GMN is smaller (2.10 v.s. 2.48). This is reasonable as GMN has explicitly encoded the constraints as opposed to EGNN and EGNNReg that learn to remember constraints by training.
7

---

Published as a conference paper at ICLR 2022
Table 2: Generalization across different systems, All models are trained in the (3,2,1) scenario.
[Train]
=500
Train
1500
3,2,1
2,4,0
1,0,3
Average
3,2,1
2,4,0
1,0,3
Average
GNN
7.58
8.06
8.37
8.00
4.73
4.98
5.58
5.10
EGNN
4.67
3.42
4.40
4.16
2.54
2.75
3,49
2.93
EGNNReg
7.01
4.49
6.62
6.04
2.62
2.62
4.29
3.18
GMN
2.48
2.53
3.28
2.76
2.10
2.18
2.65
2.31
Data efficiency. Fig. 4 depicts the prediction errors on (3,2,1) when the training size varies. is observed that GMN acts steadily, justifying its benefit in data efficiency. Once again, both EGNN and EGNNReg deliver much worse performance when the training size is small, and they approach GMN when the training dataset is enlarged sufficiently. In physics, it is important to ensure the physics rules that are discovered by relatively small number of experiments to be general enough for explaining universal phenomena. Hence, data efficiency, as a key advantage of GMN, comes as an important require- ment for learning to model physical systems. Besides, by the comparison between GNN and other equivariant models, it is seen that equivariance is a crucial point for performance guarantee.
78
GIN
500
1006
1230
2003
2008
3000
Figure 4: Data efficiency.
The generalization capability across different systems. It interesting check how the models perform when trained on one system but tested on others. Hence, Table Sablee[te tests the generalization from (3,2,1) to (2,4,0) and (1,0,3). Interestingly, the performances of all models on new systems are comparable with the original environment. We conjecture that this ability could be attributed to the usage of GNN in capturing the combination diversity of the objects. As before, GMN performs best.
Ablation studies. 1. The func- tion (22 is designed to be equivari- ant in Eq. 12 To justify this ne- cessity, we replace (2 with an MLP of the same size, and report the er- rors in Table 3 from which we con- firm that removing equivariance in- curs detriment to the performance. 2. The default setting of GMN in Eq. iis is leveraging unshared acceler- ation inference (p(()) for sticks and (p2(f x ) for hinges. Here we investigate different cases by assigning the identical form of p22 for sticks and hinges with shared 1-1 or unshared parameters. Applying identical (p2(f! 2 o (1-1) mostly outperforms the case by us- ing (2f)), probably owing to the better expressivity of the former version. Yet, both cases are worse than our design, implying that the dynamics of sticks and hinges should be modeled dis- tinctly. 3. We have introduced normalization term in Eq. [2] Table jablosjddemonstraats that eliminating this term leads to divergence during training, possibly owing to the numerical instability in the for- ward/backward propagation. 4. We have also implemented GMN-L that replaces the hand-crafted FK in GMN with a learnable black-box equivariant function. GMN-L outperforms EGNN in var- ious settings, verifying the validity of using our proposed equivariant message passing layer and the object-level message q in Eq. 6 Yet. GMN-L still yields minor gap with GMN, implying the benefit of involving domain knowledge. The full results are deferred to Appendix[]
Table 3: Ablations. "O.E" denotes numerical over-flow.
Train
500
Train| = 1500
3
5,3,3
3,2,1
5,3,3
GMN GMN-L
2.48
4.08
2.10
2.86
3.19
4.34
2.28
3.03
w/o Equivariance
3.74
4.41
2.46
3.29
92 with (f.. ti U
1-1 1-1
2.86
4.15
2.20
3.00
22 with (fl, x Us shared
3.86
4.25
2.30
3.10
P with only ff
3.10
4.39
2.34
4.19
& with only fl, shared w/o Normalization
2.91
4.94
2.39
3.45
3.15
O.F.
O.F.
O.F.
4.2 APPLICATIONS ON REAL-WORLD DATASETS
This subsection introduces how to apply GMN to practical applications including MD17 (Chmiela al., (2017) and CMU Motion Capture (CMU 2003). is not required to manually derive the entire kinematics for these complex systems; instead, each input system is decomposed as set of particles and sticks (e.g., the circles in Fig. 5).. which can be directly processed by our current formulation of GMN without any modification. The core is modeling partial cength-constraigg of the input system with disjoint sticks. The full details of the kinematics decomposition trick are in Appendix C
8

---

Published as a conference paper at ICLR 2022
Table 4: Prediction error 1) on MDI7 dataset. Results averaged across 3 runs.
Aspirin
Benzene
Ethanol
Malonaldehyde
Naphthalene
Salicylic
Toluene
Uracil
RF
10.94
103.72+12
4.64+0m
13.93-0m
0,50+am
1.23=0.01
=0.93 0.64+001
TFN
12.37
58.48:19
13.62-00
0.49:
1,03
10.8 0.84
SE
11.12=006
68.11-a77
4.74
13,89-01:
0,52+0om
1.13-0m
10.88=006 0.79+002
EGNN
14,4
62,40-05s
4.642am
13.64_001
0.47amm
1.02aam
11.
11.78gom
0 0.64gno1
EGNNReg
13.82=01
61.68-037
6.0620m
13,49-00
0.63+am
1.68-0m
11 11.05=a001
0 2.66=001
GMN
10.14
4.83
-3.11-0m
0.40
=091
10.22:20ce
2000
GMN-L
9.761am
$4.17
4.63.am
12.82_nm
0.41
0.88
10:45
0.59
MD17. We adopt MD17 (Chmiela et al., 2017) which involves the trajectories of eight molecules generated via molecular dynamics simulation. Our goal here is to pre- dict the future positions of the atoms given the current system state. We observe that the lengths of chemical bonds remain very stable during the simulation, making it reasonable to model the bonds sticks. The complete implementation details are deferred to Appendix E As
Aspirin
Benzene
Ethanol
Molonaldehyde
Naphthalene
Salicylic
Toluene
Uracil
presented in Table 4 GMN outperforms other competi- tive equivariant models on 7 of the 8 molecules. Particu-
Figure Molecules in MD17.
larly, on molecules with complex structures (e.g., Aspirin, Benzene, and Salicylic), the improvement of GMN is more significant, showcasing the benefit of constraint modeling on the bonds. Yet, we also observe that the constraint-aware models (GMN and EGNNReg) perform worse than others on Ethanol, possibly because Ethanol is a relatively small molecule with simple structure, where considering the bond constraints possibly makes less benefit but instead hinders the learning. Sur- prisingly, GMN-L showcases very competitive performance on this dataset. It surpasses GMN on 4 of the molecules, exhibiting that learnable FK works in practice even it does not involve domain knowledge of the constraint into kinematics modeling.
CMU Motion Capture. We use the motion data from the CMU Motion Capture Database (CMU) 2003). which contains the trajectory of human motion in various scenarios. Different parts of the hu- man body could be treated as hard rigid-body constraints. We focus on the walking motion of single object (subject #35) containing 23 trials, similar to Kipf et al. (2018). As depicted in Table 5] GMN outperforms other models by a large margin, verifying the efficacy of our equivariant constraint module on modeling complex rigid bodies. We further provide visualization in Fig. 6 where
GMN predicts the motion accurately while EGNN yields larger error. Again, GMN-L, although is inferior to GMN, is better than other methods remarkably.
Table 5: Prediction error (x 10 3) on motion capture. Results averaged across runs.
GNN
TFN
-E(3)-Tr.
RF
EGNN
EGNNReg
GMN
GMN-L
67.3
66.9-27
60.9
197.0
59.1
59.5
43.9
50.9
38
18
Figure 6: Left to Right: initial position, GMN, EGNN (all in bluc). Ground truths are in red.
CONCLUSION
In this paper, we propose Graph Mechanics Networks (GMN) that are capable of characterising constrained systems of interacting objects. The core is making use of generalized coordinates, by which the constraints are implicitly and exactly encapsulated the forward kinematics. To enable Euclidean symmetry, we have developed general form of equivariant functions to simulate the interaction forces and backward dynamics, whose expressivity is theoretically justified. For the simulated systems with particles, sticks and hinges, GMN outperforms existing methods regarding prediction error, constraint satisfaction and data efficiency. Moreover, the evaluations on two real- world datasets support the generalization ability of GMN towards complex systems.

---

Published as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
The complete proof of the theorems is provided in Appendix A The hyper-paramelers and other experimental details are provided Appendix
Our code is available at:
cww/hanw/cacccccccccccocccaaaa.a GMN
7 ETHICS STATEMENT
The research in this paper does NOT involve any human subject, and our dataset not related to any issue of privacy and can be used publicly. All authors of this paper follow the ICLR Code of Ethics httpss//cccccc/abbblcoooooffffit
ACKNOWLEDGMENTS
This work was jointly supported by the following projects: the Scientific Innovation 2030 Major Project for New Generation of AI under Grant NO. 2020AAA0107300, Ministry of Science and Technology of the People's Republic of China; the National Natural Science Foundation of China (No.62006137) 137); Tencent AI Lab Rhino-Bird Visiting Scholars Program (VS2022TEG001); Beijing Academy of Artificial Intelligence (BAAI).
REFERENCES
Peter W Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. arXiv preprint
arXiv:1612.00222, 2016.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields. Science udvances, 3(5)ee 2017.
CMU. Carnegie-mellon motion capture database. 2003. URL
cmu edu
George Cybenko. Approximation by superpositions a sigmoidal function. Mathematics of control, signals and systems, 2(4):303-314, 1989.
Mingyu Ding, Zhenfang Chen, Tao Du, Ping Luo, Josh Tenenbaum, and Chuang Gan. Dynamic visual reasoning by learning differentiable physics models from video and language. NeurIPS, 34, 2021.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolu- tional neural networks for equivariance to lie groups on arbitrary continuous data. International Conference Machine Learning, pp. 1176-1176. PMLR, 2020a.
Mare Finzi, Ke Alexander Wang, and Andrew Gordon Wilson. Simplifying hamiltonian la- grangian neural networks via explicit constraints. arXiv preprint arXiv:2010.13550 2020b.
Fabian B Fuchs, Daniel E Worrall, Volker Fischer, and Max Welling. ())-transformrss: roto- translation equivariant attention networks. arXiv preprint arXiv:2006. 10503, 2020.
Chuang Gan, Jeremy Schwartz, Seth Alter, Martin Schrimpf, James Traer, Julian De Freitas, Jonas Kubilius, Abhishek Bhandwaldar, Nick Haber, Megumi Sano, et al. Threedworld: A platform for interactive multi-modal physical simulation. NeurIPS, 2021.
Kurt Horiik. Approximation capabilities of multilayer feedforward networks. Neural networks, (2):251-257, 1991.
Michael J Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. Lietransformer: equivariant self-attention for lie groups. In International Confer- on Machine Learning, pp. 4533-4543 PMLR, 2021.
10

---

Published as a conference paper at ICLR 2022
Thomas Kipf, Ethan Fetaya, Kuan-Chieh Wang, Max Welling. and Richard Zemel. Neural relational inference for interacting systems. arXiv preprint arXiv:1802.046877 2018.
Jonas Köhler, Leon Klein, and Frank Noć. Equivariant flows: sampling configurations for multi- body systems with symmetric energies. arXiv preprint 191000111, 2019.
Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B Tenenbaum, and Antonio Torralba. Learning par- ticle dynamics for manipulating rigid bodies, deformable objects, and fluids. arXiv preprint arXiv:1810.01166, 2018.
Michael Lutter, Christian Ritter, and Jan Peters. Deep lagrangian networks: Using physics as model prior for deep learning. arXiv preprint arXhv:1907.044900 2019.
Karolis Martinkus, Aurelien Lucchi, and Nathanaël Perraudin. Scalable graph networks for particle simulations. AAAI, 2021.
Mises and Hilda Pollaczek- Geiringer. Praktische verfahren der gleichungsauflösung. ZAMM- Journal of Applied Mathematics and Mechanics/Zeitschrfft für Angewandte Mathematik und Mechanik, 9(1):58-77, 1929.
Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li Fei-Fei, Joshua B Tenenbaum, and Daniel LK Yamins. Flexible neural representation for physics prediction. arXiv preprint arXiv: 100.18000004444 2018
Tobias Pfaff, Meire Fortunato, Alvaro based simulation with graph networks. arXiv preprint arXiv:2010.03409, 2020.
and Peter W Battaglia. Learning mesh-
Alvaro Sanchez-Gonzale, Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integrators. arXiv preprint
2019.
Alvaro Gonzalzz Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning simulate complex physics with graph networks. In International Confer- ence on Machine Learning, pp. 8459-8466 PMLR, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling, E(n) equivariant graph neural net- works. arXiv preprint arXiv:2102.09844, 2021.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular confor- mation generation. arXiv preprint arXiv:2105, 03902, 2021.
Joshua B T'nnenbaum, Charles Kemp, Thomas I. Griffiths, and Noah D) Goodman. How to grow mind: Statistics, structure, and abstraction. science, (6022) 2011.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and translation-equiivrrant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08229, 2018.
Benjamin Ummenhofer, Lukas Prantl, Nils Thuerey, and Vladlen Koltun. Lagrangian fluid simu- lation with continuous convolutions. In International Conference Learning Representations, 2019.
Zonghan WI, Shirui Pan, Fengwen Chen, Guodong Long, Chengqi Zhang, and S Yu Philip. A comprehensive survey on graph neural networks. IEEE transactions on neural networks and learning systems, 32(1):4422 24, 2020.
Shuqi Yang, Xingzhe He, and Bo Zhu. Learning physical constraints with neural projections. arXiv preprint arXiv:2006. 12745, 2020.
11

---

Published as a conference paper at ICLR 2022
A THE PROOF OF THEOREM
1
In the following, we define the orthogonal group as ((d) {0 E Rdxd o 00 =a] Prior to providing the proof for Theorem I we first list two necessary lemmas below.
Lemma 1. For any two matrices Z1, Z2 Raxm, we have this equivalence: 30 e O(d), OZ, Z2 + 2 2 Z Zz.
Proof. We only need to prove the "=" direction, as "-" holds clearly. Suppose the SVD decom- position of Z as Z1 V with the left-singular matrix U1 E O(d), the singular diag- onal matrix S1 E Rd and the right-singular matrix V12 E ((m). Since Z Zı Z Zz, then there must exists certain SVD decomposition of Z₂ that shares the same singular matrix and right-singular matrix with Z,, implying that Z2 = U2S1 V₁",, where U2 E ((d). Hence, VT U,U, Z1 which concludes the proof owing to the orthog-
Z₂ = U2S,V U2U1 onality U2U1
Lemma 2. The function f: Rdxm - Rm is invariant on O(d), namely, S(OZ) = 1(2), YO E
O(d), ZZ E Rdxm, if and only if there exists function 9 : Rmxmm - Rmm satisfying f(Z) = g(Z Z).
Proof. The sufficiency is obvious. We now prove the necessity. We define the equivalence class [Zo] = 12 30 E (d). oz = Zo}. Since f is invariant to the orthogonal transformation, means f is actually function the equivalence class, i.e., f(Z)
f(Zz). On the other hand, according to Lemma we have [Z = [Zg] + Z Z = Z Zz, which implies the one- to-one correspondence between [Z] and Z Z; hence, there must exist function f' leading to
Z (Z 2), and f' is continuous in terms of any invariant metric such as the norm I Z ||
Overall, ((Z) - s((Z)) (( Z)) = g Z Z).
We are now ready to prove Theorem Ithat is copied below for better readability. Theorem 1. If d and the row rank of Z is full, i.e rank( Z)) d, then for any continuous anthogonality--quivrrie ((Z, < forrr arbitrarily small error
function Z(Z h). there must exist an MLP ow satisfying II
Proof. Without loss of generality, we will omit the non-vector term h, which does not change the story but let our proof more concise. Because Z is of full row-rank, the columns of an arbitrary function ( can be represented as a linear combination of the columns of Z; in other words, there must exist function π R -> Rmxm' giving rise to (ZZ) Z ()) Considering the orthogonality-cquivvriann we derive the property of π (Z) as:
*(0Z) (Z(
OZZ(OZZ) OZT(Z),
(14) Since d < and the row-rank of Z is full, we perform the compact SVD decomposition on Z, Sz 0, and Vz E RRaxd satisfying
+ (T(ZZ)= (Z)
namely, Z- U.S2VZ I where Uz E O(d), Radd
V Vz= I., Eq 14 becomes:
Zm(OZ) = ZT(Z),
(OZ) UnS2V π (Z),
= Sz S2 (OZ) Sz SZZV π (Z)
(15)
Vectore Z Z) are respectively the eigenvalues and cigenvectors of Z and are clearly invariant to the orthogonal transformation of Z. Their values can be numerically approximated by iterative programs, such the power method (Mises & Pollaczek- Geiringer 1929), thus can be treated as the continuous functions of Z. Let us define gz(Z) = Sz 2 ( Z). Then, we keep deriving Eq. by:
Given that Sz Eigen(2 2 2) and Vz
SZVZ क (OZ) = S² K (Z),
Soz VOzz (OZ) = V2 K (Z)),
7 g (Z)
(16)
12

---

Published as a conference paper at ICLR 2022
According to Lemma 2 for a certain function g. By checking the formulation of (2) as demonstrated before, we arrive
the function 9 satisfies Eq. 16 if and only it written as g' g( Z
( Zz)
U2S2 a (Z),
(Z)
Uuzg(Z (ea Z),
UUS2VZ Vz g( Z),
SV 9(Z Z),
Zn(ZZ' Z).
(17)
Here the function 17 can be approximated by MLP whose universality has been justified by (Cybenko 1989
Horni 1991). The conclusion of Theorem is proved.
Corollary 1. Assume 2 d, and also the entries Z are drawn independently from a distribution that is absolutely continuous with respect to the Lebesgue measure in IR. Then, almost surely, the conclusion of Theorem TneeoremTholee hoolss
Proof. This is straightforward. When rank Z) < d, the columns of Z are located in subspace of Fd (for example line or a plane in the 3D space), whose measure is zero. Therefore, the probability for making rank( Z) d is 1, and we almost surely have the same conclusion Theorem |
Corollary 2. For any continuous aaatogoonltt--eqiivrrr function Φ h) whose output is lo- cated in the linear subspace expanded by the columns of Z. the conclusion of Theorem holds universally.
Proof. According to the definition of (Z),, we still obtain (Z = (Z))). Let us assume > m (otherwise directly obtain Theorem D then the full SVD decomposition of Z is Z UzS2VZ with Uz E (d). Sz E ddx and Vz E (m). But here, Sz is not strictly positive.
S
Suppose Sz
where S. >
We retain that SzV2 n (Z) = g( Z) by imitating the
0
proof in Eq. 14 16
Analogous to Eq. 17 we derive,
((Z)
(T(Z),
Uzg(Z Z),
U2S2 V2 Vz (S.² .0) g( Z),
ZVz (ST 0) g(Z Z),
Zn(Z Z)
(18)
which concludes the proof.
B THE DYNAMICS ANALYSES FOR STICKS AND HINGES
The analytical forms of the dynamics for sticks and hinges can be found from mechanics book. Here we derive the formulas on our own to make our paper more celf-contained. For simplicity, we consider all particles to be of the equal mass and the sticks of no mass.
Dynamics analysis of sticks. In Fig. 7 (Let), we assume the forces acting on particles 1 and 2 are separately fi and fz. By following the theorem of the motion of the center of mass, the acceleration of the center is given by
fi + f2
g0
(19)
13

---

Published as a conference paper at ICLR 2022
MoL
f2
Vot
Xoz
f2
d
Aob
Xg2
f1
00
Figure 7: Left: Dynamics of sticks. Right: Dynamics of hinges.
The rotation accelerations of particles and around the center are calculated by
M X fi + Lo2 X fz
0o1
Oo2 J
(20)
m 20 1 ml 2022|
where M defines the total torque and J is the moments of inertia.
Dynamics analysis of hinges. The analysis for hinges is more complicated than sticks. Fig. (Right), the forces particles 0, and 2 are fo. fi and fz. By using Newton's 's second law,
f
ao + 01 + a2
(21)
m
where the aggregated force is f = fo + fi + fz. In addition, the kinematics relations between the three particles show that
a1 ao 0o1 x01 Por x Pol.
(22)
02 = ao 0o x 202 + Ou2 x Vo2.
(23)
where Oo and Oo denote the speed and acceleration of the rotation angle of particle 1 around 0, and Voi is the corresponding linear velocity; the symbols oo.. oo2 and Vo2 defined similarly. Moreover, the relative rotation acceleration of particle to O is caused by the external force fi and the inertia force -mag acting 1, which derives that
00:
Mol
Fo1
(fi map)
(24)
Jo
Analogously,
002
M0z
Lu2
(f2 ma,)
(25)
Jo2
m Ø02 2
After substituting Eq. 24 into Eq. 22 and Eq. 25 into Eq. 23 and then rearranging Eq. 21 we have
ag (I + ep1en + 6020
(26)
where f 0p1 x Po1 Oor x (I cor e (I e02e fa with ep and eg
represent the unit vectors along 2o1 and Mo2, respectively. Wee have applied a trick to simplify the derivation by observing Ro x X
(I eor ( Besides, I eai e - e02e invertible,
| 2an
hence Eq. 26 is always meaningful.
Substituting Eq. back into Eq. 24 and 25. derives the values ou and 002. Note that in $ 3.3 employ the denotation of generalize lizd coordinates by qo. qu and qo which indeed share the same meaning with the Cartesian coordinates un, Un and ag of particle O.
C KINEMATICS DECOMPOSITION
We manually build the forward kinematics (Eq. 99) the stick (and hinge), which relies on the domain knowledge of the underlying physics. However, for complex systems, longer require
14

---

Published as a conference paper at ICLR 2022
Table 6: Prediction error (x10-2) on various types of systems. The first column "p, 8, h" denotes the scenario with p isolated particles, sticks and hinges. Models are trained with 500 samples.
GMN
EGNN
EGNNReg
GNN
TFN
SE(3)-Tr.
RF
Linear
1,2,0
1.84
2.81
2.94
5.33
11.54
5.54
3.50
8.23
2,0,1
2.02
2.27
2.66
5.01
9.87
5.14
3.07
7.55
2,4,0
2.34
3.59
3.87
8.05
11.30
9.22
5.37
10.10
0,5,0
2.54
4.13
4.29
8.63
11.92
9.83
5.94
10.48
7,0,1
2.39
2.66
3,41
7.05
10.67
8.38
4.66
9.71
1.0,3
3.21
4.56
5.14
8.32
11.62
9.57
5.91
9.90
3,2,1
2.48
4.67
7.01
7.58
11.66
8.95
5,25
9.76
4,8,0
3.69
4.79
7.09
9.65
12.05
11.21
7.59
11.45
0,10.0
2.92
4.75
5.03
9.83
13.43
11.42
7.59
11.36
8,0,4
3.37
4.17
5.32
9.49
11.72
11.12
7.51
11.44
2,0,6
4.06
5.06
5.58
10.13
12.13
11.74
8.15
11.61
5,3,3
4.08
4.59
6.31
9.77
12.23
11.59
7.73
11.62
to derive the kinematics of the entire system. Instead, we propose kinematics decomposition, simple yet effective trick that can decompose each input system (an arbitrary graph) into particles example, for each molecule, select certain bonds as
and sticks. Taking the MD17 dataset sticks (the circles in Fig. 5 and the remaining atoms as isolated particles; this way, we obtain set of particles and sticks. Note that different sticks not allowed intersect; otherwise, it will generate two values for the intersecting particle of two sticks and cause ambiguity if these two values are distinct. Although this kind kinematics decomposition will only maintain partial constraints, it greatly enlarges the application scope of our current formulation Eq. any revision. More importantly, our experiments verify that GMN by this formulation sufficient to surpass other methods on complex systems like molecules. On CMU Motion Capture, since the motion graph contains no circle and is of the tree-like structure, is tractable to derive the exact forward kinematics by recursive kinematics computation from the root node. Yet, we still encourage the usage the above kinematics decomposition for its easy implementation and compatibility with GMN.
(5 without
D FULL ALGORITHMIC DETAILS
In the main body of the paper, for better readability, we first introduce the general pipeline of our method in [.1 and then present the implementation details taking the angels into account in Here, we combine them into one singe algorithmic flowchart in Alg.
3.3
E
MORE EXPERIMENTAL DETAILS AND RESULTS
Hlyper-parameter and baselines. For GNN, RF, EGNN, EGNNReg, and GMN, we empirically find that the following hyper-parameters generally work well, and use them across all experimental evaluations: batch size 200, Adam optimizer with learning rate 0.0005, hidden dim 64, and weight decay 1 x 10-10, All models are evaluated with four layers. 3E(3)-Transformer and TFN do not perform well on our datasets, potentially due to the challenge of highly complex and constrained systems. Consequently, we tune the hyper-parameters and adopt the following configuration: batch size 100, learning rate 0.001, hidden dim 64, representation degrees 3 and weight decay 1 x 10-8 Models trained for 600 epochs datasets. EGNNReg is a variant of EGNN that explicitly adds the constraint error into its training loss by regularization factor of X. In our experiments, also treat À as hyper-parametee and choose À that yields the best performance within the range [0.01, 0.1].
the simulation dataset, and 500 epochs on the real-world
Detailed experimental setup MD17. We randomly split the dataset into train/validationneest sets containing 500/2000/2000 frame pairs respectively. We choose = 5000 as the span between the input and prediction frames, and the difference in positions as the input velocity. The hyper-
15

---

Published as a conference paper at ICLR 2022
Algorithm 1 Graph Mechanics Networks (GMNs)
Input: Initial states of all particles {S (a), " and features {4]} N. Learnable equivari- ant functions (P, P2:, 1), v Layer number L.
Compute the generalized coordinates of all structural objects {(gh Lh he).)))) for layer 1 = 1 to I. do for particle
to N do Calculate the interaction force f and feature hf for each particle by:
N . ,1 ees))
$ 1 Σ
pi (ar)a
(27)
end for for object lo 1 to K do Inference the generalized Cartesian acceleration d% via:
q = Σ pp((f))
(for sticks)
(28)
1EO
d Σ
92 a v. (for hinges)
(29)
iEO
Derive the generalized angle acceleration:
0x: Eieo, Tsi X ft
Vi Os, (for sticks)
(30)
EiEo, |
0z. Ibi x (f: qn) | 1|
Vi E On.
(for hinges)
(31)
Update the positions and velocities as follows:
d v Σ M
D + t
(32)
VEO2
d = dh
1-1 Th.
(33)
0/.
h
+ 0 Vi On.
(34)
iEO2
Perform the forward kinematics for each particle in Oμ:
a gh. + rot(O(a) ا- Vi O₁,
(35)
v dh + 0% x aty Vi E O
(36)
end for end for
Output: The predicted states of all particles 1 1
parameters of all models are kept the same as the synthetic dataset. We masked out the hydrogen atoms, focusing the prediction of large atoms. We further augment the original molecular graph with 2-hop neighbors, and concatenate the hop index with atom numbers the connected atoms as well as the edge type indicator the edge feature, similar to Shi et al. (202I). We use the norm of velocity concatenated with the atom number as the node feature. We randomly select bonds without commonly coonneee-cconeetted atoms sticks, and the rest of atoms isolated particles. Although by this means not all the bond lengths are preserved, our experiment does illustrate that this is simple but effective strategy of applying GMN complex systems like molecules.
Detailed experimental setup on CMU Motion Capture We first split the data into train/val/test sets containing 1/6/6 trials respectively. We then sample from the trials to get 200/600/600 frame
16

---

Published as a conference paper at ICLR 2022
Figure 8: Left: initial position(s). Middle: the prediction(s) of GMN (in blue). Right: the predic- tion(s) of EGNN (in blue). Ground truth is marked in red. Better viewed by zooming in.
2
4,
Figure 9:
Left: initial position (in blue). Middle: the prediction of GMN (in blue). Right: the prediction of EGNN (in blue). Ground truths are marked in red. Better viewed by zooming in.
pairs with 7 = 30. We use the norm of velocity node feature. We also augment the edges with 2-hop neighbors, and use the edge type indicator as edge feature. As for GMN, we sample 6 key bones of human body (as specified in Appendix C as sticks, resulting in system with 19 isolated particles and 6 sticks. Empirically, we select the edges connecting nodes (0, 11), (2, 3), (7, 8), (12, 13), (17, 18), and (24, 25) as sticks on the motion capture dataset, which indeed are the key parts of human body like the arms and legs.
More visualizations. We visualize the prediction outcomes by GMN and EGNN in Fig. g on the simulation dataset. GMN is found to be able to track the ground-truth trajectories accurately. whereas EGNN yields clear position errors and particularly breaks the constraints for sticks and hinges. These results are consistent with the performance in Table I Fig. provides extra visual- ization on Motion Capture. Similarly, GMN yields more accurate prediction than EGNN.
More experimental results. In Table 6 we provide comprehensive performance comparison of different models under more scenarios with 500 training samples. Clearly, GMN outperforms other provide the learning curve for GMN,
models all object combinations involved. Besides, EGNN, and EGNNReg on the simulation dataset (3,2,1) in Fig. I0 GMN yields lower training loss and testing loss than EGNN and EGNNReg, benefiting from its constraint modeling.
17

---

Published as a conference paper at ICLR 2022
ONN (Trak)
0335
CANN (
EONN (Tra
CGINN (Twe]
BOINRegg (Trai)
6.20
BONNReg (Test)
0:15
0.00
Epoch
Figure 10: Learning curves on (3,2,1) with 500 training samples.
F MORE ABLATIONS
Hinge treated as two sticks. We treat 0-1, 0-2 as two sticks, and apply the stick FK respectively. Afterwards, translate the two sticks such that the Os coincide at the midpoint of their predicted positions. From Table 7 we find that this strategy performs worse than the binge-modeled GMN, since splitting the hinge into two sticks overlooks the kinematics at the connected point. Yet and still, it performs better than EGNN, again verifying the benefit of our proposed constraint modeling.
Table Ablation on hinge FK.
Train| 500
TTain 1500
3,2,
2,0,6
5,3,3
3,2,1
2,0,6
5,3,3
EGNN
4.67
5.06
4.59
2.54
3.50
3,42
EGNNReg
7.01
5.58
6.31
2.62
3.61
3.07
GMN (Stick only)
3.02
4.32
4.21
2.37
3.30
2.88
GMN
2.48
4.06
4.08
2.10
3.22
2.86
Charges as node or edge feature. In the experiment we by default assign Pgy C.ty as the edge feature (denoted as "Edge C"). Here we instead concatenate c; to the node feature of particle i, and set all Eij = 0 (denoted as "Node + C"). In Table we observe that these alternatives on charges make very limited difference performance, and indeed the models charges from node features, which is truly the case as depicted in Eq. 2 and
learn the interaction of
Table 8: Comparison of charge-assigning strategies.
Node + C
Edge + C 3,2,
1,2,0
2,0,1
3,2,
0,10,0
5,3,3
1,2,0
2,0,1
0,10,0
5,3,3
EGNN
2.89
2.28
4.25
4.80
4.50
2.81
2.27
4.67
4.75
4.59
EGNNReg
3.17
2.74
8.20
5.01
6.64
2.94
2.66
7.01
5.03
6.31
GMN
1.89
2.01
2.63
3.07
4.02
1.84
2.02
2.48
2.92
4.08
G CONSTRAINT SATISFACTION
The constraint error is computed as the total change in the lengths of sticks and hinges between the input and output, averaged per trajectory. Specifically, for hinges, the two edges are both considered.
Constraint error.
Results. One vital feature of GMN is that maintains the geometrical constraints exactly and inherently. To show this, Table 9 records the corresponding constraint errors of several typical models. The results do verify our claim that GMN always outputs near-zero zero errors (all below le-4). Although EGNNReg that augments EGNN with regulation helps in reducing the constraint errors, is data-driven and limited by the number of training samples; further, the constraints are pursued softly, making it defective for the applications where hard constraints indispensable.
18

---

Published as a conference paper at ICLR 2022
Table 9: Constraint enror various types of systems.
|Trii]| 500
Train = 1500
1,2,0
2,0,1
3,2,
0,10,0
5,3,3
1,2,0
2,0,1
3,2,
0,10,0
5,3,3
GNN
0.200
0.386
0.492
0.154
0.468
0.225
0.426
0.779
0.251
0.772
EGNN
0.220
0.370
0.714
0.248
0.760
0.217
0.317
0.521
0.139
0.596
EGNNReg
0.172
0.146
0.232
0.198
0.241
0.159
0.053
0.091
0.097
0.075
GMN
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
0.000
H
MORE DISCUSSIONS ON GENERALIZATION
In Table we compare the generalization capability of GMN with other methods. Here, we denote and GMN trained
GMN trained on (3,2,1) and tested across different scenarios as and tested on the same dataset as GMN-Original. is observed from Table TO that GMN has strong generalization capability, since the transfer performance is very close to the original setting.
Table 10: Comparison of GMN in the transfer and original settings.
|Train|
500
Train
1500
3,2,1
2,4,0
1,0,3
Average
3,2,1
2,4,0
1,0,3
Average
GMN-Transfer
2.48
2.53
3.28
2.76
2.10
2.18
2.65
2.31
GMN-Original
2.48
2.34
3.21
2.68
2.10
2.01
2.44
2.18
I LEARNABLE FK
is indeed instrumental to discuss whether a learnable black-box function, which requires less do- main knowledge, could also yield competitive performance, and if our hand-crafted FK still shows advantage over the learnable counterpart. To answer these questions, we replace the hand-crafted part (Eq.
9 well the Euler angle computations in Sec. 3.1. with the following equations: p(g gs
1-1 f!), x!" + v/. where is the equivariant message passing layer we propose in Sec. S.2. By this design, the parameterized FK preserves its equivariant property (and the theoretical universality). and compared to EGNN, it additionally leverages the information from the object-level generalized coordinates q We denote this variant of GMN as GMN-L. More- over, since the parameterized FK inevitably loses the constraint-preservinn property compared with the exact FK, therefore we also augment it with explicit constraint regularization, akin to what we did EGNNReg. We hence denote this variant as GMN-LReg.
of ((h D
We evaluate the performance of GMN-L, GMN-LReg and compare them with GMN with exact FK as well as EGNN and EGNNReg Table TI We interestingly find that GMN-L consistently outperforms EGNN in various settings (as well as the regularized version). which again verifies both the validity of our proposed equivariant message passing layer and the efficacy of leveraging object- level message (i.e., %% for the inference of FK. At the same time, GMN-L and GMN-LReg yield minor gap with GMN, showing the evidence that hard-coding the constraints replaces a nontrivial amount of learning complexity.
Table 11: Comparison with GMN-L and GMN-LReg.
Train
500
Train
1500
1,2,0
2,0,1
3,2,
0,10,0
5,3,3
1,2,0
2,0,1
3,2,
0,10,0
5,3,3
EGNN
2.81
2.27
4.67
4.75
4.59
2.59
1.86
2.54
2.79
3.25
EGNNReg
2.94
2.66
7.01
5.03
6.31
2.74
1.58
2.62
3.03
3.07
GMN-L
2.32
2.09
3.19
3.88
4.34
1.93
1.56
2.28
2.72
3.03
GMN-LReg
2.52
2.23
3.34
3.67
4.31
1.91
1.88
2.49
2.61
3.00
GMN
1.84
2.02
2.48
2.92
4.08
1.68
1.47
2.10
2.32
2.86
19

---

Published as a conference paper at ICLR 2022
J ROBUSTNESS TO ERRORS IN THE PHYSICAL PRIOR OF CONSTRAINTS
is interesting to test the robustness of our model w.r.t. the noisy constraints. This scenario would sometimes arise in real-world scenarios where we might not be certain about the exact connectivity of the rigid body, and thus would involve slight errors in domain expertise. Since our paper focuses on the distance constraint other than the angle constraint, the following investigations will only involve noise into the stick connectivity. We design three random perturbation operations on the input rigid body prior: (Join) randomly selecting isolated particles and joining them as if there is stick connecting; 2. (Split) randomly selecting an existing stick and splitting it as two isolated particles; 3. (Join Split) conducting operation and at the same time; and 4. (Change in Length) randomly adding Gaussian noise (O, 0.1L) to the length of stick, where L is its original length. Note that the operation is conducted independently for every training sample each time it fed into the network.
We observe that these perturbations, although somehow hinder the performance, in general do not jeopardize the performance too much (difference in MSE S 0.30), indicating that GMN is not sensitive to slight errors of the input physical prior of the constraints and it is still able learn some degree of given the wrong constraints.
We summarize the results in Table 12
Table 12: Robustness test in various scenarios.
Train
500
Train
1500
3,2,
5,3,3
8,6,0
3,2,1
5,3,3
8,6,0
GMN GMN w/ Join
2.48
4.08
2.84
2.10
2.86
2.22
2.59
4.27
2.98
2.22
3.16
2.37
GMN w/ Split
2.57
4.11
2.95
2.16
3.13
2.26
GMN w/ Join Split
2.63
4.16
3.01
2.31
3.01
2.34
GMN w/ Change in Length
2.75
4.36
3.11
2.15
3.15
2.41
K MORE DISCUSSIONS ON THE DECOMPOSITION
is possible to decompose into bigger objects rather than just sticks. As comparison, we further adopt the hinge-wise decomposition (i.e., decompose the system into particles and hinges), and investigate the performance on both MD17 and Motion Capture. We denote this model as GMN-LH, where stands for hinges. The results are depicted in Table [3 and Table 14 On MD17, GMN- L. yields a little bit worse performance than GMN-L on several molecules, while giving desirable results on Ethanol and Benzene. On Motion Capture, GMN-LH outperforms GMN-L by small gap. while is still worse than GMN. By default, we still encourage to perform the decomposition via sticks as sticks are actually the basic building blocks of hinges and other larger rigid objects.
Table 13: Prediction error (x10-3) 3) on MD17 dataset. Results averaged across runs.
Aspirin Benzene Ethanol Malonaldehyde Naphthalene Salicylic Toluene
Uracil
EGNN
14.41 2015 62.40 4.64
13. 23.61
0.47
1.02
11.7820.7
0.64
GMN
10.14:000 48.12.0.00
4,83
13. 1 1000
2.40
0.91
10.22:0.0s
0.59.001
GMN-L
9.764011
54.17.00.0
4.03
12082
0.41 10.01
10.45 0.08
0.5 0.59.00 1000
GMN-LH
10.25 +aos 52.02+
4.62
12.83
0.41 =0.0
1.03
10.8 +0.14
0.59 4o09
Table 14: Prediction error (x 10 3) on motion capture. Results averaged across runs.
EGNN
GMN
GMN-I.
GMN-LH
59.1
43.9
50.9
48.7
20


Reviews:
[{"review_id": "lEP41HMA0S-", "reviewer": "Reviewer_F8Jp", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths**\n\nThe presentation is clear and professional.\n\nThe task is well defined and has some physical relevance.\n\nThe supplementary materials are extensive and include a quality code repository that is well documented.\n\n**Weaknesses**\n\nThere is a lack of novelty and scope in the work. The following problems are fundamental to the contribution and it is not clear to me how the authors could remedy these faults within the rebuttal period unless they would like to flatly refute the points:\n\n1. Most of the core ideas presented here are captured and applied more widely, including to real-world scenarios such as the case of robotics mentioned in the introduction of this work, in ​​Deep Lagrangian Networks: Using Physics as Model Prior for Deep Learning Lutter et al. ICLR 2019. (This reference is missing.)\n\n2. The scope of the theoretical and empirical contributions of the paper are limited to very small systems of sticks and hinges. There are existing works in this area that scale to more complicated rigid systems without the need for hard coded physics (Sanchez-Gonzalez et al. ICML 2018), to large-scale systems of thousands of particles (Sanchez-Gonzalez et al. ICML 2020), and to systems with constraints (Pfaff ICLR 2021). The contributions are thus quite limited in the context of the achievements of the field (and in particular when the work of Lutter et al. is considered.)\n\n3. The general formulation of an equivariant message passing layer given in Equation 11 does appear novel, although the actual use case with $Z = (f,x,v) = (ma, x, v)$ is very close to the extension given by ​​Satorras et al. (EGNN) in their equation 7 that shows how to extend their method to using velocities. It is trivial to extend their method to include an equivariant treatment of accelerations ($v^{l+1}_i = v^{l}_i + a^{l+1}_i$). It is also in the spirit of the work presented here to include such a physical inductive bias (i.e. that the rate of change of the velocity $v$ is set to be exactly the acceleration $a$ and so on).\n\nAll the tables should include errors. I appreciate that the results suggest that the GMN is the dominant model in this system but it is still necessary to indicate the variation in the recorded performances.\n\nIt is also not correct to state that ‘GMN meets constraints spontaneously.’ The constraint satisfaction has been hand crafted into the model whereas the compared models are required to learn the constraints. This is nearly a minor point but it does make an incorrect claim.\n\nThe self-containment of the paper, as described in Appendix B with reference to the derivation of the dynamics of sticks and hinges, whilst admirable, is not relevant to the novelty or quality of the scientific contribution of the work.\n\n**Minor comments**\n\nThe following notes are very minor and may be subjective and as such I have not factored them in to my score for the paper. I would not update the score in either direction if the authors ignore or act on these notes, they're just things that I spotted as I went through the paper.\n- It's unclear to me what the first sentence of the abstract is trying to say. 'Prevailing yet challenging' is a little confusing to me, is it not the case that challenging topics are also prevalent?\n- 'mainly stem from that...' in the second sentence of the abstract needs to be rephrased\n- 'just like what physicists used to do' in the introduction is a little childish, though perhaps we should wonder what it is that physicists are now getting up to\n- 'Coulombian force' is an unusual variation on 'Coulomb force'\n- Page 4 'to inference the acceleration' should be 'to infer'\n\n\n### Rebuttal Update\nI have raised my score (3->5) in light of the additional experiments and lowered my confidence (3->2) as I have not been able to engage as thoroughly as the other reviewers who are apparently satisfied with the paper.", "summary_of_the_paper": "This paper develops a model that augmented interaction networks with mechanical constraints and develops equivariant message passing schemes to ensure physical realism in the model outputs. The work focuses on a set of toy systems that contains particles connected by sticks and hinges and the developments are made for these systems in particular. There are some results on generalised equivariant message passing and the universality of the proposed method. ", "summary_of_the_review": "The ideas presented are physically well motivated but only incrementally advance on the state of the field by combining many existing ideas, and the application scope is very narrow so it is not the case that the paper derives a significant contribution by showing that these techniques taken together are greater than the sum of their parts etc. The presentation is very professional and their is some novelty in extending the equivariant message passing scheme of Satorras et al., but these are not strong enough contributions to justify acceptance alone.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "dX_stVQQEKR", "reviewer": "Reviewer_4PY9", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "# Strenghts \n\n- The introduction does a good job of motivating the need for the proposed work. \n- The description of EGNN and the preliminary material more generally provide a good understanding of the relevant prior work. \n- The model essentially relies on an object-centric view and the dynamics of each object are treated unitarily. This could also be thought of as a novel instance of object-centric-based learning. \n- The theoretical analysis around the newly proposed equivariant layer is nice and useful. \n- Since the EGNN layer is a particular case of the proposed one, the derived result also has theoretical implications for this prior work and improves our understanding of its limitations. \n- The evaluation equips all the baselines with minimal information about the presence of the rigid objects (i.e. the edge-type indicator functions). This ensures a certain level of fairness of the evaluation methodology. \n- I like that the evaluation is nicely factorized into different settings (e.g. small vs large training sets and results across different types of systems). This provides further insights into how the model performs in different regimes. \n- Table 2 provides a nice empirical confirmation that the model works as expected and that the constraints are satisfied. \n- The generalization of the proposed models to novel system configurations is tested in Table 3 and the results look good. \n- The ablation study confirms the intuition / theoretical motivations provided for various design decisions described in earlier parts of the paper. \n\n# Weaknesses\n\n-  The evaluation only looks at a single simulated dataset. While these simulations are nice and insightful, I would have expected some real-world evaluation as well. Even the original work of Kipf considered a few real-world datasets. For instance, the authors could have looked at some motion capture data where different parts of the human body could be treated as rigid-body constraints. \n- Related to the point above, the forward kinematics part of the model relies on having a perfect understanding of the nature of the constraint and relies completely on domain knowledge. This raises two issues: 1) What happens when the constraint is so complicated we might not be able to derive some forward kinematics equations like for the simple hinges and sticks? 2) This additional domain knowledge provides very important information that the baselines do not necessarily have. For these two reasons, it would have been useful to also consider a model where Equation (9) is learned by an abstract learnable function. This would have provided a model less reliant on very specific domain knowledge and more centered on the inductive bias of object-centric modeling.\n- The considered constraints are relatively simple (i.e. small rigid bodies) and do not really test how the model can cope with more complex constraints. For instance, the sum in Equation (7) might increase too much in magnitude as the object is formed of many more parts. This naturally brings the question of whether the model can cope in practice with very large rigid objects? \n- None of the reported results have standard errors around them. The experiments should be run at least across multiple initializations in order to assess the statistical significance of the results.  \n\n", "summary_of_the_paper": "The paper proposes a new graph neural network model for predicting the dynamics of n-body systems. The two main innovations of this model are:\n1. Ability to handle systems with rigid-body constraints. \n2. A new 3D transformation equivariant layer with universality properties.  ", "summary_of_the_review": "Balancing the weaknesses and strengths from the list above, the paper is, in my view, at the threshold between acceptance and rejection. Since there is no neutral score I can select, I am voting for weak reject for now, but I am open to changing my score during the discussion period. \n\nEDIT: Updated my score to 8. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "5kidu5TlH9U", "reviewer": "Reviewer_vbzj", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nS1. The GMN model delivers exactly what it promises: exact constraint preservation, with E(n) equivariance. \n\nS2. Claim C2, seems like a good idea in general. In fact, leaving hard constraints aside, it seems like a nice alternative to model in Eq. (7) of the E(n) paper, to be able to directly feed the velocity vector to the message function, and perhaps give more expressive power.\n\nS2. The model shows expected advantages in terms of data efficiency over the baselines.\n\nS3. The model shows expected out of distribution generalization properties.\n\nWeaknesses:\n\nW1. The model requires specifying the forward kinematics function for the constraint (to go from generalized positions and velocities to cartesian positions and velocities) as a differentiable function, which in many cases is of a big ask. Furthermore, the approach for writing this constraints does not seem very compositional, for example two separate analyses (forward functions, sets of generalized coordinates, etc) are used to distinguish the hinge and the stick constraints, when in practice a hinge is pretty much two sticks where one of the particles is shared. Would it not be possible to somehow share some of these mechanisms between the hinge and the stick (e.g. just writing the hinge as two sticks in sequence?).\n\nW2. The work focuses a lot on the specific combination of adding exact constraint preservation to the E(N) equivariant model. I believe in most cases the E(n) equivariant model is often used for domains where hard mechanical constraints are not that common (e.g. molecule property prediction, or force prediction). I think it is great to extend the E(N) model this way, but I am not 100% sure how many practitioners interested on the E(n) model will actually end up using the constrain part.\n\nW3. On the other hand there are other baselines from the literature that have looked more explicitly into mechanical systems with hard constraints, which are not currently referenced in the paper. For example, https://arxiv.org/pdf/1810.01566.pdf models a fluid with a floating rigid box made of 64 particles by making a whole prediction for the canonical coordinates of the whole box, and then using the forward kinematic model of the box to obtain the cartesian positions, with an approach that is very similar to equations (6,7,8,9) in this work, so this should probably be referenced. Another common strategy for learning constraints without having to specify the forward kinematics function is to add noise to the training dataset, and ask the model to correct the noise (mujoco rigid systems: https://arxiv.org/pdf/1806.01242.pdf, particle systems: https://arxiv.org/abs/2002.09405, meshes https://arxiv.org/abs/2010.03409), so the model learns to correct constraint violations, and keep constraint violations low during a long rollout, so this should probably be used as a baseline (giving the length of the sticks as an additional static feature). I believe the approach proposed in this paper would probably still do better than training with noise, but it is also much more complex in terms of assumptions that it makes, so it would be an interesting trade off to show. Note that in these type of domains full rotational equivariance is usually not as important since gravity often breaks the symmetry for at least one of the axes.\n\nOther minor comments/observations:\n\nO1. \"first recording data and then inducting formulas\". If feel the this sentence is a bit extraneous to the paper contributions, but if you want to make that point, probably worth referencing this: https://arxiv.org/abs/2006.11287\n\nO2. \"However, all above approaches have ignored the symmetry in physics\". There is a large subset of models that partially implement symmetries by enforcing translation equivariance but not rotation equivariance, which should probably be mentioned and referenced in the related work (e.g. https://openreview.net/pdf?id=B1lDoJSYDH, https://arxiv.org/abs/2002.09405, https://arxiv.org/abs/2010.03409).\n\nO3. \"Euclidean transformation (translation or rotation) of the input states\". I think \"reflection\" should be part of that list too, since the paper uses E(n) assumptions for the proofs (otherwise this would involve the SE(n) group, and I think the appendix proof would not hold).\n\nO4. I think equation (2) is technically incorrect, since in the E(n) paper the function used to update the hidden node embedding (phi_h in e.q. 6 from E(n) paper), is applied after the sum aggregation, which is not possible in Eq(2). Similarly, the two outputs in eq (5) are confusing, since the actual functional form of phi for vector and non vector outputs will end up having different functional form (as per the last sentence of section 3.3, which again does not clarify whether rho_w is shared or not between the vector and scalars outputs).\n\nO5. Corollary 1. \"assume the entries of Z are drawn independently\", actually this is quite a strong assumption, there may be many cases in systems where forces, velocities and relative positions are actually parallel to each other (e.g. an object falling in a uniform field with initial zero velocity), so the rank of the matrix would be 1. So probably worth skipping Corollary 1 and just going to Corollary 2.\n\nO6. \"the generalized coordinates include the states of particle 0 denoted as ... R^3  the rotation Euler angles of stick 01 ... R^3 ... and 02 .. R3\". Actually, if the particle 0 has only 3 degrees of freedom this means they are point particles (e.g. no solid rotation), so would not this mean then sticks 01 and 02, would only have 2 generalized coordinates each (e.g. latitude and longitude), instead of 3?\n\nO7. \"The edge feature is provided by a concatenation of the product of charges\" Did the model had any trouble learning to perform this product if instead just giving the charge as a node feature of the inputs?\n\nO8. \"EGNN deliver much worse performance when the training size is small\". Do you have an intuition of why GMN is more data efficient than EGNN, considering both models are similarly equivariant?\n\nO9. Table 3. \"are comparable with the original environment\" would be possible to add an extra row with the performance GMN trained on those datasets (e.g. what the test performance would be if the model actually saw data from that distribution during training?).\n\nO10. Table 4. Is there an interpretation of why the shared model works worse for the phi_2 that receives the three inputs than the unshared one, but in the case of phi_2 receiving just one input is the other way around?\n\n\nSmall typos/grammar:\n\"Considering constraints, by no means, is not easy\" weird grammar.\n\"How the size of the training set influences.\" --> Data efficiency\n\"We have introduced an normalization\" (typo)\n\n", "summary_of_the_paper": "The authors propose Graph Mechanics Networks (GMN), a E(n) equivariant model that explicitly models systems with rigid constraints, with exact constraint preservation by means of using generalized coordinates. In my opinion the main contributions are as follows:\n\nC1. Establish math for writing models that operate on cartesian coordinates, make predictions in very generalized coordinates and then converted back to cartesian. This is an approach that has been used in previous work, but is usually hidden as an implementation detail.\n\nC2. Extend the E(n) (Satorras et al.) equivariant message function to take and output a variable number of vectors with a proof demonstrating that the approach is flexible enough to actually approximate any E(n) equivariant function. I think this is a novel low level modeling contribution.\n\nC3. Extent the E(n) (Satorras et al.) model to build an E(n) equivariant model for exact constraint preservation. Although a bit niche, this is the first model of that nature.", "summary_of_the_review": "The paper has a reasonable balance of novel contributions along different axes. I am not sure any one of them is sufficiently novel/interesting to meet the ICLR bar on its own, but put together and following corresponding revisions, I think the could make make a bigger ICLR-worthy contribution. My only concern is that it may be hard to reach the right audience, since I suspect many readers would be interested on each of the contributions individually, but not sure many will be interested on all 3 of them simultaneously. I would perhaps recommend including \"E(n)\" or \"equivariant\" in the title of the paper to reach a larger audience.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "1Jpa1oz1WXz", "reviewer": "Reviewer_Dcng", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nThe paper is decently clearly written (modulo some typos/phrasings noted below, but nothing that impedes the understanding of the reader), provides all necessary proofs and makes 2 main contributions\n\n- formulating an algorithm which incorporates Generalized Coordinates and inverse/forward kinematics\n- proposes an extension of the EGNN framework from working on single vectors to matrices, including some proofs about the equivariance and expressivity.\n\n\nIncorporating the constraints clearly helps generalization on  the experiments\n\nWeaknesses:\n\n- My main criticism is with regards to lack of novelty/generality: one has to manually derive the forward/inverse kinematics of the constrained system, which requires knowledge of the underlying physics already...at which point one might want to go one step further and derive some approximation as well? Alternatively, I would have expected the incorporation of the generalized coordinates to be done in a more black box fashion, i.e. encoding the knowledge that there are *some* generalized coordinates in the architecture and letting the models derive *which* we are talking about (this could possibly be attempted with invertible neural networks for the kinematics and some form of constraint-seeking algorithm that tries to minimize the total number of dimensions used in the generalized coordinate latent for example?)\n- only studied on a single setting, but this is a minor nitpick as it's similar to other previous works\n- reporting the constraint violation in the main body is only marginally meaningful as of course a system which enforces the constraints *by design* will achieve 0 constraint violation and that part in particular is not learned. \n\nMinor nitpick:\n- the paper mentions rotation and translation in the intro but then focuses mainly on rotation? doesn't change much for the final results but it confused me a bit \n\n\nTypos/phrasings  (not accounted for in my review, just as feedback for the authors):\n\n- page 2: \"have achieved desired performance on the N-body system...lacks of constraint\" is meant to  say \"have achieved desirable performance ....lack constraints\"?\n- page 3 \"is proportion to the its\" \"is proportional to the\" ?\n- page 4: \"despite the desired performance\" => \"desireable\" again?\n- page 5: \"can be generalized to the function with multiple input vectors\" => to functions with multiple input vectors? ", "summary_of_the_paper": "The paper focuses on the formulation of a graph neural network approach towards the learning of update rules of constrained systems, extending prior work (Interaction entwork, EGNN) by learning the updates of constrained components in the generalized coordinates expressing those constraints and incorporating forward and inverse kinematics in the algorithm.\nComparisons are made with Linear, Basic GNN, TensorForce Networks, SE3-Transformers, Radial Field and EGNNs by evaluating them on a constrained version of the N-body problem introduced by Kipf et al. 2018 and it is shown that the proposed method better adheres to constraints and generalizes better on the studied problem.", "summary_of_the_review": "The paper presents and incremental improvement over previous works and was done cleanly, so a marginal accept. If it had been exploring a fully learned generalized coordinate transformation that uses an inductive bias but does not need to hand design the coordinates and kinematics I'd give a higher score.\n\n://EDIT: updated my score", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The manuscript develops a new kind of graph neural network (a Graph Mechanics Network; GMN) that is particularly well suited to representing and making predictions about physical mechanics systems (and data with similar structure). It does so by developing a way to build geometric constraints implicitly and naturally into the forward kinematics of the network, while still allowing for effective learning from data. The manuscript proves some essential properties of the new architecture and runs experiments both with simulated particles, hinges, sticks (and their combination), as well as with motion capture data. 
Reviewers were generally impressed by the writing and clarity of the work, as well as the main results. In addition, in those cases where reviewers thought that the experiments were lacking, the authors delivered effective new experiments to address those concerns (e.g. looking  at mocap and molecular datasets). One reviewer initially scores the manuscript as a Reject/3 on the basis of concerns about novelty and the scope of the theoretical and experimental contributions of the paper. However, they adjust their score 3->5 based on the rebuttal presented by the authors (including new experiments). The reviewer also downgrades their certainty (from 3->2) on the basis of the engagement from reviewers offering higher scores.
Overall, the manuscript presents a promising contribution to the graph networks literature and I agree with the general consensus in favour of publication.
Decision: Accept (Poster)