Roto-translated Local Coordinate Frames For Interacting Dynamical Systems

1 Introduction . Modelling interacting dynamical systems –systems of interacting objects with highly non-linear and time-dependent behaviour– with neural networks is attracting a significant amount of interest [ 27 , 3 , 19 ] for its potential to learn long-term behaviours directly from observations . A large class of these systems consists of objects in the physical space , for instance pedestrians in a traffic scene [ 31 , 37 ] or colliding subatomic particles [ 5 ] . These systems can be formalized as geometric graphs , in which the nodes describe the physical coordinates of the objects among other features . Kipf et al . [ 27 ] introduced Neural Relational Inference ( NRI ) to learn geometric graph dynamical systems using the variational autoencoding framework [ 25 , 38 ] . Following [ 27 ] , dynamic NRI [ 19 ] advocated sequential latent variable models to encode time-transient behaviours . Both approaches and the majority of learning algorithms for dynamical systems assume an arbitrary global coordinate system to encode time-transient interactions and model complex behaviours . In this work , we posit that taking into account the relative nature of dynamics is key to accurate modelling of interacting dynamical systems . Represented by geometric graphs , learning algorithms of dynamical systems subscribe themselves to the Newtonian space . The absolute notion of Newtonian space and mechanics , however , determines that there exist infinite inertial frames that connect with each other by a rotation and translation . Each of these inertial frames is equivalent in that they can all serve as global coordinate frames and , thus , an arbitrary choice is made . Notwithstanding this arbitrariness , the dynamics of the system are invariant to the choice of a global coordinate frame up to a rotation and translation , in what is also known as Galilean invariance . Put otherwise , geometric graphs of interacting dynamical systems often exhibit symmetries that if left to their own devices lead models to subpar learning . Inspired by the notion of Galilean invariance , we focus on inducing roto-translation invariance in interacting dynamical systems and their geometric graphs to sustain the effects of underlying 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . pathological symmetries . Symmetries , invariances and equivariances have attracted an increased interest with learning algorithms in the late years [ 11 , 12 , 53 , 54 , 46 , 17 ] . The reason is that with the increasing complexity of new tasks and data , exploiting the symmetries improves sample efficiency [ 17 , 41 ] by requiring fewer data points and gradient updates . To date the majority of works on exploiting symmetries or data augmentations are with static data [ 11 , 44 , 54 ] . We argue and show that invariance in representations of dynamic data is just as important , if not more , as it is critical in accounting for the inevitable increased pattern complexity and non-stationarity . We induce roto-translation invariant representations in graphs by local coordinate frames . Each local coordinate frame is centered at a node-object in the geometric graph and rotated to match its angular position –yaw , pitch , and roll . Since all intermediate operations are performed on the local coordinate frames , the graph neural network is roto-translation invariant and the final transformed output is roto-translation equivariant . We obtain equivariance to global roto-translations by an inverse rotation that transforms the predictions back to the global coordinates . We make the following three contributions . First , we introduce canonicalized roto-translated local coordinate frames for interacting dynamical systems formalized in geometric graphs . Second , by operating solely on these coordinate frames , we enable roto-translation invariant edge prediction and roto-translation equivariant trajectory forecasting . Third , we present a novel methodology for natural anisotropic continuous filters based on relative linear and angular positions of neighboring objects in the canonicalized local coordinate frames . We continue with a brief introduction of relevant background and then the description of our method . We present related work and finish with experiments and ablation studies on a number of settings , including modelling pedestrians in 2D traffic scenes , 3D particles colliding , and 3D human motion capture systems . 2 Background . 2.1 Interacting dynamical systems and geometric graphs . An interacting dynamical system consists of i = 1 , . . . , Nt objects , whose position p = ( px , py , pz ) ⊤ and velocity u = dpdt = ( ux , uy , uz ) ⊤ in the Euclidean space are recorded over time t. The state of the i-th object at timestep t is described by xti = [ p t i , u t i ] , adopting for clarity a column vector notation and using [ · , · ] to denote vector concatenation . Over the past few years , a natural way that has emerged for organizing interacting dynamical systems is by geometric graphs [ 3 , 27 , 19 ] through space and time , G = { Gt } Tt=1 , where Gt = ( Vt , Et ) is the snapshot of the graph at timestep t. The nodes Vt = { vt1 , . . . , v t Nt } of the graph correspond to the objects in the dynamical system , with vti corresponding to the state of the i-th object , x t i . The edges Et ⊆ { etj , i = ( vtj , v t i ) | ( vtj , v t i ) ∈ Vt × Vt } of the graph , correspond to the interactions from node-object j to node-object i . We use N ( i ) to denote the graph neighbours of node vi . In the absence of domain knowledge about how objects connect , for instance the links between atoms in molecules , the graph is fully connected . Explicit inference of the graph structure can be achieved by using latent edges ztj , i corresponding to the edges e t j , i . Graph neural networks [ 42 , 32 , 18 ] exchange messages between neighbors and update the vertex and edge embeddings per layer , commonly referred to as message passing h ( l ) j , i = f ( l ) e ( [ h ( l−1 ) i , h ( l−1 ) j , i , h ( l−1 ) j ] ) ( 1 ) h ( l ) i = f ( l ) v h ( l−1 ) i , □ j∈N ( i ) h ( l ) j , i , ( 2 ) where h ( l ) i is the embedding of node vi at layer l and h ( l ) j , i is the embedding of edge ej , i at layer l. fe , fv denote differentiable functions such as MLPs and □ denotes a permutation invariant function , commonly a summation or an average . Many graph neural networks rely on isotropic filters [ 26 ] , although various ways [ 47 , 36 ] to circumvent this constraint have also been explored . 2.2 Neural relational inference . Kipf et al . [ 27 ] proposed neural relational inference ( NRI ) , a variational autoencoding inference model [ 25 , 38 ] that explicitly infers the graph structure over a discrete latent graph and simultaneously learns the dynamical system . Given the input trajectories the encoder learns to infer interactions as latent edges zj , i ∈ [ 0 , 1 ] K , sampled from a concrete distribution [ 34 , 23 ] with K edge types . The decoder receives the inferred interactions and the past trajectories , and learns the dynamical system , pθ ( x|z ) = ∏T t=1 pθ ( x t+1|x1 : t , z ) . The encoder is a regular graph neural network without explicitly taking time into account and learns to infer latent interactions by maximizing the evidence lower bound for the next time step . The decoder is another graph neural network that either assumes Markovian dynamics pθ ( xt+1|x1 : t , z ) = pθ ( xt+1|xt , z ) or is recurrent through time . As the prior and encoder in NRI assume static interactions ( e.g . whether forces between charged particles are attractive or repulsive ) , dynamic NRI [ 19 ] replaces them with a sequential relation prior based on past states , pϕ ( z|x ) = ∏T t=1 pϕ ( zt|x1 : t , z1 : t−1 ) , and an approximate relation posterior based on both the past and future states . The decoder is also reformulated as pθ ( x|z ) = ∏T t=1 pθ ( x t+1|x1 : t , z1 : t ) , taking into account the dynamic nature of interactions . 2.3 Invariance and equivariance . Last , we give a very brief introduction to invariance and equivariance . A function f : X → Y is equivariant [ 54 ] under a group of transformations if every transformation π ∈ Π of the input x ∈ X can be associated with a transformation ψ ∈ Ψ of the output , ψ [ f ( x ) ] = f ( π [ x ] ) . A special case is invariance , where Ψ = { I } , the identity transformation , f ( x ) = f ( π [ x ] ) . In this work , we are interested in translation , i.e. , f ( x ) + τ = f ( x+ τ ) with the translation vector τ , and rotation invariance/equivariance , i.e. , Qf ( x ) = f ( Qx ) using the rotation matrix Q . 3 Roto-translation invariance with local coordinate frames . In this section we present our method , termed LoCS ( Local Coordinate frameS ) . We start with the derivation of roto-translated local coordinate frames and continue with the formulation of graph networks and continuous anisotropic filters operating in these frames . 3.1 Local coordinate frames . Starting from the spatio-temporal graph G , we focus for clarity on pairs of node-objects in the same time step , xti , x t j . In the real world , objects are not point particles and have a spatial extension . Central to our method is the use of the angular positions ω = ( θ , ϕ , ψ ) ⊤ , otherwise known as yaw , pitch and roll , that describe the orientation of a rigid body with respect to the axes of the coordinate system . We , thus , augment the states xti with the angular positions , using v t i = [ p t i , ω t i , u t i ] to denote the augmented state that captures the angular position as well as the linear position and velocity . Our method draws inspiration from Galilean invariance and inertial frames of reference that capture the relative locations and motions of all the objects in a system . We introduce Nt local coordinate frames , one per object in the system . For the i-th object , a local coordinate frame is one in which both the linear and the angular position lie on the origin . By the adoption of the local coordinates with respect to object-centric frames , the behaviors between objects will not depend on the arbitrary positions of objects in the absolute Newtonian space . In other words , the local coordinates offer invariance to global translations and rotations and do not bias the learning algorithm . Our goal , thereby , is to compute the relative local coordinates of all objects j = 1 , . . . , Nt , while iterating over the i-th reference object . The transformation from global to local coordinate systems is in two steps . Per target node , we first translate the origin to match its linear position by a translation transformation . Since velocities and angular positions are translation invariant , we only need to perform the translation to the linear positions . This gives us the relative positions rtj , i = [ ptj − pti ] . Then , we canonicalize the local coordinate frame to match the target object ’ s orientation by a rotation transformation , described by the rotation matrix Q ( ωi ) . The coordinates of all other objects are analogously transformed given the i-th local coordinate frame . The rotations are performed independently and equivalently to the state components of the j-th object , namely the relative –due to the translation transformation performed first– linear positions , the angular positions and the velocities . A schematic overview of the proposed transformation in a 2D setting is presented in fig . 1 . Using tensor operations , we compactly write down the transformed state as : R̃ ( ω ) = Q ( ω ) ⊕Q ( ω ) ⊕Q ( ω ) ( 3 ) vtj|i = R̃ t⊤ i [ rtj , i , ω t j , u t j ] ( 4 ) where ⊕ denotes the direct sum operator that concatenates two matrices along the diagonal resulting in a block diagonal matrix and R̃t⊤i = R̃ ⊤ ( ωti ) to reduce notation clutter . The rotation matrix in eq . ( 3 ) has three entries . Each Q ( ω ) independently transforms the relative linear positions , the angular positions and the velocities . Local coordinate frames in 2D We start with the simpler case where the dynamical system resides in the 2D Euclidean space , for instance pedestrians in a traffic scene . In 2 dimensions , the angular position is a scalar value , namely the yaw angle θ . Thus , the rotation matrix for a target node vi is : Q ( ωi ) = Q ( θi ) = ( cos θi − sin θi sin θi cos θi ) ( 5 ) Local coordinate frames in 3D When in the 3D Newtonian space the chain of transformations is the same ; first , we translate the origin to match each target node ’ s linear position , and then we rotate the local coordinate frames to match each object ’ s orientation . The translation transformation is identical to the 2D case . However , the rotation transformation is more involved . For the 3D case , we must decompose Q ( ω ) into 3 chained elemental rotations , described by the matrices Qz ( θ ) , Qy ( ϕ ) and Qx ( ψ ) . Qz ( θ ) describes a rotation around the z-axis by an angle θ , Qy ( ϕ ) describes a rotation around the y-axis by an angle ϕ and Qx ( ψ ) describes a rotation around the x-axis by an angle ψ . We use the following convention that dictates the order of rotations , Q ( ω ) = Qz ( θ ) Qy ( ϕ ) Qx ( ψ ) . Each elemental rotation matrix has similar structure to eq . ( 5 ) . We provide the complete description of all rotation matrices in appendix A.2 . After the computation of the rotation matrix , the states are transformed identically to eq . ( 4 ) . The local coordinate frames are invariant with respect to global translations and rotations , either in 2 or 3 dimensions . We provide the detailed derivations of Q ( ω ) for the 2D and 3D case in appendices A.1 and A.2 , respectively , and a detailed proof in appendix A.3 .

Summary: This paper presents an approach to  enable roto-translation invariant edge prediction and roto-translation equivariant trajectory forecasting. Rotation and translation invariants are achieved by transforming global coordinate into local coordinate frames of each object in the graph neural network. Experiments are done on forecasting trajectories of charged particles, traffics and human motions.