Constrained Graph Mechanics Networks

1 INTRODUCTION Representing and reasoning about the relations and dynamics of a group of interacting objects is among the core aspects of human intelligence ( Tenenbaum et al. , 2011 ) . As a motivating example , we consider the Nbody system ( Kipf et al. , 2018 ) where the movement of a single charged particle is attracted or repelled by other charged particles . Physicists have revealed that this process can be modeled by Newton ’ s laws along with the Coulomb force . One may wonder , however , if we can teach a machine to redis- cover the underlying physics by solely observing the particles ’ states . This thinking has actually inspired the study of learning to model interacting systems , which now is a prevailing topic in machine learning ( Battaglia et al. , 2016 ; Thomas et al. , 2018 ; Köhler et al. , 2019 ; Sanchez-Gonzalez et al. , 2019 ; Fuchs et al. , 2020 ; Martinkus et al. , 2021 ; Satorras et al. , 2021 ) . ∗Equal contributions : Wenbing Huang and Jiaqi Han ; B Corresponding authors : Yu Rong and Fuchun Sun . †This work is done when Jiaqi Han works as an intern in Tencent AI Lab . Learning to model interacting systems is challenging . First , the systems are combinatorially complex , on account of that objects can be composed in combinatorially many possible arrangements ( Battaglia et al. , 2016 ) . This challenge , to some extent , can be addressed by making use of Graph Neural Networks ( GNNs ) ( Wu et al. , 2020 ) . By regarding objects as nodes and interactions as edges , GNNs extract information via message passing , which is able to characterize arbitrarily ordered objects and combinatorial relations . The second challenge is related to an important symmetry in physics : the model we use should be equivariant to any Euclidean transformation ( translation/reflection/rotation ) of the input . This complies with the fact that physics rules keep unchanged regardless of the reference coordinate system . Several works ( Fuchs et al. , 2020 ; Satorras et al. , 2021 ) have investigated equivariance upon GNNs and exhibited remarkable benefits on N-body . Another challenge , though less explored , is that the systems could be geometrically constrained . Geometric constraints arise in common practical systems , for example , in robotics where the joints of mechanical arms should be linked one by one , or in biochemistry where the atoms of molecules are connected by chemical bonds . When modeling these constrained systems , it is crucial to enforce the model to output legal predictions . For instance , the lengths/angles of chemical bonds in a molecular closely determine its chemical property which will be changed dramatically if the structural constraints are broken . As mentioned above , equivariant GNNs ( Fuchs et al. , 2020 ; Satorras et al. , 2021 ) have achieved desirable performance on the N-body system—this system , nevertheless , lacks of constraint . Considering constraints is not easy , as the dynamics of all elements within a constrained system evolve in a joint and complex manner . Unfortunately , to the best of our knowledge , there is no research ( particularly among equivariant GNN methods ) that learns to model the dynamical systems of multiple interacting objects under geometrical constraints . In this paper , we propose Graph Mechanics Network ( GMN ) that can tackle the above three challenges simultaneously : I . To cope with the combinatorial complexity , GMN takes advantage of graph models to encode the states of and interactions between objects , analogous to previous approaches . II . For the constraint satisfaction , GMN resorts to generalized coordinates , a well known notion in conventional mechanics to model the dynamics of structural objects . Here , a structural object ( such as the stick and hinge in Fig . 1 ) is defined as a set of multiple rigidly connected particles . In GMN , the constraints are implicitly encoded in the forward kinematics that describes the Cartesian states as the function of generalized coordinates . This strategy can inherently maintain the constraints and requires no external regulation . III . GMN is equivariant . In GMN , we need to compute the interaction forces in the Cartesian space ( see Eq . 5 ) , and infer the accelerations of the generalized coordinates based on the inverse dynamics ( see Eq . 6 ) , both of which are learned by a general form of equivariant functions ( see Eq . 11 ) . Notably , the proposed equivariant formulation is more efficient to compute compared to the previous versions based on spherical harmonics ( Thomas et al. , 2018 ; Fuchs et al. , 2020 ) , or more general than the one developed by EGNN ( Satorras et al. , 2021 ) . More importantly , we have theoretically discussed when and how our formulation can universally approximate any equivariant function . To evaluate the advantages of GMN , we have constructed a simulated dataset composed of three types of objects : particles , sticks and hinges , which is a complex case of the N-body system ( Kipf et al. , 2018 ) and can be used as building blocks for common systems . Under various scenarios in terms of different ratios of object types and different numbers of training data , we empirically verify the superiority of GMN compared to state-of-the-art models in prediction accuracy , constraint satisfaction and data efficiency . In addition , we also test the effectiveness of GMN on two real-world datasets : MD17 ( Chmiela et al. , 2017 ) and CMU Motion Capture ( CMU , 2003 ) . 2 RELATED WORK . Learning to simulate complex physical systems has been shown to greatly benefit from using GNNs . Interaction Network ( IN ) proposed by Battaglia et al . ( 2016 ) is the first attempt for this purpose , and it can be deemed as a special kind of GNNs to learn how the system interacts and how the states of particles evolve . Later researches have extended IN in different aspects : HRN ( Mrowca et al. , 2018 ) utilizes hierarchical graph convolution for tackling objects of various geometrical shapes and materials , NRI ( Kipf et al. , 2018 ) further explicitly infers the interactions with the help of a variational auto-encoder , and Hamiltonian graph networks ( Sanchez-Gonzalez et al. , 2019 ) equip GNNs with ordinary differential equations and Hamiltonian mechanics for energy conservation . However , all above approaches have ignored the symmetry in physics and the GNN models they use are not Euclidean equivariant . There is a subset of models ( Ummenhofer et al. , 2019 ; Sanchez-Gonzalez et al. , 2020 ; Pfaff et al. , 2020 ) that partially implement symmetries , but they only enforce translation equivariance but not rotation equivariance . On the other hand , it is nontrivial to enforce rotation equivariance . Tensor-Field networks ( Thomas et al. , 2018 ) uses filters built from spherical harmonics to allow 3D rotation equivariance , and this idea has been developed by SE ( 3 ) Transformer ( Fuchs et al. , 2020 ) that further takes the attention mechanism into account . Anther class of works ( Finzi et al. , 2020a ; Hutchinson et al. , 2021 ) resorts to the Lie convolution for the equivariance on any Lie group , based on lifting and sampling . Recently , EGNN ( Satorras et al. , 2021 ) has proposed a simple yet effective form of equivariant message passing on graphs , which does not require computationally expensive higher-order representations while still achieving better performance on N-body . As interpreted before , ensuring geometrical constraints is crucial for many practical systems , which , nevertheless , is seldom investigated in aforementioned works . Although several attempts ( Yang et al. , 2020 ; Finzi et al. , 2020b ) have been proposed for learning to enforce constraints , they explicitly augment the training loss with soft Lagrangian regulation and thus are completely data-driven and have no guarantee of generalization for limited training data . DeLaN ( Lutter et al. , 2019 ) also employs generalized coordinates to describe the kinematics of the rigid object . Nevertheless , it only target on the physical process of a single object , other than the complex systems with multiple rigid and structural objects that are the focus of this paper . In DPI-Net ( Li et al. , 2018 ) , the BoxBath task does share the similar strategy to us by first predicting the canonical coordinates of the box and then using the forward kinematic model to obtain the Cartesian positions . However , the passing messages in DPI-Net are scalars other than directional vectors ( positions , velocities , and accelerations ) used in our work . After all , both DeLaN and DPI-Net never study equivariant models . 3 GRAPH MECHANICS NEURAL NETWORK . We begin with introducing the N-body system ( Kipf et al. , 2018 ) , as illustrated in Fig . 1 ( a ) . This system consists of N interacting particles { Pi } Ni=1 of the same mass , and the kinematics states of each particle are defined as Si = ( xi , vi ) , where xi , vi ∈ R3 are the position and velocity vectors , respectively . There could be certain non-vector information of each particle ( such as charge ) , which is represented by a c-channel feature hi ∈ Rc 1 . Suppose the system we study is conservative and the dynamics is driven by the interaction force fij ∈ R3 between any pair of particles i and j . According to Newton ’ s second law , the acceleration of particle i , ai ∈ R3 is proportional to the aggregated force from other particles ∑ j 6=i fij . All symbols will be specified with a superscript t for temporal denotations , e.g. , xti indicating the position of particle i at time t. In this paper , we are mainly concerned with the prediction task : we need to seek out a function φ ( { ( S0i , h0i ) } Ni=1 ) given the initial states { S0i } Ni=1 and features { h0i } Ni=1 to forecast the future states { STi } Ni=1 at time T . As presented in Introduction , two kinds of inductive biases have been explored previously . The first one is to apply the graph structure to capture the distribution of particle states and their interactions ( Battaglia et al. , 2016 ; Kipf et al. , 2018 ) , where , particularly , the particle states S0i ( along with h0i ) are as node features and the interaction forces f 0 ij as edge messages . In this way , the transition function φ boils down to a GNN model . The second inductive bias is that φ should be equivariant to any translation/reflection/rotation of the input states . By saying equivariance , we imply φ ( { ( g · S0i , h0i ) } Ni=1 ) = g · φ ( { ( S0i , h0i ) } Ni=1 ) , ∀g ∈ O ( 3 ) , ∀S0i , ∀h0i . ( 1 ) Here , O ( 3 ) defines the 3D orthogonal group ( Fuchs et al. , 2020 ) that consists of translation , reflection and rotation transformations ; g · S0i denotes to perform transformation g on the states S0i , and it is instantiated as R ( x0i + b ) for the position and Rv 0 i for the velocity , where R ∈ R3×3 is the orthogonal matrix and b ∈ R3 is the translation vector . Several works ( Thomas et al. , 2018 ; Köhler et al. , 2019 ; Fuchs et al. , 2020 ; Satorras et al. , 2021 ) have investigated both inductive biases , among which EGNN ( Satorras et al. , 2021 ) has achieved promising performance on N-body . The typical process in EGNN iterates the following steps : 1Although hi could be of multiple dimensions , we still call it as a non-vector value since the dimension of hi is distinct from that of the kinematics vector , for example xi , the latter of which has intrinsic geometry and should obey translation/rotation equivariance . This is also why we do not denote hi in bold . ali , h l i = ∑ j ϕegnn ( x l−1 ji , h l−1 i , h l−1 j , eji ) , ( 2 ) vli = ψ ( h l−1 i ) v l−1 i + a l i , ( 3 ) xli = x l−1 i + v l i , ( 4 ) where the superscript l denotes the l-th layer ; the acceleration ali returned by the message aggregation of ϕegnn in Eq . 2 is adopted for the update of the velocity vli in Eq . 3 ( multiplied by a scalar ψ ( hl−1i ) ∈ R ) , followed by the renovation of the position xli in Eq . 4 . The formu- lation of ϕegnn is physically reasonable , since the interaction ( actually the Coulomb force ) between particles i and j truly depends on their relative position xl−1ji = x l−1 i − x l−1 j , node features h l−1 i and hl−1j , and edge feature eji . To enable equivariant message passing , EGNN has developed a specific form of ϕegnn to let ali be equivariant and h l i be invariant in terms of the input x l−1 ji , which will be presented in Eq . 10 . Notice that the computations for f li and h l i are actually by two different functions , and we have abbreviated them into one in Eq . 2 ( and also Eq . 5 later ) , since they share the same inputs ; besides , their parameters are shared following EGNN .

Summary: The paper proposes a new graph neural network model for predicting the dynamics of n-body systems. The two main innovations of this model are: 1. Ability to handle systems with rigid-body constraints.  2. A new 3D transformation equivariant layer with universality properties.