L EARNING S YMMETRIC R EPRESENTATIONS FOR E QUIVARIANT W ORLD M ODELS
Anonymous authors
Paper under double-blind review
A BSTRACT
Incorporating symmetries leads to highly data-efﬁcient and generalizable world models.
However, characterizing how underlying symmetries manifest in the input space is often difﬁcult.
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.
Speciﬁcally, 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.
I NTRODUCTION
Symmetry has proved to be a powerful inductive bias for improving generalization in supervised and unsupervised learning.
A symmetry group deﬁnes 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 efﬁcient, more sample efﬁcient, 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 difﬁcult 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 efﬁciency 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.  R ELATED W ORK
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) deﬁne disentangled representations based on symmetry, with latent factors considered disentangled if they are independently transformed by commuting subgroups.
Within this deﬁnition, 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 S n permutation symmetry. We learn symmetric representations for groups G and explicitly enforce G -equivariance constraints to latent transition networks.
B ACKGROUND
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 ( g 1 g 2 , s )
= a ( g 1 , a ( g 2 , s ))
for g 1 , g 2 ∈ G, s ∈ S .
A real representation of the group G is a linear group action, given by a map ρ : G → GL n ( 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, ρ ( g 1 ◦ g 2 ) =
ρ ( g 1 ) ρ ( g 2 ) .
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.
Our goal is to learn equivariant world models where the underlying symmetry group of the input space is difﬁcult to compute. Our method gives a general template that can be ﬁt to different equivariant neural networks, symmetries, and data types. We describe the general template and the implementation in examples.
4. M ODEL O VERVIEW
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 (cid:48)
=
T ( s, a ) given state s and action a .
Our goal is to learn T from tuples ( s, a, s (cid:48) ) collected from ofﬂine 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 T Z : 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 (cid:48)
=
T ( s, a ) , the model will then automatically generalize to gs (cid:48)
=
T ( gs, ga ) , enabling improved generalization and sample efﬁciency.
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 difﬁcult 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.
Let ( s, a, s (cid:48) ) be a ground truth transition triplet and s (cid:48)(cid:48) be an incorrect next state s (cid:48)(cid:48)
(cid:54) =
s (cid:48) .
Let z
=
E ( S ( s )) , z (cid:48) = E ( S ( s )) , and z (cid:48)(cid:48) =
E ( S ( s )) , then
[Equation]
Minimizing this loss pushes T ( z, a ) towards z (cid:48) and away from the incorrect sample z (cid:48)(cid:48) .
Symmetric Embedding Network S
We use non-equivariant CNNs for all environments, but the speciﬁc architecture varies (see Appendix D.2).
For example, in the object-centric environments, the symmetry group is G = C 4 × S 5 , 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 S 5 by permuting the 5-dimensional axis and of C 4 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 = C 4 × S 5 , the equivariant encoder E is shared over all 5 objects and uses group convolution over C 4 (Cohen & Welling, 2016a), thus achieving C 4 × S 5 -equivariance.
The transition function T is implemented as a GNN (message-passing neural network) with edge and node networks which use C 4 -convolutions for their linear layers.
Since GNNs are S 5 equivariant by deﬁnition and the linear layers within the GNNs are C 4 -equivariant, this is C 4 × S 5 -equivariant.
For other implementations, see Table 1 and Appendix D.2. [Figure]
For these environments, we consider a case where the action space
A is the same as the symmetry group G = A = SO(3) , similar to Quessard et al. We choose to model the latent space Z using SO(3) as well, which is not a linear group representation, but a set with a group action.
Symmetric Embedding In the case of G = 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 ∈ R 3 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) (cid:110) ( R 3 , +) -equivariant method which works over point clouds.
Here z ∈ SO(3) and a ∈ SO(3) .
We consider z as 3 points in R 3 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
T Z ( z, a ) =
az which is also equivariant,
[Equation]
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 .
4. G ENERALIZING OVER THE MDP A CTION S PACE
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 (cid:48)
⊂
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 (cid:48)
=
A , i.e. every MDP action is Gtransformed version of one in A (cid:48) .
We assume that S is approximately equivariant after training, which, as T and E are constrained to be equivariant, implies T S ( s, a )
=
T ( E ( S ( s )) , a ) is equivariant.
We also assume low error for T on the restricted action set ( s 1 , a (cid:48) , s 2 ) ∈ S × A (cid:48)
× S .
Then given ( s 1 , a, s 2 ) ∈ S ×
A × S , there exists g ∈ G such that a = ρ A ( g ) · a (cid:48) .
Let s (cid:48)
i =
ρ S ( g − 1 ) s i .
Then
{+T S ( s, a ) ≈ T S ( ρ S ( g ) · s (cid:48) 1 , ρ A ( g ) · a (cid:48) ) ≈ ρ S ( g ) · T S ( s (cid:48) 1 , a (cid:48) ) ≈ ρ S ( g ) · s (cid:48) 
 ≈ s 2 .+}
If performance is good for A (cid:48) and S has low equivariance error, then performance will be good for A .
E XPERIMENTS 5. S ETUP
We choose ﬁve environments with varying symmetries to evaluate our models.
The ﬁrst three environments 2D Shapes, 3D Blocks, and Rush Hour are grid-worlds with ﬁve 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-speciﬁed 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 deﬁne a ρ S acting on the 2D image space which approximates the true group action.
5. M ODEL ARCHITECTURES AND TRAINING
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 (cid:48) ) 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 speciﬁc implementation details are provided in Appendix D.2.
5. M ETRICS
In order to evaluate the performance of our model in latent space without reconstruction, we use standard ranking metrics from Kipf et al.
(2020), modiﬁed versions of these metrics to adapt to continuous state spaces, and two metrics for evaluating the equivariance of the learned model. 5.3. A CCURACY M ETRICS
Hits, Hard Hits, and MRR
The evaluation samples are ranked according to the pairwise Ldistance 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 (cid:48) n of states that are close to the true next state s (cid:48)
(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. [Table] for the ρ deﬁnitions.
5.3. E QUIVARIANCE M ETRICS
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 (cid:48) ) 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.
Speciﬁcally, the equivariance error of the symmetric embedding is calculated as
[Equation]
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 deﬁned.
We therefore propose a proxy for the equivariance error using invariant distances.
For a pair of input states s, s (cid:48) , an equivariant model f will have the same distances (cid:107) f ( s ) − f
( s (cid:48) )
(cid:107) and (cid:107) f
( gs ) − f ( gs (cid:48) )
(cid:107) 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,
[Equation]
The distance invariance error is computed as
[Equation]
We evaluate both the symmetric embedding network ( DIE(S) ) and entire model ( DIE(model) ).
5. M ODEL PERFORMANCE COMPARISON
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 [Table]
[Table] 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 ﬁgures 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 (cid:48) 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.
5. G ENERALIZATION FROM LIMITED ACTIONS
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 modiﬁcations to the episode lengths. For Reacher, the action space represents
[Table] [Table] [Table] 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 ).
C ONCLUSION AND F UTURE W ORK
We demonstrate a ﬂexible 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 difﬁcult to isolate symmetry.
E THICS S TATEMENT
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 beneﬁt or harm, our method enables both options.
R EPRODUCIBILITY S TATEMENT
We will open-source our code, including all models and data generation scripts, thus allowing all experiments to be fully reproduced. A O UTLINE
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 deﬁnition in Section E. B S ETUP : E QUIVARIANT W ORLD M ODELS
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 =
(cid:104)S , A , T, R, γ (cid:105) , 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 reﬂections and rotations.
Ravindran (2004) study group symmetry in MDPs as a special case of MDP homomorphisms.
Symmetry in MDPs.
Symmetry in MDPs is deﬁned 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 ﬁnding 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 =
(cid:104)S , A , T, R, γ (cid:105) to another M = (cid:104)S , A , T, R, γ (cid:105) which needs to preserve the transition and reward structure (Ravindran, 2004).
The mapping h consists of a tuple of surjective maps
h = (cid:104) φ, { α s | s ∈ S}(cid:105) , 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 (cid:48) )
| φ ( s ) , α s ( a )) =
(cid:80) s (cid:48)(cid:48) ∈
φ − 1 ( φ ( s (cid:48) ))
T ( s (cid:48)(cid:48) | s, a ) , (2) and the reward function is also preserved R ( φ ( s ) , α s ( a ))
=
R ( s, a ) , for all s, s (cid:48) ∈ S and for all a ∈ A .
An MDP isomorphism from an MDP M to itself is call an automorphism of M .
The collection of all automorphisms of M along with the composition of homomorphisms is the automorphism group of M , denoted Aut( M ) .
We speciﬁcally care about a subgroup of G ⊆
Aut( M ) which is usually easily identiﬁable from environments a priori and thus we can design appropriate equivariant network architectures to respect it, such as C 4 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 deﬁnition, when an MDP M 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:
[Equation]
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 ))
(cid:55)→ ρ 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 L g : S → S and state-dependent action transformation K sg : 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 :
[Equation]
for all g ∈ G , s, s (cid:48) ∈ S and a ∈ A .
The learned latent embedding z for all states in the evaluation set for Reacher is shown in Figureand 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.
D T RAINING DETAILS D. E NVIRONMENTS
2D Shapes & 3D Blocks There are ﬁve 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 signiﬁcance of equivariance.
Reacher This environment makes a small modiﬁcation to the original MuJoCo environment
Reacher-v2 .
As we do not consider rewards, we ﬁx the goal position to the position [0 . 2 , 0 . 2] so [Figure] 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
(cid:55)→ az by
[Figure]
D. M ODEL 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 ﬁlters for every layer and use 8 ﬁlters for Reacher and 16 ﬁlters 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 ﬁrst 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 (cid:48) of the same dimension as the input z .
We use one-hot encoding for discrete actions. [Figure]
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 ∈ R 3 and then performing Gram-Schmidt orthogonalization.
Only two vectors are necessary since orthogonality and orientation determine the third, after producing two orthonormal vectors u (cid:48) , v (cid:48) , the third vector w (cid:48) is uniquely determined by the property that it completes a positively-oriented orthonormal frame and can be computed by cross product.
In summary, [Equation] D. D ATASETS AND H YPERPARAMETERS
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 ﬁnd 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.
E G ROUP R EPRESENTATIONS
We explain the notation and deﬁnitions of the different representations of the groups considered in the paper and displayed in Table 1.
The ρ std representation of C 4 or D 4 on R 2 is by 2-by-2 rotation and reﬂection matrices.
The ρ std representation of S 5 permutes the standard basis of R 5 .
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 ﬁxes R as ρ triv ( g ) · x = x .
For D 4 , ρ ﬂip ( g ) = ± 1 is a representation on R depending only on if g contains a reﬂection.
Given representations ( ρ 1 , R n 1 ) and ( ρ 2 , R n 2 ) of Gand G 2 , ( ρ 1 (cid:2) ρ 2 )( g 1 , g 2 )( v ⊗ w ) =
g 1 v ⊗ g 2 w gives a representation on G 1 × G 2 on R n 1 ⊗ R n 2 .