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 goal of this paper is to learn world models - where the network must predict actions from images - in an efficient manner.  This is done by exploiting symmetries in the data.   Results are shown for a number of world model benchmarks, showing improvement over standard methods.