Learning Symmetric Representations for Equivariant World Models
Authors: Jung Yeon Park
Ondrej Biza
Linfeng Zhao
Jan-Willem van de Meent
Robin Walters
Venue: ICLR 2022 Submitted (2022)
Keywords: equivariant
symmetry
contrastive loss
world models
transition
representation theory
generalization

Abstract: Encoding known symmetries into world models can improve generalization. However, identifying how latent symmetries manifest in the input space can be difficult. As an example, rotations of objects are equivariant with respect to their orientation, but extracting this orientation from an image is difficult in absence of supervision. In this paper, we use equivariant transition models as an inductive bias to learn symmetric latent representations in a self-supervised manner. This allows us to train non-equivariant networks to encode input data, for which the underlying symmetry may be non-obvious, into a latent space where symmetries may be used to reason about outcomes of actions in a data-efficient manner. Our method is agnostic to the type of latent symmetry; we demonstrate its usefulness over $C_4 \times S_5$ using $G$-convolutions and GNNs, over $D_4 \ltimes (\mathbb{R}^2,+)$ using $E(2)$-steerable CNNs, and over $\mathrm{SO}(3)$ using tensor field networks. In all three cases, we demonstrate improvements relative to both fully-equivariant and non-equivariant baselines. 

Under review as a conference paper at ICLR 2022
LEARNING SYMMETRIC REPRESENTATIONS FOR EQUIVARIANT WORLD MODELS
Anonymous authors
Paper under double-blind review
ABSTRACT
Incorporating symmetries leads highly data-efficient and generalizable world models. However, characterizing how underlying symmetries manifest in the input space is often difficult. We provide a method map an input space (c.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 equívariant neural networks to operate previously accessible domains. We evaluate the effectiveness of this approach in domains with distinct forms of underlying symmetry. many cases, we demonstrate improvements relative to both fully-equivariant and non-equivariant baselines.
INTRODUCTION
Symmetry has proved be 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 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 that it not always obvious how symmetry group acts pairs of images in Figure [ On the left, we have a pair of MNIST digits, for which a 2D rotations in pixel space should induce corresponding rotation in feature space. Here it is possible achieve state-of tftt--ff- accuracy using an ((2)-equivariant network (Weiler Cesa 2019). By contrast, exploiting the underlying symmetry is much more challenging for the pair of images 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 particular rotation.
input data set. As example, let us consider the two
CAPVANA
CARVANA
Figure 1: On MNIST, the rotation is easy to compute, allowing for the use of equivariant models. The rotation difficult compute for the car, making it harder to apply equivariant methods (Carvanal 2017)
In this paper, we consider the task of learning symmetric representations of data domains where transformations cannot be hard-coded, which is to say that the group action 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 characterize, into 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.

---

Under review as a conference paper at ICLR 2022
As concrete instantiation of this idea, 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 contrastive loss. The underlying intuition behind this approach is that the symmetry group of the transition model acts as inductive bias that guides the embedding network possible. At the same time, incorporating symmetries into the transition dynamics has the potential to improve both data efficiency and out-of-distribution generalization.
representation that is as equivariant as
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 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 follows:
We introduce a meta-architee archttrrci for equivariant world models with symmetric embedding networks, which can be trained end-to-end using triples from replay buffer by minimizing a contrastive loss.
We demonstrate that this meta- aretttercie can be used to learn world models with a variety of equivariances in a self-supervisee manner.
Morcover, we show that models that have been trained using only a subset of all input actions can generalize 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 various groups across 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 al. 2018), or convolution in the Fourier domain (Esteves al. 2017).
They have been applied to many data types such as gridded data (Weiler & Cesal 2019). spherical data (Cohen et al. 2018), point clouds (Dym & Maron (020), 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 (2044) 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 al. 12021) similarly learn basis for Lie algebra generating symmetry group while simultane- ously learning parameters for the equivariant convolution over this symmetry group. Benton et al. (2020) propose 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 group element, whereas our transition model is given by an equivariant neural network.
2

---

Under review as a conference paper at ICLR 2022
Structured Latent World Models World models learn state representations by ignoring unneces- sary information unrelated to predicting environment dynamics. Such models are frequently used for dineedimmnsionnl image inputs, and usually employ (1) reconstruction loss (Ha & Schmidhuber 2018 Watter et al., 2015; Hafner et al.) 2019 (020) or (2) constrastive loss. Contrastive loss is known 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 al. 2020a), thus we use it for training our models. For example, Kipp al.
(2020) learn object factored representations for structured world modeling with GNNs, which respect
S, 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 set G together with composition map o: G x G + G. The group must contain an identity 1 E G and each element g E G must be invertible. An action of the group G element of the group to a permutation Ty E Perm(S) of the elements of S. Composition of group elements is compatible with the action such that a(g) 8) = a(g,, a(g2, 8)) for 91, G. E S. A real representation of the group G is a linear group action, given by a map p: G - GL,, (R) which maps each element of G to an invertible n X n matrix. The multiplication table of these matrices must match that of the abstract group elements under composition. That is, p(91 o g2) = ((g))
set S is map a: G - Perm(S) mapping each
See Hall (0033) 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 px and py acting on X and Y respectively, we say that f is equivariant if, for all x E X, g E G, we have f(xx (g) x) py (g) fxx.. This means that if the input 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 om-llinarities.
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 x A + S which outputs the next state s' -Tss given state and action a. Our goal is learn T' from tuples (s, a, s') collected from offline data. As the space S x 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 symmetries with group representations ps and PAA. We assume that (,4 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 (ps(g)) s, p사(g) a) ps(g) T(s, a). We wish to enforce this neural network model for T. trained to predict s' = T(s, a), the model will then automatically generalize to gs' T(gs, ga), enabling improved generalization and sample efficiency.
3

---

Under review as a conference paper at ICLR 2022
Symmetric Fabbeding
St.
Eneader
Transtionn Model
S
Ut
E
zt
Loss
T
t++++
z++1
a+ Force Joint 1 of Foree Joint 2 of
G-equivariant
Figure 2: Diagram of the model architecture of our G-equivariant world model. The features in red have an explicit G-action p. The networks red are G-equivariant. The example Reacher input has G = D2 symmetry. The MDP actions have G-representation type Ptlp meaning they reversed in sign by reflections and unaltered by rotations. The Symmetrie Embedding Network a CNN and the Encoder and Transition model are E(2)-CNNS with fiber group D4.
Current methods for constructing model classes of equivariant neural networks require that ps and PA are known. In our case, however, although we assume the symmetry group G and group action PA. we assume that ps is explicitly known. In many environments, is a pixel-level input and the transformation Ps may be very difficult to describe. Thus, we cannot directly enforce symmetry for T using an equivariant neural network.
Meta-Architecture To learn equivariant world model without access to ps, learn symmetric abstract state mapping from states 8 to abstract states z in a space Z with an explicit action pz of the symmetry group G. We then learn transition model in latent space where we 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: + V. This is a non-equivariant neural network acts as feature extractor and maps images into space y with simpler explicit symmetry group action Py. 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: x A + Z. We explicitly enforce E and 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, 8') be a ground truth transition triplet and s" be an incorrect next state 8" * s. Let E(S(s)), z' E(S(s)), and z" E(S(s)), then
((s,a,8',8") = ([T(2, a) 2'1l omax(3 [IT(2, a) (l)
Minimizing this loss pushes T(2, a) towards z' and away from the incorrect sample =".
Symmetrie Embedding Network S We use non-equivariant CNNs for all environments, but the specific architecture varies (see Appendix D.2h. For example, in the object-centric environments, the symmetry group is G = C4 x Sg. the cyclic group order acting by T/ 7/2-rotations and the permutation group the 5 objects. The output has shape [B, C, 5, 4) and carries an action Py of S3 by permuting the 5-dimensional axis and C4 by cyclically permuting the 4-dimensional axis.
Equivariant Encoder F 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 Sr. the equivariant encoder E is shared over all objects and uses group convolution over C4 (Cohen & Welling 2016a), thus achieving C4 x Sæ-equivariance. The transition function T is neural network) with edge and node networks which use Ci-convolutions for their linear layers. Since GNNs are Sg equivariant by definition and the linear layers within the GNNs are C4 Cy-qquivariant, this is C4 x Ss-equivariant. For other implementations, Table and Appendix D.2
implemented as a GNN (message-paseee
4

---

Under review as a conference paper at ICLR 2022
4.2 (3( STRUCTURED SYMMETRIC EMBEDDINGS
S S A
g (2))
S
T
Rod88
(8.a) [sss.o0n" "t
g- (s) = (gs)
T
noa( na((ana)))
Figure 3:
of the transition function for SO(3) object manipulation.
For these environments, we consider case where the action space A is the same the symmetry group G = A = 50(3), similar to Quessard et al. (020). We choose model the latent space 2 using (3) as well, which is not linear group representation, but set with group action.
the case of G (3)3) symmetry, we expect the symmetric embedding network to detect the pose of object = in 3D. We omit the equivariant encoder (by setting E = id) and instead use a two-part Symmetric Embedding Network that directly encodes y = 2 using a down- wwwnssamplnng CNN whose output is passed an MLP, and converted to an element of SO(3).
Symmetric Embedding
To force the output of the symmetric embedding network y to be an element of SO(3), we have the last layer output vectors u, R and perform Gram-Schmidt orthogonalization to construct 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.
We implement an equivariant transition model using Tensor Field Networks (Thomas al. 2018 Geiger et al. 2020). This is an SO(3) (IR,+ -eqquivaaan method which works over point clouds. Here 2 E SO(3) and a E SO(3). We consider z as 3 points in R&& 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 -0epee-epressentaation (A(g) a gag
Transition
The MDP action a SO(3) is rotation matrix and the latent state : E SO(3) is positively-oriente orthogonal coordinate frame. Though Z = A = SO(3), these different semantics lead to differing G SO(3) actions with pz (2) g z but pz (a) gag If 2 is correctly learned, then the ground truth latent transition function can be represented by simple matrix multiplication (z(2, a) = az which is also equivariant,
Tz(pz(g) (4())))))
(gog (g=>) gaz = (2
This method, which we label MatMul, is similar to the latent transition model used in Quessard et al. 22000) except in our framework the groundtruth a e SO(3) is provided to aid learning 2.
4.3 GENERALIZING OVER THE MDP ACTION SPACE
Although the state does not have a known group action ps. the MDP action does have known PA- 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'' C 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 PC A' A, i.e. every MDP action is G- transformed version of one in A'. We assume that S is approximately equivariant after training, which, as T and are constrained to be equivariant, implies Ts(x, a) T(E(((s)), a) equivariant. We also assume low error for I on the restricted action set (51, a', 82) E E x A' x S. Then given (#1, a, s2) E S x x S, there exists E such that a = pa(g) a'. Let s = ps(g 1)s, Then Ts(s, a) (s(ps(g) 21 (9( d') ps(g) Ts( a') R ps(g) 8/2 82-
5

---

Under review as a conference paper at ICLR 2022
If performance is good for A' and S has low equivariance error, then performance will be good for A.
5 EXPERIMENTS
5.1 SETUP
We choose five environments with varying symmetries 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. 2020j. Rush Hour is a variant of 2D Shapes where objects can move relative to their orientation. We consider symmetry to 4/2 rotations and object permutations. We also evaluate in a continuous control domain, the Reache
MuJoCo environment, with rotational and translational symmetries. The last domain is of rotating 3D teapot, where action is an element of (3()). We consider two action spaces: a small discrete action space (S) with rotations of 20 in SO(3) and 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.I
We compare three types of models: (a) non-equivariant model with no enforced symmetry, (b) fully-equivariant model with a mis-specified symmetry where ps is a simple transformation of the pixels, and (c) method. For 3D Teapot, we forgo the fully equivariant baseline as it is hard to define a Ps acting on the 2D image space which approximates the true group action.
5.2 MODEL 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 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 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 Table [
objects for the grid world environments
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, states and thus overlap between training and evaluation datasets is highly unlikely.
have either combinatorially large state space (with objects) continuous
As equivariant networks have more parameters than the non-equivariant counterparts, reduce the number of hidden dimensions accordingly to keep the number of parameters approximately constant for all models. The Adam (Kingma & Ba 2014) optimizer used for all experiments. All other specific implementation details are provided Appendix D.2
5.3 METRICS
In order to evaluate the performance of model in latent space without reconstruction, 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.
5.3.1 ACCURACY METRICS
Hits, Hard Hits, and MRR The evaluation samples are ranked according to the pairwise L distance of the predicted next states and the true next states (both are encoded in latent space). Hits at Rank h (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 variant of Hits at Rank & (HH k) where generate negative samples s/, of states that are close to the true next state (see Appendix for more details) and count the number of times that the distance to the positive sample was lower than the distance 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.

---

Under review as a conference paper at ICLR 2022
Environment
2D Shapes & 3D Blocks
Rush Hour
Reacher
3D Teapot
Observation
50x50x3
50x50x3
128x
64x64x1
(df d R (joint forces)
Action
up.right,down,left
fwd.Jeft.hack.rght
50(3)
C4 x St ( pot.; obj. perm.)
C4 x Ss ( rot.; obj. perm.)
D₂ (R", +) & rot: flip; trans.)
Symmetry
SO(3)
(Pood @ Prog. RR) BiPaa, R²)
Z-rep: pz
(Pad. RR2) 8 (Paid, RP)
(prg B4 & puiv
(matrix mult.)
A-rep: PA
(Preg. R.) (Przds RD)
(Peite RJ*
(pmp. Ry
gag
(conjugation)
2-layer CNN (2D)
Non-Equ. Extractor
2-layer CNN
7-layer CNN
4 conv, FC layers
4-layer CNN (3D)
F((2)-covv,
Equ. Encoder
MLP C.-conv
MLP Cl-conv
1d. (Nonc)
Dg-FC layers
GNN + C. -ony Cohen Welling 10
MLP Tensor Field
Equ. Transition
GNN Ca-conv
MLP CI2)-CN
Geiger et al. 2020 mult.
Weiler & Cesa 2019
Scurselli et 2008
Table The symmetry and implementation for each domain. See Appendix for the p definitions.
5.3.2 EQUIVARIANCE METRICS
Equivariance Error EE In order to analyze the equivariance of each model, we generate 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 ps(g) cannot be computed using g and 8, our synthetic datasets allow us to render both & and ps(g) during generation. Specifically, the equivariance error of the symmetric embedding is calculated
EE-
Es.g (g) Ss) (ps s)]]
Distance Invariance Error DIE The above equivariance error always be applied the sym- metric embedding network when its output space is spatial and we can manually perform group actions on the outputs. However it cannot be applied the latent space Z in the case of non-equivariant models since the group action on the latent space pz cannot be meaningfully defined.
We therefore propose a proxy for the equivariance error using invariant distances. For a pair of input states s', an equivariant model f will have the same distances [s] - [[s)) and (gs) f (gs')) I assuming the action of G is norm preserving as is for all transformations considered in the paper.
Due to the linearity of the action, [g(s)) f(gs') = [[gf(s) - gs('') (f - f(s ))ll =
[I(f(s) f(s'))ll- The distance invariance error is computed
DIE = E.,''s -g [sllf(s) f(s') -If(gs))) f(g'')lll] III]
We evaluate both the symmetric embedding network (DIE(S)) and entire model (DIE(model)).
5.4 MODEL PERFORMANCE COMPARISON
The results are shown in Tables 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 higher H I on Rush Hour but has slightly lower MRR than either the fully cquivariant and our model on Reacher. Surprisingly, the fully equivariant model performs very well even when the group action the input space Ps is not correct. Due the skewed perspective, can see that the simple pixel-level transformation maps training data to -iitt----ittiittionn images which never seen by the model. We hypothesize that equivariance does not hamper its performance
7

---

Under review as a conference paper at ICLR 2022
Hits@) I
MRR
DIE(S)
DIE (model)
Model
EE(S)
(10 step. %)
(10 step, %)
(x10-
(10 step. x10-
None
=4.3
99.0
a.80
3.30
3.85.12.0
3D Blocks
Full
99.9
0.82=0.5
3.36
5.54+4.8
Ours
99.9
100.0.0
0.86.10.4
3.3211.
3.16.115
None
95.9
97. 97.9
0.37=0.07
1.17
26.0=7.13
Rush Hour
Full
92.0=a.1
95.9
0.00=0.00
0.00
0.05
Ours
93.3=a.7
96.5#t2.0
0.26=0.0m
4. 4.12
10.0=4.45
Table 2: Model Performance on 3D Blocks and Rush Hour.
H@10 10
MRR
DIE(S)
DIE (model)
HHR
EB(S)
Model
BB(S)
stop,
(1 step, %)
(1 step, %)
(x10 ?)
(1 step)
None (S)
0.1
2.34
MatMul (S)
31.6s1.2
2.39:000
None
100
88.3
1.26
4.53
0.56
TPN (8)
39.7a1.
2.22.001
Reacher
Full
100+0.0
95.321.
1.19
3.51
0.39
7.4
2,409
MaiMul (L)
100
0.05:00
Ours
10010.0
94.1
1.29.10.
4.05.1077
0.52.001
TFN (L)
4.9.0.7
2,41:00
Reacher
3D Teapot
Table 3: Model performance on Reacher (left) and 3D Teapot (right) environments. For the 3D Teapot models, None is the non-equivariant model, Matmul is the matrix multiplication model and Equiv the version using Tensor Field Networks. (S) denotes the small discrete action space with 6 rotations of 2x (3) and (L) denotes any rotation in SO(3). 30
on training data, but only constrains its extrapolation capabilities to ontto--distribution samples. In 3D Teapot, 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 TEN fails with large actions (L). We hypothesize that the difference is caused by the contrastive loss function; small actions lead to local minima where all states get mapped to close-by by latent states. For the equivariance metrics, our model outperforms the -quivarriia model in all domains on the case for accuracy metrics, the fully equivariant model performs surprisingly well with possibly the same reason outlined above.
while performs similarly on EE(S) and DIE(S). As was
Visualization of latent embeddings We visualize the latent embedding for our model to qualita- tively analyze what kind of representations are learned, All figures are provided in the Appendix for space. Figure
all the learned embeddings for Reacher for all observations in the evaluation set and shows sample transition in both pixel and latent space. The encoded current state is highlighted red and the encoded next state is highlighted 2' highlighted in orange. which we factor into irreducible representations (irreps) before visualizing (see Hall (2003)). Some irreducible representations are I-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.5 GENERALIZATION FROM LIMITED ACTIONS
We now train on 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 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 similar modifications to the episode lengths. For Reacher, the action space represents
8

---

Under review as a conference paper at ICLR 2022
H@l (10 step, %)
MRR (I step, %)
DIE(S)
DIE (model)
EE(S)
(x10*)
(10 step)
CNN Ours/Full
2.8+
5.3+0.4
0
0=0
0.19+0.05
100.0.00
99.9
0.10
0.:
0.1
Table Generalization results for 2D Shapes trained using only the up action and evaluated on all actions. Due to the simplicity of the environment, simple CNN turns out to be equivariant, so both the baseline CNN and Ours/Full have an equivariant symmetric embedding network.
H@I (10 step, %)
MRR (10 step, %)
DIE(S) (x10-2)
DIE (model) (10 step, x x10-3)
EF(S)
None
12.3
61.8213
0.98
3.64
181
Full
83.7
86.0.31
0.81
3.32.1.11
14.8 3.65.11.6
Ours
99.9 10.0
100
0.96
5.0.1..7
Table 5:
Generalization results for 3D Blocks with limited actions. The training set only contains the up, right and down actions; the evaluation set contains all four actions.
H@ 10
MRR
DIE(S) (x10-
DIE (model) (1 step, x10-
EE(S)
(I step, %)
(I step, %)
None
86.5+3.0
50.6
1.22+#011
6.54+1.9
6.95#1.5
Full
89.4+11
61.8+13
1.18+0.1
3.62+0.8
4.87
Ours
90.8144.5
59.4
1.28
5.45:10.9
4.95
Table 6: Reacher with limited actions. The models were trained on data where the second joint is constrained to be positive and evaluated on unconstrained data.
joint actuation forces E [-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 6 show results for 2D Shapes, 3D Blocks, and Reacher respectively. We see that our 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@ 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-o--distribution samples. these limited actions experiments, the fully equivariant model cannot extrapolate correctly and achieves lower performance than our model.
method successfully generalize
Figure 5 shows embeddings for all states in the evaluation dataset for our model and the non- cquivariant model trained on only the up action. Our model shows a clear x 5 grid, while the non-equivariant model learns degenerate solution (possibly encoding only the row index x)).
CONCLUSION AND FUTURE WORK
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 variety of domains and equivariant neural network architectures. Our methods confer some of the advantages of equivariant neural networks in situations where they did previously apply, such 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 isolate symmetry.

---

Under review as a conference paper at ICLR 2022
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 be fully reproduced.
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mare-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. arXiv preprint arXiv: (tti::/990600222000 2019.
Brandon Anderson, Truong-Son Hy, and Risi Kondor. Cormorant: Covariant molecular neural networks. arXiv preprint 1eXivv1906.04000 2019.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer normalization, 2016.
Erkao Bao and Lingi Song. Equivariant neural networks and equivarification. arXiv preprint 17X1v:1906..0777,, 2019.
Gregory Benton, Mare Finzi, Pavel Izmailov, and Andrew Gordon Wilson. Learning invariances in neural networks. arXiv preprint arXiv:2010.1188 2020.
Alexander Bogatskiy, Brandon Anderson, Jan Offermann, Marwah Roussi, David Miller, and Risi Kondor. Lorentz group equivariant neural network for particle physics. In International Conference on Machine Learning. pp. 992-1002. PMLR, 2020.
Carvana. Carvana Image Masking Challenge, 2017. URL htttp::/kasgle..com
carvana maskinmmskkngg challll
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Taco S. Cohen and Max Welling. Group equivariant convolutional networks. International 2016a.
conference on machine learning (ICML), pp.
Taco S. Cohen and Max Welling. Stecrable CNNs. arXiv preprint arXiv:1612.08449 2016b.
Taco S. Cohen, Mario Geiger, Jonas K ohler, and Max Welling. Spherical CNNs. In International Conference on Learning Representations (ICLR), 2018.
Taco S. Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolu- tional networks and the icosahedral CNN. In Proceedings of the 36th International Conference on Machine Learning (ICML), volume 97, pp. 1321-1330, 2019.
Nima Dehmamy, Robin Walters, Yanchen Liu, Dashun Wang, and Rose Yu. Automatic symmetry discovery with lie algebra convolutional network. arXiv preprint arXiv:2109.07103, 2021.
Nadav Dym and Haggai Maron. On the universality of rotation equivariant point cloud networks. arXiv preprint 0rX1::2110002449, 2020.
Carlos Esteves, Christine Alanc-Blanchette, Xiaowei Zhou, and Kostas Daniilidis. Polar transformer networks. arXiv preprint arXiv: 1709.01889 2017.
Luca Falorsi, Pim de Haan, Tim R Davidson. Nicola De Cao, Maurice Weiler, Patrick Forré, and Taco Cohen. Explorations in homeomorphic variational auto-encoding. arXiv preprint arXiv: 1807 04689, 2018.
10

---

Under review as a conference paper at ICLR 2022
Mario Geiger, Tess Smidt, Alby M., Benjamin Kurt Miller, Wouter Boomsma, Bradley Dice, Kos- tiantyn Lapchevskyi. Maurice Weiler, Michat Tyszkiewicz, Simon Batzner, Martin Uhrin, Jes Frellsen, Nuri Jung, Sophia Sanbom, Josh Rackers, and Michael Bailey. Euclidean neural networks: e3nn, 2020. URL htt.::tttt:///ttmm org/ 0 5281 /znodo 5292912
David Ha and Jürgen Schmidhuber. World models. arXiv preprini arXiv: 10122, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193,, 2020.
Brian C Hall. Lie groups, Lie algebras, and representations: an elementary introduction, volume 10. Springer, 2003.
Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerehner. Towards definition of disentangled representations, 2018.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In Francis Bach and David Blei (eds.). Proceedings of the 32nd International Conference Machine Learning, volume 37 of Proceedings of Machine Learning Research, pp. 448-456 Lille, France, 07-09 2015. PMLR.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint 1411::141.669800 2014.
Thomas N. Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020.
Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution neural networks to the action of compact groups. In Proceedings of the 35th International Conference on Machine Learning (ICML), volume 80, pp. 2747-2755, 2018
Haggai Maron, Or Litany, Gal Chechik, and Ethan Fetaya. On learning sets of symmetric elements. In International Conference on Machine Learning, pp. 6734-6744. PMLR, 2020.
Shravan Matthur Narayanamurthy and Balaraman Ravindran. On the hardness of finding symmetries in markov decision processes. In Proceedingy of the 25th international conference on Machine learning, pp. 688 695, 2008.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.0378, 2018.
Robin Quessard, Thomas Barrett, and William Clements. Learning disentangled representations and group structure of dynamical environments. H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 19727-19737 Curran Associates, Inc., 2020. URL http::/://///cceeeeeeeeeee URttttpp::/pooeeeeeeoo wtwp:eeweddis.......eeeeeeeeeeee neurips........... cc/
pdf
Balaraman Ravindran. An algebraic approach abstraction in reinforcement learning. University of Massachusetts Amherst, 2004.
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. JEEE transactions on neural networks, 20(1):61-80, 2008.
Aravind Srinivas, Michael I Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint arXiv:2004.00 040000400436,.....00000010 2020.
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang. Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-an and translation -qquvavainnn neural networks for 3d point clouds. arXiv preprint arXiv: arXiv: 8888000022229 2018

---

Under review as a conference paper at ICLR 2022
Elise van der Pol, Thomas Kipf, Frans A Olichoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. arXiv preprint arXiv:2002.11963 2020a.
Elise van der Pol, Daniel E Worrall, Herke van Hoof, Frans A Olichoek, and Max Welling, Mdp homo- morphic networks: Group symmetries in reinforcement learning. arXiv preprint arXiv:2006. 16908, 2020b.
Robin Walters, Jinxi Li, and Rose Yu. Trajectory prediction using equivariant continuous convolution. arXiv preprint arXiv:2010. 13344, 2020.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Gareett (eds.), Advances in Neural Information Processing Systems, volume 28. Curran Associates, Inc., 2015. URL UtL///tttps///rrcceeii.oo neurip....e..s.....p............ cc/
Maurice Weiler and Gabriele Cesa. General E(2)-qquivariant steerable CNNs. In Advances in Neural Information Processing Systems (NeurIPS), pp. 14334-1423 2019.
Maurice Weiler, Fred Hamprecht, and Martin Storath. Learning steerable filters for rotation equivariant CNNs. Computer Vision and Pattern Recognition (CVPR), 2018.
Daniel E Worrall, Stephan J Garbin, Daniyar Turmukhambetov, and Gabriel J Brostow. Harmonic networks: Deep translation and rotation equivariance. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5028-5037, 2017.
Allan Zhou, Tom Knowles, and Chelsea Finn. Meta-learning symmetries by reparameterization. arXiv preprint arXiv:2007.02933, 2020.
Martin Zinkevich and Tucker Balch. Symmetry in markov decision processes and its implications for single agent and multi agent learning. In Proceedings of the 18th International Conference on Machine Learning. Citeseer, 2001.
12

---

Under review as a conference paper at ICLR 2022
A OUTLINE
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 Section D. and D.2 We further explain the notation and definition in Section E
B SETUP: EQUIVARIANT WORLD MODELS
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 AS, R,7), with state space S, action space A, (deterministic) transition function I S x A + S, reward function R : S x A + R. and discount factor y [0, 1].
Symmetry can appear in MDPs naturally (Zinkevich & Balch 2001 Narayanamurthy & Ravin dran 2008 Ravindran 2004 van der Pol et al.) | (020b), which we can exploit using equivariant networks. For example, van der Pol et al. (2020b) study geometric transformations, such as reflec- tions and rotations. [Ravindran (0044) 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
(u) is an MDP homomorphism h -M - M that maps M to itself and thus preserves its structure. Zinkevich & Balch (2001) show the invariance value function for an MDP with symmetry. Narayanamurthy & Ravindran (2008) prove that finding exact symmetry in MDPs is graph isomorphism complete.
Ravindran 2004) provide comprehensive overview of using MDP homomorphisms for state abstraction and study symmetry in MDPs as 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 on policy learning, which needs to preserve both transition and reward structure and thus has optimal value equivalence (Ravindran) 2004).
MDP homomorphism h M M is mapping from one MDP M (S, A, T, 7) to another M (S, AS R, which needs to preserve the transition and reward structure (Ravindran) 2004 The mapping consists of tuple of surjective maps h = (ò, {a, E S}), where o S + S is the state mapping and ag A + A is the state- dependent action mapping. The mappings are constructed to satisfy the following conditions: (1) the transition function is preserved T (ø (s') d(s), (a()) = E<rea' T (s" s, a), (2) and the
More formally,
reward function is also preserved R (d(s), (a)) = R(s, a), for all 8, S and for all a E A.
An MDP isomorphism from 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 (ut.
We specifically care about a subgroup of G C Aut(M) which is usually easily identifiable from environments priori and thus 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 x A +
design appropriate equivariant network architectures to
Equivariant transition. By definition, when MDP M has symmetry, any state-action pair (s, a) and its transformed counterpart (ps(g) s, PA(g) a) are mapped to the same abstract state-action pair by h (M) (ó(s), as (a)) (ø(gs), ags (ga)), for all s E S, A,g G. Therefore, the transition function T x A + S should be G-equivariant:
T(ps(g) s, (g())) a) = ps(g) a(s,
(I)
for all s S,a E A,g.G.
13

---

Under review as a conference paper at ICLR 2022
State-dependent action transformation. Note that the group operation acting action space A depends on state, since G actually acts on the product space x A: (g. (s, + A (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. (022b) also has a joint group action on state and action space, which is denoted state transtormation L.g S S and state-dependent action transformation K2 : A + A. Table I in van der Pol ct al. (2020b) outlines state and action transformations for their environments, and all of action transl formations are not state-dependent.
Similarly in our case, geometric transformations are usually acting globally on the environments S x A. thus states and actions are transformed accordingly. We use the factorized form and omit the state-dependency PA(g; 8) of action transformation P.(g)), since the action transformations do not depend on states PA 8) = PA(g) for all g E G. E.S
Learning transition with equivariant networks. In this paper we are mainly interested in learn- ing transition functions which are equivariant under symmetry transformations and can be high- dimensional.
We apply the idea of learning equivariant transition models in the latent space 2, where 2 is the space of symmetric representations, on various environments with different symmetry groups G. We assume do not explicitly know Ps since 8 is high-dimensional. We factorize the group representation on state and action S x A as latent state transformation pz E(s) and (A s) a. In the deterministic case, the transition model can be modeled by G-equivariant networks in latent state Z and action space A:
pz(g) T(E(s),)) T(pz(s) ((s). a),
for all g E G, 8,8' E and E A.
C LEARNED LATENT REPRESENTATIONS
The learned latent embedding for all states in the evaluation set for Reacher is shown in Figure and the embeddings are factored into irreducible representations. The 2-dimensional representations show 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 latent space traversal with smooth rotations, showing that the learned representations correctly encode the correct symmetries. Figure/ojhor 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 TRAINING DETAILS
D. ENVIRONMENTS
2D Shapes & 3D Blocks There are five objects arranged in a 5 x 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 time, unless it is blocked by the boundaries by another object. Observations 50 x 50 x 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 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 Rea c- As we do not consider rewards, we fix the goal position to the position [0.2, 0.2) so
14

---

Under review as a conference paper at ICLR 2022
(a) 8, s
(b) = E Z.
Figure Learned symmetric embeddings for Reacher: pixel observation s (left top) and next observation s' (left bottom), latent representation of the evaluation set (right). The representation type of PDa,reg which we factor into irreducible representations before visualizing (see HHlll (2003)). All encoded samples in the evaluation set are shown and the encoded current observation is colored red and the encoded next observation is colored orange. There clear circular pattern that match joint rotations.
Angle: 0"
60
120"
180
240
300
360
Figure 5: Latent space traversal for matrix multiplication (Matmul) transition model in Teapot. Top row: Ground truth rotation of a teapot. Bottom row: Projection in the latent space of the Matmul model, which is regularized to consist only of valid rotation matrices. The model, by construction, does not have fixed reference frame; hence, the two sequences are offset by the learned latent reference frame.
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 x 500 RGB image to 128 128 3. The previous and current frames are then stacked an observation to encode velocities. The default camera position gives slightly skewed perspective, see Table [
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 coordinate frame 2 E 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 by
15

---

Under review as a conference paper at ICLR 2022
-10
2
115
-.0
1
225
-3.0
-3.5
0.8
1.0
1.2
-3
(a)
(b)
Figure 6: 2D Shapes: learned embeddings for all states in the evaluation when trained on only the up action. Our model (left) is able to generalize well and learns the correct underlying x 5 grid. The non-equivariant model (right) learns degenerate solution
E SO(3) where a is rotation of 7200330 around the x-wy-or z-axis. In the large action space, actions may be any element (O(3).
D.2 MODEL ARCHITECTURES
Symmetric embedding network S For all models and environments except for 3D Teapot, CNNs with BatchNorm loffe & Szegedy (2015) and ReLU activations between each convolutional layer. For 3D Teapot, the symmetric embedding network maps directly the latent space use 4 convolutional layers followed by 3 fully connected layers. The output a 3 x 3 rotation matrix. The number of layers for each environment is given in Table For the non-equivariant symmetric embedding networks, we use 32 convolutional filters for every layer and use filters for Reacher and 16 filters for all other environments.
Encoder The object-oriented environments use 3-layer MLPs with 512 hidden units for the non-equivariant networks and 256 for the equivariant counterparts. There 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 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 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 [
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 equivariant version. Actions are concatenated to the latent and are input into the transition models which then outputs z' of the same dimension as the input z. We use one-hot encoding for discrete actions.
a MI.P with 512 hidden units for the non-equivariant version and 256 for the
(a) Blocks
(b) Reacher
Figure 7: Original observations and their G-transformed versions
16

---

Under review as a conference paper at ICLR 2022
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-orientee orthonormal frame. This is achieved by having the network output two vector u, R 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' v)u'
n' u/|lu||s w' = u' x x
v' | (u' vyu'll
y lu' e' 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 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 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 learning rate
10-1, batch size 1024, 100 epochs, and the hinge margin 7 = We find that these hyperparameters work well for all other environments, except that Reacher uses batch size of 256 and mixed precision training was used for both non-equivariant, fully-equivariant, and method, in order keep the batch size relatively high for stable contrastive learning. Most experiments were run on a single Nvidia RTX 20801i except for 3D Cubes which used a single Nvidia P100 12GB.
E
GROUP REPRESENTATIONS
We explain the notation and definitions of the different representations of the groups considered in the paper and displayed in Table
The Astd representation of C4 D4 R2 is by 2-by-2 rotation and reflection matrices. The Pstd representation of Sş permutes the standard basis of R The regular representation Preg of G' permutes the basis element of R|G| according to the multiplication table of G. The trivial representation G fixes as Ptriv (g) depending only on if contains reflection. Given representations (p). R"1 and (p2. R²*) of G and G2, (01 X p2) (g) (v s w) = g1v 3 gou gives a representation on G1 G2 on R 6 R22
x. For D4, (Mip(g) +1 representation
R
17


Reviews:
[{"review_id": "vq1_Pe938no", "reviewer": "Reviewer_hasv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found this to be a very difficult paper to read and understand.  In part it is a new area for me, but also I feel the authors are writing to others who live in their own world rather than a broader community of researchers.\n\nFigure 2 is helpful, but I could not understand what is the purpose of the symmetric embedding network S as opposed to the encoder E?  Why two networks instead of one?\n\nAlso, as far as I understand S, E and T are just standard neural networks, there is nothing about SO(3) per se or the mathematics of transformations engineered into them.  It appears the group theory mainly comes into play in how the networks are trained, as opposed to the computations they perform per se.  If so, then to me this seems a shame because I think where group theory could help us the most is in moving away from standard MLP models to more powerful and interesting computational frameworks.\n\nThe tables 2-6 showing model performance are also very hard to understand - extremely cryptic.   A bit more information about what these abbreviations mean and what is being measured would be helpful.\n", "summary_of_the_paper": "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.\n", "summary_of_the_review": "An interesting idea but a difficult paper to read and understand.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "p9jKRIU3d5", "reviewer": "Reviewer_LG1G", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I found this paper incredibly dense in parts, and the presentation quite confusing. The paper mixes a variety of interesting but sometimes disparate topics - equivariance as it applies to networks, the concept of symmetry as it applies to group actions, and equivariant transitions. In a sense the authors might be trying to sell too much here, and as a result the actual contributions are quite hard to tease out. The quality of writing is also a concern. To the extent that I follow the arguments, the main strength is the richness of the problems considered here - equivariance is indeed a hot topic, and the general goal of trying to recover representations that capture aspects of symmetries of the underlying phenomena (e.g. in the latent space) when from observations such as visual images, these symmetries are not easy to tease out (e.g. the rotating car example in Figure 1.). The experimental results do appear to be promising, although the use of some of the metrics (e.g. DIE) are only relevant for the restricted case of the Euclidean groups or its sub-groups.\n\nWeaknesses of the article include the quality of presentation which could be vastly improved, e.g, the related work section is rather disorganized - there doesn't seem to be a common thread. Simplifying the claims and accordingly the presentation, would add focus to the paper. There is also a rather a strong limitation of their assumption that for the task at hand the relevant groups are known. For particular cases there could be much more straightforward solutions. For example, Quessard et al. in the article \"Learning Group Structure and Disentangled Representations of Dynamical Environments\" (NeurIPS 2020) offer a much cleaner solution where a group G acts on pixel space, in such a way that the operations can be captured by matrix multiplication. The basic formulation in the Quessard et al. paper could certainly be applied to the teapot task it seems, one doesn't need the complexity of the equivariant transitions in Section 4.2 of the present paper to do this.\n", "summary_of_the_paper": "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).", "summary_of_the_review": "A lot is packed in to the paper, making the message a bit confusing, and the main contribution not so easy to tease out. Further, it isn't clear to me what the theoretical guarantees are and it seems that a fair amount has to be known to apply the methods and to effect the equivariant transitions. To me it seems that simpler solutions do exist, at least for some cases, but perhaps I am missing something here. I also found the symmetric embedding steps in the Meta-Architecture to be somewhat engineered. The authors are indeed grappling with some interesting and very relevant questions, but separating the theory from what is engineered is not so easy (for this reviewer at least).", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "gHlmOulw3Wv", "reviewer": "Reviewer_Uu3z", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nThe goal of the paper in using equivariance in latent space models in RL is interesting and motivated. Another positive point about the paper is that it evaluates its proposal on diverse problems. However, I have questions and concerns about the claims of contribution, the basic setup, theoretical presentation, and experimental results of this paper.\n\nContributions\n------------------\nThe paper claims that “the idea of learning symmetric embedding networks has, to our knowledge, not previously been proposed or demonstrated.” and this is discussed as one of the main contributions. There have been several works in this area in recent years, including using symmetries in latent space for dynamical environments; e.g., \n\n- Higgins, Irina, et al. \"Towards a definition of disentangled representations.\" arXiv preprint arXiv:1812.02230 (2018).\n- Quessard, Robin, Thomas D. Barrett, and William R. Clements. \"Learning group structure and disentangled representations of dynamical environments.\" arXiv preprint arXiv:2002.06991 (2020).\n\n\nThe proposed model\n---------------------------\nIn the general case, the symmetry transformation of the action should be state-dependent. This is noted in prior work such as the one cited below. The theoretical setup used in the paper is therefore inadequate for symmetric environments. Generally, there is no discussion on prior works on symmetric MDPs and the way this work fits in.\n\n- van der Pol, E., Worrall, D., van Hoof, H., Oliehoek, F., & Welling, M. (2020). MDP homomorphic networks: Group symmetries in reinforcement learning. Advances in Neural Information Processing Systems, 33.\n\nAlso, the motivation for having an equivariant encoder \"after\" a non-equivariant one is not clear, since the embedding will not be equivariant. This is especially so given that the equivariance of the representation is not enforced in any other way. My understanding is that the loss function used here only ensures that the transition in the latent space is consistent with the state space - i.e., the loss is exactly the same as that of Kipf et al (2019), am I correct?\n\nExperiments\n---------------\nWhile the diversity in the choice of environments is great and the evaluation metric is reasonable, it is difficult to read much from the results. There are also some baselines that are difficult to justify: why use an equivariant network with a particular assumption of linear action on the input space where we know the group action is non-linear? This is further complicated by the fact that such a baseline is producing good results. With the exception of visualizations in the appendix (which are supportive), the results are mixed, borderline, or difficult to interpret. \n\n", "summary_of_the_paper": "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.\n", "summary_of_the_review": "see above", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "LzXIvr9VCzO", "reviewer": "Reviewer_5ro3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The presentation can be improved in several ways. I agree that the paper studies various models, symmetry assumptions, and different controlled environments so it is a challenge to make it self-contained. Still, there are several things that I found hard to understand:\n\n- 'We consider the case of MDPs with symmetry as in van der Pol et al. (2020)'. Maybe I missed something but this model is not motivated anywhere.  is it obvious for computational purposes? it is not obvious to me why study/build on this model? The background section gives some motivation behind Kipf et al. and structured contrastive models but Van der Pol et. al 2020 is mentioned here for the first time.\nHow central is the MDP assumption in validation of the main idea? Is it reasonable to ask to demonstrate the main idea on other structured spaces?\n\n-'We require that ρA(G) · A′ = A. That is, every MDP action is a G-transformed version of one in A′ .' this should be mentioned at the beginning of Model overview, Section 4.1.\n\nThe experiments are not completely convincing. In particular, fully equivariant networks that assume the correct transformation $\\rho_S$  is a simple transformation of the pixels performs very well in all tests and thus, somewhat contradicts the main motivation of the paper.\n\nThe paper is riddled with typos, some mentioned below: \n-Future work will include applying *out* method  - The MDP action *is* input directly to T. - 'each element $g \\in G$ *must* invertible with respect to composition' .. We demonstrate that this meta-architecture can be *use* to learn world models with a variety. We then learn a transition model in latent space where *can* enforce symmetry", "summary_of_the_paper": "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. \n\nThe 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. ", "summary_of_the_review": "I give marginally above as the experimental setup seems rigorous. The authors are quite transparent about the lack of theoretical novelty when they state they are not proposing a new equivariant neural network design and instead, they build upon previous work to demonstrate working in new domains with unspecified group actions. I will revisit the score depending on other reviews. I can not comment on the utility of the idea only in MDP structured space but in general, the problem is relevent.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a new method for learning symmetric representations for equivariant world models. All reviewers recognized the interesting results in the paper. The reviewers have raised some concerns, which were not addressed well yet after the rebuttal. For example, Reviewers LG1G and Uu3z mentioned about the limitation of using the group for a task and the generality of the approach. Reviewer 5ro3 mentioned about the lack of novelty. Though they gave 6, they were quite neutral about the paper acceptance. Eventually, after a second round of deiscussions, we had to make this difficult decision: The current form of this paper is not ready for publications.
Decision: Reject