Equivariant Representation Learning via Class-Pose Decomposition

Giovanni Luca Marchetti* Gustaf Tegnér* Anastasiia Varava Danica Kragic

School of Electrical Engineering and Computer Science, KTH Royal Institute of
Technology Stockholm, Sweden

###### Abstract

We introduce a general method for learning representations that are
equivariant to symmetries of data. Our central idea is to decompose the latent
space into an invariant factor and the symmetry group itself. The components
semantically correspond to intrinsic data classes and poses respectively. The
learner is trained on a loss encouraging equivariance based on supervision
from relative symmetry information. The approach is motivated by theoretical
results from group theory and guarantees representations that are lossless,
interpretable and disentangled. We provide an empirical investigation via
experiments involving datasets with a variety of symmetries. Results show that
our representations capture the geometry of data and outperform other
equivariant representation learning frameworks.

##  1 INTRODUCTION

For an intelligent agent aiming to understand the world and to operate
therein, it is crucial to construct rich _representations_ reflecting the
intrinsic structures of the perceived data ([2]). A variety of self-supervised
approaches have been proposed to address the problem of representation
learning such as (variational) auto-encoders ([21, 1]) and contrastive
learning methods ([6, 20]). These approaches are applicable to a wide range of
scenarios due to their generality but often fail to capture fundamental
aspects hidden in data. For example, it has been recently shown that
_disentangling_ intrinsic factors of variation requires specific biases or
some form of supervision ([24]). This raises the need for representation
learning paradigms leveraging upon specific structures carried by data.

Figure 1: An illustration of our equivariant representation decomposing
intrinsic class and pose of data.

A fundamental geometric structure of datasets consists of their _symmetries_
([16]). Such structure arises in several practical scenarios. Consider for
example images depicting rigid objects in different poses. In this case the
symmetries are rigid transformations (translations and rotations) that _act_
on datapoints by transforming the pose of the depicted object. Symmetries not
only capture the geometry of the pose but additionally preserve the object’s
shape, partitioning the dataset into intrinsic _invariant_ classes. The joint
information of shape and pose describes the data completely and is recoverable
from the symmetry structure alone. Another example of symmetries arises in the
context of an agent exploring an environment. Actions performed by a mobile
robot can be interpreted as changes of frame i.e., symmetries transforming the
data the agent perceives (see Figure 2). Assuming the agent is capable of
odometry (measurement of its own movement), such symmetries are collectable
and available for learning. All this motivates the design of representations
that rely on symmetries and behave coherently with respect to them – a
property known as _equivariance_ ([3, 18]). ††footnotetext: *Equal
contribution.

In this work we introduce a general framework for equivariant representation
learning. Our central idea is to encode _class and pose separately_ by
decomposing the latent space into an invariant factor $\mathcal{E}$ and a
symmetry component $G$ (see Figure 1). The pose component $G$ extracts
geometry from data while the class component is interpretable and necessary
for a lossless representation. We then train a representation learner
$\varphi:\ \mathcal{X}\rightarrow\mathcal{E}\times G$ via a loss encouraging
equivariance relying on supervision from relative symmetries between
datapoints. Our methodology is based on a theoretical result guaranteeing that
under mild assumptions an ideal learner achieves isomorphic representations by
being trained on equivariance alone. Another advantage of our framework in the
presence of multiple symmetry factors is that each of them can be varied
independently by acting on the pose component. This realizes disentanglement
in the sense of [18] which, as mentioned, would not be possible without the
information carried by symmetries.

Figure 2: Actions performed by a mobile agent can be seen as symmetries of the
perceived data.

We rely on the abstract language of (Lie) _group theory_ in order to formalize
symmetries and equivariance. As a consequence, our framework is general and
applicable to arbitrary groups of symmetries and to data of arbitrary nature.
This is in contrast with previous works on equivariance often focusing on
specific scenarios. For example, a number of works focus on Euclidean
representations and linear or affine symmetry groups ([13, 38, 32]) while
others enforce equivariance via group _convolutions_ ([8, 7]). The latter are
only applicable when data consists of signals over a base space $P$ (e.g., a
pixel plane or a voxel grid) and symmetries are induced by the ones of $P$,
which is limiting and does not lead to interpretable and structured
representations. On the other hand, some recently introduced frameworks aim to
jointly learn the equivariant representation together with the latent
dynamics/symmetries ([22, 31]). Although this has the advantage that the group
of symmetries is not assumed to be known a priori, the obtained representation
is again unstructured, uninterpretable, and comes with no theoretical
guarantees.

We empirically investigate our framework on image datasets with a variety of
symmetries including translations, dilations and rotations. We moreover
provide both qualitative and quantitative comparisons with competing
equivariant representation learning frameworks. Results show that our
representations exhibit more structure and outperform the baselines in terms
of latent symmetry prediction. Moreover, we show how the preservation of
geometry of our framework can be applied to a _mapping_ task: for data
collected by a mobile agent our representation can be used to extract maps of
multiple environments simultaneously. We provide a Python implementation
together with data and code for all the experiments at a publicly available
repository 111https://github.com/equivariant-ml/equivariant-representation-
learning. In summary, our contributions include:

  * •

A method for learning equivariant representations separating intrinsic data
classes from poses.

  * •

A general mathematical formalism based on group theory, which ideally
guarantees lossless and disentangled representations.

  * •

An empirical investigation via a set of experiments involving various group
actions, together with applications to scene mapping through visual data.

##  2 THE MATHEMATICS OF SYMMETRIES

We now introduce the necessary mathematical background on symmetries and
equivariance. The modern axiomatization of symmetries relies on their
algebraic structure i.e., composition and inversion. The properties of those
operations are captured by the abstract concept of a _group_ ([33]).

######  Definition 2.1.

A group is a set $G$ equipped with a _composition map_ $G\times G\rightarrow
G$ denoted by $(g,h)\mapsto gh$, an _inversion map_ $G\rightarrow G$ denoted
by $g\mapsto g^{-1}$, and a distinguished _identity element_ $1\in G$ such
that for all $g,h,k\in G$:

_Associativity_ | _Inversion_ | _Identity_
---|---|---
$g(hk)=(gh)k$ |  $g^{-1}g=gg^{-1}=1$ | $g1=1g=g$

Examples of groups include the permutations of a set and the general linear
group $\textnormal{GL}(n)$ of $n\times n$ invertible real matrices, both
equipped with usual composition and inversion of functions. An interesting
subgroup of the latter is the special orthogonal group
$\textnormal{SO}(n)=\\{A\in\textnormal{GL}(n)\ |\ AA^{T}=1,\
\textnormal{det}(A)=1\\}$, which consists of linear orientation-preserving
isometries of the Euclidean space. An example of a commutative group (i.e.,
such that $gh=hg$ for all $g,h\in G$) is $\mathbb{R}^{n}$ equipped with vector
sum as composition.

The idea of a space $\mathcal{X}$ having $G$ as a group of symmetries is
abstracted by the notion of group _action_.

######  Definition 2.2.

An action by a group $G$ on a set $\mathcal{X}$ is a map
$G\times\mathcal{X}\rightarrow\mathcal{X}$ denoted by $(g,x)\mapsto g\cdot x$,
satisfying for all $g,h\in G,\ x\in\mathcal{X}$:

_Associativity_ |  | _Identity_
---|---|---
$g\cdot(h\cdot x)=(gh)\cdot x$ |  | $1\cdot x=x$

In general, the following actions can be defined for arbitrary groups: $G$
acts on any set _trivially_ by $g\cdot x=x$, and $G$ acts on itself seen as a
set via (left) _multiplication_ by $g\cdot h=gh$. A further example of group
action is $\textnormal{GL}(n)$ acting on $\mathbb{R}^{n}$ by matrix
multiplication.

Maps which preserve symmetries are called _equivariant_ and will constitute
the fundamental notion of our representation learning framework.

######  Definition 2.3.

A map $\varphi\colon\ \mathcal{X}\rightarrow\mathcal{Z}$ between sets acted
upon by $G$ is called _equivariant_ if $\varphi(g\cdot x)=g\cdot\varphi(x)$
for all $g\in G,x\in\mathcal{X}$. It is called _invariant_ if moreover $G$
acts trivially on $\mathcal{Z}$ or, explicitly, if $\varphi(g\cdot
x)=\varphi(x)$. It is called _isomorphism_ if it is bijective.

Now, group actions induce classes in $\mathcal{X}$ called _orbits_ by
identifying points related by a symmetry.

######  Definition 2.4.

Consider the equivalence relation on $\mathcal{X}$ given by deeming $x$ and
$y$ equivalent if $y=g\cdot x$ for some $g\in G$. The induced equivalence
classes are called _orbits_ , and the set of orbits is denoted by
$\mathcal{X}/G$.

Figure 3: Orbits of a (free) group action represent intrinsic classes of data.
Each orbit is isomorphic to the symmetry group $G$ itself.

For example the orbits of the trivial action are singletons, while the
multiplication action has a single orbit. Data-theoretically, an orbit may be
seen as an invariant, maximal class of data induced by the symmetry structure.
In the example of rigid objects acted upon by translations and rotations,
orbits indeed correspond to shapes (see Figure 3).

It is intuitive to assume that a nontrivial symmetry $g\not=1\in G$ has to
produce a change in data. If no difference is perceived, one might indeed
consider the given transformation as trivial. We can thus assume that no point
in $\mathcal{X}$ is fixed by an element of $G$ different from the identity or,
in other words, $g\cdot x\not=x$ for $g\not=1$. Such actions are deemed as
_free_ and will be the ones relevant to the present work.

######  Assumption 2.1.

The action by $G$ on $\mathcal{X}$ is free.

The following is the core theoretical result motivating our representation
learning framework, which we will discuss in the following section. The result
guarantees a general decomposition into a trivial and a multiplicative action
and describes all the equivariant isomorphisms of such a decomposition.

######  Proposition 2.1.

The following holds:

  * •

There is an equivariant isomorphism

$\mathcal{X}\simeq(\mathcal{X}/G)\times G$ (1)

where $G$ acts trivially on the orbits and via multiplication on itself, i.e.,
$g\cdot(e,h)=(e,gh)$ for $g,h\in G$, $e\in\mathcal{X}/G$. In other words, each
orbit can be identified equivariantly with the group itself.

  * •

Any equivariant map $\varphi:\ (\mathcal{X}/G)\times
G\rightarrow(\mathcal{X}/G)\times G$ is a right multiplication on each orbit
i.e., for each orbit $O\in\mathcal{X}/G$ there is an $h_{O}\in G$ such that
$\varphi(O,g)=(O^{\prime},gh_{O})$ for all $g\in G$. In particular, if
$\varphi$ induces a bijection on orbits then it is an isomorphism.

We refer to the Appendix for a proof. The first part of the statement can be
interpreted in plain words as a _decomposition of classes from poses_ for any
free group action. According to this terminology, a pose is abstractly an
element of an arbitrary group $G$ while a class is an orbit. The intuition
behind the second part of the statement is that any equivariant map performs
an orbit-dependent ‘change of frame’ in the sense that elements of an orbit
$O$ get composed on the right by a symmetry $h_{O}$ depending on $O$. This
will imply that our representations can differ from ground-truth ones only by
such change of frames and will in fact guarantee isomorphic representations
for our framework.

##  3 METHOD

###  3.1 General Equivariant Representation Learning

In the context of _representation learning_ the goal of the model is to learn
a map deemed ‘representation’ $\varphi\colon\
\mathcal{X}\rightarrow\mathcal{Z}$ from the data space $\mathcal{X}$ to a
_latent space_. The learner optimizes a loss $\mathcal{L}\colon\
\mathcal{M}\rightarrow\mathbb{R}$ over parameters $\theta\in\mathcal{M}$ of
the map $\varphi=\varphi_{\theta}$. The so-obtained representation can be
deployed in downstream applications for an improved performance with respect
to operating in the original data space.

The central assumption of _equivariant_ representation learning is that data
carries symmetries which the representation has to preserve. As discussed in
Section 2, this means that a group of symmetries $G$ acts on both
$\mathcal{X}$ and $\mathcal{Z}$ and that the representation $\varphi$ is
encouraged to be equivariant via the loss. While the action on $\mathcal{Z}$
is designed as part of the model, the action on $\mathcal{X}$ is unknown in
general and has to be conveyed by data. Concretely, the dataset consists of
triples $(x,g,y)$ with $x\in\mathcal{X},g\in G$ and $y=g\cdot x$. The group
element $g$ carries symmetry information which is relative between $x$ and
$y$. Equivariance is then naturally encouraged via a loss in the form:

$\mathcal{L}(\theta;x,g,y)=d\left(\varphi_{\theta}(y),\
g\cdot\varphi_{\theta}(x)\right).$ (2)

Here $d\colon\ \mathcal{Z}\times\mathcal{Z}\rightarrow\mathbb{R}_{\geq 0}$ is
a similarity function on $\mathcal{Z}$. Generally speaking, $d$ does not
necessarily need to satisfy the axioms for a distance but we at least require
it to be _positive definite_ i.e., $d(z,z^{\prime})=0$ iff $z=z^{\prime}$.
Note that we assume the group together with its algebraic structure to be
known a priori and not inferred during the learning process. Its action over
the latent space is defined in advance and constitutes the primary inductive
bias for equivariant representation learning.

###  3.2 Learning to Decompose Class and Pose

Motivated by Proposition 2.1, we propose to set the latent space as:

$\mathcal{Z}=\underbrace{\mathcal{E}}_{Class}\times\underbrace{G}_{Pose}$ (3)

with $G$ acting trivially on $\mathcal{E}$ and via multiplication on itself.
Here, $\mathcal{E}$ is any set which is meant to represent classes of the
encoded data. Since there is in general no prior information about the action
by symmetries on $\mathcal{X}$ and its orbits, $\mathcal{E}$ has to be set
beforehand. Assuming $\mathcal{E}$ has enough capacity to contain
$\mathcal{X}/G$, Proposition 2.1 shows that an isomorphic equivariant data
representation is possible in $\mathcal{Z}$. By fixing (positive definite)
similarity functions $d_{\mathcal{E}}$ and $d_{G}$ on $\mathcal{E}$ and $G$
respectively, we obtain a joint latent similarity function
$d(z,z^{\prime})=d_{\mathcal{E}}(z_{\mathcal{E}},z^{\prime}_{\mathcal{E}})+d_{G}(z_{G},z^{\prime}_{G})$,
where the subscripts denote the corresponding components. When
$G\subseteq\textnormal{GL}(n)$ is a group of matrices, a typical choice for
$d_{G}$ is the (squared) Frobenius distance i.e., the Euclidean distance for
matrices seen as flattened vectors. The equivariance loss
$\mathcal{L}(\theta;x,g,y)$ in Equation 2 then reads:

$\underbrace{d_{\mathcal{E}}\left(\varphi^{\mathcal{E}}(y),\
\varphi^{\mathcal{E}}(x)\right)}_{Invariant}\
+\underbrace{d_{G}\left(\varphi^{G}(y),\
g\varphi^{G}(x)\right).}_{Multiplication-Equivariant}$ (4)

Here we denoted the components of the representation map by
$\varphi=(\varphi^{\mathcal{E}},\varphi^{G})$ and omitted the parameter
$\theta$ for simplicity. To spell things out, $\varphi^{\mathcal{E}}$
encourages data from the same orbit to lie close in $\mathcal{E}$ (i.e.,
$\varphi^{\mathcal{E}}$ is ideally invariant) while $\varphi^{G}$ aims for
equivariance with respect to multiplication on the pose component $G$.

If $\varphi_{\mathcal{E}}$ is injective then Proposition 2.1 guarantees
lossless (i.e., isomorphic) representations, which we summarize in the
following corollary:

######  Corollary.

Suppose that $\varphi_{\mathcal{E}}$ is injective. Then
$\mathcal{L}(\theta;x,g,y)=0$ for all $x,g,y=g\cdot x$ if and only if
$\varphi_{\theta}$ is an equivariant isomorphism.

In order to force injectivity, we propose a typical solution from contrastive
learning literature ([6]) encouraging latent points to spread apart. To this
end, we opt for the standard _InfoNCE loss_ ([28]), although other choices are
possible. This means that we replace the term
$d_{\mathcal{E}}\left(\varphi^{\mathcal{E}}(y),\
\varphi^{\mathcal{E}}(x)\right)$ in Equation 4 with

$\frac{1}{\tau}d_{\mathcal{E}}\left(\varphi^{\mathcal{E}}(y),\
\varphi^{\mathcal{E}}(x)\right)+\mathbb{E}_{x^{\prime}}\left[e^{-\frac{1}{\tau}d_{\mathcal{E}}\left(\varphi^{\mathcal{E}}(x^{\prime}),\
\varphi^{\mathcal{E}}(x)\right)}\right].$ (5)

The hyperparameter $\tau\in\mathbb{R}_{>0}$ (‘temperature’) controls the
amount of latent entropy. Following [28, 6], we set the class component as a
sphere $\mathcal{E}=\mathbb{S}^{m}$ by normalizing the output of
$\varphi_{\mathcal{E}}$. This allows to deploy the cosine dissimilarity
$d_{\mathcal{E}}(z,z^{\prime})=-\textnormal{cos}(\angle zz^{\prime})=-z\cdot
z^{\prime}$ and is known to lead to improved performances due to the
compactness of $\mathcal{E}$ ([36]).

###  3.3 Parametrizing via the Exponential Map

The output space of usual machine learning models such as deep neural networks
is Euclidean. Our latent space (Equation 3) contains $G$ as a factor, which
might be non-Euclidean as in the case of $G=\textnormal{SO}(n)$. In order to
implement our representation learner $\varphi$ it is thus necessary to
parametrize the group $G$. To this end, we assume that $G$ is a differentiable
manifold (with differentiable composition and inversion maps) i.e., that $G$
is a _Lie group_. One can then define the _Lie algebra_ $\mathfrak{g}$ of $G$
as the tangent space to $G$ at $1$.

We propose to rely on the _exponential map_ $\mathfrak{g}\rightarrow G$,
denoted by $v\mapsto e^{v}$, to parametrize $G$. This means that $\varphi$
outputs an element $v$ of $\mathfrak{g}$ that gets mapped into $G$ as $e^{v}$.
Although the exponential map can be defined for general Lie groups by solving
an appropriate ordinary differential equation, we focus on the case
$G\subseteq\textnormal{GL}(n)$. The Lie algebra $\mathfrak{g}$ is then
contained in the space of $n\times n$ matrices and the exponential map amounts
to the matrix Taylor expansion $e^{v}=\sum_{k\geq 0}v^{k}/k!$. For specific
groups the latter can be simplified via simple closed formulas. For example,
the exponential map of $\mathbb{R}^{n}$ is the identity while for
$\textnormal{SO}(3)$ it can be efficiently computed via the Rodrigues’ formula
([23]).

###  3.4 Relation to Disentanglement

Our equivariant representation learning framework is related to the popular
notion of _disentanglement_ ([2, 17]). Intuitively, in a disentangled
representation a variation of a distinguished aspect in the data is reflected
by a change of a single component in the latent space. Although there is no
common agreement on a rigorous formulation of the notion ([24]), a proposal
has been addressed in [18]. The presence of multiple dynamic aspects in the
data is formalized as an action on $\mathcal{X}$ by a decomposed group

$G=G_{1}\times\cdots\times G_{n}$ (6)

where each of the factors $G_{i}$ is responsible for the variation of a single
aspect. A representation $\varphi\colon\ \mathcal{X}\rightarrow\mathcal{Z}$ is
then defined to be disentangled if (i) there is a decomposition
$\mathcal{Z}=\mathcal{Z}_{1}\times\cdots\times\mathcal{Z}_{n}$ where each
$\mathcal{Z}_{i}$ is acted upon trivially by the factors $G_{j}$ with
$j\not=i$ and (ii) $\varphi$ is equivariant.

Our latent space (Equation 3) automatically yields to disentanglement in this
sense. Indeed, in the case of a group as in Equation 6 we set
$\mathcal{Z}_{i}=G_{i}$. In order to deal with the remaining factor
$\mathcal{Z}_{0}=\mathcal{E}$, a copy of the trivial group $G_{0}=\\{1\\}$ can
be added to $G$ without altering it up to isomorphism. The group $G\simeq
G_{0}\times\cdots\times G_{n}$ acts on $\mathcal{Z}=\mathcal{E}\times
G=\mathcal{Z}_{0}\times\cdots\times\mathcal{Z}_{n}$ as required for a
disentangled latent space. In conclusion, our work fits in the line of
research aiming to infer disentangled representation via indirect and weak
forms of supervision ([25]), of which symmetry structures are an example.

##  4 RELATED WORK

Equivariant Representation Learning. Models relying on symmetry and
equivariance have been studied in the context of representation learning.
These models are typically trained on variations of the equivariance loss
(Equation 2) and are designed for specific groups $G$ and actions on the
latent space $\mathcal{Z}$. The pioneering _Transforming Autoencoders_ ([19])
learn to represent image data translated by $G=\mathbb{R}^{2}$ in the pixel
plane, with $\mathcal{Z}$ consisting of several copies of $G$ (‘capsules’)
acting on itself. Although such models are capable of learning isomorphic
representations, the orbits are not explicitly modeled in the latent space. In
contrast, our invariant component $\mathcal{E}$ is an interpretable
alternative to multiple capsules making orbits recoverable from the
representation. A series of other works represent data in a latent space
modelled via the group of symmetries $G$ ([12, 35, 40]). These works however
either do not reserve additional components dedicated to orbits, obtaining a
representation that forgets the intrinsic classes of data, or address specific
groups i.e., $\textnormal{SO}(3)$, the torus
$\textnormal{SO}(2)\times\textnormal{SO}(2)$ and product of cyclic groups
respectively. _Affine Equivariant Autoencoders_ ([13]) deal with affine
transformations of the pixel-plane (shearing an image, for example) and
implement a latent action through a hand-crafted map $t\colon\
G\rightarrow\mathcal{Z}=\mathbb{R}^{n}$. Groups of rotations
$\textnormal{SO}(n)$ linearly acting on a Euclidean latent space
$\mathcal{Z}=\mathbb{R}^{n}$ are explored in [38, 32]. Since rotating a vector
around itself has no effect, linear actions are not free (for $n\geq 3$),
which makes isomorphic representations impossible. _Equivariant Neural
Rendering_ ([11]) proposes a latent voxel grid on which $\textnormal{SO}(3)$
acts approximately by rotating and interpolating values. In contrast, our
latent group action is exact and thus induces no loss of information. We
provide an empirical comparison to both linear Euclidean actions and
Equivariant Neural Rendering in Section 5. Lastly, [37] have recently proposed
to learn equivariant representations by splitting the latent space into an
invariant component and an equivariant one, which bears similarity to our
framework. Differently from us, however, the model is trained via an
equivariance loss in the data space $\mathcal{X}$ and thus requires the group
actions over data to be known a priori. As previously discussed, this is
limiting and often unrealistic in practice.

Convolutional Networks. Convolutional layers in neural networks ([8, 7])
satisfy equivariance a priori with respect to transformations of the pixel
plane. They were originally introduced for (discretized) translations and
later extended to more general groups ([9, 26, 10]). However, they require
data and group actions in a specific form. Abstracty speaking, data need to
consist of vector fields over a base space (images seen as RGB fields over the
pixel plane, for example) acted upon by $G$, which does not hold in general.
Examples of symmetries not in this form are changes in perspective of first-
person images of a scene and rotations of rigid objects on an image. Our model
is instead applicable to arbitrary (Lie) group actions and infers equivariance
in a data-driven manner. Moreover, equivariance through $G$-convolutions alone
is hardly suitable for representation learning as the output dimension
coincides with the input one. Dimensionality reduction techniques deployed
together with convolutions such as max-pooling or fully-connected layers
disrupt equivariance completely. The latent space in our framework is instead
compressed and is ideally isomorphic to the data space $\mathcal{X}$
(Proposition 2.1).

Table 1: Datasets involved in our experiments, with the corresponding group of symmetries and number of orbits. | Sprites | Shapes | Multi-Sprites | Chairs | Apartments
---|---|---|---|---|---
$\mathcal{X}$ |  |  |  |  |
$G$ | $\mathbb{R}^{3}$ | $\mathbb{R}^{3}$ | $\mathbb{R}^{6}$ | $\textnormal{SO}(3)$ | $\mathbb{R}^{2}\times\textnormal{SO}(2)$
$|\mathcal{X}/G|\ $ | $3$ | $4$ | $27$ | $3$ | $2$

World Models. Analogously to group actions, Markov Decision Processes (MDPs)
from reinforcement learning and control theory involve a possibly stochastic
interaction $\mathcal{A}\times\mathcal{X}\rightarrow\mathcal{X}$ with an
environment $\mathcal{X}$ via a set $\mathcal{A}$ of moves. In general, no
algrbraic structure (such as a group composition) is assumed on $\mathcal{A}$.
In this context, a representation equivariant with respect to the action is
referred to as _World Model_ ([14, 22, 29]) or _Markov Decision Process
Homomorphism_ (MDPH) ([31]). MDPHs are usually deployed as pre-trained
representations for downstream tasks or trained jointly with the agent for
exploration purposes ([30]). However, the latent action
$\mathcal{A}\times\mathcal{Z}\rightarrow\mathcal{Z}$ of an MDPH is learned
since no prior knowledge is assumed around $\mathcal{A}$ or the environment.
This implies that the resulting representation is unstructured and
uninterpretable. We instead assume that $G=\mathcal{A}$ is a group acting
(freely) on $\mathcal{X}$, which enables us to define a geometrically-
structured and disentangled latent space that guarantees isomorphic
equivariant representations. We provide an empirical comparison to MDPHs in
Section 5.

##  5 EXPERIMENTS

###  5.1 Dataset Description

Our empirical investigation aims to assess our framework via both qualitative
and quantitative analysis on datasets with a variety of symmetries. To this
end we deploy five datasets summarized in Table 1: three with translational
symmetry extracted from dSprites and 3DShapes ([27, 4]), one with rotational
symmetry extracted from ShapeNet ([5]) and one simulating a mobile agent
exploring apartments and collecting first-person views. The latter is
extracted from Gibson ([39]) and generated via the Habitat simulator ([34]).
Datapoints are triples $(x,g,y)$ where $x,y$ are $64\times 64$ images, $g\in
G$ and $y=g\cdot x$. We refer to the Appendix for a more detailed description
of the datasets.

###  5.2 Baselines and Implementation Details

We compare our method with the following models designed for learning
equivariant representations:

MDP Homomorphisms (MDPH) from [31], [22]: a framework where the representation
$\varphi:\ \mathcal{X}\rightarrow\mathcal{Z}$ is learnt jointly with the
latent action $T:\ G\times\mathcal{Z}\rightarrow\mathcal{Z}$. The two models
are trained with the equivariance loss $\mathbb{E}_{x,g,y=g\cdot
x}[d(\varphi(y),\ T(g,\varphi(x)))]$ (cf. Equation 2). In order to avoid
trivial solutions, an additional ‘hinge’ loss term
$\mathbb{E}_{x,x^{\prime}}[\max\\{0,\ 1-d(x,x^{\prime})\\}]$ is optimized that
encourages encodings to spread apart in the latent space. This is analogous to
(the denominator of) the InfoNCE loss (Equation 5) which we rely upon to avoid
orbit collapse in $\mathcal{E}$. Differently from us, an MDPH does not assume
any prior knowledge on $G$ nor any algebraic structure on the latter. However,
this comes at the cost of training an additional model $T$ and losing the
structures and guarantees provided by our framework.

Linear: a model with $\mathcal{Z}=\mathbb{R}^{3}$ on which
$\textnormal{SO}(3)$ acts by matrix multiplication. Such a latent space has
been employed in previous works ([38, 32]). The model is trained with the same
loss as MDPH i.e., equivariance loss together with the additional hinge term
avoiding collapses such as $\varphi=0$. Note that the action on $\mathcal{Z}$
is no longer free (even away from $0$) since rotating a vector around itself
has no effect. Differently from our method, the model is thus forced to lose
information in order to learn an equivariant representation.

Equivariant Neural Renderer (ENR) from [11]: a model with a tensorial latent
space $\mathcal{Z}=\mathbb{R}^{C\times D\times H\times W}$, thought as a copy
of $\mathbb{R}^{C}$ for each point in a $D\times H\times W$ grid in
$\mathbb{R}^{3}$. The group $\textnormal{SO}(3)$ _approximately_ acts on
$\mathcal{Z}$ by rotating the grid and interpolating the obtained values in
$\mathbb{R}^{C}$. The model is trained trained jointly with a decoder $\psi:\
\mathcal{Z}\rightarrow\mathcal{X}$ and optimizes variation of the equivariance
loss incorporating reconstruction: $\mathbb{E}_{x,g,y=g\cdot
x}[d_{\mathcal{X}}(y,\ \psi(g\cdot\varphi(x)))]$. We set $d_{\mathcal{X}}$ as
the standard binary cross-entropy metric for (normalized) images. Although the
action on $\mathcal{Z}$ is free, the latent discretisation and consequent
interpolation make the model only approximately equivariant.

Figure 4: Left: visualization of encodings through $\varphi$ from the Sprites,
Chairs and Apartments datasets. The images display the projection to the
annotated components of $\mathcal{Z}$ and data are colored by their ground-
truth class. Each latent orbit from Apartments is compared to the view from
the top of the corresponding scene. Right: same visualization for the baseline
models MDPH and Linear on the Chairs dataset.

We implement the equivariant representation learner $\varphi$ as a ResNet-$18$
([15]), which is a deep convolutional neural network with residual
connections. We train our models for $100$ epochs through stochastic gradient
descent by means of the Adam optimizer with learning rate $10^{-3}$ and batch
size $16$. The distance $d_{G}$ is set as the squared Euclidean one for
$G=\mathbb{R}^{n}$ and for $G=\textnormal{SO}(2)\subseteq\mathbb{R}^{2}$,
while it is set to the squared Frobenius one for $G=\textnormal{SO}(3)$. The
invariant component consists of a sphere
$\mathcal{E}=\mathbb{S}^{7}\subseteq\mathbb{R}^{8}$ (see Section 3)
parametrized by the normalized output of $8$ neurons in the last layer of
$\varphi$. All the models from Section 5.2 implement the same architecture
(ResNet-$18$). ENR moreover implements $3$D convolutional layers around the
latent space as suggested in the original work ([11]). The latent action model
$T$ for MDPH is implemented as a two-layer deep neural network ($128$ neurons
per layer) with ReLu activation functions. For MDPH we set
$\textnormal{dim}(\mathcal{Z})=8+\textnormal{dim}(G)$, which coincides with
the output dimensionality of our model.

Table 2: Hit-rate (mean and std over $3$ runs) on test trajectories of increasing length. Dataset | Model | 1 Step | 10 Steps | 20 Steps
---|---|---|---|---
Sprites | Ours | ${\bf 1.00}_{\pm 0.00}$ | ${\bf 1.00}_{\pm 0.00}$ | ${\bf 1.00}_{\pm 0.00}$
MDPH | ${\bf 1.00}_{\pm 0.00}$ | ${\bf 1.00}_{\pm 0.00}$ | $0.98_{\pm 0.02}$
Shapes | Ours | ${\bf 1.00}_{\pm 0.00}$ | ${\bf 1.00}_{\pm 0.00}$ | ${\bf 1.00}_{\pm 0.00}$
MDPH | ${\bf 1.00}_{\pm 0.00}$ | $0.99_{\pm 0.01}$ | $0.96_{\pm 0.04}$
Multi-Sprites | Ours | ${\bf 1.00}_{\pm 0.00}$ | ${\bf 0.93}_{\pm 0.03}$ | ${\bf 0.93}_{\pm 0.03}$
MDPH | ${\bf 1.00}_{\pm 0.00}$ | $0.28_{\pm 0.06}$ | $0.11_{\pm 0.01}$
Chairs | Ours | ${\bf 0.98}_{\pm 0.01}$ | ${\bf 0.94}_{\pm 0.01}$ | ${\bf 0.94}_{\pm 0.01}$
Linear | $0.89_{\pm 0.08}$ | $0.87_{\pm 0.10}$ | $0.87_{\pm 0.10}$
MDPH | ${\bf 0.98}_{\pm 0.00}$ | $0.88_{\pm 0.07}$ | $0.78_{\pm 0.13}$
ENR | ${\bf 0.98}_{\pm 0.00}$ | $0.91_{\pm 0.01}$ | $0.82_{\pm 0.05}$
Apartments | Ours | ${\bf 0.99}_{\pm 0.00}$ | ${\bf 0.99}_{\pm 0.00}$ | ${\bf 0.99}_{\pm 0.00}$
MDPH | $0.98_{\pm 0.02}$ | $0.94_{\pm 0.03}$ | $0.86_{\pm 0.05}$

###  5.3 Visualizations of the Representation

In this section we present visualizations of the latent space of our model
(Equation 3), showcasing its geometric benefits. The preservation of
symmetries coming from equivariance enables indeed to transfer the intrinsic
geometry of data explicitly to the representation. Moreover, the invariant
component $\mathcal{E}$ separates the orbits of the group action, allowing to
distinguish the intrinsic classes of data in the latent space. Finally, the
representation from our model automatically disentangles factors of the group
as discussed in Section 3.4.

Figure 4 (left) presents visualizations of encodings through $\varphi$ for the
datasets Sprites, Chairs and Apartments. For each dataset we display the
projection to one component of $\mathcal{E}$ as well as a relevant component
of the group $G$. Specifically, for Sprites we display the component
$\mathbb{R}^{2}\subseteq G=\mathbb{R}^{3}$ corresponding to translations in
the pixel plane, for Chairs we display a circle $\textnormal{SO}(2)\subseteq
G=\textnormal{SO}(3)$ corresponding to one Euler angle while for Apartments we
display the component $\mathbb{R}^{2}\subseteq
G=\mathbb{R}^{2}\times\textnormal{SO}(2)$ corresponding to translations in the
physical world. For Apartments, we additionally compare representation of each
of the two apartments with the ground-truth view from the top.

As can be seen, in all cases the model correctly separates the orbits in
$\mathcal{E}$ through self-supervision alone. Since the orbits are isomorphic
to the group $G$ itself, the model moreover preserves the geometry of each
orbit separately. For Sprites, this means that (the displayed component of)
each orbit is an isometric copy of the pixel-plane, with disentangled
horizontal and vertical translations. (Figure 4, top-left). For Apartments,
this similarly means that each orbit exhibits an isometric copy of the real-
world scene. One can recover a map of each of the explored scenes by, for
example, estimating the density of data in $\mathcal{Z}$ (Figure 4, bottom-
right) and further use the model $\varphi$ to localize the agent within such
map. Our equivariant representation thus solves a _localization and mapping
task_ in a self-supervised manner and of multiple scenes simultaneously.

As a qualitative comparison, Figure 4 (right) includes visualizations for the
models MDPH (trained with $\textnormal{dim}(\mathcal{Z})=3$) and Linear on the
Chairs dataset. As can be seen, the latent space of MDPH is unstructured: the
geometry of $\mathcal{X}$ is not preserved and classes are not separated. This
is because the latent action of MDPH is learned end-to-end and is thus
uninterpretable and unconstrained a priori. For Linear the classes are
organized as spheres in $\mathcal{Z}$, which are are the orbits of the latent
action by $G=\textnormal{SO}(n)$. Such orbits are not isomorphic to $G$ (one
Euler angle is missing) since the action is not free. This means that
$\mathcal{Z}$ loses information and does not represent the dataset faithfully.

###  5.4 Performance Comparison

In this section we numerically compare our method to the equivariant
representation learning frameworks described at the beginning of Section 5. We
evaluate the models through _hit-rate_ , which is a standard score that allows
to compare equivariant representations with different latent space geometries
([22]). Given a test triple $(x,g,y=g\cdot x)$, we say that ‘$x$ hits $y$’ if
$\varphi(y)$ is the nearest neighbour in $\mathcal{Z}$ of $g\cdot\varphi(x)$
among a random batch of encodings $\\{\varphi(x)\\}_{x\in\mathcal{B}}$. For a
test set, the hit-rate is then defined as the number of times $x$ hits $y$
divided by the test set size. We set the number of aforementioned random
encodings to $|\mathcal{B}|=32$. For each model, the nearest neighbour is
computed with respect to the same latent metric $d$ as the one used for
training. In order to test the performance of the models when acted upon
multiple times in a row, we generate test sets where $g$ is a trajectory i.e.,
it is factorized as $g=g_{1}g_{2}\cdots g_{T}$ for $T\in\\{1,10,20\\}$. Hit-
rate is then computed after sequentially acting by the $g_{i}$’s in the latent
space. This captures the accumulation of errors in the equivariant
representation and thus evaluates the performance for long-term predictions.
All the test sets consist of $10\%$ of the corresponding dataset.

The results are presented in Table 2. As can be seen, all the models perform
nearly perfectly on single-step predictions with the exception of Linear
($89\%$ hit-rate). For the latter the latent group action is not free, which
prevents learning a lossless equivariant representation and thus degrades the
quality of predictions. On longer trajectories, however, our model outperforms
the baselines by an increasing margin. MDPH accumulates errors due to the lack
of structure in its latent space: its latent action is learned, which does not
guarantee stability with respect to composition of multiple symmetries. The
degradation of performance for MDPH is particularly evident in the case of
Multi-Sprites ($11\%$ hit-rate on $20$ steps), which is probably due to the
large number of orbits ($27$) and the consequent complexity of the prediction
task. Our model is instead robust even in presence of many orbits ($93\%$ hit-
rate on Multi-Sprites) due to the dedicated invariant component $\mathcal{E}$
in its latent space.

When the latent space is equipped with a group action, stability on long
trajectories follows from _associativity_ of the group composition and the
action (see Definition 2.1 and 2.2). This is evident from the results for the
Chairs dataset, where our model and Linear outperform MDPH on longer
trajectories ($94\%$ and $87\%$ against $78\%$ hit-rate on $20$ steps) and
exhibit a stable hit-rate as the number of steps increases. Even though ENR
carries a latent group action, it still accumulates errors ($82\%$ hit-rate on
$20$ steps) due to the discretization of the its latent space i.e., the latent
grid acted upon by $\textnormal{SO}(3)$. The consequent interpolation makes
the latent action only approximately associative, causing errors to accumulate
on long trajectories.

##  6 CONCLUSIONS, LIMITATIONS AND FUTURE WORK

In this work we addressed the problem of learning equivariant representations
by decomposing the latent space into a group component and an invariant one.
We theoretically showed that our representations are lossless, disentangled
and preserve the geometry of data. We empirically validated our approach on a
variety of groups, compared it to other equivariant representation learning
frameworks and discussed applications to the problem of scene mapping.

Our formalism builds on the assumption that the group of symmetries is known a
priori and not inferred from data. This is viable in applications to robotics,
but can be problematic in other domains. If a data feature is not taken into
account among symmetries, it will formally define distinct orbits. For
example, the eventual change in texture for images of rigid objects has to be
part of the symmetries in order to still maintain shapes as the only intrinsic
classes. A framework where the group structure is learned might be a feasible,
although less interpretable alternative to prior symmetry knowledge that
constitutes an interesting line of future investigation.

##  7 Acknowledgements

This work was supported by the Swedish Research Council, Knut and Alice
Wallenberg Foundation and the European Research Council (ERC-BIRD-884807).

## References

  * [1] Pierre Baldi  “Autoencoders, unsupervised learning, and deep architectures”  In _ICML workshop on unsupervised and transfer learning_ , 2012  JMLR
  * [2] Yoshua Bengio, Aaron Courville and Pascal Vincent  “Representation Learning: A Review and New Perspectives”  In _IEEE Transactions on Pattern Analysis and Machine Intelligence_ 35.8, 2013, pp. 1798–1828  DOI: 10.1109/TPAMI.2013.50
  * [3] Yoshua Bengio et al.  “A meta-transfer objective for learning to disentangle causal mechanisms”  In _arXiv preprint arXiv:1901.10912_ , 2019
  * [4] Chris Burgess and Hyunjik Kim  “3D Shapes Dataset”, https://github.com/deepmind/3d-shapes/, 2018
  * [5] Angel X. Chang et al.  “Shapenet: An information-rich 3d model repository”  In _arXiv preprint arXiv:1512.03012_ , 2015
  * [6] Ting Chen, Simon Kornblith, Mohammad Norouzi and Geoffrey Hinton  “A simple framework for contrastive learning of visual representations”  In _ICML_ , 2020  PMLR
  * [7] Taco S. Cohen, Mario Geiger and Maurice Weiler  “A General Theory of Equivariant CNNs on Homogeneous Spaces”  In _NeurIPS_ , 2019
  * [8] Taco S. Cohen and Max Welling  “Group Equivariant Convolutional Networks”  In _ICML_ , 2016  PMLR
  * [9] Taco S. Cohen, Mario Geiger, Mario Köhler and Max Welling  “Spherical CNNs”  In _ICLR_ , 2018
  * [10] Taco S. Cohen, Maurice Weiler, Berkay Kicanaoglu and Max Welling  “Gauge equivariant convolutional networks and the icosahedral CNN”  In _ICML_ , 2019  PMLR
  * [11] Emilien Dupont et al.  “Equivariant neural rendering”  In _ICML_ , 2020  PMLR
  * [12] Luca Falorsi et al.  “Explorations in Homeomorphic Variational Auto-Encoding”  In _arXiv preprint arXiv:1807.04689_ , 2018
  * [13] Xifeng Guo, En Zhu, Xinwang Liu and Jianping Yin  “Affine Equivariant Autoencoder.”  In _IJCAI_ , 2019
  * [14] David Ha and Jürgen Schmidhuber  “World models”  In _arXiv preprint arXiv:1803.10122_ , 2018
  * [15] Kaiming He, Xiangyu Zhang, Shaoqing Ren and Jian Sun  “Deep Residual Learning for Image Recognition”  In _CVPR_ , 2016  IEEE
  * [16] Irina Higgins, Sébastien Racanière and Danilo Rezende  “Symmetry-Based Representations for Artificial and Biological General Intelligence”  In _Frontiers in Computational Neuroscience_ Frontiers, 2022, pp. 28
  * [17] Irina Higgins et al.  “beta-VAE: Learning Basic Visual Concepts with a Constrained Variational Framework”  In _ICLR_ , 2017
  * [18] Irina Higgins et al.  “Towards a definition of disentangled representations”  In _arXiv preprint arXiv:1812.02230_ , 2018
  * [19] Geoffrey E. Hinton, Alex Krizhevsky and Sida D. Wang  “Transforming auto-encoders”  In _International conference on artificial neural networks_ , 2011, pp. 44–51  Springer
  * [20] Ashish Jaiswal et al.  “A survey on contrastive self-supervised learning”  In _Technologies_ 9.1 MDPI, 2020, pp. 2
  * [21] Diederik P. Kingma and Max Welling  “Auto-encoding variational bayes”  In _ICLR_ , 2014
  * [22] Thomas Kipf, Elise Pol and Max Welling  “Contrastive learning of structured world models”  In _ICLR_ , 2020
  * [23] Kuo Kan Liang  “Efficient conversion from rotating matrix to rotation axis and angle by extending Rodrigues’ formula”  In _arXiv preprint arXiv:1810.02999_ , 2018
  * [24] Francesco Locatello et al.  “Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations”  In _ICML_ , 2019  PMLR
  * [25] Francesco Locatello et al.  “Weakly-supervised disentanglement without compromises”  In _International Conference on Machine Learning_ , 2020, pp. 6348–6359  PMLR
  * [26] Haggai Maron, Or Litany, Gal Chechik and Ethan Fetaya  “On Learning Sets of Symmetric Elements”  In _ICML_ , 2020  PMLR
  * [27] Loic Matthey, Irina Higgins, Demis Hassabis and Alexander Lerchner  “dSprites: Disentanglement testing Sprites dataset”, https://github.com/deepmind/dsprites-dataset/, 2017
  * [28] Aaron Oord, Yazhe Li and Oriol Vinyals  “Representation learning with contrastive predictive coding”  In _arXiv preprint arXiv:1807.03748_ , 2018
  * [29] Jung Yeon Park et al.  “Learning Symmetric Embeddings for Equivariant World Models”  In _ICML_ , 2022
  * [30] Deepak Pathak, Pulkit Agrawal, Alexei A Efros and Trevor Darrell  “Curiosity-driven exploration by self-supervised prediction”  In _ICML_ , 2017  PMLR
  * [31] Elise Pol, Thomas Kipf, Frans A Oliehoek and Max Welling  “Plannable approximations to MDP homomorphisms: Equivariance under actions”  In _AAMAS_ , 2020
  * [32] Robin Quessard, Thomas Barrett and William Clements  “Learning Group Structure and Disentangled Representations of Dynamical Environments”  In _NeurIPS_ , 2020
  * [33] Joseph J. Rotman  “An introduction to the theory of groups”  Springer Science & Business Media, 2012
  * [34] Manolis Savva et al.  “Habitat: A Platform for Embodied AI Research”  In _ICCV_ , 2019  IEEE
  * [35] Loek Tonnaer et al.  “Quantifying and learning linear symmetry-based disentanglement”  In _ICML_ , 2020
  * [36] Tongzhou Wang and Phillip Isola  “Understanding contrastive representation learning through alignment and uniformity on the hypersphere”  In _ICML_ , 2020  PMLR
  * [37] Robin Winter et al.  “Unsupervised Learning of Group Invariant and Equivariant Representations”  In _NeurIPS_ , 2022
  * [38] Daniel E. Worrall, Stephan J. Garbin, Daniyar Turmukhambetov and Gabriel J. Brostow  “Interpretable transformations with encoder-decoder networks”  In _ICCV_ , 2017  IEEE
  * [39] Fei Xia et al.  “Gibson Env: real-world perception for embodied agents”  In _CVPR_ , 2018  IEEE
  * [40] Tao Yang et al.  “Towards building a group-based unsupervised representation disentanglement framework”  In _arXiv preprint arXiv:2102.10303_ , 2021

Equivariant Representation Learning via Class-Pose Decomposition:
Supplementary Materials

##  Appendix A Proofs of Theoretical Results

######  Proposition.

The following holds:

  * •

There is an equivariant isomorphism

$\mathcal{X}\simeq(\mathcal{X}/G)\times G$ (7)

where $G$ acts trivially on the orbits and via multiplication on itself, i.e.,
$g\cdot(e,h)=(e,gh)$ for $g,h\in G$, $e\in\mathcal{X}/G$. In other words, each
orbit can be identified equivariantly with the group itself.

  * •

Any equivariant map $\varphi:\ (\mathcal{X}/G)\times
G\rightarrow(\mathcal{X}/G)\times G$ is a right multiplication on each orbit
i.e., for each orbit $O\in\mathcal{X}/G$ there is an $h_{O}\in G$ such that
$\varphi(O,g)=(O^{\prime},gh_{O})$ for all $g\in G$. In particular, if
$\varphi$ induces a bijection on orbits then it is an isomorphism.

###### Proof.

We start by proving the first statement. Choose a system of representatives
$\mathcal{S}\subseteq\mathcal{X}$ for orbits, that is $\mathcal{S}$ contains
exactly one element for each class. Consider the map $f\colon\
(\mathcal{X}/G)\times G\rightarrow\mathcal{X}$ given, for $s\in\mathcal{S}$
and $g\in G$, by $f([s],g)=g\cdot s,$ where $[s]$ denotes the orbit of $s$. It
is straightforward to check that $f$ is indeed equivariant. Now if
$f([s],g)=f([t],h)$ for $s,t\in\mathcal{S}$ and $g,h\in G$ then $g\cdot
s=h\cdot t$ and $s,t$ are thus in the same orbit, which implies $s=t$ because
of uniqueness of representatives. But then $h\cdot s=g\cdot s$ and,
equivalently, $g^{-1}h\cdot s=s$, from which we deduce $g=h$ since the action
is free. This shows that $f$ is injective. Finally, for $x\in\mathcal{X}$, one
can write $x=g\cdot s$ for the representative $s$ of the orbit of $x$, which
means that $x=f([s],g)$. That is, $f$ is surjective and thus also bijective,
which concludes the proof of the first statement.

We now prove the second statement. Consider an equivariant map $\varphi:\
(\mathcal{X}/G)\times G\rightarrow(\mathcal{X}/G)\times G$. For each orbit
$O\in(\mathcal{X}/G)$ denote by $h_{O}\in G$ the element such that
$\varphi(O,1)=(O^{\prime},h_{O})$. Then by equivariance
$\varphi(O,g)=\varphi(O,g1)=(O^{\prime},gh_{O})$, as desired. ∎

##  Appendix B Description of Datasets

In our experiments we deploy the following datasets, which are also summarized
in Table 1:

Sprites: extracted from _dSprites_ ([27]). It consists of grayscale images
depicting three sprites (heart, square, ellipse) translating and dilating in
the pixel plane. The group of symmetries is $G=\mathbb{R}^{3}$: a factor
$\mathbb{R}^{2}$ translates the sprites in the pixel plane while the last copy
of $\mathbb{R}$, which is isomorphic via exponentiation to $\mathbb{R}_{>0}$
equipped with multiplication, acts through dilations. The dataset size is
$3\times 10^{4}$ and there are three orbits, each corresponding to a sprite.

Shapes: extracted from _3DShapes_ ([4]). It consists of colored images
depicting four objects (cube, cylinder, sphere, pill) on a background divided
into wall, floor and sky. Again, $G=\mathbb{R}^{3}$ but with the action given
by color shifting: each of the factors $\mathbb{R}$ acts by changing the color
of the corresponding scene component among object, wall and floor. The dataset
size is $4\times 10^{4}$ and there are four orbits, each corresponding to a
shape.

Multi-Sprites: obtained from Sprites by overlapping images of three colored
sprites (with fixed scale). The group of symmetries is
$G=\mathbb{R}^{6}=\mathbb{R}^{2}\times\mathbb{R}^{2}\times\mathbb{R}^{2}$,
each of whose factors $\mathbb{R}^{2}$ translates one of the three sprites in
the pixel plane. The added colors endow the sprites with an implicit ordering,
which is necessary for the action to be well-defined. The dataset size is
$3\times 10^{4}$. Since the scene is composed by three possibly repeating
sprites, there are $3^{3}=27$ orbits corresponding to the different
configurations.

Chairs: extracted from _ShapeNet_ ([5]). It consists of colored images
depicting three types of chair from different angles. The group of symmetries
is $G=\textnormal{SO}(3)$, which rotates the depicted chair. The dataset size
is $3\times 10^{4}$ and there are three orbits, each corresponding to a type
of chair.

Apartments: extracted from _Gibson_ ([39]) and generated via the _Habitat_
simulator ([34]). It consists of colored images of first-person renderings of
two apartments (‘Elmira’ and ‘Convoy’). The data simulate the visual
perception of an agent such as a mobile robot exploring the two apartments and
collecting images and symmetries. The latter belong to the group of two-
dimensional orientation-preserving Euclidean isometries
$G=\mathbb{R}^{2}\times\textnormal{SO}(2)$ and coincide with the possible
moves (translations and rotations) by part of the agent. One can realistically
imagine the agent perceiving the symmetries through some form of odometry
i.e., measurement of movement. Note that the action by $G$ on $\mathcal{X}$ is
_partially_ defined since $g\cdot x$ is not always possible because of
obstacles. As long as the agent is able to reach any part of each of the
apartments, the latter still coincide with the two orbits of the group action.
The dataset size $2\times 10^{4}$.

All the datasets consist of triples $(x,g,y)$ where $x,y$ are $64\times 64$
images, $g\in G$ and $y=g\cdot x$.