Learning Physical Dynamics with Subequivariant Graph Neural Networks
Abstract
Graph Neural Networks (GNNs) have become a prevailing tool for learning physical dynamics.
However, they still encounter several challenges: 1) Physical laws abide by symmetry, which is a vital inductive bias accounting for model generalization and should be incorporated into the model design.
Existing simulators either consider insufficient symmetry, or enforce excessive equivariance in practice when symmetry is partially broken by gravity.
2) Objects in the physical world possess diverse shapes, sizes, and properties, which should be appropriately processed by the model.
To tackle these difficulties, we propose a novel backbone, Subequivariant Graph Neural Network , which 1) relaxes equivariance to subequivariance by considering external fields like gravity, where the universal approximation ability holds theoretically; 2) introduces a new subequivariant object-aware message passing for learning physical interactions between multiple objects of various shapes in the particle-based representation; 3) operates in a hierarchical fashion, allowing for modeling long-range and complex interactions.
Our model achieves on average over 3% enhancement in contact prediction accuracy across 8 scenarios on Physion and 2 × lower rollout MSE on RigidFall compared with state-of-the-art GNN simulators, while exhibiting strong generalization and data efficiency.
Code and videos are available at our project page: https://hanjq17.github.io/SGNN/.
Introduction
Learning and predicting the complicated physical dynamics and interactions between objects are vital to many tasks including physical reasoning [8, 38, 7, 19], scene understanding [3] and model-based planning and control [20, 23, 24, 37, 1].
Several learning-based differentiable simulators [31, 30, 22, 27] have been proposed, and they obtain promising achievements in simulating various kinds of interacting objects including rigid and fluids.
The success mainly relies on the prevalence of graph neural networks [14, 30, 15], which are desirable tools for physical simulation, by modeling particles as nodes, 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 to right 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 bias has 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 [30] and DPI [22] would fail to learn the real dynamics that can generalize to all directions.
In the example of dominoes, if the data in training are positioned from left to right, GNS can predict well during testing when the dominoes also fall in the same direction, but performs poorly if the scene is rotated horizontally (see demo in the Supplementary Video).
This observation implies that GNS overfits training samples without learning the true dynamics that obey the symmetry, limiting its generalization.
With this consideration, there have been a number of works, named geometrically equivariant graph neural networks [16], that leverage symmetry as an inductive bias in learning to simulate.
These models 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 [32] and GMN [17] are exerted E(3)- equivariance, the full symmetry of the 3D Euclidean space.
Such constraint is so strong that it penalizes all directions in the 3D space, which cannot be applied to scenarios with external fields, like gravity.
The existence of gravity breaks the symmetry in the vertical direction, reducing E(3) to its subgroup.
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 scale of 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 weakness of current GNN-based methods: From a practical point of view, they seldom explicitly involve the geometric object information into message passing.
Moreover, when modeling multiple interacting objects of different shapes, the interactions between or within objects are usually different.
For example, the former aims to exchange momentum or energy across objects, while the latter usually accounts for the geometrical constraint.
It is thus important to involve objectness to distinguish interactions between particles within objects from those across objects.
In this work, we propose Subequivariant Graph Neural Networks (SGNN) that consist of several features to tackle the above challenges.
We relax equivariance to subequivariance and design subequivariant functions to physical scenarios with the existence of gravity, and excitingly, we have proved that our designed form inherits the approximation universality.
We formulate a novel subequivariant object-aware message passing framework for modeling dynamics and interactions between objects of varying shapes.
We incorporate the subequivariant message passing into a hierarchical 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 strong generalization 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 neural networks as physical simulators of dynamical systems [2, 26, 18].
Graph Network Simulator (GNS) proposed by [30] has been showcased a simple yet powerful tool in simulating large systems in particle-based representation of rigid and fluids by dynamically constructing interaction graph and performing multiple steps of information propagation.
DPI [22] adds one-level of hierarchy to the rigid and predicts the rigid transformation via generalized coordinates, with applications to manipulation.
There are also works that incorporate strong physical priors like Hamiltonian mechanics [29, 39], geometrical constraints [28] and shapes [27].
Despite the promising empirical results, these works have not given sufficient considerations on symmetry especially when external force like gravity presents, leaving room for enhancing their generalization to unseen testing data.
Geometrically equivariant graph neural networks.
Equivariant graph neural networks [16] are a family of GNNs that are specifically designed to meet the constraint of certain symmetry, mostly involving translations, rotations, and/or reflections in Euclidean space.
This goal is approached by several measures, including solving group convolution with irreducible representation [34, 12] or leveraging invariant scalarization [35] like taking the inner product [32, 17].
Our approach also belongs to the scalarization family together with EGNN [32] and GMN [17].
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-channel version by taking into consideration a stack of multiple vectors.
These equivariant GNNs operate on particle graphs or point clouds, while we additionally consider a combination of both particle- and object-level information for physical simulation.
More importantly, they are assumed a full Euclidean symmetry with strong equivariance constraint, while we elaborate how to relax the constraint in the existence of external fields like gravity by leveraging subequivariance.
Equivariance on subgroups.
E( n )-Steerable CNNs [36, 6] develop convolutional kernels that meet equivariance on E(3)/E(2) and their subgroups by leveraging restricted representation.
EMLP [10] obtains equivariance on arbitrary matrix groups.
However, these approaches require specifying the particular group/subgroup to perform equivariance, while our goal here is to relax the equivariance on subgroups of E(3) by considering external force field, having more physical implications.
Moreover, these works rely on computationally expensive operations like irreducible representation or solving group constraints, while our formulation resorts to scalarization, which is easy to implement, efficient to 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 for learning 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 some state information, including 1.
the position ⃗ x ( t ) i ∈ R 3 and the velocity ⃗ v ( t ) i ∈ R 3 , which are both directional 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 distance between two particles is smaller than a threshold r , namely, E ( t ) =
{ ( i, j ) :
∥ ⃗ x ( t ) i − ⃗ x ( t ) j ∥ 2 < r } .
We denote 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 at time t , which can be favorably modeled by GNNs, i.e. , ⃗ x ( t +1)
i
[Equation]
Since the prediction at different time shares the same model, we will henceforth omit the temporal superscript t for all variables for brevity.
The conventional GNN simulators φ GNN [22, 30, 26] 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 function and node update function, respectively.
The prediction is obtained by conducting several iterations of message passing.
Nevertheless, such form does not guarantee the desirable symmetry context with common choices of ϕ and ψ , and meanwhile the object information has not been elaborated, leaving room for a more exquisite message passing scheme for learning complex physical dynamics.
Recent works such as EGNN [32] and GMN [17] have actually considered E(3)-equivariance in the design of the functions ϕ and ψ to pursue symmetry.
However, they are not applicable to the case when symmetry is partially violated by gravity, which motivates our proposal in § 3.2.
Prior to introducing our 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).
𝒁 " =
[ 𝒙 " , 𝒗 " ] 𝒁 ) =
[ 𝒙 ) , 𝒗 ) ]
EGNN
GMN 𝒁 ") =
𝒁 " ⊖ 𝑪 , " || 𝒁 ) ⊖ 𝑪 , )
|| 𝒁 " ⊖ 𝒁 ) 𝒎 ") =
[ 𝒁 ") , 𝒈 ] 𝜎 ([ 𝒁 ") , 𝒈 ] T [ 𝒁 ") , 𝒈 ] ) 𝒁 ")
=
𝒁 " ⊖𝒁 ) 𝒎 ")
= 𝒁 ") 𝜎 ( 𝒁 ")4 𝒁 ") )
SGNN 𝒎 ")
= 𝒙 " −𝒙 ) 𝜎 (|| 𝒙 " −𝒙 ) || 6 )
Object 𝑜 ( 𝑖)
Object 𝑜 ( 𝑗) GNS, DPI 𝒎 ")
= 𝜎 ( 𝒙 " −𝒙 ) , 𝒗 " , 𝒗 ) ) 𝑪 , () )
𝒁 " 𝒁 ) 𝒙 " , 𝒗 " 𝒙 ) , 𝒗 ) 𝑪 , (" ) 𝒈 𝒈 𝒙 ) , 𝒗 ) 𝒙 " , 𝒗 "
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 for translation where t ∈ R 3 .
Note that for the input of f , we have added the the right-arrow superscript on ⃗ Z to distinguish it from the scalar h that is unaffected by the transformation, akin to our notations for the position ⃗ x i , velocity ⃗ v i , and attribute h i .
It is non-trivial to derive equivariant function particularly for orthogonal transformations.
GMN [17] 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 resultant invariant 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 a matrix multiplication of the basis ⃗ Z with V .
By joining the analyses in GMN along with [35], we have the universality of the formulation in Eq.
(3), which is provided in Appendix A.
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 is initialized as [ ⃗ x i , ⃗ v i ] , 1 ≤ i ≤ N by involving both the position and velocity.
We additionally initialize the object features by pooling its particles: for object k , we have ⃗ C k =
[Equation]
{ i : o [Equation] i )=
k } h i , 1 ≤ k ≤ M .
We also introduce a binary operation “ ⊖ ” which yields ⃗ Z i ⊖ ⃗ Z j =
[ ⃗ Z i − ⃗ Z j , ⃗ v i , ⃗ v j ] .
It transforms the translation-equivariant vectors to be invariant and stacks those invariant vectors together, resulting in a translation-invariant representation while enriching the input vectors with more channels.
The direction of gravity is set to be along the vertical axis.
In a high-level overview, we denote our message passing as the function φ that updates each particle given 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 neighbors of node i , and both ϕ ⃗ g and ψ ⃗ g are subequivariant and will be defined in detail in Eq.
In detail, we first derive the multi-channel vector input ⃗ Z ij for the interaction between particle i and j in Eq.
(5) by stacking the three terms along the channel dimension, including the relative information of the particle with respect to its belonged object: ⃗ Z i ⊖ ⃗ C o
( i ) and ⃗ Z i ⊖
⃗ C
o ( i ) , and the relative information between particles
⃗ Z i ⊖ ⃗ Z j .
By this means, we arrive at a translation-invariant representation ⃗ Z ij with rich object-aware geometric information.
For the invariant features we simply concatenate them in Eq.
Subsequently, in Eq.
(7) we feed the interactions ⃗ Z ij and h ij into the subequivariant message function ϕ ⃗ g , yielding the vector message ⃗ M ij and scalar message m ij 2 .
Finally, Eq.
(8) first performs message aggregation and then updates the states by another subequivariant function ψ ⃗ g , obtaining the eventual result.
With the updated ⃗ Z ′ i and h
′ i , one can readily obtain the output with equivariance or invariance 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 in taking the invariant inner product before feeding the vectors into the MLP.
Such design shares a similar spirit in nature with existing works including EGNN [32] and GMN [17].
The interaction modeling 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 dynamics between 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-object correlation ⃗
Z i ⊖ ⃗ C
o ( i ) in our SOMP.
Besides, GNS [30] and DPI [22] only enforces translation equivariance 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 will illustrate 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 object information can also be constantly updated during message passing; 3.
The expressivity of SOMP is enhanced over EGNN and GMN with object information considered.
We provide detailed theoretical comparisons 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.
Since the full symmetry is violated by gravity ⃗ g ∈ R 3 , and the dynamics of the system will naturally preserve a gravitational acceleration in the vertical direction.
By this means, the orthogonal symmetry is no longer maintained in every direction but only restricted to the subgroup O ⃗ g (3) :
= { O ∈ O (3) |
O ⃗ g = ⃗ g } , that is, the rotations/reflections around the gravitational axis.
We term such a reduction of equivariance as a novel 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 augment the 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 of the proposed SOMP is depicted in Fig.
8 in Appendix A.3.
We leverage our designed subequivariant function with such desirable properties for ϕ ⃗ g and ψ ⃗ g in our subequivariant object-aware message passing.
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 scenes are complicated, possibly involving contact, collision, and friction amongst multiple objects.
In light of this, we further incorporate our message passing into a multi-stage hierarchical modeling framework.
One of our interesting findings here is the edge separation.
This is indeed motivated by the consideration that the interactions between or within objects are usually different, similar to
[Figure]
automorphism graph networks [9, 25, 33] where the neighboring objects formulate the isomorphism group of edges.
The former serves as a bridge of exchanging momentum and energy via interaction forces, while the latter usually accounts for maintaining the rigid or plastic constraints.
Therefore, it would be beneficial to disentangle these two kinds of interactions in message passing, which also distinguishes our hierarchical modeling from DPI [22].
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]
11) 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 considered as a node and the original object information in Eq.
(5-8) is omitted (this is why we denote the second input 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 calculation of “ ⊖ ” 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) , and similarly, c k ∥
c l = Mean-Pool ( i,j ) ∈E inter (cid:0)
h ′ i
∥
h ′ j (cid:1) .
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 [21].
Physion [4] is a large scale dataset created by the ThreeDWorld simulator [13], which consists of eight different scenarios, including Dominoes,
Contain, Collide, Drop, Roll, Link, Support, and Drape.
These scenarios possess diverse physical scenes with complicated object interactions, serving as a challenging benchmark for evaluating dynamics models.
Particularly, Drape also involves simulating deformable objects.
RigidFall [21] is a simulation dataset whose scenes involve several cubes falling and colliding under varying magnitude of gravitational acceleration.
[Table]
4.1 Baselines
We compare SGNN with several SOTA particle-based GNN simulators, including two non-equivariant models GNS [30] and DPI [22], as well as two equivariant models EGNN [32] and GMN [17].
Notably, the data preprocessing provided by Physion [4] for particle-based methods does not consider the 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) in the trajectories.
This motivates us to consider three evaluation protocols: 1.
Use the original training and testing set with directional bias.
Train on the original set and test on the randomly-rotated testing set.
Train on the randomly-rotated training set (akin to applying data augmentation of rotations [5]), and test on the randomly-rotated testing set.
Here all rotations are restricted to the ones around 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 B. Note that EGNNs, GMNs, and our SGNN always produce exactly the same result in the three scenarios due to their equivariance (or subequivariance).
For RigidFall, such bias is not 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] for the 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 will contact during the model rollout.
Rollout MSE, the mean squared error between model rollout and the ground truth.
We reuse the training protocol and hyper-parameters suggested in Physion for baselines and our model to ensure a fair comparison.
In detail, all MLPs are initialized withprojection 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 use 4 iterations in each message passing of our model.
We run RANSAC [11] during testing on all models to help enforce the rigid constraint as suggested by Physion.
More experimental details including a comparison of computational complexity are deferred to Appendix B.
[Figure]
Results.
We present the average contact prediction accuracy with std in Table 1, and the rollout MSE in Fig.
The detailed MSE on all 8 scenarios is deferred to Appendix C.2 due to space limit.
We interpret our results by answering the following questions.
Q1.
Is maintaining the symmetry important to models that learn to simulate?
For the non-equivariant models GNS and DPI, their performance usually 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 from 78.6% to 53.3% on Dominoes when comparing GNS* with GNS.
Such observation also holds for DPI and on more scenarios like Collide and Support.
Using data augmentation would help in relieving the issue, but is still hard to recover the performance.
The accuracy on Dominoes yielded by GNS-Rot is 74.7%, still worse than 78.6% by GNS*.
Our SGNN instead leverages subequivariance to generalize to all horizontal directions, and is always guaranteed to produce the same prediction regardless of any rotations around gravity, and there is also no need to apply any data augmentation of rotations during 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 very competitive on Collide.
The improvements are very significant (>5% accuracy enhancement over the second best) on scenarios including Dominoes, Contain, and Link, where there are multiple objects contacting, colliding, and interacting.
This observation is also supported by the rollout MSE curves 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, enforcing such 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 around 60%, while SGNN achieves promising results with subequivariant object-aware message passing.
Generalization across scenarios.
We further evaluate the generalization capability of our model across various scenarios.
To fulfill this goal,
we leverage models trained in certain scenarios
[Figure]
the training/testing to test them in others.
We summarize the results in Fig.
5, where the models are trained in the scenarios indexed by the rows and evaluated 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 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 dynamics and interactions of various objects.
[Table]
Ablation studies.
Here we conduct several ablations to inspect how our proposed components contribute to the overall performance, and the results are given in Table 2.
Subequivariance.
We replace our subequivariant formulation of ϕ ⃗ g and ψ ⃗ g by the E( 3 )-equivariant counterpart without gravity.
It is clear that, in this way, the model is restricted by strong equivariance constraints, which leads to a significant degradation in performance.
This verifies the necessity of relaxing the full-equivariant model with our introduced method for physical scenarios with gravity.
Moreover, we extend the Steerable E(2)-CNNs [36, 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 object information from message passing.
Without this necessary information, the model might fail to capture useful geometric vectors like ⃗ x i − ⃗ x o ( i ) , which is closely related to physical quantities like torque.
Hierarchy.
We compare our model with its counterpart without hierarchy.
The flat model encounters an average of 4% drop in the prediction accuracy, showcasing the advantages of leveraging the particle- and object-level message passing for modeling complex object interactions in physical scenes.
Edge separation.
We employ different sets of edges in the inter-object and inner-object message passing.
We argue that it relieves the learning complexity of the message-passing with different types of relations.
It is observed that removing edge separation leads to detriment in the accuracy in various scenarios.
Generalization toward other rotations.
To further investigate whether incorporating our subequivariance truly helps the model to learn the effect of gravity, we apply a rotation around a non-gravity
[Figure]
axis, resulting in scenarios in Fig.
6 where dominoes are placed on an incline while gravity still points downwards vertically.
Interestingly,
SGNN well generalizes to this novel scenario and reasonably simulates the effect of gravity.
Particularly, the domino at the bottom starts to slide 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 model consistently yields accurate predictions across various scenarios.
4.3 Evaluation on RigidFall
[Table]
Experimental setup.
We use the code provided by [22].
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 training set ranging from 200 to 5000.
More details are in Appendix B.
Results.
The results are presented in Table 3 and visualized in Fig.
Our model consistently yields the best predictions regardless of the time step t and the size of training set.
Notably, the rollout
[Figure]
Discussion
Limitations and Future Work.
Our model relies on the particle-based representation of the physical systems, which, in real-world scenarios, might be difficult to obtain, or usually with noise.
Our model does not explicitly consider self-contact.
Future directions include augmenting our dynamics model with strong physical priors like Hamiltonian [29, 39], or combining it with groundings of visual prior [21] with application to more sophisticated tasks like physical reasoning or robot manipulation.
Conclusion.
We propose Subequivariant Graph Neural Networks for modeling physical dynamics of multiple interacting objects.
We inject appropriate symmetry into a hierarchical message passing framework, 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 to systems 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.
Acknowledgments and Disclosure of Funding
We thank the anonymous reviewers for their constructive suggestions.
This project was in part supported by MIT-IBM Watson AI Lab, Amazon Research Award Mitsubishi Electric.
Dr. Wenbing Huang was supported by the following projects: the National Natural Science Foundation of China (No.62006137); Guoqiang Research Institute General Project, Tsinghua University (No.
2021GQG1012);
Beijing Outstanding Young Scientist Program (No. BJJWZYJH012019100020098). 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] 2.
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 experiments multiple times)?
[Yes] See Table 1.
(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. 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] 5.
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]