Learning Symmetric Representations for Equivariant World Models

1 INTRODUCTION . Symmetry has proved to be a powerful inductive bias for improving generalization in supervised and unsupervised learning . A symmetry group defines an equivalence class of inputs in terms of a set of transformations that can be performed on this input , along with corresponding transformations for the output . The last years have seen many proposed equivariant models that incorporate symmetries into deep neural networks ( Cohen & Welling , 2016a ; b ; Cohen et al. , 2019 ; Weiler & Cesa , 2019 ; Weiler et al. , 2018 ; Kondor & Trivedi , 2018 ; Bao & Song , 2019 ; Worrall et al. , 2017 ) . This results in models that are often more parameter efficient , more sample efficient , and safer to use by behaving more consistently in new environments . However , a major impediment to the wider application of equivariant models is that it is not always obvious how a symmetry group acts on an input data set . As an example , let us consider the two pairs of images in Figure 1 . On the left , we have a pair of MNIST digits , for which a 2D rotations in pixel space should induce a corresponding rotation in feature space . Here it is possible to achieve state-of-the-art accuracy using an E ( 2 ) -equivariant network ( Weiler & Cesa , 2019 ) . By contrast , exploiting the underlying symmetry is much more challenging for the pair of images on the right , which show the same three-dimensional object in two orientations . For these images , there is also an underlying symmetry group of rotations , but it is not easy to characterize the transformation in pixel space that is associated with a particular rotation . In this paper , we consider the task of learning symmetric representations of data in domains where transformations can not be hard-coded , which is to say that the group action on the data is not known . We propose training a standard network to learn a mapping from an input space , for which the group action is difficult to characterize , into a latent space , for which the action is known . This mapping , which we refer to as a symmetric embedding network , can then be composed with any equivariant network for downstream predictions . As a concrete instantiation of this idea , we focus on learning world models , i.e . models that encode the effects of actions in the state space of an environment . We combine the symmetric embedding network with an equivariant transition network , which are trained end-to-end by minimizing a contrastive loss . The underlying intuition behind this approach is that the symmetry group of the transition model acts as an inductive bias that guides the embedding network to a representation that is as equivariant as possible . At the same time , incorporating symmetries into the transition dynamics has the potential to improve both data efficiency and out-of-distribution generalization . The idea of learning symmetric embeddings using a mix of standard and equivariant networks has , to our knowledge , not previously been proposed or demonstrated . However , note that we are not proposing a new equivariant neural network design . In fact , our approach is useful precisely because it can be paired with any existing equivariant neural network to extend its applicability to new domains with unknown group actions . We apply our method to 5 different domains , 3 different symmetry groups , and 3 different equivariant architectures . We summarize our contributions as follows : • We introduce a meta-architecture for equivariant world models with symmetric embedding networks , which can be trained end-to-end using triples from a replay buffer by minimizing a contrastive loss . • We demonstrate that this meta-architecture can be used to learn world models with a variety of equivariances in a self-supervised manner . • Moreover , we show that models that have been trained using only a subset of all input actions can generalize to unseen input actions at test time . 2 RELATED WORK . Equivariant Neural Networks A multitude of equivariant neural networks have been devised to impose various symmetry groups with respect to various groups across a variety of data types . They require that the group G is known and the the group action on input , output , and hidden spaces is explicitly constructed . Examples include G-convolution ( Cohen & Welling , 2016a ) , G-steerable convolution ( Cohen & Welling , 2016b ; Weiler & Cesa , 2019 ) , tensor product and Clebsh-Gordon decomposition ( Thomas et al. , 2018 ) , or convolution in the Fourier domain ( Esteves et al. , 2017 ) . They have been applied to many data types such as gridded data ( Weiler & Cesa , 2019 ) , spherical data ( Cohen et al. , 2018 ) , point clouds ( Dym & Maron , 2020 ) , and sets ( Maron et al. , 2020 ) . They have found applications in many domains including molecular dynamics ( Anderson et al. , 2019 ) , particle physics ( Bogatskiy et al. , 2020 ) , and trajectory prediction ( Walters et al. , 2020 ) . In particular , Ravindran ( 2004 ) consider symmetry in the context of Markov Decision Processes ( MDPs ) and van der Pol et al . ( 2020b ) construct equivariant policy networks for policy learning . Our work also considers MDP with symmetry but focuses on learning equivariant world models ( see Appendix B ) . Learning Symmetry Our work occupies a middle ground between equivariant neural networks in which the group and its representations are known and symmetry discovery models . Symmetry discovery methods attempt to learn both the group and actions from data . For example , Zhou et al . ( 2020 ) learn equivariance by learning a parameter sharing scheme using meta-learning . Dehmamy et al . ( 2021 ) similarly learn a basis for a Lie algebra generating a symmetry group while simultaneously learning parameters for the equivariant convolution over this symmetry group . Benton et al . ( 2020 ) propose an adaptive data augmentation scheme , where they learn which group of spatial transformations best supports data augmentation . Higgins et al . ( 2018 ) define disentangled representations based on symmetry , with latent factors considered disentangled if they are independently transformed by commuting subgroups . Within this definition , Quessard et al . ( 2020 ) learn the underlying symmetry group by interacting with the environment , where the action space is a group of symmetry transformations . Except for the teapot task , we handle the more general case where the action space may be different from the symmetry group . Their latent transition is given by multiplication with a group element , whereas our transition model is given by an equivariant neural network . Structured Latent World Models World models learn state representations by ignoring unnecessary information unrelated to predicting environment dynamics . Such models are frequently used for high-dimensional image inputs , and usually employ ( 1 ) reconstruction loss ( Ha & Schmidhuber , 2018 ; Watter et al. , 2015 ; Hafner et al. , 2019 ; 2020 ) or ( 2 ) constrastive loss . Contrastive loss is known to be less computationally costly and can produce good representations in learning from high-dimensional inputs ( Oord et al. , 2018 ; Anand et al. , 2019 ; Chen et al. , 2020 ; Srinivas et al. , 2020 ; van der Pol et al. , 2020a ) , thus we use it for training our models . For example , Kipf et al . ( 2020 ) learn object factored representations for structured world modeling with GNNs , which respect Sn permutation symmetry . We learn symmetric representations for groups G and explicitly enforce G-equivariance constraints to latent transition networks . 3 BACKGROUND . We provide some background on symmetry groups and highlight the difference between abstract symmetry groups and their concrete representations . Here we assume we know the abstract group , but only some of its relevant representations . Groups , Actions , and Representations A symmetry group consists of a set G together with a composition map ◦ : G×G→ G. The group must contain an identity 1 ∈ G and each element g ∈ G must be invertible . An action of the group G on a set S is a map a : G→ Perm ( S ) mapping each element of the group g to a permutation πg ∈ Perm ( S ) of the elements of S. Composition of group elements is compatible with the action such that a ( g1g2 , s ) = a ( g1 , a ( g2 , s ) ) for g1 , g2 ∈ G , s ∈ S. A real representation of the group G is a linear group action , given by a map ρ : G→ GLn ( R ) which maps each element of G to an invertible n × n matrix . The multiplication table of these matrices must match that of the abstract group elements under composition . That is , ρ ( g1 ◦ g2 ) = ρ ( g1 ) ρ ( g2 ) . See Hall ( 2003 ) for additional background on groups and their representations . Equivariant and Invariant Functions Given a function f : X → Y between vector spaces X and Y and a group G equipped with representations ρX and ρY acting on X and Y respectively , we say that f is equivariant if , for all x ∈ X , g ∈ G , we have f ( ρX ( g ) · x ) = ρY ( g ) · f ( x ) . This means that if the input is transformed by g the output will be transformed correspondingly . The composition of equivariant functions is equivariant . Thus we can model equivariant functions using equivariant neural networks which alternate equivariant linear layers and equivariant non-linearities . 4 SYMMETRIC EMBEDDINGS FOR EQUIVARIANT WORLD MODELS . Our goal is to learn equivariant world models where the underlying symmetry group of the input space is difficult to compute . Our method gives a general template that can be fit to different equivariant neural networks , symmetries , and data types . We describe the general template and the implementation in examples . 4.1 MODEL OVERVIEW . Equivariant World Models Let S be the state space and A be the action space of an environment . We consider a deterministic transition function T : S × A → S which outputs the next state s′ = T ( s , a ) given state s and action a . Our goal is to learn T from tuples ( s , a , s′ ) collected from offline data . As the space S × A is combinatorially large , we wish to learn a compact state representation of the state and an accurate model of T that can generalize to unseen transitions . This is accomplished by learning a state abstraction map S → Z and then learning transitions in latent space TZ : Z ×A → Z . To learn world models from fewer samples , we exploit inherent symmetries of the environment . Let G be a group of symmetries with group representations ρS and ρA . We assume that ρA ( g ) is independent of state for simplicity as many domains do not require state-dependent actions . Additional details on the symmetric MDPs and the setup are provided in the Appendix . The transition function T is equivariant if T ( ρS ( g ) · s , ρA ( g ) · a ) = ρS ( g ) · T ( s , a ) . We wish to enforce this on a neural network model for T . If trained to predict s′ = T ( s , a ) , the model will then automatically generalize to gs′ = T ( gs , ga ) , enabling improved generalization and sample efficiency . Current methods for constructing model classes of equivariant neural networks require that ρS and ρA are known . In our case , however , although we assume the symmetry group G and group action ρA , we assume that ρS is not explicitly known . In many environments , s is a pixel-level input and the transformation ρS may be very difficult to describe . Thus , we can not directly enforce symmetry for T using an equivariant neural network . Meta-Architecture To learn an equivariant world model without access to ρS , we learn a symmetric abstract state mapping from states s to abstract states z in a space Z with an explicit action ρZ of the symmetry group G. We then learn a transition model in latent space where we can enforce symmetry using an equivariant neural network . We learn the symmetric abstract state mapping in two parts . First we map the pixel-space state to an intermediate space Y using a symmetric embedding S : S → Y . This is a non-equivariant neural network acts as a feature extractor and maps images into a space Y with a simpler explicit symmetry group action ρY . We then map the intermediate space to the lower-dimensional latent space using an equivariant encoder E : Y → Z . The encoder discards features that are unnecessary to predict dynamics and reduces dimensionality . Lastly , we compute the transition in latent space , T : Z ×A → Z . We explicitly enforce E and T to be equivariant using equivariant neural networks . For training these networks , we employ the same self-supervised contrastive loss in Kipf et al . ( 2020 ) . Let ( s , a , s′ ) be a ground truth transition triplet and s′′ be an incorrect next state s′′ 6= s′ . Let z = E ( S ( s ) ) , z′ = E ( S ( s ) ) , and z′′ = E ( S ( s ) ) , then L ( s , a , s′ , s′′ ) = ‖T ( z , a ) − z′‖+ αmax ( β − ‖T ( z , a ) − E ( z′′ ) ‖ , 0 ) . Minimizing this loss pushes T ( z , a ) towards z′ and away from the incorrect sample z′′ . Symmetric Embedding Network S We use non-equivariant CNNs for all environments , but the specific architecture varies ( see Appendix D.2 ) . For example , in the object-centric environments , the symmetry group is G = C4 × S5 , the cyclic group of order 4 acting by π/2-rotations and the permutation group on the 5 objects . The output y has shape [ B , C , 5 , 4 ] and carries an action ρY of S5 by permuting the 5-dimensional axis and of C4 by cyclically permuting the 4-dimensional axis . Equivariant Encoder E and Transition Model T The encoder and transition model are both implemented using G-equivariant neural networks . In the object-centric environments with G = C4 × S5 , the equivariant encoder E is shared over all 5 objects and uses group convolution over C4 ( Cohen & Welling , 2016a ) , thus achieving C4 × S5-equivariance . The transition function T is implemented as a GNN ( message-passing neural network ) with edge and node networks which use C4-convolutions for their linear layers . Since GNNs are S5 equivariant by definition and the linear layers within the GNNs are C4-equivariant , this is C4 × S5-equivariant . For other implementations , see Table 1 and Appendix D.2 . 4.2 SO ( 3 ) -STRUCTURED SYMMETRIC EMBEDDINGS Symmetric Embedding In the case ofG = SO ( 3 ) symmetry , we expect the symmetric embedding network to detect the pose of object z in 3D . We omit the equivariant encoder ( by setting E = id ) and instead use a two-part Symmetric Embedding Network that directly encodes y = z using a down-sampling CNN whose output is passed to an MLP , and converted to an element of SO ( 3 ) . To force the output of the symmetric embedding network y to be an element of SO ( 3 ) , we have the last layer output 2 vectors u , v ∈ R3 and perform Gram-Schmidt orthogonalization to construct a positively oriented orthonormal frame ( see Appendix ) . This method is also used by Falorsi et al . ( 2018 ) , who conclude it produces less topological distortion than alternatives . Transition We implement an equivariant transition model using Tensor Field Networks ( Thomas et al. , 2018 ; Geiger et al. , 2020 ) . This is an SO ( 3 ) n ( R3 , + ) -equivariant method which works over point clouds . Here z ∈ SO ( 3 ) and a ∈ SO ( 3 ) . We consider z as 3 points in R3 and add the origin to get a 4 point cloud . We embed the actions as features over these 4-points . The MDP action a is then set as a feature over these 4 points , which has SO ( 3 ) -representation ρA ( g ) · a = gag−1 . The MDP action a ∈ SO ( 3 ) is a rotation matrix and the latent state z ∈ SO ( 3 ) is a positively-oriented orthogonal coordinate frame . Though Z = A = SO ( 3 ) , these different semantics lead to differing G = SO ( 3 ) actions with ρZ ( g ) ( z ) = g · z but ρZ ( g ) ( a ) = gag−1 . If z is correctly learned , then the ground truth latent transition function can be represented by a simple matrix multiplication TZ ( z , a ) = az which is also equivariant , TZ ( ρZ ( g ) ( z ) , ρA ( g ) ( a ) ) = ( gag −1 ) ( gz ) = gaz = ρZ ( g ) TZ ( z , a ) . This method , which we label MatMul , is similar to the latent transition model used in Quessard et al . ( 2020 ) , except in our framework the groundtruth a ∈ SO ( 3 ) is provided to aid learning z .

Summary: The paper proposes to learn latent representation and transition models that are equivariant to symmetry transformations of states and actions. The proposed “meta-architecture” consists of an encoder followed by an equivariant encoder and an equivariant transition model. The model is trained using (state, action, state) triples, where contrastive loss is used to prevent the collapse. The paper evaluates the proposed model in four different settings. In each case, a different choice of symmetry group and equivariant architecture is used. The quality of the learned latent representation is evaluated by comparing the relative distance of the predicted next state embedding to the true next state embedding, in comparison to negative samples.