Year: 2021

Abstract:
Incorporating symmetries leads to highly data-efficient and generalizable world models. However, characterizing how underlying symmetries manifest in the input space is often difficult. We provide a method to map an input space (e.g. images) where we do not know the effect of transformations onto a feature space that transforms in a known manner under these operations. Specifically, we consider equivariant transition models as an inductive bias for learning the encoder. Our method allows existing equivariant neural networks to operate in previously inaccessible domains. We evaluate the effectiveness of this approach in domains with 3 distinct forms of underlying symmetry. In many cases, we demonstrate improvements relative to both fully-equivariant and non-equivariant baselines.

heading: 1 INTRODUCTION
text: 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 cannot 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 BACKGROUND
text: 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.

heading: 4 SYMMETRIC EMBEDDINGS FOR EQUIVARIANT WORLD MODELS
text: 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.

heading: 4.1 MODEL OVERVIEW
text: 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 cannot 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.

heading: 4.3 GENERALIZING OVER THE MDP ACTION SPACE
text: Although the state does not have a known group action ρS , the MDP action does have known ρA. In the domains we consider, although the state is high-dimensional and has non explicit symmetry, the action is low dimensional and has clear symmetry. The MDP action is passed directly to T and thus bypasses the non-equivariant part S of the neural network. Since the neural network is explicitly equivariant with respect to the MDP action, it is thus feasible to train the neural network using only a proper subset A′ ⊂ A of the action space, and then test on the entire A. This may be useful in domains in which data collection is costly. Since the samples from S are still i.i.d., the non-equivariant neural network S is still able to learn well.
In generalization experiments, we require that ρA(G) · A′ = A, i.e. every MDP action is Gtransformed version of one in A′. We assume that S is approximately equivariant after training, which, as T andE are constrained to be equivariant, implies TS(s, a) = T (E(S(s)), a) is equivariant. We also assume low error for T on the restricted action set (s1, a′, s2) ∈ S × A′ × S. Then given (s1, a, s2) ∈ S ×A× S , there exists g ∈ G such that a = ρA(g) · a′. Let s′i = ρS(g−1)si. Then
TS(s, a) ≈ TS(ρS(g) · s′1, ρA(g) · a′) ≈ ρS(g) · TS(s′1, a′) ≈ ρS(g) · s′2 ≈ s2.
If performance is good for A′ and S has low equivariance error, then performance will be good for A.

heading: 5 EXPERIMENTS

heading: 5.1 SETUP
text: We choose five environments with varying symmetries to evaluate our models. The first three environments 2D Shapes, 3D Blocks, and Rush Hour are grid-worlds with five moving objects, based on (Kipf et al., 2020). Rush Hour is a variant of 2D Shapes where objects can move relative to their orientation. We consider symmetry to π/2 rotations and object permutations. We also evaluate in a continuous control domain, the Reacher-v2 MuJoCo environment, with rotational and translational symmetries. The last domain is of a rotating 3D teapot, where an action is an element of SO(3). We consider two action spaces: a small discrete action space (S) with 6 rotations of 2π30 in SO(3) and a large continuous action space (L) of any rotation in SO(3). All environments use images as observed states. Additional details are given in the Appendix D.1.
We compare three types of models: (a) a non-equivariant model with no enforced symmetry, (b) a fully-equivariant model with a mis-specified symmetry where ρS is a simple transformation of the pixels, and (c) our method. For 3D Teapot, we forgo the fully equivariant baseline as it is hard to define a ρS acting on the 2D image space which approximates the true group action.

heading: 5.2 MODEL ARCHITECTURES AND TRAINING
text: As each environment contains different symmetries, the model architecture is customized for every environment keeping the meta-architecture the same. We use object-oriented structured models which factorize the latent state space and latent action space over objects for the grid world environments of 2D shapes, 3D blocks, and Rush Hour. Though the objects and actions are factorized, the world model must account for the pairwise interactions between objects (e.g. actions to move one object can be blocked by another object). The encoder E is shared over all objects and a GNN is used for the transition function T . The Reacher and 3D Teapot environments do not consider objects and thus we do not use model components that consider permutations. A summary of the environments, different symmetries, representation types, and model architectures are given in Table 1.
We do not consider the reward as our focus is on constructing accurate latent representations and their dynamics. A random policy was used to create training and evaluation datasets of (s, a, s′) tuples. For all environments, we have either a combinatorially large state space (with objects) or continuous states and thus overlap between training and evaluation datasets is highly unlikely.
As equivariant networks have more parameters than the non-equivariant counterparts, we reduce the number of hidden dimensions accordingly to keep the number of parameters approximately constant for all models. The Adam (Kingma & Ba, 2014) optimizer was used for all experiments. All other specific implementation details are provided in Appendix D.2.

heading: 5.3 METRICS
text: In order to evaluate the performance of our model in latent space without reconstruction, we use standard ranking metrics from Kipf et al. (2020), modified versions of these metrics to adapt to continuous state spaces, and two metrics for evaluating the equivariance of the learned model.

heading: 5.3.1 ACCURACY METRICS
text: Hits, Hard Hits, and MRR The evaluation samples are ranked according to the pairwise L2 distance of the predicted next states and the true next states (both are encoded in latent space). Hits at Rank k (H@k) measures the average percentage of time that the predicted next state is within k-nearest neighbors of the encoded true next state. The mean reciprocal rank (MRR) is the average inverse rank. We also consider a variant of Hits at Rank k (HH@k) where we generate negative samples s′n of states that are close to the true next state s
′ (see Appendix for more details) and count the number of times that the distance to the positive sample was lower than the distance t o the negative samples. This is a harder version of H@k as the model must distinguish between close negative samples and the true positive sample in latent space.

heading: 5.3.2 EQUIVARIANCE METRICS
text: Equivariance Error EE In order to analyze the equivariance of each model, we generate a version of the evaluation dataset where one element of the symmetry group acts on the tuple (s, a, s′) and calculate the true equivariance error for the embedding network. Although by assumption ρS(g) · s cannot be computed using g and s, our synthetic datasets allow us to render both s and ρS(g) · s during generation. Specifically, the equivariance error of the symmetric embedding is calculated as
EE = Es,g [|ρY(g) · S(s)− S(ρS(g) · s)|] .
Distance Invariance Error DIE The above equivariance error can always be applied to the symmetric embedding network when its output space is spatial and we can manually perform group actions on the outputs. However it cannot be applied to the latent space Z in the case of non-equivariant models since the group action on the latent space ρZ cannot be meaningfully defined.
We therefore propose a proxy for the equivariance error using invariant distances. For a pair of input states s, s′, an equivariant model f will have the same distances ‖f(s)−f(s′)‖ and ‖f(gs)−f(gs′)‖ assuming the action of G is norm preserving as it is for all transformations considered in the paper. Due to the linearity of the action, ‖f(gs) − f(gs′)‖ = ‖gf(s) − gf(s′)‖ = ‖g(f(s) − f(s′))‖ = ‖(f(s)− f(s′))‖. The distance invariance error is computed as
DIE = Es,s′,g [|‖f(s)− f(s′)‖ − ‖f(gs)− f(gs′)‖|] . We evaluate both the symmetric embedding network (DIE(S)) and entire model (DIE(model)).

heading: 5.4 MODEL PERFORMANCE COMPARISON
text: The results are shown in Tables 2,3. In general, the ranking metrics (Hits and MRR) show that all three models are accurate on the 3D blocks, Rush Hour, and Reacher environments. The non-equivariant model achieves a higher H@1 on Rush Hour but has a slightly lower MRR than either the fully equivariant and our model on Reacher. Surprisingly, the fully equivariant model performs very well even when the group action on the input space ρS is not correct. Due to the skewed perspective, we can see that the simple pixel-level transformation maps training data to out-of-distribution images which are never seen by the model. We hypothesize that equivariance does not hamper its performance
on training data, but only constrains its extrapolation capabilities to out-of-distribution samples. In 3D Teapot, we observe that our Tensor Field Networks (TFN) and Matrix Multiplication (MatMul) models outperform the non-equivariant model, but have different performance for different action spaces. With small actions (S), Tensor Field Networks (TFN) and Matrix Multiplication (MatMul) perform similarly, while TFN fails with large actions (L). We hypothesize that the difference is caused by the contrastive loss function; small actions can lead to local minima where all states get mapped to close-by latent states. For the equivariance metrics, our model outperforms the non-equivariant model in all domains on DIE(model), while it performs similarly on EE(S) and DIE(S). As was the case for accuracy metrics, the fully equivariant model performs surprisingly well with possibly the same reason outlined above.
Visualization of latent embeddings We visualize the latent embedding z for our model to qualitatively analyze what kind of representations are learned. All figures are provided in the Appendix for space. Figure 4 plots all the learned embeddings for Reacher for all observations in the evaluation set and shows a sample transition in both pixel and latent space. The encoded current state z is highlighted in red and the encoded next state is highlighted z′ is highlighted in orange. which we factor into irreducible representations (irreps) before visualizing (see Hall (2003)). Some irreducible representations are 1-dimensional and are plotted as a line. The 2-dimensional irreps show a clear circular pattern, match the joint rotations of the environment.
Figure 5 shows the traversal of rotations in pixel and latent space for 3D Teapot. The latent space can choose its own base coordinate frame and thus is oriented downwards. We can clearly see that the effective rotations relative to the objects’ orientation perfectly align, demonstrating that the learned embeddings correctly encode 3D poses and rotations.

heading: 5.5 GENERALIZATION FROM LIMITED ACTIONS
text: We now train on a limited subset of actions and evaluate on datasets generated with the full action space. These experiments aim to verify that our model, even where all components are not designed to be equivariant, can learn a good equivariant representation which can generalize to unseen actions.
We perform experiments on the 2D Shapes, 3D Blocks, and Reacher domains. For 2D Shapes, the training data only contains ‘up’ actions and we set the number of episodes to 100,000 with length 1 to avoid any distribution shifts in the data (e.g. performing up continuously will produce many transitions where all blocks are blocked by the boundaries). For 3D Blocks, we omit the left action in training and use similar modifications to the episode lengths. For Reacher, the action space represents
joint actuation forces ∈ [−1, 1] for each of the two joints. We restrict the range of the force for the second joint to be positive, meaning that the second arm rotates in only one direction.
Tables 4,5,6 show results for 2D Shapes, 3D Blocks, and Reacher respectively. We see that our method can successfully generalize over unseen actions compared to both the non-equivariant and fully equivariant baselines. The non-equivariant baseline in particular performs poorly on all domains, achieving only 2.8% on Hits@1 and 5.5% on MRR for 2D Shapes. The fully equivariant model performs worse than our method for 3D Blocks and achieves a similar performance on Reacher. As the fully equivariant model performs well when trained on all actions but does not perform as well in these generalization experiments, these results lend support to our hypothesis that the incorrect pixel-level equivariance bias limits its extrapolation abilities to out-of-distribution samples. In these limited actions experiments, the fully equivariant model cannot extrapolate correctly and achieves lower performance than our model.
Figure 6 shows embeddings for all states in the evaluation dataset for our model and the nonequivariant model trained on only the up action. Our model shows a clear 5 × 5 grid, while the non-equivariant model learns a degenerate solution (possibly encoding only the row index x).

heading: 6 CONCLUSION AND FUTURE WORK
text: We demonstrate a flexible method which can be used to extend equivariant neural networks to domains with known symmetry types, but transformation properties which cannot be easily explicitly described. We apply our method across a variety of domains and equivariant neural network architectures. Our methods confer some of the advantages of equivariant neural networks in situations where they did not previously apply, such as generalization to data outside the training distribution. Future work will include applying our method to tasks besides world models and using our method to develop disentangled and more interpretable features in domains with known but difficult to isolate symmetry.
ETHICS STATEMENT
Our paper does directly address domains with privacy or safety concerns. However, our method can be used in robotics applications to train robots using fewer data samples. To the extent that robotics technology can be used for benefit or harm, our method enables both options.
REPRODUCIBILITY STATEMENT
We will open-source our code, including all models and data generation scripts, thus allowing all experiments to be fully reproduced.

heading: A OUTLINE
text: Our appendix is organized as follows. First, in Section B, we provide an additional formal setup of the problem. The visualization of learned latent symmetric representations is presented in Section C, followed by the details of training environments and network architectures in Section D.1 and D.2. We further explain the notation and definition in Section E.

heading: B SETUP: EQUIVARIANT WORLD MODELS
text: In this section, we provide a technical background for building equivariant world models, which we use in learning symmetric representations.
We model our interactive environments as Markov decision processes. A (deterministic) Markov decision process (MDP) is a 5-tuple M = 〈S,A, T,R, γ〉, with state space S, action space A, (deterministic) transition function T : S ×A → S, reward function R : S ×A → R, and discount factor γ ∈ [0, 1]. Symmetry can appear in MDPs naturally (Zinkevich & Balch, 2001; Narayanamurthy & Ravindran, 2008; Ravindran, 2004; van der Pol et al., 2020b), which we can exploit using equivariant networks. For example, van der Pol et al. (2020b) study geometric transformations, such as reflections and rotations. Ravindran (2004) study group symmetry in MDPs as a special case of MDP homomorphisms.
Symmetry in MDPs. Symmetry in MDPs is defined by the automorphism group Aut(M) of an MDP, where an automorphism g ∈ Aut(M) is an MDP homomorphism h :M→M that maps M to itself and thus preserves its structure. Zinkevich & Balch (2001) show the invariance of value function for an MDP with symmetry. Narayanamurthy & Ravindran (2008) prove that finding exact symmetry in MDPs is graph isomorphism complete.
Ravindran (2004) provide a comprehensive overview of using MDP homomorphisms for state abstraction and study symmetry in MDPs as a special case. A more recent work by van der Pol et al. (2020b) builds upon the notion of MDP homomorhpism induced by group symmetry and uses it in an inverse way. They assume knowledge of MDP homomorphism induced by symmetry group is known and exploit it. Different from us, their focus is on policy learning, which needs to preserve both transition and reward structure and thus has optimal value equivalence (Ravindran, 2004).
More formally, an MDP homomorphism h : M → M is a mapping from one MDP M = 〈S,A, T,R, γ〉 to another M = 〈S,A, T ,R, γ〉 which needs to preserve the transition and reward structure (Ravindran, 2004). The mapping h consists of a tuple of surjective maps h = 〈φ, {αs | s ∈ S}〉, where φ : S → S is the state mapping and αs : A → A is the statedependent action mapping. The mappings are constructed to satisfy the following conditions: (1) the transition function is preserved T (φ (s′) | φ(s), αs(a)) = ∑ s′′∈φ−1(φ(s′)) T (s
′′ | s, a), (2) and the reward function is also preserved R (φ(s), αs(a)) = R(s, a), for all s, s′ ∈ S and for all a ∈ A. An MDP isomorphism from an MDPM to itself is call an automorphism ofM. The collection of all automorphisms ofM along with the composition of homomorphisms is the automorphism group ofM, denoted Aut(M). We specifically care about a subgroup of G ⊆ Aut(M) which is usually easily identifiable from environments a priori and thus we can design appropriate equivariant network architectures to respect it, such as C4 rotation symmetry of objects. Additionally, while MDP homomorphisms pose constraints to the transition and reward function, we only care about the transition function, especially the deterministic case T : S ×A → S.
Equivariant transition. By definition, when an MDPM has symmetry, any state-action pair (s, a) and its transformed counterpart (ρS(g) · s, ρA(g) · a) are mapped to the same abstract state-action pair by h ∈ Aut(M): (φ(s), αs(a)) = (φ(gs), αgs(ga)), for all s ∈ S, a ∈ A, g ∈ G. Therefore, the transition function T : S ×A → S should be G-equivariant:
T (ρS(g) · s, ρA(g) · a) = ρS(g) · T (s, a), (1) for all s ∈ S, a ∈ A, g ∈ G.
State-dependent action transformation. Note that the group operation acting on action space A depends on state, since G actually acts on the product space S ×A: (g, (s, a)) 7→ ρS×A(g) · (s, a). However, in most cases, including all of our environments, the action transformation ρA does not depend on state. As a bibliographical note, the formulation in van der Pol et al. (2020b) also has a joint group action on state and action space, which is denoted as state transformation Lg : S → S and state-dependent action transformation Ksg : A → A. Table 1 in van der Pol et al. (2020b) outlines state and action transformations for their environments, and all of action transformations are not state-dependent.
Similarly in our case, geometric transformations are usually acting globally on the environments S ×A, thus states and actions are transformed accordingly. We use the factorized form and omit the state-dependency ρA(g; s) of action transformation ρA(g), since the action transformations do not depend on states ρA(g; s) = ρA(g) for all g ∈ G, s ∈ S.
Learning transition with equivariant networks. In this paper we are mainly interested in learning transition functions which are equivariant under symmetry transformations and can be highdimensional.
We apply the idea of learning equivariant transition models in the latent space Z , where Z is the space of symmetric representations, on various environments with different symmetry groups G. We assume we do not explicitly know ρS since S is high-dimensional. We factorize the group representation on state and action S ×A as latent state transformation ρZ(g) ·E(s) and ρA(g; s) · a. In the deterministic case, the transition model can be modeled by G-equivariant networks in latent state Z and action space A:
ρZ(g) · T (E(s), a) = T (ρZ(g) · E(s), ρA(g) · a), (2)
for all g ∈ G, s, s′ ∈ S and a ∈ A.

heading: C LEARNED LATENT REPRESENTATIONS
text: The learned latent embedding z for all states in the evaluation set for Reacher is shown in Figure 4 and the embeddings are factored into irreducible representations. The 2-dimensional representations show a circular pattern, mimicking the rotation of joints. Figure 5 shows a sample observation in pixel space (ground truth) in the top row and its encoded latent embedding in the bottom row. The latent space can have a different orientation than the ground truth. Applying rotations to the low-dimensional embedding yields a latent space traversal with smooth rotations, showing that the learned representations correctly encode the correct symmetries. Figure 6 show learned embeddings for all states in the evaluation dataset for our model (left) and the non-equivariant model (right) when trained on only the up action.

heading: D TRAINING DETAILS
text: D.1 ENVIRONMENTS
2D Shapes & 3D Blocks There are five objects are arranged in a 5× 5 grid and each object can occupy a single cell. Actions are the 4 cardinal directions for each object and an action moves one object at a time, unless it is blocked by the boundaries or by another object. Observations are 50 × 50 × 3 RGB images for both 2D Shapes and 3D Blocks, with pixel values normalized to [0, 1]. The observations in 2D shapes are top down views of the grid and each object has a different color-shape combination. For 3D Blocks, the observations are rendered isometrically with a skewed perspective and each block has a z-height, introducing partial occlusions to the image.
Rush Hour We create a variant of 2D Shapes called Rush Hour. Each object has an orientation and the action is relative to the object’s orientation: {forward, backward, left, right}. This increases the importance of rotational orientation in the environment increasing the significance of equivariance.
Reacher This environment makes a small modification to the original MuJoCo environment Reacher-v2. As we do not consider rewards, we fix the goal position to the position [0.2, 0.2] so
that features related to the goal are ignored. Instead of using the 11-dimensional state, we use pixel observations as images and preprocess them by cropping slightly and downsampling the original 500× 500× 3 RGB image to 128× 128× 3. The previous and current frames are then stacked as an observation to encode velocities. The default camera position gives a slightly skewed perspective, see Table 1.
3D Teapot The 3D teapot enviroment contains images of the Utah teapot model rendered into the 64× 64 grayscale images. The teapot varies in pose which can be described by a coordinate frame z ∈ SO(3). We consider both a small (S) and large (L) action space for this environment. In the small action space, 6 actions may be taken corresponding to multiplication of the pose z 7→ az by
a ∈ SO(3) where a is a rotation of π2π/30 around the x−,y−,or z-axis. In the large action space, actions may be any element a ∈ SO(3).
D.2 MODEL ARCHITECTURES
Symmetric embedding network S For all models and environments except for 3D Teapot, we use CNNs with BatchNorm Ioffe & Szegedy (2015) and ReLU activations between each convolutional layer. For 3D Teapot, the symmetric embedding network maps directly to the latent z space so we use 4 convolutional layers followed by 3 fully connected layers. The output is a 3× 3 rotation matrix. The number of layers for each environment is given in Table 1. For the non-equivariant symmetric embedding networks, we use 32 convolutional filters for every layer and use 8 filters for Reacher and 16 filters for all other environments.
Encoder E The object-oriented environments use 3-layer MLPs with 512 hidden units for the non-equivariant networks and 256 for the equivariant counterparts. There is a ReLU activation after the first and second layers and a LayerNorm (Ba et al., 2016) after the second layer. For Reacher, we use 3 convolutional layers followed by 3 fully connected layers. The 3D Teapot does not have an explicit encoder, i.e. it is the identity function. The output of the non-equivariant encoder is a 2-dimensional vector for 2D Shapes, 3D Blocks, and Rush Hour and a 4-dimensional vector for Reacher. The output of the equivariant encoders for each environment is listed in Table 1.
Transition T The object-oriented environments use GNN transition models where the edge and node networks have the same structure as the encoder (3-layer MLPs). For Reacher and 3D Teapot, the transition model T is a MLP with 512 hidden units for the non-equivariant version and 256 for the equivariant version. Actions are concatenated to the latent z and are input into the transition models which then outputs a z′ of the same dimension as the input z. We use one-hot encoding for discrete actions.
Gram-Schmidt embedding for Teapot transition model In the case of the teapot domain, the transition model is constrained to output an element of SO(3) representing a positively-oriented orthonormal frame. This is achieved by having the network output two vector u, v ∈ R3 and then performing Gram-Schmidt orthogonalization. Only two vectors are necessary since orthogonality and orientation determine the third, after producing two orthonormal vectors u′, v′, the third vector w′ is uniquely determined by the property that it completes a positively-oriented orthonormal frame and can be computed by cross product. In summary,
u′ = u/‖u‖, v′ = v − (u ′ · v)u′
‖v − (u′ · v)u′‖ ,
w′ = u′ × v′, y = [u′ v′ w′].
D.3 DATASETS AND HYPERPARAMETERS
For training, we use 1000 episodes of length 100 as training data for the grid world environments (2D shapes, 3D blocks, Rush Hour), 2000 episodes of length 10 for Reacher, and 100,000 episodes of length 1 for the 3D teapot. For Reacher, the starting state is restricted to a subset of the whole state space, so we perform warm starts with 50 random actions in order to generate more diverse data. The evaluation datasets are generated with different seeds from the training data to ensure that transitions are different.
For the object-oriented environments, we follow the hyperparameters used in (Kipf et al., 2020): a learning rate of 5 × 10−4, batch size of 1024, 100 epochs, and the hinge margin γ = 1. We find that these hyperparameters work well for all other environments, except that Reacher uses a batch size of 256 and mixed precision training was used for both non-equivariant, fully-equivariant, and our method, in order to keep the batch size relatively high for stable contrastive learning. Most experiments were run on a single Nvidia RTX 2080Ti except for 3D Cubes which used a single Nvidia P100 12GB.

heading: E GROUP REPRESENTATIONS
text: We explain the notation and definitions of the different representations of the groups considered in the paper and displayed in Table 1.
The ρstd representation of C4 or D4 on R2 is by 2-by-2 rotation and reflection matrices. The ρstd representation of S5 permutes the standard basis of R5. The regular representation ρreg of G permutes the basis element of R|G| according to the multiplication table of G. The trivial representation of G fixes R as ρtriv(g) · x = x. For D4, ρflip(g) = ±1 is a representation on R depending only on if g contains a reflection. Given representations (ρ1,Rn1) and (ρ2,Rn2) of G1 and G2, (ρ1 ρ2)(g1, g2)(v ⊗ w) = g1v ⊗ g2w gives a representation on G1 ×G2 on Rn1 ⊗ Rn2 .

Review summaries:
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. 

This article assumes knowledge of an underlying group to learn symmetric latent representations, by utilizing equivariant transition models. The paper claims that the basic setup can be applied to a variety of existing equivariant networks (or non-equivariant ones) to improve performance as well as data efficiency. Experimental results are shown on 3 tasks, which are somewhat hand-crafted, in that for each of these tasks the allowed groups are assumed to be known a priori. Some of these experimental results are intriguing but the results themselves beg the question of why one needs all the complex machinery, at least in comparison with other recent efforts which admit a more direct mechanism (e.g. with regard to SO(3) and the teapot rotation task - see comments below).

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. 

This paper proposes to learn an equivariant world model without access to group representation on input space $\rho_S$. To this end, it learns a symmetric abstract state mapping from states $s$ to abstract states $z$ in a space $Z$ with an explicit action $\rho_Z$ of the symmetry group G. It then learns a transition model in latent space that can enforce symmetry using an equivariant neural network.   The problem of learning the equivariant model without assuming any group representation on input space is important and motivated well in the paper. The paper evaluates the model and compares it extensively on various benchmarks.