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, 32, 1].
Several learning-based differentiable simulators [27, 26, 20, 23]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 [11, 26, 12], 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 [26] 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 [13], 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 [28] and GMN [14] 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 theefficacy of our SGNN for learning physical dynamics on a large-scale challenging dataset Physionwith 8 different scenarios, and a 3-object RigidFall dataset.
Experimental results show that our modelis 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 graph neuralnetworks as physical simulators of dynamical systems.
Interaction Network (IN) [2] is the pioneerin modeling the dynamics of evolving systems represented by objects and particles.
[22] introducesa hierarchical representation of objects to efficiently perform message passing on different objects.
NRI [16] leverages variational auto-encoders to infer the latent relationship between objects from theevolution of the system.
Graph Network Simulator (GNS) proposed by [26] has been showcased asimple yet powerful tool in simulating large systems in particle-based representation of rigids andfluids by dynamically constructing interaction graph and performing multiple steps of informationpropagation.
DPI [20] adds one-level of hierarchy to the rigids and predicts the rigid transformationvia generalized coordinates, with applications to manipulation and control.
There are also works thatincorporate strong physical priors like Hamiltonian mechanics [25, 33], geometrical constraints [24]and shapes [23].
Despite the promising empirical results, these works have not given sufficientconsiderations on symmetry, a vital factor for a simulator to generalize to unseen testing data.
Geometrically equivariant graph neural networks.
Equivariant graph neural networks [13] 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 [29, 10], liftingand sampling via Lie algebra [8, 15], or leveraging invariant scalarization [30] like taking the innerproduct [28, 14].
Our approach also belongs to the scalarization family together with EGNN [28]and GMN [14].
Specifically, EGNN models the interaction with the invariant distance as input,computed as an inner product of the relative positions, and GMN generalizes to a multi-channelversion by taking into consideration a stack of multiple vectors.
These equivariant GNNs operate onparticle graphs or point clouds, while we additionally consider a combination of both particle- andobject-level information for physical simulation.
More importantly, they are assumed a full Euclideansymmetry with strong equivariance constraint, while we elaborate how to relax the constraint in the existence of external fields like gravity by leveraging subequivariance.
Besides, [31] and [6] aim todevelop equivariant models given any E(3)/E(2) and their subgroups, while our focus lies in how theexternal force field relaxes the equivariance from an E(3)-equivariant model.
Moreoever, they aresteerable CNNs and not explored on GNNs especially challenging tasks like physical simulation.
Subequivariant Graph Neural Networks3.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 i } , E ( t ) (cid:17) .
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, 26, 22] 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 [28] and GMN [14] 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 [14]proposes a multichannel scalarization form, which yields,
[Equation]
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)) σ :
𝒁 " =
[ 𝒙 " , 𝒗 " ] 𝒁 ) =
[ 𝒙 ) , 𝒗 ) ]
EGNN
GMN 𝒁 ") =
𝒁 " ⊖ 𝑪 , " || 𝒁 ) ⊖ 𝑪 , )
|| 𝒁 " ⊖ 𝒁 ) 𝒎 ") =
[ 𝒁 ") , 𝒈 ] 𝜎 ([ 𝒁 ") , 𝒈 ] T [ 𝒁 ") , 𝒈 ] ) 𝒁 ")
=
𝒁 " ⊖𝒁 ) 𝒎 ")
= 𝒁 ") 𝜎 ( 𝒁 ")4 𝒁 ") )
SGNN 𝒎 ")
= 𝒙 " −𝒙 ) 𝜎 (|| 𝒙 " −𝒙 ) || 6 )
Object 𝑜 ( 𝑖)
Object 𝑜 ( 𝑗) GNS, DPI 𝒎 ")
= 𝜎 ( 𝒙 " −𝒙 ) , 𝒗 " , 𝒗 ) ) 𝑪 , () )
𝒁 " 𝒁 ) 𝒙 " , 𝒗 " 𝒙 ) , 𝒗 ) 𝑪 , (" ) 𝒈 𝒈 𝒙 ) , 𝒗 ) 𝒙 " , 𝒗 "
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 [30], 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 )
=
{ j : ( 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 a similar spirit in nature with existing works including EGNN [28] and GMN [14].
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 ]akin to our setup.
Nevertheless, GMN is not object-aware, lacking the necessary object information,while in our message-passing we explicitly involve the particle-object correlation ⃗
Z i ⊖ ⃗ C
o ( i ) .
Anotherline of work including GNS [26] and DPI [20] only enforces translation equivariance by directlyfeeding [ ⃗ x i − ⃗ x j , ⃗ v i , ⃗ v j ] into an MLP without taking the inner product.
These approaches fail topreserve the O( 3 )-equivariance and consequently have weaker generalization as we will illustrate inour experiments.
Fig.
1 summarizes and compares the message-passing schemes of GNS, EGNN,
GMN and our SOMP.
We provide detailed comparisons in the theoretical perspective in Appendix byshowing that EGNN and GMN are indeed special cases of SGNN.
We now provide 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]
, (9) 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 the appendix.
We thus leverage our designedsubequivariant function with such desirable properties for ϕ ⃗ g and ψ ⃗ g in our subequivariant objectaware message passing introduced above.
We immediately have the following theorem guaranteeingthe 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 motivated bythe consideration that the interactions between or within objects are usually different.
The formerserves as a bridge of exchanging momentum and energy via interaction forces, while the latterusually accounts for maintaining the rigid or plastic constraints.
Therefore, it would be beneficialto disentangle these two kinds of interactions by different types of 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 ) } .
Object-level message passing.
Given the renewed particle-level information produced by the
first stage, we are now ready to construct the object-level message passing graph.
The object-level
message passing is proceeded as
[Equation]
[Figure]
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 [26] and DPI [20], as well as two equivariant models EGNN [28] and GMN [14].
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.
Details are in Appendix.
Note that EGNNs, GMNs, and our SGNN always produce exactly the sameresult in the three scenarios due to their equivariance (or subequivariance).
For RigidFall, such bias isnot observed since the cubes are falling vertically driven by gravity.
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.
[Table]
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.
We use4 iterations in each message passing of our model.
We run RANSAC [9] during testing on all modelsto help enforce the rigid constraint as suggested by Physion.
More experimental details including acomparison 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 curvesin 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 these physical scenes
[Figure]
with gravity involved.
For instance, in the challenging scenarios Dominoes, Contain, and Link, theaccuracy of EGNN and GMN only stays around 60%, while SGNN achieves promising results withsubequivariant 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 summa- rize the results in Fig.
5, where themodels are trained in the scenarios indexed by the rows and evaluated onthose indexed by the columns.
It isobserved that our model exhibits significantly stronger generalization than
GNS, the best-performed baseline on
Physion.
For example, our model, although trained with the dynamics of
Contain or Link, still yields a high accuracy (>78%) when tested on Drop,showcasing that our model is more advantageous in learning the dynamicsand 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 [31, 6]to GNNs as an alternative to obtain equivariance on the O(2) subgroup (more details are results arein Appendix C.5).
Although better than EGNN, the adapted Steerable E(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 to
[Table]
capture 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 model
[Table]
encounters 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-objectmessage passing.
We argue that this design relieves the learning complexity of the message-passingwith different types of relations.
This is observed that removing edge separation consistently leads todetriment in the accuracy in various scenarios.
Qualitative comparisons.
We provide visualization samples on Physion in Fig.
More cases areshown by our Supplementary Video.
Our model consistently yields accurate predictions.
4.3 Evaluation on RigidFall
Experimental setup.
We use the code provided by [20].
We record the rollout MSE at t = 20 and 40 .
To compare the data efficiency of different models, we consider four cases with the size of trainingset ranging from 200 to 5000.
More details are in Appendix.
Results.
The results are presented in Table 3.
Our model consistently yields the best predictionsregardless of the time step t and the number of training samples.
Notably, the rollout error stillmaintains a very low value even if only 200 training samples are provided, which verifies the strongdata-efficiency and generalization of SGNN.
Consistent observations are also visualized in Fig.
[Figure]
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 [25, 33], 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 toits 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 withthe 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.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental 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 experi-
ments multiple times)?
[Yes] See Table 1, 3.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] In Appendix.
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, if applicable?
[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]