Learning Physical Dynamics with Subequivariant Graph Neural Networks
Anonymous Author(s)
AffiliationAddress email
Abstract
Graph Neural Networks (GNNs) have become a prevailing tool for learning physical dynamics.
However, they still encounter several challenges: 1) Physical lawsabide by symmetry, which is a vital inductive bias accounting for model generalization and should be incorporated into the model design.
Existing simulators eitherconsider insufficient symmetry, or enforce excessive equivariance in practice whensymmetry is partially broken by gravity.
2) Objects in the physical world possessdiverse shapes, sizes, and properties, which should be appropriately processed bythe model.
To tackle these difficulties, we propose a novel backbone, Subequivariant Graph Neural Network , which 1) relaxes equivariance to subequivariance byconsidering external fields like gravity, where the universal approximation abilityholds theoretically; 2) introduces a new subequivariant object-aware message passing for learning physical interactions between multiple objects of various shapes inthe particle-based representation; 3) operates in a hierarchical fashion, allowing formodeling long-range and complex interactions.
Our model achieves on averageover 3% enhancement in contact prediction accuracy across 8 scenarios on Physion and 2 × lower rollout MSE on RigidFall compared with state-of-the-art GNNsimulators, while exhibiting strong generalization and data efficiency.
Introduction
Learning and predicting the complicated physical dynamics and interactions between objects arevital to many tasks including physical reasoning [7, 17], scene understanding [3] and model-basedplanning and control [18, 21, 34, 1].
Several learning-based differentiable simulators [28, 27, 20, 24]have been proposed, and they obtain promising achievements in simulating various kinds of interactingobjects including rigid and fluids.
The success mainly relies on the prevalence of graph neuralnetworks [12, 27, 13], which are desirable tools for physical simulation, by modeling particles asnodes, physical relations as edges, and their interactions as the message passing thereon.
Notably, physical world exhibits certain symmetries.
For example, the way dominoes fall from left toright will exactly be preserved when the system is rotated horizontally to another direction.
Inherently,humans reason about the dynamics with the preservation of such symmetry, and this inductive biashas been endorsed by the physical laws that also abide by the symmetries in our 3D world.
Yet,regardless of this inductive bias, current differentiable simulators like GNS [27] and DPI [20] wouldfail to learn the real dynamics that can generalize to all directions.
In the example of dominoes, ifthe data in training are positioned from left to right, GNS can predict well during testing when thedominoes also fall in the same direction, but performs poorly if the scene is rotated horizontally (seedemo in the Supplementary Video).
This observation implies that GNS overfits training sampleswithout learning the true dynamics that obey the symmetry, limiting its generalization.
With this consideration, there have been a number of works, named geometrically equivariant graphneural networks [14], that leverage symmetry as an inductive bias in learning to simulate.
Thesemodels are designed such that their outputs will rotate/translate/reflect in the same way as the inputs,hence retaining the symmetry.
However, models like EGNN [29] and GMN [15] are exerted E(3)-equivariance, the full symmetry of the 3D Euclidean space.
Such constraint is so strong that itpenalizes all directions in the 3D space, which cannot be applied to scenarios with external fields, likegravity.
The existence of gravity breaks the symmetry in the vertical direction, reducing E(3) to itssubgroup.
We formally characterize this phenomenon of equivariance relaxation as subequivariance .
In general, simulating physics on datasets like Physion [4] are highly challenging, due to the scaleof the system (on average thousands of particles per system), the diversity of the interactions ( e.g. ,collision, friction, gravity), as well as multiple shapes, materials, or even rigidness of the objects.
These factors require unique efforts in designing the simulators, which brings another weaknessof current GNN-based methods: From a practical point of view, they seldom explicitly involve thegeometric object information into message passing.
Moreover, when modeling multiple interactingobjects of different shapes, the interactions between or within objects are usually different.
Forexample, the former aims to exchange momentum or energy across objects, while the latter usuallyaccounts for the geometrical constraint.
It is thus important to involve objectness to distinguishinteractions between particles within objects from those across objects.
In this work, we propose Subequivariant Graph Neural Networks (SGNN) that consist of severalfeatures to tackle the above challenges.
We relax equivariance to subequivariance and designsubequivariant functions to physical scenarios with the existence of gravity, and excitingly, we haveproved that our designed form inherits the approximation universality.
We formulate a novelsubequivariant object-aware message passing framework for modeling dynamics and interactionsbetween objects of varying shapes.
We incorporate the subequivariant message passing into ahierarchical model to deal with long-range and complicated object interactions
We demonstrate the efficacy of our
SGNN for learning physical dynamics on a large-scale challenging dataset Physion with 8 different scenarios , and a 3-object RigidFall dataset
Experimental results show that our model is capable of yielding more accurate dynamics prediction is highly data-efficient and has stronggeneralization compared with the state-of-the-art learning-based differentiable physical simulators.
Related Work
GNN-based physical dynamics simulators.
There have been many works that employ graphneural networks as physical simulators of dynamical systems [2, 23, 16].
Graph Network Simulator(GNS) proposed by [27] has been showcased a simple yet powerful tool in simulating large systemsin particle-based representation of rigid and fluids by dynamically constructing interaction graphand performing multiple steps of information propagation.
DPI [20] adds one-level of hierarchyto the rigid and predicts the rigid transformation via generalized coordinates, with applicationsto manipulation.
There are also works that incorporate strong physical priors like Hamiltonianmechanics [26, 35], geometrical constraints [25] and shapes [24].
Despite the promising empiricalresults, these works have not given sufficient considerations on symmetry especially when externalforce like gravity presents, leaving room for enhancing their generalization to unseen testing data.
Geometrically equivariant graph neural networks.
Equivariant graph neural networks [14] area family of GNNs that are specifically designed to meet the constraint of certain symmetry, mostlyinvolving translations, rotations, and/or reflections in Euclidean space.
This goal is approached byseveral measures, including solving group convolution with irreducible representation [31, 11] orleveraging invariant scalarization [32] like taking the inner product [29, 15].
Our approach alsobelongs to the scalarization family together with EGNN [29] and GMN [15].
Specifically, EGNNmodels the interaction with the invariant distance as input, computed as an inner product of therelative positions, and GMN generalizes to a multi-channel version by taking into consideration
a
Equivariance on subgroups.
E( n )-Steerable CNNs [33, 6] develop convolutional kernels that meetequivariance on E(3)/E(2) and their subgroups by leveraging restricted representation.
EMLP [9]obtains equivariance on arbitrary matrix groups.
However, these approaches require specifying theparticular group/subgroup to perform equivariance, while our goal here is to relax the equivariance onsubgroups of E(3) by considering external force field, having more physical implications.
Moreover,these works rely on computationally expensive operations like irreducible representation or solvinggroup constraints, while our formulation resorts to scalarization, which is easy to implement, efficientto compute, and also comes with necessary universality guarantee.
Subequivariant Graph Neural Networks
3.1 Background
GNN-based simulators.
We start by introducing the mechanism of GNN-based simulators forlearning physical dynamics.
We consider the particle-based representation of a physical system with
N particles consisting of M objects.
At time t , each particle within the system possesses somestate information, including 1.
the position ⃗ x ( t ) i ∈ R 3 and the velocity ⃗ v ( t ) i ∈ R 3 , which are bothdirectional vectors; 2.
some attributes h i ∈
R n without geometric context, such as the rigidness; 3.the dynamic spatial connections with other particles, where an edge will be constructed if the distancebetween two particles is smaller than a threshold r , namely, E ( t )
= { ( i, j ) :
∥ ⃗ x ( t ) i − ⃗ x ( t ) j ∥ 2 < r } .
Wedenote the object category of particle i as o
( i ) ∈
Z satisfying o ( i ) =
k if particle i belongs to object k .
The goal here is to predict the position of the next step ⃗ x
( t +1)
i given the above system information attime t , which can be favorably modeled by GNNs, i.e. , ⃗ x ( t +1)
i =
φ GNN (cid:16) { ⃗ x ( t ) i } , { ⃗ v ( t ) i } , { h
[Equation]
Since the prediction at different time shares the same model, we will henceforth omit the temporalsuperscript t for all variables for brevity.
The conventional GNN simulators φ GNN [20, 27, 23] offer a favorable solution by leveraging messagepassing on the interaction graph, which computes:
[Equation]
where N ( i ) =
{ j : ( i, j ) ∈ E} is the neighbors of node i , and ϕ , ψ are the edge message functionand node update function, respectively.
The prediction is obtained by conducting several iterations ofmessage passing.
Nevertheless, such form does not guarantee the desirable symmetry context withcommon choices of ϕ and ψ , and meanwhile the object information has not been elaborated, leavingroom for a more exquisite message passing scheme for learning complex physical dynamics.
Recentworks such as EGNN [29] and GMN [15] have actually considered E(3)-equivariance in the designof the functions ϕ and ψ to pursue symmetry.
However, they are not applicable to the case whensymmetry is partially violated by gravity, which motivates our proposal in § 3.2.
Prior to introducingour work, we first provide necessary preliminaries related to equivariance.
Equivariance.
In this paper, we basically focus on equivariance in terms of E(3) transformations:translation, rotation, and reflection.
The formal definition is provided below.
Definition 1 (Equivariance and Invariance) .
We call that the function f : R 3 × m × R n → R 3 × m ′is E(3)-equivariant, if for any transformation g in E(3), f ( g · ⃗ Z , h ) =
g · f ( ⃗ Z , h ) , ∀ ⃗ Z ∈ R 3 × m ,∀
h ∈ R n .
Similarly, f is invariant if f ( g · ⃗ Z , h ) =
f ( ⃗ Z , h ) , ∀ g ∈ E(3).
In Definition 1, the group action · is instantiated as g · ⃗ Z :=
O ⃗ Z for the orthogonal transformation(rotation and reflection) where O ∈ O (3) :
= { O ∈ R 3 × 3 | O ⊤ O = I } , and g · ⃗ Z := ⃗ Z + t fortranslation where t ∈ R 3 .
Note that for the input of f , we have added the the right-arrow superscripton ⃗ Z to distinguish it from the scalar h that is unaffected by the transformation, akin to our notationsfor the position ⃗ x i , velocity ⃗ v i , and attribute h i .
It is non-trivial to derive equivariant function particularly for orthogonal transformations.
GMN [15]proposes a multichannel scalarization form, which yields,
[Equation]
𝒁 " =
[ 𝒙 " , 𝒗 " ] 𝒁 ) =
[ 𝒙 ) , 𝒗 ) ]
EGNN
GMN 𝒁 ") =
𝒁 " ⊖ 𝑪 , " || 𝒁 ) ⊖ 𝑪 , )
|| 𝒁 " ⊖ 𝒁 ) 𝒎 ") =
[ 𝒁 ") , 𝒈 ] 𝜎 ([ 𝒁 ") , 𝒈 ] T [ 𝒁 ") , 𝒈 ] ) 𝒁 ")
=
𝒁 " ⊖𝒁 ) 𝒎 ")
= 𝒁 ") 𝜎 ( 𝒁 ")4 𝒁 ") )
SGNN 𝒎 ")
= 𝒙 " −𝒙 ) 𝜎 (|| 𝒙 " −𝒙 ) || 6 )
Object 𝑜 ( 𝑖)
Object 𝑜 ( 𝑗) GNS, DPI 𝒎 ")
= 𝜎 ( 𝒙 " −𝒙 ) , 𝒗 " , 𝒗 ) ) 𝑪 , () )
𝒁 " 𝒁 ) 𝒙 " , 𝒗 " 𝒙 ) , 𝒗 ) 𝑪 , (" ) 𝒈 𝒈 𝒙 ) , 𝒗 ) 𝒙 " , 𝒗 " where the inner product ⃗ Z ⊤ ⃗ Z ∈ R m × m is firstly computed and concatenated with h , the resultantinvariant term is then transformed by a function (usually a Multi-Layer Perceptron (MLP)) σ :
R m × m + n (cid:55)→ R m × m ′ producing V ∈ R m × m ′ , and the directional output is acquired by taking amatrix multiplication of the basis ⃗ Z with V .
By joining the analyses in GMN along with [32], wehave the universality of the formulation in Eq.
(3), which is provided in Appendix.
3.2 Subequivariant Object-aware Message Passing
To simultaneously benefit from symmetry and object-aware information, we develop Subequivariant
Object-aware Message Passing (SOMP).
It characterizes information aggregation between multichannel vectors and scalars, while taking into account both the particle and object information.
We use⃗
Z i ∈ R 3 × m to indicate a stack of m multi-channel 3D vectors; particularly, ⃗ Z i =
[ ⃗ x i , ⃗ v i ] , 1 ≤ i ≤ Nby involving both the position and velocity.
We additionally initialize the object features by poolingits particles: for object k , we have ( ⃗ C k , c k ) =
1 |{ i : o ( i )=
k }| (cid:80) { i : o ( i )=
k } ( ⃗ Z i , h i ) , 1 ≤ k ≤ M .
Wealso introduce a binary operation “ ⊖ ” which yields ⃗ Z i ⊖ ⃗ Z j =
[ ⃗ Z i − ⃗ Z j , ⃗ v i , ⃗ v j ] .
This operation notonly produces a translation-invariant representation but also enriches the input stack of vectors withmore channels.
The direction of gravity is set to be along the third axis, namely ⃗ g
=
In a high-level overview, we denote our message passing as the function φ that updates each particlegiven the input states of all particles, objects, and graph connectivity:
[Equation]
Specifically, φ is unfolded as the following message passing and aggregation computations:
[Equation]
where ∥ is the concatenation along the channel dimension, N ( i )
= {
[Equation] [Equation]
i, j ) ∈ E} is the neighborsof node i , and both ϕ ⃗ g and ψ ⃗ g are subequivariant and will be defined in detail in Eq.
In detail, wefirst derive the multi-channel vector input ⃗ Z ij for the interaction between particle i and j in Eq.
5by stacking the three terms along the channel dimension, including the relative information of theparticle with respect to its belonged object: ⃗ Z i ⊖ ⃗ C
o ( i ) and ⃗ Z i ⊖
⃗ C
o ( i ) , and the relative informationbetween particles
⃗ Z i ⊖ ⃗ Z j .
By this means, we arrive at a translation-invariant representation ⃗ Z ijwith rich object-aware geometric information.
For the invariant features we simply concatenate themin Eq.
Subsequently, in Eq.
7 we feed the interactions ⃗ Z ij and h ij into the subequivariant messagefunction ϕ ⃗ g , yielding the vector message ⃗ M ij and scalar message m ij 1 .
Finally, Eq.
8 first performsmessage aggregation and then updates the states by another subequivariant function ψ ⃗ g , obtaining theeventual result.
With the updated ⃗ Z ′ i and h
′ i , one can readily obtain the output with equivariance orinvariance as desired, which, in our case, implies ⃗ Z ′ i =
[ ⃗ x ′ i ] by setting m ′
= 1 .
Comparison with existing GNN-based simulators.
The core of satisfying equivariance lies intaking the invariant inner product before feeding the vectors into the MLP.
Such design shares asimilar spirit in nature with existing works including EGNN [29] and GMN [15].
The interactionmodeling of EGNN only considers the relative position ⃗ x i − ⃗ x j and its inner product ∥
⃗ x i − ⃗ x j ∥ 2 ,which might have limited expressivity while tackling complex physical interactions and dynamicsbetween objects.
GMN further extends to a multi-channel interaction ⃗ Z i ⊖ ⃗ Z j where ⃗ Z i =
[ ⃗ x i , ⃗ v i ] .
Nevertheless, it lacks the necessary object information, while we explicitly involve the particle-objectcorrelation ⃗ Z i ⊖ ⃗ C
o ( i ) in our SOMP.
Besides, GNS [27] and DPI [20] only enforces translationequivariance by directly feeding [ ⃗ x i − ⃗ x j , ⃗ v i , ⃗ v j ] into an MLP without taking the inner product.
They fail to preserve the O( 3 )-equivariance and consequently have weaker generalization as we willillustrate in our experiments.
Fig.
1 summarizes and compares the message-passing schemes of GNS,
EGNN, GMN and our SOMP.
Particularly, our design of SOMP takes careful considerations: 1.
Equivariance is still permitted with both geometric and scalar object features involved; 2.
The objectinformation can also be constantly updated during message passing; 3.
The expressivity of SOMP isenhanced over EGNN and GMN with object information considered.
We provide detailed theoreticalcomparisons in Appendix A.2 by showing that EGNN and GMN are indeed special cases of SGNN.
We now present the detailed formulations of ϕ ⃗ g in Eq.
7 and ψ ⃗ g in Eq. 8.
Since the full symmetry isviolated by gravity ⃗ g ∈ R 3 , and the dynamics of the system will naturally preserve a gravitationalacceleration in the vertical direction.
By this means, the orthogonal symmetry is no longer maintainedin every direction but only restricted to the subgroup O ⃗ g (3) :
= { O ∈ O (3) |
O ⃗ g = ⃗ g } , that is, therotations/reflections around the gravitational axis.
We term such a reduction of equivariance as anovel notion: subequivariance .
To reflect this special symmetry, we augment Eq.
3 by:
[Equation]
where σ : R ( m +1) × ( m +1)
→ R ( m +1)
× m
′ is an MLP.
Compared with Eq.
(3), here we just augmentthe directional input with ⃗ g .
Interestingly, such a simple augmentation is universally expressive:
Theorem 1.
Let f ⃗ g
( ⃗ Z , h ) be defined by Eq.
Then, f ⃗ g is O ⃗ g (3) -equivariant.
More importantly,
For any O ⃗ g (3) -equivariant function ˆ f ( ⃗ Z , h ) , there always exists an MLP σ satisfying ∥ ˆ f − f ⃗ g ∥ <
ϵfor arbitrarily small positive value ϵ .
The proof is non-straightforward and deferred to Appendix A.1.
The detailed architectural view ofthe proposed SOMP is depicted in Fig.
8 in Appendix A.3.
We leverage our designed subequivariantfunction with such desirable properties for ϕ ⃗ g and ψ ⃗ g in our subequivariant object-aware messagepassing.
We immediately have the following theorem guaranteeing the validity of our design.
Theorem 2.
The message passing φ (Eq. 4) is O ⃗ g (3) -equivariant.
3.3 Application to Physical Scenes with Hierarchical Modeling
This subsection introduce the entire architecture of our subequivariant GNN.
Many physical scenesare complicated, possibly involving contact, collision, and friction amongst multiple objects.
Inlight of this, we further incorporate our message passing into a multi-stage hierarchical modelingframework.
One of our interesting findings here is the edge separation.
This is indeed motivatedby the consideration that the interactions between or within objects are usually different, similar toautomorphism graph networks [8, 22, 30] where the neighboring objects formulate the isomorphismgroup of edges.
The former serves as a bridge of exchanging momentum and energy via interactionforces, while the latter usually accounts for maintaining the rigid or plastic constraints.
Therefore, itwould be beneficial to disentangle these two kinds of interactions in message passing, which alsodistinguishes our hierarchical modeling from DPI [20].
The overall flowchart is provided in Fig.
To be specific, we alternate between the following three stages that implement φ in Eq.
(4) distinctly.
Particle-level inter-object message passing.
We start
by
modeling the local interactions from the
particle-level, where we
only
involve
those between different
objects
That
is
[Equation]
where E inter = { ( i, j ) : ( i, j ) ∈ E , o ( i )
̸
= o ( j ) } .
[Figure]
[Equation]
[Figure]
Object-level message passing.
Given the renewed particle-level information produced by thefirst stage, we are now ready to construct the object-level message passing graph.
The object-levelmessage passing is proceeded as
[Equation]
) where E obj comprises the edges between interacted objects, indicating E obj = { ( k, l ) : ∃ ( i, j ) ∈
E inter , o ( i ) =
k, o ( j ) =
l } .
As a specific instantiation of φ in Eq.
(4), here each object is consideredas a node and the original object information in Eq.
(5-8) is omitted (this is why we denote the secondinput of φ 2 as the empty set ∅ ).
Furthermore, we leverage a pooling of the updated particle-level information from the first stage as the object-level interactions, or formally, we re-define the calculationof “ ⊖ ” for two different objects indexed by k, l as ⃗ C k ⊖ ⃗ C l = Mean-Pool ( i,j ) ∈E inter (cid:16)
⃗ Z ′ i ⊖ ⃗ Z ′ j (cid:17) , andsimilarly,
c k ∥
c l = Mean-Pool ( i,j ) ∈E inter (cid:0)
h ′ i
∥
h ′ j (cid:1) .3.
Particle-level inner-object message passing.
Finally, given the updated object state information,we carry out the message passing for the particles only along the inner-object edges E inner =
{ ( i, j ) :(
i, j ) ∈ E , o ( i ) =
o ( j ) } , namely,
[Equation]
We use { ⃗ Z ′′ i } Ni =1 as the proposal of the positions in the next time step, akin to the way in § 3.2.
Experiments
We conduct evaluations on Physion [4] and RigidFall [19].
Physion [4] is a large scale datasetconsisting of eight different scenarios, including Dominoes, Contain, Collide, Drop, Roll, Link,
Support, and Drape.
These scenarios possess diverse physical scenes with complicated objectinteractions, serving as a challenging benchmark for evaluating dynamics models.
Particularly, Drapealso involves simulating deformable objects.
RigidFall [19] is a simulation dataset whose scenesinvolve several cubes falling and colliding under varying magnitude of gravitational acceleration.
4.1 Baselines
We compare SGNN with several SOTA particle-based GNN simulators, including two non-equivariantmodels GNS [27] and DPI [20], as well as two equivariant models EGNN [29] and GMN [15].
Notably, the data preprocessing provided by Physion [4] for particle-based methods does not considerthe camera-angle for the rotation of 3D coordinates, leading to the situation that in some scenarios, e.g. ,
Dominoes, Link, Support, and Collide, there exist a directional bias (from left to right horizontally) inthe trajectories.
This motivates us to consider three evaluation protocols: 1.
Use the original trainingand testing set with directional bias.
Train on the original set and test on the randomly-rotatedtesting set.
Train on the randomly-rotated training set (akin to applying data augmentation ofrotations [5]), and test on the randomly-rotated testing set.
Here all rotations are restricted to the onesaround gravity.
We dub the results obtained from the three cases GNS*, GNS, and GNS-Rot for GNS,and similarly for DPI.
For EGNN and GMN, we also propose to adapt them to be subequivariant(EGNN-S and GMN-S) by adding an extra vector along the gravity in their update of velocities by gravity
[Table]
4.2 Evaluation on Physion
Experimental setup.
We strictly follow the training and evaluation protocol proposed in [4] forthe particle-based simulators.
We employ two evaluation metrics: 1.
Contact prediction accuracy.
This metric, as also adopted by [4], evaluates whether two targeted objects labeled in the dataset willcontact during the model rollout.
Rollout MSE, the mean squared error between model rolloutand the ground truth.
We reuse the training protocol and hyper-parameters suggested in Physionfor baselines and our model to ensure a fair comparison.
In detail, all MLPs are initialized with 3projection layers and a hidden dimension of 200.
The networks are trained with an Adam optimizer,using an initial learning rate 0.0001 and an early-stopping of 10 epochs on the validation loss.
Weuse 4 iterations in each message passing of our model.
We run RANSAC [10] during testing onall models to help enforce the rigid constraint as suggested by Physion.
More experimental detailsincluding a comparison of computational complexity are deferred to Appendix.
Results.
We present the average contact prediction accuracy with std in Table 1, and the rollout MSEin Fig. 3.
The detailed MSE on all 8 scenarios is deferred to Appendix due to space limit.
We interpretour results by answering the following questions.
Q1.
Is maintaining the symmetry important tomodels that learn to simulate?
For the non-equivariant models GNS and DPI, their performanceusually encounters a significant drop if they are trained on the original data but tested on the rotated,particularly on the scenarios with bias in directions.
For example, the accuracy of GNS drops from78.6% to 53.3% on Dominoes when comparing GNS* with GNS.
Such observation also holds for DPIand on more scenarios like Collide and Support.
Using data augmentation would help in relieving theissue, but is still hard to recover the performance.
The accuracy on Dominoes yielded by GNS-Rot is74.7%, still worse than 78.6% by GNS*.
Our SGNN instead leverages subequivariance to generalizeto all horizontal directions, and is always guaranteed to produce the same prediction regardless ofany rotations around gravity, and there is also no need to apply any data augmentation of rotationsduring training.
Interestingly, EGNN-S and GMN-S offers improvements over EGNN and GMN,implying the necessity of subequivariance, but are still outperformed by SGNN by a large margin.
Q2.
Does SGNN learn the physics and simulate the dynamics better than other models?
As displayed in
Table 1, SGNN achieves the highest prediction accuracy on 7 out of the 8 scenarios, and is also verycompetitive on Collide.
The improvements are very significant (>5% accuracy enhancement over thesecond best) on scenarios including Dominoes, Contain, and Link, where there are multiple objectscontacting, colliding, and interacting.
This observation is also supported by the rollout MSE curves
[Figure]
in Fig.
3, showing our SGNN consistently yielding lower prediction error along the trajectories.
Q3.
How does SGNN perform compared with other equivariant models?
For EGNN and GMN, enforcingsuch a strong constraint of E( 3 )-equivariance leads to inferior performance in scenes with gravity.
For instance, in Dominoes, Contain, and Link, the accuracy of EGNN and GMN only stays around60%, while SGNN achieves promising results with subequivariant object-aware message passing.
Generalization across scenarios.
We further evaluate the generalization capability of our modelacross various scenarios.
To fulfill this goal,
we leverage models trained in certain scenarios
[Figure]
to test them in others.
We summarize the results in Fig.
5, where the models are trainedin the scenarios indexed by the rows andevaluated on those indexed by the columns.
It is observed that our model exhibits significantly stronger generalization than GNS,the best-performed baseline on Physion.
For example, our model, although trainedwith the dynamics of Contain or Link, stillyields a high accuracy (>78%) when testedon Drop, showcasing that our model ismore advantageous in learning the dynamics and interactions of various objects.
[Table]
Ablation studies.
Here we conduct several ablations to inspect how our proposed componentscontribute to the overall performance, and theresults are given in Table 2. 1.
Subequivariance.
We replace our subequivariant formulation of ϕ ⃗ gand ψ ⃗ g by the E( 3 )-equivariant counterpart without gravity.
It is clear that, in this way, the modelis restricted by strong equivariance constraints,which leads to a significant degradation in performance.
This verifies the necessity of relaxingthe full-equivariant model with our introducedmethod for physical scenarios with gravity.
Moreover, we extend the Steerable E(2)-CNNs [33, 6] to
GNNs as an alternative to obtain equivariance on the SE(2) subgroup (more details are results are in
Appendix C.5).
Although better than EGNN, the adapted Steerable SE(2)-GNN is still inferior to
SGNN, implying our physics-inspired SOMP is more advantageous on simulating physical dynamics.2.
Object-aware message passing.
We replace ⃗ C k and c k in Eq.
10 by zeros, eliminating the objectinformation from message passing.
Without this necessary information, the model might fail tocapture useful geometric vectors like ⃗ x i − ⃗ x o ( i ) , which is closely related to physical quantities liketorque.
Hierarchy.
We compare our model with its counterpart without hierarchy.
The flat modelencounters an average of 4% drop in the prediction accuracy, showcasing the advantages of leveragingthe particle- and object-level message passing for modeling complex object interactions in physicalscenes.
Edge separation.
We employ different sets of edges in the inter-object and inner-object
[Table]
message passing.
We argue that it relieves the learning complexity of the message-passing with
[Table]
different types of relations.
It is observed that removing edge separation leads to detriment in theaccuracy in various scenarios
[Table]
Generalization toward other rotations.
To further investigate whether incorporating our subequivariance truly helps the model to learn the effect of gravity
[Table]
we apply a rotation around a non-gravity
[Figure]
axis, resulting in scenarios in Fig.
6 where dominoes are placed on an incline while gravitystill points downwards vertically.
Interestingly,
SGNN well generalizes to this novel scenarioand reasonably simulates the effect of gravity.
Particularly, the domino at the bottom starts toslide down along the table driven by gravity.
More visualizations are provided in the video.
Qualitative comparisons.
We provide visualization samples on Physion in Fig.
More cases are shown by our Supplementary Video.
Our modelconsistently yields accurate predictions across various scenarios.
4.3 Evaluation on RigidFall
GroundTruth
Ours
DPI
Experimental setup.
We use the code providedby [20].
We record the rollout MSE at t =20 and 40 .
To compare the data efficiency ofdifferent models, we consider four cases withthe size of training set ranging from 200 to 5000.
More details are in Appendix.
Results.
The results are presented in Table 3and visualized in Fig.
Our model consistentlyyields the best predictions regardless of the timestep t and the size of training set.
Notably, therollout error still maintains very low even if only200 training samples are provided, verifying the strong data-efficiency and generalization of SGNN.
Discussion
Limitations and Future Work.
Our model relies on the particle-based representation of the physicalsystems, which, in real-world scenarios, might be difficult to obtain, or usually with noise.
Our modeldoes not explicitly consider self-contact.
Future directions include augmenting our dynamics modelwith strong physical priors like Hamiltonian [26, 35], or combining it with groundings of visualprior [19] with application to more sophisticated tasks like physical reasoning or robot manipulation.
Conclusion.
We propose Subequivariant Graph Neural Networks for modeling physical dynamicsof multiple interacting objects.
We inject appropriate symmetry into a hierarchical message passingframework, and takes into account both particle- and object-level state messages.
In particular,subequivariance is a novel concept for characterizing the relaxation of equivariance from a group to its subgroup.
We show that relaxation from full equivariance to subequivariance can be applied tosystems with gravity involved.
Experiments verify that SGNN accurately captures the dynamics with the guarantee of the desirable symmetry, and exhibits strong generalization with high data-efficiency
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] Our contributions are stated in the introduction.
(b) Did you describe the limitations of your work?
[Yes] In Sec. 5.
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] In Appendix.
(b) Did you include complete proofs of all theoretical results?
[Yes] In Appendix.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main ex-
perimental results (either in the supplemental material or as a URL)?
[Yes] See our
supplementary material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See Appendix.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] See Table 1, 3.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] In Appendix.4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] In Sec. 4.
(b) Did you mention the license of the assets?
[Yes] In Appendix.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, ifapplicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]