EqR: Equivariant Representations for Data-Efficient Reinforcement Learning
Authors: Arnab Kumar Mondal
Vineet Jain
Kaleem Siddiqi
Siamak Ravanbakhsh
Venue: ICLR 2022 Submitted (2022)
Keywords: Equivariance
Invariance
Representation learning
Reinforcement learning
Symmetric MDPs
MDP homomorphism
Lie parameterization.

Abstract: We study different notions of equivariance as an inductive bias in Reinforcement Learning (RL) and propose new mechanisms for recovering representations that are equivariant to both an agent’s action, and symmetry transformations of the state-action pairs. Whereas prior work on exploiting symmetries in deep RL can only incorporate predefined linear transformations, our approach allows for non-linear symmetry transformations of state-action pairs to be learned from the data itself. This is achieved through an equivariant Lie algebraic parameterization of state and action encodings, equivariant latent transition models, and the use of symmetry-based losses. We demonstrate the advantages of our learned equivariant representations for Atari games, in a data-efficient setting limited to 100k steps of interactions with the environment. Our method, which we call Equivariant representations for RL (EqR), outperforms many previous methods in a similar setting by achieving a median human-normalized score of 0.418, and surpassing human-level performance on 8 out of the 26 games.

Under review as a conference paper at ICLR 2022
EQR: EQUIVARIANT REPRESENTATIONS FOR PATA-EFFICIENT REINFORCEMENT LEARNING
Anonymous authors Paper under double-bling review
ABSTRACT
We study different notions of equivariance as an inductive bias Reinforcement Learning (RL) and propose new mechanisms for recovering representations that are equivariant to both an agent's s action, and symmetry transformations of the state-action pairs. Whereas prior work exploiting symmetries in deep RL can only incorporate predefined linear transformations, our approach allows for non- linear symmetry transformations of state-action pairs to be learned from the data itself. This is achieved through an equivariant Lie algebraic parameterization state and action encodings, cquivariant latent transition models, and the use of symmetry-basee based losses. We demonstrate the advantages our learned equivari- ant representations for Atari games, in data-efficient setting limited to 100K steps of interactions with the environment. Our method, which we call Equiv- ariant representations for RL (EqR), outperforms other comparable methods statistically reliable evaluation metrics.
INTRODUCTION
The recent success of deep reinforcement learning (François-Lavet et al., 2018) in applications to games such as Atari (Mnih et al., 2015), Go (Silver al., 2016) and Poker (Brown & Sandholm, 2019), to applications in robotics (Levine et al., 2016) and autonomous navigation (Bellemare et al., 2020) has demonstrated its promise as the framework of choice for sequential decision making. However, the use of reward as the only signal for representation learning with high dimensional states and actions leads to tremendous data inefficiency. Notably, almost all success stories of RI. rely on vast amounts of data or simulations with a huge computational overload.
More data-efficient representation learning (Bengio et al., 2013) requires stronger inductive biases, though the search for a general yet strong inductive bias is still under way. One general approach to place a central role on transformations of the data, where invariance and equivariance to set of transformations imposes strong conditions the learned representations. This viewpoint is particularly appealing in RL, where the agent in control of some of these transformations through its own actions. Fig. 1 illustrates this concept using the example of 2D pendulum. Moreover,
g.s
g.s
g.s
g.s
(b)
An illustration typical symmetries a pendulum, and corresponding transformations of state and action for a group equivariant transition model: (a) shows how reflection of the agent's state results in a permutation of the action, denoted by a (b) shows how rotation of the agent's 's state results in invariance of the action in the absence of gravity. The state transitions example), which can captured by our symmetry transformation- based transition model. Note rotational symmetry can hold even when gravity is present. In this case, symmetry transformations include rotations (and reflections) that preserve the Hamiltonian. Such non-lincar energy-preee peeggy-preservig transformations of statc-actions in the pixel space can become linear in the embedding space.
Figure 1:
be modeled group actions (2D rotations in this

---

Under review as a conference paper at ICLR 2022
transformations naturally lead to notion of disentanglement in the representations (Higgins et al., 2018), potentially enabling better out-of-distribution generalization (Higgins et al., 2017: Thomas et al., 2017). The recent success of self-super learning approaches that rely a (predefined) set of transformations (Chen et al., 2020; Zbontar et al., 2021), and also within the context of RI. (Yarats et al., 2021; Laskin al., 2020b), further highlights the importance of transformations in data-efficient representation learning.
Motivated by these observations, this work develops broader perspective on the notion of equiv- ariant representation learning within RL. In particular, we integrate equivariance under the agent's action and equivariance under the symmetries of the environment into single latent variable model that is equivariant to an priori unknown group of non-linear transformations of state-action pairs. In contrast to the traditional approach of using symmetric Markov Decision Processes (MDPs), we argue for modeling the larger group of state-action symmetries (separate from reward symmetries), and show how to parameterize the latent embeddings of states and actions to make the representa- tions equivariant to continuous transformations of the environment resulting from agent's s action. We benchmark our approach, which we call Equivariant representations for RL (E) on the 26 games in the Atari 100K benchmark (Kaiser et al., 2019), where outperform other comparable methods using reliable evaluation metrics (Agarwal et al., 2021). Our approach, however, is not restricted to this domain. It is applicable in any setting where the transformations that an agent undergoes be expressed using matrix Lie groups, including autonomous driving, navigation, and robotics.
2 RELATED WORK
The use of transformations, be it in autaentgmenttiion or self-supervision, has become common ingredient in recent representation learning methods for deep RL. However, theoretical work on symmetry in RL goes back to Zinkevich & Balch (2001) and Ravindran & Barto (2001), both of which use symmetric MDPs. more recent use of this formalism is in van der Pol et al. (2020b); Mondal et al. (2020), where policy networks. with built-in equivariance, are shown to improve data- efficiency. Closely related notions, that motivated the early work on symmetric MDPs, are model minimization (Ravindran & Barto, 2002), state abstraction (Ravindran & Barto, 2003; Li et al., 2006), MDP homomorphism (Ravindran & Barto, 2004) and lax bisimulations (Taylor, 2008). In particular, MDP homomorphism, which requires equivariance under an agent's action, encompasses the general idea of model-based reinforcement learning. As examples, a latent MDP that matches the state dynamics and the reward distribution of the environment is learned in (van der Pol al., 2020a; Gelada et al., 2019).
Other works in RL that are relevant using a learned model of the environment. While some methods such as SimPLe (Kaiser et al., 2019), learn this transition model the pixel level, the majority of methods use latent space model. The latent space is either learned using reconstruction (Hafner et al., 2019a;b), or through self-supervision and contrastive methods (CURL, Laskin et al., 2020b). However, there is evidence that the improvement in sample efficiency is largely due to image augmentation, as seen in Laskin al. (2020a) and DrQ (Yarats al., 2021). Using a reconstruction- based method also inefficient because similar to pixel level models, one needs to learn potentially irrelevant details. The fact that variations of model-free algorithms such as Data-Efficient Rainbow (DER) (van Hasselt et al., 2019) and OTRainbow (Kielak, 2019) can achieve a similar performance to reconstruction-basee methods without explicit representation learning components confirms this intuition. More recently SPR (Schwarzer et al., 2021) shows that data augmentation and improvements in Rainbow combined with particular forms of self-supervision can significantly improve the sample efficiency, producing state-of-the-are results in sample-efficient representation learning in RL.
objective are those that attempt to increase data-efficiency
BACKGROUND
3.1 GROUPS AND THEIR REPRESENTATIONS
A group G = {g] is a set, equipped with an associative binary operation, such that the set is closed under this operation, and each element g has unique inverse, such that their composition gives the identity g g= e. Any subset g' s 9 that is closed under binary operation the groups forms subgroup. A group can act on set X by transforming its elements X through
2

---

Under review as a conference paper at ICLR 2022
bijection. We use a: G X X to denote the group action, and for brevity replace a(g.z) with g. moving forward. The action captures some of the structure of ς due to two constraints the identity element acts trivially e
I; and composition of actions is equal to action of the composition, i.e, (gg') x = g (g' x),, Vg.g' E G. X is then called a G-set. Any g-action partitions X into orbits 2 = {g x g 9}. and we denote the set orbits under G-action using X/g. G-action is transitive iff its action results in single orbit.
Parameterizing Lie Groups In this work, we assume G is (any sub sub-group of) a classical Lie group over R. These are the groups that can be represented using invertible matrices in other words, we consider groups whose action linear on real vector spaces. We use (()) to denote a linear representation of G, and Py IR I for the action (a.k.a. the representation) of g E G. Two other greek letters T and λ. are also used for this purpose. Many such Lie groups are identifiable by their Lie algebra g (i)() This connection enables a simple parameterization p(g) using set Y of linear bases for their Lie algebra- that is Pg = exp gg. E("),, where exp(Y) Σj=o Jr is the matrix exponential. We refer to this parameterization later in Section 4. Such linear representations the form of invertible matrices can be used for both continuous transformations (e.g.. 3D rotation) and finite groups (x90" rotations).
3.2 MDP HOMOMORPIISS AND SYMMETRIC MDPs
We define an MDP as the 4-tuple M -- A, R, T) where S and A are respectively the set of states and actions, R S A E is the reward function, and T : A x - R is the state transition function.? For two MDPs M = (8, A, R T) and M (8, A, R2 T), MDP homomorphism can be defined as tuple H (hs, ha) where ha S $ is the state mapping and ha :Sx Å is the state dependent action mapping. These two mappings satisfy the following invariance and equivariance conditions:
1) Invariance of the reward:
R(hs(s), ha (s, a)) ) = R(s,a) Vs,aeSxA & 2) Equivariance of the deterministic transition model under the agent's action:
(1)
(shs(s), (s, a)) = hs(TTss,)) a)) Vs, E x
(2)
probabilistic variation of the above equation for stochastic MDP (Bloem-Reddy & Teh, 2020) is:
T(ho(s') hs(s), as) s"e[T'
Σ T(s" s,a) Vs, E E
(3)
where [s']nr hi
-2 (ha(s') the equivalence class of s' under hg.
In related literature, MDP homomorphism is often used for minimization of the MDP, because the optimal policy of M can be lifted to obtain the optimal counterparts for M.
Symmetric MDPs The automorphism group 9M = Aut(M) of an MDP identifies the set of symmetry transformations of state-actions that preserve the reward and the transition dynamics:
R(s,a) = (s,g-(s,()))
Vg E A S (5)
A (4)
T''-s'l (s, and g gss = T(g (s,a)) Vg 8,8
We refer to a reward function A that satisfies Eq. (4) as Gye-invariant reward function and deterministic transition function T' that satisties Eq. (5) a Guu-qvvivriaan transition function. Note that this is distinct notion from invariance and equivariance under agent's action in the context of MDP homomorphism. Here, the action refers the action of symmetry group, while in MDP homomorphism, the equivariance is to the action of the agent. We use group action or G-action to make this distinction clear when necessary.
For symmetric MDP that satisfies both Eq. (4) and Eq. (5), both the optimal action-value and optimal policy functions become invariant under GM action (Ravindran & Barto, 2001) that is,
Q(s,a) a) Qg (8,a) a)) and a(a, s) <g s)) Vg 9M 8, 8x A.
(6)
1 'This relation is bijective for "simply connected"' Lie groups
We ignore the discount factor for brevity.
3

---

Under review as a conference paper at ICLR 2022
The connection of symmetric MDPs to MDP homomorphism is due to the fact that symmetries can be used to define a homomorphism H: : .M + M by collapsing the state-actions that form an orbit under GM- Formally,the collapsed MDP MM (8,A, R, T) is defined by $ S/GM: A A/Gm, R((s, a) = R(s, a) and TM (8, a))x))))) = T(s' 18,). This results in symmetry-based model minimization of symmetric MDPs.
4 DESIDERATA FOR SYMMETRY- BASED REPRESENTATION IN RL
Separating Transition and Reward Symmetries One important choice is between using the symmetry group of the MDP (G)) versus the symmetry group of state-transitions (Gr), where Gr is the group of transformations of state-action pairs that leads to equivariant deterministic transitions, as given by Equation 5. The former is a subgroup of the latter 9M s 9r, i.c, the symmetries of a transition model contains the symmetries of the MDP. In fact it easy see that 9M- Gr n GR. where Gn is the group of transformations of state-action pairs that preserve the one step-reward and only satisfy Equation 4. We observe that working with a larger symmetry group Gr has two benefits: 1) it creates a stronger inductive bias for the model, because in many real-world settings can involve a range of symmetries in transitions that are not present in the reward. For example, an agent's navigation of 2D map often has the symmetry of the Euclidean group, while the reward (e.g., riving at a particular location) breaks this symmetry; 2) Separate modeling of transition symmetries facilitates transfer to new tasks, where the reward is changing.
Invariance/Equivariance in moode-frreemmodel-basee RL If the objective is to carry out model- free RL, Eq. (6) motivates the need to learn action-valle functions, the policies that are invariant to symmetries of the MDP (GM). For a deterministic policy, the invariance of Eq. (6) becomes an equivariance constraint: 9 (s)) T 8). As it essentially leads to model minimization, van der Pol et al. (2020b): Mondal et al. (2020) use this idea to improve sample efficiency when the groups actions the agent's action space are known permutations. However, if our objective just to learn a wymmetry-baseed model of the environment (i.e., transition and reward functions), Eq. (5) suggests that we need to learn a -queqivvrrann transition function.
Symmetries in Latent Transition Model While it is possible to learn the state transition model in the observation space that is equivariant to the agent's action, for high-dimensional inputs this could be quite challenging since the model has to learn details of the environment that are irrele- vant to the RL agent. Using state and action embeddings enables learning of the transition model in the latent space. Indeed the constraint momorphism Section 3.2. Working in the latent space has an additional benefit when it comes to symmetries: we can assume that the G action on the latent state-ace sate-action pairs is linear through p(g) despite having non-linear transformations in the observation space.
the model and the embedding is that of the MDP ho-
This figure demonstrates the relationship between two types of cquivariance latent variable modeling for an MDP with symmetric transition function. Green arrows (ver- tical plane) identify a diagram for transition models in an MDP homomorphism. A model T and state embedding fune- tion hg that are equivariant under agent's action makes this diagram commute. Red arrows (horizontal plane) identify the commutativity diagram for a symmetric transition function of an MDP in the latent space. Here the state-action embedding (8, a) is produced through the symmetry transformation of another state-action embedding (s.
Figure 2:
ston typrevtryy
Action transtioo
hg hg
Greep tovvdtio
7
From the fact that symmetries of states G3 S G is a subgroup of the state-action or transition symme- try, follows that Pg p()) can be divided into two parts: 1) T E 7(Gs) the group representation acting on the state embedding, and; 2) ng e k(G), the group representation for state-dependent action embedding."
'This is because p()) can be seen as a representation that is induced by the representation r(Gx) of its subgroup p = Indg, T.
4

---

Under review as a conference paper at ICLR 2022
At this point combine the requirement for MDP homomorphism Eq. (2), with that of the G-equivariant transition model Eq. (5) of a symmetric MDP. The result is the following two constraints in symmetric latent variable model Vs, e 3 x A and g g (see Figure 2): (7)
T(hg(s), (s, has h (T(s
TyT (a = (Tgps)))))(())())))))))))))))))) kihha(sa))
(8)
Matrix Embedding of States and Actions While the design choice of this subsection is not nec- essary, we see that it can significantly simplify the constraints on a symmetry based model. We propose use group representations for our state, and state action embeddings his S 7(Gs) and hA : S + x()). This choice assumes that & action on state and state-action pairs is fransitive, so that each state, and state-action pair can be mapped to (or at least one) group member. To empha- size this in our notation, we use k(s) instead of ha(*) (8) and similarly use T(8, a) instead of ha (s. a) for state, and netendependent action embedding respectively. This choice of embedding has several benefits: First, the learned embeddings are automatically equivariant to symmetry transformations of the state, and state-actions:
() - (s) (T)(ss)
and <(g- - =(gg((s,)))xxxx<(,,
Vs,uESxA.gEE.
This means that the symmetries of the state-action pairs are preserved and now take linear form in the latent space. Note that while the embeddings are automatically equivariant, they may be equiv- ariant to irrelevant non-linear transformations of the input. Word modeling constraints ensure the relevance of these non-linear transformations that are captured by the group equivariant embeddings above. Moreover, this embedding enables the following simple choice for the transition model
T(T(s), x(s, a)== x(s,a)T(s)
(10)
which simply transforms the state-embedding T(s) through the linear group action state- dependent action encoding s( a). Using this transition model, the action equivariance constraint of Eq. (7), and equuivariaee constraint of Eq. (8) simplifies to: Vs, a E SxA A given a state transition triplet {s, (s,,,s')
T(5') = =(),)(()))T($
(11)
Tg)(s,0)T(S) (nn((s))=nn(((s rT
(12)
In practice our model seeks to satisfy these two constraints via the direct minimization of appropriate loss functions, as will be discussed Section 5.
Decomposition of the Latent Space The decomposition G G1 x... subgroups disentangle the factors of variation in the dataset (Higgins al., 2018). This gives us a way to represent the latent embedding space as a direct product each factor varies independently by actions of subgroup of & on the latent embedding. Intuitively such a symmetry-based disentanglement provides an effective inductive bias particularly when there modularity that temporally coherent changes in the environment (sparse) subset of factors. We impose this decomposed structure in the form of a direct sum for state representation and the state- dete-deepnndnt action representation that T(s) @ Th( and k(s, a().a) D+ &&g (s, a,,,, where & e {l...., K} and (91.... .9)... Moreover the representation of the symmetry group G acting on the state embedding and the state-dependent action embedding is decomposed Tg = @ Tax and N @* Hi where 9k e GI: for (1,..., K) and G. Combining this block structure with the Lie parameterization of Section 3.1 we get
x G into a direct product of
K subgroups of G, where
due the change of
To()) = ① exxp (IAALO()))
<o(Ta(s), = eexp (F &
(13)
where we use any standard neural network implement the ag and 3e functions above 5, As we can backpropagate through this function, the network parameters 0, can be learned end to end. The
"s noted by Caselles-Dupré et al. (2019), simply having product structure in the latent space does not guarantee disentanglement, and further constraints required. In this work, do not impose any additional constraints disentanglement.
'We denote both, the neural networks which maps to group representations, and network parameters by lowercase greek alphabets.
5

---

Under review as a conference paper at ICLR 2022
choice of the subgroup depends the symmetries of the RL environment, and it only affects the set of bases (E(+); ); in Eq. (13). For example, in Atari games, the screen often has multiple objects undergoing 2-D translations and rotations, and one can use blocks of the 2-D Special Euclidean (2) that comprise translation and rotations of Euclidean space (EE), For more realistic 3D environments, such as those of interest in robotics, self-driving cars and third person games, one can use SE(3), which is the group of 3-D translations and rotations. Also, in theory, we only need specify group that "contains" the group of interest as a subgroup. For example, if our state- actions only have 90º rotational symmetry, we may use more general group for the representation (e.g., (2)). The embedding function can define a homomorphism into the relevant subgroup.
5
LOSS FUNCTIONS
We consider a standard RL setup where the agent interacts with its environments in episodes and we have access ((st, ar Tt, St+1 St+1)) where 8t is the state, at is the action taken by the agent, Trr is the reward received and 8t-1 is the observed next state at timestep t. Below we describe three loss functions that encode the equivariance//ivariance constraints of Eqs. (1), (11) and (12).
Action Equivariant Transition Loss Eq. (11) Given triplets (st, at, St+1) from our dataset we simply apply loss function 1 such as square loss that penalizes the difference between two sides:
LAET (0, Φ) = - (To(#t+1), (st, ai (s)))
(14)
The choice of the embedding space and the latent transition function ensure that state embeddings are transformed by linear group action of the action embeddings. Minimization of LAET encourages
these symmetry transformations to capture state transitions resulting from the agent's 's action.
Group Equivariant Transition Loss Eq. (12) For this we need By in addition to (st, ar), where can be any state (at different time step in the same or a different episode.) We find the group transformation that maps to s' the latent space using Tg- Te(s))Ta(s)) Using this we
can rewrite Equation Eq. (12) as
Te Te( ro(s)) Ki (o (o) (s at) (s()))
(15)
Since the state- reae-deenndent action encoding (st for the pair (st, at) is also produced by a neural network, the only missing part in the equation above is Kig the -epte-dep action transformation. We use a neural network Pw : Tg - Kÿ to infer it from state transformation Tg.
Example 1. To get an intuition for what this network is doing consider the example of a pendulum without gravity, with rotation and reflection symmetry O(2) as shown in Fig. 1, where inputs to the networks (s) are image sequences and the (ideal) embeddings To(s), Fo, a) are x-y coordinates plus angular velocity and torque respectively. If we rotate the pendulum using a rotation matrix To we expect state-dependent action embedding to remain the same since the effect of torque remains similar after rotation. However, if we transform the pendulum by reflection around the vertical axis, we expect that the effect of torque will be negated. Pw parameterizes this dependence.
A loss function l could then measure the difference between left and the right hand side in the equation above
(ww,) (To(() (To - nn ax) (oo (o:)) (st, a) ())((()())))))
(16)
Action Invariant Reward Loss - Eq. (1) While LAET and LGET enforce the equivariance of the latent transition model to an agent's 's action and the symmetry group, they do not encode information of the reward in the state representations. In order for the latent model to be homomorphic to the underlying MDP of the environment we match the reward at every state embedding using a reward predictor network ry (s) + R. We measure the difference between the predicted reward and the actual reward at time step t 1:
In(v,0.0 0)
r++1)2
(17)
'In practice use the normalized square loss (Y,Y))
Y'
1

---

Under review as a conference paper at ICLR 2022
6 APPLICATION TO MODEL-FREE RL
Following previous success of using transition models for representation learning in model-free RL. (Gelada et al., 2019; Schwarzer al., 2021) add the losses discussed above the Temporal Difference (TD) error in Deep Q-learning. In practice, we need to make three modifications to our model/loss.
Target Network trivial solution to both equivariance enforcing losses of Section is to encode all states and actions using an identity matrix. This problem in different contexts is known as the problem of collapse in representation learning. While using the reward signal helps in avoiding the collapse it is often not sufficient specially in sparse reward settings. Following Schwarzer et al. (2021), we use target network to encode state Se-1 and Sy in Eq. (16) in which the network parameters do receive gradient and moreover are copied from the online network. We explicitly drop the subscripts to differentiate the target from the online network in this section (e.g., + T).
Projection Head for Transition Losses Strict enforcement of symmetry constraints by model can be overly restrictive when the environment has non-symmetric components, or when the our transition model is too simplistic. For this reason. following the previous work, we enforce the losses on a learnable projection of the state embedding. That is before application of the loss & in Eqs. (14) and (16) we pass the embedding through a projection head.
M-step prediction Following the success of (Schwarzer et al., 2021) because of long-term state embedding predictions, we predict state embeddings and rewards for M-steps.
6.1 PUTTING IT ALL TOGETHER
Considering M consecutive state-actions {stt+ M. and it 3t = (o()), we predict the state embeddings and the rewards of next M steps:
Ptrm Ko (ft-m-1, 0t+m-1 12tm-1111--11111 and fit F Vm E {1...., M}
Note that we are using å for M-step model prediction the embedding to distinguish them from the latent embedding r. and the embedding produced by the target network J ((ttmm). The same applies to the M-step predicted reward f and observed reward r. We then project these embedding using a projection head Pę to produce êetm= Pc (f+em) and žerm p( Using this notation, our final expressions for LAFT and LR are:
2ttm
Zt-m
M Σ (frmm - Ttmm) m=1
LAFT
Σ
and LR=
(18)
12ttmll2
m=1
For LIET we need (8t, 81+1) and another state 8. From their embedding using the notation above we get Tg Tyri; -1" the linear transformation between them, and n (Tg), the state- dependent action transformation. Now for Ty. we obtain the predicted next state from Fy as Ty+1 <,)((æ,, (:) āy Pw (x )x((,, az) and from 2z41 as fv41 TyLre1. Before penalizing the difference between these embeddings, we project them to 2e+1 (&v.1) and Ør.1 b(ār-1) using projection head by, to get the final expression for LGET:
Dr+1
7r41
LGET
(19)
1
Q-learning We pass the representation de to Q-learning head d to learn policies based on the output of the Q-value estimator. The Q-value estimator is learnt by minimizing:
LDQN (g, 0) (gg (s:). oz) (re max (T(-()))
(20)
We use the data efficient adaptation of Rainbow (van Hasselt et al., 2019: Hessel et al., 2018) which combines many improvements over the original DQN(Mnih et al., 2013) such as Distributional (Dabben et al., 2018 8), Dueling DQN (Wang et al., 2016), Double DQN (Van Hasselt et al., 2016). The total loss optimized by our model is:
L = LDQN + An + AaLGET NaLAET
(21)
7

---

Under review as a conference paper at ICLR 2022
where A1. 22 and A3 hyper-parameters. Motivated by the performance improvements due to augmentation reported in recent literature (Yarats et al., 2021; Schwarzer et al., 2021), we also augment our states by shifting and changing the pixel intensity before encoding them. Fig. 5 in Appendix B.2 shows a detailed schematic of our model. We provide the algorithm for our model in Appendix B.4 and details of network architectures in Appendix B.3.
7 EXPERIMENTS
We test our method on suite of 2D Atari games, which is a popular benchmark used in RL. The full Atari suite consists of 57 games with typically 50 million environment steps. We use the sample- efficient Atari suite introduced by Kaiser et al. (2019), which consists of 26 games with only 100,000 environment steps of training data available. connected Lie subgroup blocks including General Linear G2) Special Euclidean S2(2), and Translation T(2), see Appendix A for more details. Unless stated otherwise, our EqR model uses SE))) subgroup blocks, with K = 12 blocks and M = 5 steps during training. LAET is always used to train EqR inorder to model the non-linear transfomations in the state space resulting from agent's s actions as linear group actions in the latent space. LGET, which makes transition model equivariant with respect to the symmetry transformation of state-actions, and LR are optional. We build our implementation
our experiments, we use three types of simple
top of SPR's (Schwarzer al., 2021), which is based on r pyy (Stooke & Abbeel, 2019) and Py PyTorch (Paszke et al., 2019). We use the same underlying RL algorithm and hyperparameters used by SPR for fair comparision.
Evaluation Metries We compute the average episodic return (the 'game score') the end of train- ing and normalize it with respect to human scores, as is standard practice. The human- normalized agent score random score
score (HNS) is given by human score radomm score" Since there is considerable variance across different runs, the mean and the median are not very reliable metries. Instead, Agarwal et al. (2021) pro- pose using bootstrapped confidence intervals (CI) with stratified sampling which is more suitable for small sample sizes (10 runs per game in our case). We report the Interquartile Mean (IQM), which is the mean across the middle 50% of the runs, as well as the Optimality Gap, which is the amount by which the algorithm fails to meet a minimum HNS of 1.0. We also provide performance profiles showing the fraction of runs above certain normalized score, which gives a more complete picture of the performance.
Results We use 10 seeds for every game, for every variation of our model. Figure 3 shows performance profiles for two vari- ations of our model, EqR with Ln and with LR + LAET, along with other com- parable methods. If one curve is strictly above another, the better method is said to "stochastically dominate" the other (Agarwal et al., 2021). The curves for both variations of the proposed method are almost always above the next best method, SPR (Schwarzer et al., 2021). Figure 4(a) provides results for different methods on all 26 games. The two best variations of the proposed method outperform previous methods, and the difference is statistically significant considering the CI. Table 2 in Appendix B.1 shows full results on all games, and our best model achieves super-human performance on eight games and achieves higher score than any other previous method on 13 out of the 26 games.
Zmpia
CJRL
Dig
Kcaen Butkire
tiloriaddiee with san
1.00
fractin of runs with stora
Froed on of basks with score
0.30
0.75
Humar Normalized Score (T)
Human Normaized Store
Figure 3: Performance profiles for different methods based on score distributions (left), and average score distributions (right). Shaded regions show pointwise 95% confidence bands. The higher the curve, the better the method is.
In order to better understand the effect of various modeling choices, loss functions and implemen- tation details on the performance, we now consider different variations of EqR, with the same aug- mentation as the baseline for ablation studies.
Choice of Group To understand the role of the choice of a group in the embedding space, we use our EqR model with LR. This variation of EqR is similar to DeepMDP (Gelada et al., 2019), except for the group structured latent embedding space and group action-based state transition. In order to investigate the effect of the above two group-related constructs, remove them and use an action
8

---

Under review as a conference paper at ICLR 2022
KQM
Optimality Eap
IQM
Optimality Gap
La Lest EoR. SPR
EqR. La - fazr EqR, 1a SPR
DrC
DrO CURL DER SimPLe
CURL
DER
SimPLL
0.15
0.30
0.45
0.56
0.64
0.72 0.80
0.1
0.2 0.3
0.4
0.56
0.64
0.72 0.80
Human Normalized Score
Humen Normalized Score
(a)
(b)
IOW
Optimality Gap
IDM
Octimality Gap
GL(2)
EaR, Loa EqR, FiR, l
EoR, SE(2)
EqR.
MLP
Rainbow
Rainbow
0.24 0.32 0.40 0.49
0.56
0.64
0.72
0.24 0.32 0.40 0.48
0.56 0.64 0.72
Human Normalzze Score
Human Normaizee Score
(c)
(d)
Figure 4: Plots of Interquartile Mean (IQM) and Optimality Gap computed from human-normalized scores, showing the point estimates along with 95% confidence intervals (over 10 runs for all methods, 5 runs for SimPLe). higher IQM and a lower optimality gap reflects better performance. (a) shows different methods for all 26 games, (b) shows different methods for 17 games. (c) shows the proposed model with different group choices 26 games. (d) shows the proposed model with different loss terms for all 26 games.
encoder to predict the next states directly, referring to this as MLP. We further test models with other subgroups including T(2) and GL(2). Figure 4(c) shows that adding a symmetry-base bias in the model by making the embeddings group representations and modeling the transitions as group actions is indeed helpful. The success of the model which uses SE(2) blocks might be attributed to the fact that translations and rotations are the most common types of symmetry transfor- mations present in Atari games. However, the more restrictive T(2) slightly hurts the performance, while the more general G2) performs similarly to the MLP model.
inductive
Loss functions Figure 4(d) compares the performance of EqR using SE(2) subgroup blocks with different loss components. Using EqR with the default LAET results in a considerable improvement over Rainbow with augmentation. Adding LEET improves the performance slightly, while adding only LR improves the performance even further. We hypothesize that the reward loss is playing role in both preventing representation collapse and preserving more information of the reward dis- tribution in the latent state embeddings. Adding both LGET and LR improves the performance only slightly. It might be that this prior of a equivariant transition model with respect to symmetry trans- formations of state-actions is too restrictive for some games while being beneficial for others. Based on the results in Figure 4(b), in 17 out of total of 26 games, including this loss term leads to a sta- tistically significant boost in performance. These 17 games are: 'Alien', 'BankHeist', 'BattleZone', 'oxiing,, 'ChopperCommand', 'CrazyClimber', 'DemonAttack', 'Freeway', 'Hero", 'Jamesbond', "MsPacman', 'Pong'. 'PrivateByee, 'Qbert', 'RoadRunner', 'Seaquest', 'UpNDown'. The full list of game-wor wise scores for the ablation studies are presented in Tables 3 and 4 in Appendix B.I.
8 CONCLUSION
This paper considers three major symmetry-related constructs within a coherent framework. First, there the group equivariant state and state-dependent action embedding, which we achieve through Lie parameterization. The world modeling constraints, which are discussed next, further ensure that the transformations captured by the equivariant embedding are relevant. Second, the action equivariant transition, which when combined with group equivariant embeddings, ensures that the state transitions are captured by symmetry transformations in the latent space. Third, there the group equivariant transition, which acts model itself is equivariant under symmetry transformations of state-action pairs.
additional bias and ensures that the latent transition
We provided an extensive set of experiments to evaluate the usefulness of our approach in learning state-embeddings for model-free RL. In future work we would like to explore the application of our approach in model-based RL, well as its ability to generalize across tasks. We also plan to further investigate theoretically grounded methods for combining both symmetric and asymmetric aspects of the environment in the model.

---

Under review as a conference paper at ICLR 2022
REPRODUCIBILITY STATEMENT
We have made every effort to ensure that our method reproducible. Section and Section 6 provide detailed descriptions of the loss functions and various implementation-relatee details. We present step-by-step algorithm in Appendix B.4 and have provided details of the network architecture and the list of hyperparameters used, in Sections B.3 and B.5 respectively. Section B. has individual results on each game for the proposed EqR model and its variations, to enable researchers to verify the results. Finally, we have submitted the code as part of the submission and we plan on releasing it to the public once the reviewing process is over.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare G Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in Neural Informa- tion Processing Systems, 34, 2021.
Mare G Bellemare, Salvatore Candido, Pablo Samuel Castro, Jun Gong, Marlos C Machado, Sub- hodeep Moitra, Sameera S Ponda, and Ziyu Wang. Autonomous navigation of stratospheric bal- loons using reinforcement learning. Nature, 588(7836)777-82, 2020.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence, 35(8):1798-188 18288, 2013.
Benjamin Bloem-Reddy and Yee Whye Teh. Probabilistic symmetriee and invariant neural networks. Mach. Learn. Res., 21:90-1, 2020.
Noam Brown and Tuomas Sandholm. Superhuman ai for multiplayer poker. Science, 365(6456): 885-890, 2019.
Hugo Caselles-Dupel Michael Gareia Ortiz, and David Filliat. Symmetry-basee based disentangled rep- resentation learning requires interaction with environments. Advances in Neural Information Processing Systems, 32:4606-4615, 2019.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning visual representations. International conference on machine learning. pp. 1597-1607. PMLR, 2020.
Will Dabney, Mark Rowland, Mare G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Thirty-Secone AAAI Conference on Artificial Intelligence, 2018.
Vincent François-Lavet, Peter Henderson, Riashat Islam, Mare G Bellemare, and Joelle Pineau. An introduction to deep reinforcement learning. arXiv preprint arXiv:1811.12566 2018.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In International Conference on Machine Learning, pp. 2170-2179. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.0166 2019a.
Danijar Hafner, Timothy Lillicrap, lan Fischer, Ruben Villegas, David Ha. Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2565-2565. PMLR, 2019b.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Thirty-second AAAI conference artificial intelligence, 2018.
I. Higgins, D. Amos, D. Pfau, Sébastien Racanière, Loic Matthey, Danilo Jimenez Rezende, and Alexander Lerchner. Towards a definition of disentangled representations. ArXiv, abs/1812.022 2018.
10

---

Under review as a conference paper at ICLR 2022
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. Darla: Improving zero-shot trans- fer in reinforcement learning. In International Conference Machine Learning, pp. 1480-1490. PMLR, 2017.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model based reinforcement learning for atari. In International Conference on Learning Representations, 2019.
Kacper Piotr Kielak. Do recent advancements model-based deep reinforcement learning really improve data efficiency? 2019.
Michael Laskin, Kimin Lee, Adam Stooke, Lemrll Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. arXiv preprint arXiv:2004. 14990, 2020a.
Michael Laskin, Aravind Srinivas, and Pieter Abheel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. In International Conference Machine Learning, pp. 5639- 5650. PMLR, 2020b.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. The Journal of Machine Learning Research, 17(1):13341377, 2016.
Lihong Li, Thomas J Walsh, and Michael Littman. Towards unified theory of state abstraction for mdps. ISAIM, 4:5, 2006.
Volodymyr Mnih, Koray Kavukcuoglu. David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning, nature, 51 518(7540):599-5 2015.
Arnab Kumar Mondal, Pratheeksha Nair, and Kaleem Siddiqi. Group equivariant deep reinforce- ment learning. arXiv preprint arXiv:2007.03437, 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. Advances in neural information processing systems, 32: 8026-8037, 2019.
Robin Quessard, Thomas Barrett, and William R Clements. Learning group structure and disen- tangled representations of dynamical environments. arXiv preprint arXiv:2002.0661 06991, 2020.
B. Ravindran and A. G. Barto. Symmetries and model minimization in markov decision processes. Technical report, USA, 2001.
Balaraman Ravindran and Andrew G Barto. Model minimization in hierarchical reinforcement learning. In International Symposium on Abstraction, Reformulation, and Approximation, pp. 196-211. Springer, 2002.
Balaraman Ravindran and Andrew G Barto. Smdp homomorphisms: an algebraic approach to abstraction in semi-markov decision processes. In Proceedings of the 18th international joint conference on Artificial intelligence, pp. 101 1011-1016, 2003.
Balaraman Ravindran and Andrew G. Barto. An algebraic approach to abstraction in reinforcement
learning. 2004.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive represen- tations. In International Conference on Learning Representations, 2021. URL https://openreview net http://nnnnneeeeeeeeeooooooereea f

---

Under review as a conference paper at ICLR 2022
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529 489, 2016.
Adam Stooke and Pieter Abbeel. ripyt: A research code base for deep reinforcement learning in pytorch. arXiv preprint arXiv: 19099.01000 2019.
Jonathan Taylor. Lax probabilistic bisimulation. 2008.
Valentin Thomas, Jules Pondard, Emmanuel Bengio, Marc Sarfati, Philippe Beaudoin, Marie-Jean Meurs, Joelle Pineau, Doina Precup, and Yoshua Bengio. Independently controllable features. arXiv preprint arXiv: 70000000000 01289,, 2017.
Elise van der Pol, Thomas Kipf, Frans A Olichoek, and Max Welling. Plannable approximations to mdp homomorphisms: Equivariance under actions. In Proceedings of the 19th International Conference on Autonomous Agents and MultiAgent Systems, pp. 1431-1439, 2020a.
Elise van der Pol, Daniel Worrall, Herke van Hoof, Frans Olichock, and Max Welling. Mdp homo- morphic networks: Group symmetries in reinforcement learning. Advances Neural Information Processing Systems, 33, 2020b.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. In Proceedings of the AAAI conference artificial intelligence, volume 30, 2016.
Hado van Hasselt, Matteo Hessel, and John Aslanides. When to use parametric models in re- inforcement learning? Advances in Neural Information Processing Systems, 32:14322-14333, 2019.
Ziyu Wang, Tom Schaul, Matte Hessel, Hado Hasselt, Marc Lanctot, and Nando Freitas. Dueling network architectures for deep reinforcement learning. In International conference on machine learning. pp. 1995-2003. PMLR, 2016.
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representa- tions, 2021. URL https: http::://pppeeeieeeeee
net teee/fooune forum? 1d= 6-s--cya.6s.V....
Jure Zhontar, Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. arXiv preprint 1xXi::2003.03230, 2021.
Martin Zinkevich and Tucker Balch. Symmetry in markov decision processes and its implications for single agent and multi agent learning. In In Proceedings of the 18th International Conference on Machine Learning. Citeseer, 2001.
A SUBGROUP BLOCKS AND THEIR PARAMETERIZATION
A.1 CHOICE GROUP
Atari games differ in their style of play, their objectives, the symmetry transformations of both the agent and other objects on the screen and associated symmetry transformation of the agent's 's action. But most of these games include symmetry transformations. For example, the screen often has mul- tiple objects undergoing two dimensional(2-D) translations and rotations. In this case one can use blocks of the 2-D Special Euclidean Group S2()). Each such block can capture the transformation of a particular object in the screen, including the agent. One can also use more restrictive subgroup blocks like T(2), which capture only 2D translations. For more realistic 3D environments, such as those of interest in robotics, self-driving cars and third person games, one can use SE(3), which is the group of 3-D translations and rotations. This should capture both the transformations of the objects in the environment and changes in viewpoint due to the agent's actions.
A.2 IMPLEMENTING PARAMETERIZATION
We consider three types of subgroups: GL(n) the set of all invertible linear transformations. SE(n) the set of all rotations and translations and T(n) the set of all translations in an n dimensional vector space. We provide general method to parameterize each of these, based on the type of group.
12

---

Under review as a conference paper at ICLR 2022
GL(n) As the matrix representation of GL(n) is the set of invertible matrices which has measure of is easy to parameterize it. We just generate n² parameters using a network corresponding to each element of the matrix. This gives an element from GL(n).
T(n) As T(n) just denotes translation in a n-dimensional space with group action being addition, implementing it is straightforward. We generate parameters using neural network and instead of using matrix multiplication use addition for the group action. Note that we can also use matrix representation for T(n) but it is unnecessary and inefficient.
SE(n) Unlike GL(n) and T(n), parameterizing SE(n) is bit tricky because it involves pa- rameterizing SO(n). We use homogeneous co-ordinate based representation of SE(n)
{( f). Re SO(n) and l T(n) So we need parameters for the L and another D =(n-1)
parameters for SO(n) from the neural network. As explained Section 4, we can use Lie pa- rameterization to get the elements of (0 by R ==exx exp(ex 3g E() where E(d) denote bases of the space of skew symmetric matrices and the B& the parameters of the neural network. For example, in the case of SO(2) we can use the basis ( = (-) !) Similarly, we can extend this
to (O(n) by using a basis given by D n x matrices ( V s si<jsn) whose only non-zero elements E E -1 and E(())
Although Lie parameterization gives general recipe to output representation of simple nected Lie groups like (n),,, in our implementation we use Euler parameterization because it runs faster in Pytorch. We provide the code for both. Following Quessard et al. (2020), parameter- ize each rotation matrix in SO(n) using the product rotations on D orthogonal planes in R": plane, and its non-zero elements besides the diagonal are the four values on the i,j rows and columns, which comprise the
R IT I Isicjjnn R Here RV Ruin is the rotation matrix in the
cos(O..) sin ain(O.)) We have D parameters fs. which 2D rotation matrix that is R sin(O:.)
cos 0i,))
obtain from a neural network.
The parameters in all the parameterization techniques mentioned here can be back prok-ppopagated. We summarize the number of parameters required from neural network output, representation type and the associated group actions of different subgroups in Table 1.
Table 1: Group Properties
Subgroup block type
Parameters
Representation type
Group action
General Linear GL(n)
n²
Matrix
Matrix multiplication
Special Euclidean - SE(n)
m(n-l) n n(--1)
Matrix(
Matrix multiplication
Special Orthogonal SO(n)
Matrixaxn
Matrix multiplication
2
Translation - -T())
Vectorn
Addition
13

---

Under review as a conference paper at ICLR 2022
B ATARI DETAILS
B.1 FULL RESULTS
We provide individual results on the Atari games after 100K training steps. Our results are aver- aged over 10 seeds, and the network architectures and full list of hyperparameters used to produce them are provided in Appendix B.3 and Appendix B.5.
Table compares our two best performing EqR models using SE(2) subgroup blocks with other methods.
Table 3 compares different choices of subgroup blocks with the reward loss, LR, included for all EqR models (also see Figure (a)).
Table 4 compares EqR using SE(2) subgroup blocks with different loss terms included in the training objective (see Section 5 and Figure (b)). The action equivariance transition loss, LAET always included for EqR models.
Table 2: Mean game scores on the 26 Atari games after 100K environment steps. The EqR models use (2) subgroup blocks along with an action equivariant transition loss, LATT, and are averaged over 10 seeds.
Game
Random
Human
SimPLe
DER
CURL
DrQ
SPR
EqR, Lx
EqR, La LOST
7127.7
616.9
739.9
Amidar
5,8
1719.5
88.0
188.6
142.1
102.8
176.3
140.9
138.4
Assault
527.2
Asterix
210.0
8503.3
1128.3
470.8
734.5
603.5
977.8
923.2
902.5
34.2
BattleZone
2360.0
37187.5
5184.4
10124.6
14870.0
12954.0
16651.0
13044,0
13255.0
Boxing
12.1
9.1
0.2
35.8
37.5
39.2
Breakout
1.7
30.5
16.4
1.9
4,9
16.1
17.1
17.2
16.0
Command
811.0
7387.8
1246,9
861.8
1058.5
780.3
974.8
1073.5
1142.2
Cracy Climber
10780.5
35829.4
62583.6
16185.3
12146.5
20516.5
42923.6
49399.0
52008.1
Demon Attack
152.1
1971.0
208.1
508.0
817.6
1113,4
545.2
531.4
532.1
Freeway
29.6
20.3
27.9
26.7
24.4
24.1
25.2
Frestbile
4334.7
254.7
866.8
1181.3
331.1
1821.5
1855.6
16994
Gupher
257.6
2412.5
771.0
349.5
669.3
636.3
715.2
1010.0
912.1
Hero
1027.0
30826.4
2656.6
6837.0
6279.3
3736.3
7019.2
5775.2
6118.5
Jameshond
29.0
302,8
125.3
3016
471.0
236.0
365.4
312.8
319,7
Kangaron
$2.0
3035.0
323.1
779.3
872.5
940.6
3276.4
3569.3
3296.0
Krull
1598.0
2665.5
4539.9
2851.5
4229.6
4018.1
3688.9
$614.5
5467.7
Kung Fu Master
258.5
22736.3
17257.2
14346.1
14307.8
9111..0
13192.7
18511.0
17510.9
Ms Pacman
307.5
6951.6
1480.0
1204,1
1465.5
960.5
1313.2
1317.1
1663.5
Pong
-2077
14.6
12.8
-19.3
-16.5
-8.5
-5.9
-6.0
-6.1
Private
69571.3
58.3
97.8
218.4
-13.6
124.0
76.6
88.9
Qhert
163.9
13455.0
1288.8
1152.9
1042.4
854.4
669.
773.8
814.9
Road Runner
11.5
7845.0
5640.6
9600.0
5661.0
8895,1
14220.5
13385.0
13708.8
Seaquest
68.4
42054.7
683.3
354.1
384.5
301.2
$83.1
650.3
697.9
Up N Down
533.4
11693.2
3350.3
2877.4
2955.2
3180.8
28138.5
44295.4
52118.4
Mean Human-Nomn'd
0.000
1.000
0.443
0.285
0.381
0.357
0,704
0.859
0.886
Median Norm'd
0.161
Superhuman games
N/A
2
14

---

Under review as a conference paper at ICLR 2022
Table Mcan game scores on the 26 Atari games after 100K environment steps for different choices of subgroup blocks, averaged over 10 seeds. The reward loss, LR, is included in addition to the default loss LAET-
Game
Random
Human
MLP
EqR. 73
EqR, SEa
EqR. GL2
Alien
227.8
7127.7
780.1
846.4
774.0
881.3
Amidar
1719.5
139.7
140.9
132.2
Assuut
222.4
742.0
701.5
684.0
753.8
692.3
Asterix
210.0
8503.3
973.6
1004.4
923.2
889.5
Bank
14.2
395.1
430.2
BattleZone
2360.0
37187.5
12722.6
11500.0
13044.0
13114.0
Boxing
0.1
12.1
38.0
28,9
37.5
33.4
Breakout
15.4
Chopper Command
811.0
7387.8
989.5
1028.9
1073.5
1088.1
Crazy Climber
10780.5
35829.4
43705,8
50822.1
49399.0
55018.5
Demon Attack
152.1
1971.0
518.6
544.2
5314
$10.2
Freeway
0.0
29.6
20.3
18.5
24.1
21,4
Frostbite
65.2
4334.7
1702.4
1653.8
1855.6
1797.7
Gopher
257.6
1012.5
1010.0
894.4
Hero
1027,0
30826.4
6840.0
5779.8
5775.2
5934.7
Jameshond
29.0
302.8
337.4
313.25
312.8
334.8
Kangaroo
52.0
3035.0
2994.8
2942.5
3569.3
3186,4
Krull
1598.0
2665.5
3801.5
5293.0
5614.5
5772.6
Kung Master
258.5
22736.3
13780.4
14924.2
18511.0
16002.8
Ms Pacman
307.3
6951.6
1220.8
1166.8
1317.1
1147.7
Pong
-20.7
14.6
-6.1
-11.5
-6.0
-8.2
Private Eye
24.9
69571.3
72.4
65.1
76.6
55,9
Obert
163.9
13455.0
678.4
763.2
773.8
635.2
Road Runner
11.5
7845.0
12763.2
13654.2
13385.0
12560.4
Seaquest
68.4
42054.7
656.9
647.3
650.3
633.3
Up N Down
533.4
11693.2
23130.6
58164.4
44295.4
43767.2
Mean Homan-Norm'd
0.000
1.000
0.681
0.829
0.859
0.833
Median Humun-Norm's d
0.000
1.000
0.398
0.361
0.418
0,380
# Superhaman games
N/A
7
Table 4: Mean game scores on the 26 Atari games after 100K environment steps for EqR using SE(() subgroup blocks with different loss terms included in the training objective. The action equivariance transition loss, LAET, is included for all EqR models and the scores are averaged over 10 seeds.
Random
Human
EqR
EqR, LR
EqR, LEEr
EqR, Lg
LCET
Alien
227.8
7127.7
774.0
862.5
Amidar
5.8
1719.5
134.7
140.9
135.0
138
Assault
222.4
742.0
643.1
753.8
701.3
734.3
Asterix
210.0
8503.3
824.8
923.2
864.9
902.5
Bunk Heist
14.2
753.1
407.3
395.
335.9
397.4
BattleZone
2360.0
37187.5
12805.6
13044.0
12990.4
13255,0
Boxing
0.1
12.1
32.7
37.5
34.8
39,2
Breakout
1.7
30,5
14.6
17.2
14.8
16.0
ChopperCommand
811.0
7387.8
1015.6
1073.5
934.8
1142.2
Crazy Climber
10780.5
35829.4
38483.9
49399.0
43085.6
52008.1
Demon Ataak
152.1
1971.0
$23.8
$31.4
504.6
532.1
Freeway
0.0
29.6
22.1
24.1
22.5
25.2
Frostbite
65.2
4334.7
1635.2
1855.6
1563.9
1699.4
Gopher
257.6
2412.5
695.3
1010.0
789.3
912.1
Hero
1027.0
30826.4
5763.9
5775.2
5603.8
6118.5
Jameshond
29.0
302.8
388.4
312.8
344.9
319.7
Kangaroo
$2.0
3035.0
2667.9
3569.3
2848.7
3296.0
Krull
1598.0
2665.5
4209.2
5614.5
4411.2
5467.7
Kung Fu Master
258.5
22736.3
12287.9
18511.0
16394.6
17510.9
Ms Pacman
307.3
6951.6
1141.3
1317.1
1514.7
1663.5
Pong
-20.7
14.6
-9.9
-6.5
Private Eye
24.9
69571.3
73.2
76.6
87.5
88.9
Qbert
163.9
13455.0
696,7
773.8
736.8
814,9
Road Runner
11.5
7845.0
12659.2
13385.0
13110.4
13708.8
Seaquest
68.4
42054.7
593.6
650.3
641.0
697.9
Up N Down
533.4
11693.2
29425.4
44295.4
39076.6
$2118.4
Mean Hluman-Nomits
0.000
1,000
0.682
0.859
0.749
0.886
Mediun Human-Norm's
0.000
1,000
0.337
0.418
0.377
0.398
Superhuman games
N/A
8
15

---

Under review as a conference paper at ICLR 2022
B.2
MODEL SCHEMATIC
Reward Loss
T+1
T++
7+1
N+
T
Tv
a:
9p+1
Q-learning (Rainbow)
g
T
t
98
Kg (3px 91)
( (11)))
LDON
X1
1
&
Pu
St
gg
2g
X F+1
X X+M
Pc
Pc
Action
X
2
Equivariant
2
Z
xx 4))
Loss
X
bg
LEET
De+1
Group
latent representation action dependent
Equivariant
y
transition the environment
Loss
LGET
function with parameter
copy parametric function fe
(no gradient flow)
linear group action in the
non-linear group action the original space
latent space
A schematic of the EqR model, applied to model-free RL. Green in the framework cor- responds to learning equivariance under the agent's action and red corresponds to learning equiv- ariance of the transition model with respect to symmetry transformation of the state-action. This color scheme is consistent with Figure 2. The part of the framework that corresponds to reward matching and Q-learning is shown in blue and brown respectively. The arrows in the schematic are differentiated by their heads and are described in the legend.
Figure 5:
B.3 NETWORK ARCHITECTURE
We follow the baseline RL implementation of DrQ (Yarats al., 2021) and SPR (Schwarzer al., 2021) by using the 3-layer convolutional encoder from (Mnih et al., 2015) and then use linear layer get the parameters for the Group Parameterization. The output size of this layer varies depending on the group type, the number of blocks used and the size of the group. This defines our тę. Note that the output of our encoder other neural network like the Q-head g-(-)
matrix for GL( and (E(n). We flatten before feed to
For the action encoder x(-) we use simple 1 layer MLP with batchnorm, ReLU and hidden size of 256. We concatenate the one-hot encodings of the actions with the state representations coming from To and pass it through the action encoder to get matrix representation of the group after parameterization.
For the reward predictor network ry we use a 2-layered MLP with batchnorm, ReLU and a hidden size of 256.
For the Q-head (.) we use 2-layered MLP as well.
For the projection head pc(-) we share the first layer of Q-head whereas for projection head (.) we use single layer MI.P.
16

---

Under review as a conference paper at ICLR 2022
B.4 ALGORITHM
Algorithm 1: Equivariant Representations for RL
Denote the parameters of online networks To, Kig, Pc, by as O.,
Denote the parameters of target networks т, fi, p, O,
Denote the parameters of networks Pw. 9g as d Denote the dept of the prediction as M and batch size N Initialize the replay buffer B while Training do
Collect {s, r. s'} using policy with (0.. ②) and add to the buffer B Sample minibatch of M length sequences {St:M, 4o:M: To:M) for in range(0, N) do if augmentation then
$..M (ugment(s(.v))
end
w To (s));
// state representation
2 z0
--0
for l in range(1 1, / 1) do
$ - (f 1- k- 1 // state transition by group action
z T(s);;
// target state representation // projections
& P5 p,( (ft), 2 Pc (x())):
+ +A2l TRu
Brnk
Book TR
// compute LAET at step
A (
// predict rewards // compute La at step
p-
1 -
end
j j- (0..,, N -1};
// uniformly sample an index 2/ - 7( (s)); // encode the state for that index from the batch Tg = 2
// find the group representation // next state by group action // next state by action-embeddi ng // projections // compute LEFT
8 + oj
1 Pa (T)))(Z)) d
Ø (H)), 1 - b(T));
I' - dal 11
10
10 10
- 1' ( To. #1:4)
end
1 N Σ li o, ① oplmize((O0) 1),1);
// average over minibatch // update online networks // copy weights to target networks
O.+ + O,: 0,:
end
17

---

Under review as a conference paper at ICLR 2022
B.5 HYPERPARAMETERS
In this section, we provide the full set of hyperparameters in our model. As mentioned earlier, our baseline RL algorithm closely follows SPR's (Schwarzer al., 2021) implementation of Rainbow and hence we use most of their hyperparameters setting in order to be able to compare to them. Note that the weights of LR X1, LAET X2 and LGET λ are set one whenever they are used in the model.
Table 5: Hyperparameters for ErQ (including variations) on Atari.
Parameter
Setting
Gray-scaling
True
Observation down-sampling
84 x 84
Frames stacked
4
Action repetitions
4
Reward clipping
[-1, 1]
Terminal on loss of life
True
Max frames per episode
108K
Update
Distributional Q
Dueling
True
Support of Q-distribution
51
Discount factor
0.99
Minibatch size
32
Optimizer
Adam
Optimizer: learning rate
0.0001
Optimizer: 31
0.9
Optimizer: F3
0.999
Optimizer:
0.00015
Max gradient norm
10
Priority exponent
0.5
Priority correction Exploration
0.4 1
Noisy nets
Noisy nets parameter
0.5
Training steps
100K
Evaluation trajectories
100
Min buffer size for sampling
2000
Replay period every
step
Updates per step
2
Multi-step return length Prediction depth, M
10
5
A1
I
A2
Ag
Data Augmentation
Random shifts (+4 pixels)
Intensity(scale00.05)
Parameter
Setting (T2)
Setting (SE2)
Setting (GL2)
Num Blocks (K)
32
12
12
Group Action
Addition
MatMul
MatMul
18


Reviews:
[{"review_id": "xT7RJCX491v", "reviewer": "Reviewer_pfpn", "summary": "", "questions": "", "limitations": "", "rating": 8, "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": "Overall, I think this is a quite interesting paper. It presents a very principled approach to learning better representations for RL in a data-efficient manner. Perhaps the improvement in performance is not very dramatic (and parts of this were already in earlier publications), but I think the motivation, the technical development, and the final presented technique are all of interest to the community and will be valuable contributions.\n\nI don't have many concerns about the paper but I would have liked to see some analysis of the learned representations and the transition models. Are there any interesting patterns emerging? How does the new learned state/action space look like?\n\nAlso, it was not clear to me why adding the loss term L_GET led to worse results (in median scores). It would be nice to discuss this briefly. This might be perhaps the assumed group structure is not really suitable for some of the games. \n\nSome other points\n\n- It would be nice to have colors in Figure 2. As it stands, it is a little bit hard to parse.\n- It would be nice to mention what group SE(2) is at some point (perhaps in the Experiments section).\n- L_AET is always included in training but this is only mentioned afterwards. It might be better to mention this earlier.\n\n", "summary_of_the_paper": "This paper presents a representation learning technique for data-efficient RL that makes use of invariance/equivariance of the environment wrt both state and agent actions. This work builds on two ideas: MDP homomorphism and symmetries in state-actions due to some inherent structure in the environment. MDP homomorphism refers to a mapping of states and actions of an MDP to a new space where the MDP reward structure and dynamics are preserved. Such a homomorphism can be useful if the new mapped state-action space is easier to work with. The second idea is that we may know that an environment exhibits certain regularities (i.e., may have a group structure) where we know the effect of certain operations. Then building these regularities directly into our models can then make training much more data efficient.\n\nIn this paper, the authors combine these two ideas (MDP homomorphism and symmetries in MDPs) such that state and actions are mapped to new spaces that satisfy the MDP homomorphism criteria along with the invariance/equivariance relations that result from the structure of the environment. This is achieved by minimizing a loss function that consists of 3 terms. The first term (L_AET) encourages the transition dynamics in the learned space to be the same with the original MDP. The second (L_GET) encourages the effect of group actions to be equivariant on the learned space. And finally the third one (L_R) encourages the rewards in the learned space to be the same as the original MDP.\n\nThe authors test their technique (in a Rainbow DQN agent) on Atari levels under data limited setting and show that their technique, albeit slighlty, outperforms previous techniques.", "summary_of_the_review": "Overall, I think the technique presented in this paper is well-motivated and addresses an important problem in RL. Perhaps the improvement in performance is small but I think the technical development and the final technique are valuable contributions.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "j-QYh0F05__", "reviewer": "Reviewer_dVfB", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "### Strengths\n\n- The paper is well-motivated and addresses an interesting and important problem.\n- The idea of embedding states and actions as matrices representing a group action is intriguing, and may induce interesting properties in the representation beyond those studied in this paper.\n- The paper provides sufficient background to be understood by readers with some but not extensive experience in the topic.\n\n\n### Weaknesses\n\n- **The justification for the proposed method is lacking** — section 4 provides a number of desiderata, but there is no follow-up showing empirically or theoretically that the resulting method (especially given the use of target networks and projections) will actually satisfy these desiderata.\n- **The use of Lie groups is not motivated in Section 3**. As a reader, I presumed that this class was chosen in order to justify the use of learned matrix embeddings in the proposed model, but this would be worth clarifying. Additionally, the basis of the Lie algebra is not referred to in Section 6 despite an indication that it will be used in Section 3.\n\n- **The paper suffers from a lack of clarity.** For example, in the discussion of symmetric MDPs the group action is applied to both states and state-action pairs, but it is not clear whether these represent two separate actions, or whether $g\\cdot \\langle s, a \\rangle$ has the same effect on states as $g \\cdot s$. Additionally, the notation $\\kappa (\\mathcal{G})$ and $\\tau(\\mathcal{G}_{\\mathbb{S}})$ is not clear on a first read. In general I think more explicitly stating the types of the mathematical objects used in the paper will make it much easier to parse. By default, I and probably many other RL readers default to assuming vector-valued embeddings, and it is easier to override this default when the output type of the embedding is made explicit.\n\t- Section 4 is quite long and in many cases it’s not clear whether the proposed properties are _desiderata_ or guaranteed to be satisfied by the proposed method. For example, the matrix embeddings of state actions seems to hold by construction but disentanglement presumably does not.\n- **The empirical evaluations don’t support the main goal of the paper.** Adding the group-equivariant loss doesn’t provide much improvement over the standard model-learning loss — particularly when we look at median performance. This suggests that most of the benefit over the baselines is coming from the model-learning component of the objective rather than the incorporation of symmetries. Further, the method without the symmetry-learning loss is similar to many existing representation learning methods such as DeepMDP, PSR, PBL, MuZero, LatCo, and DBC. The equivariant loss is therefore critical to the novelty of the proposed algorithm, yet it doesn’t seem to have a significant impact on overall performance. It is possible that a more careful empirical analysis on environments with more explicit symmetries may reveal that this loss is indeed useful, but the current experimental results are insufficient to show this to a satisfactory degree.\n    - The experimental results are limited. The paper only uses one evaluation scheme and doesn’t provide much insight into whether the learned model really is equivariant to the set of transformations used. \n    - Further, the baselines used in the paper don’t seem to be described or cited, which makes it difficult to determine whether a fair comparison is really being made between the different methods. The baselines I was already familiar with mostly seem to make use of model-free approaches combined with data augmentation, which leaves the model-learning component of EqR as a confounding factor. SPR, which does include a learned latent-space model, seems to attain a higher median human-normalized score than EqR with the equivariant transition loss, suggesting that at least part of the improved performance is coming from the specific parameterization of the model. \n\n### Concrete avenues for improvement\n\n- I am very intrigued by the use of matrix embeddings in the representation — typically in RL we treat representations as vectors, and it would be interesting to see whether enforcing linear actions of the transition and symmetry operators on the latent space induces particular structure. \n- It is possible that by looking at environments with more structure, the proposed method might be shown to capture the appropriate in/equivariances and therefore improve performance and sample efficiency. I recommend incorporating evaluations on e.g. Mujoco or some other environment with more explicit physically-based invariances. \n- The focus of this paper is on equivariant representations for sample efficiency; however, it seems like an arguably more promising application is in improving generalization. \n- Ensuring that symbols and their types are defined before they are used will make the paper much more readable. \n- Section 4 should be rewritten to more clearly express which of the desiderata are definitely satisfied by the proposed method, and which are not guaranteed to be satisfied. It might also be interesting to include empirical investigations into the degree to which the proposed method satisfies these desiderata. ", "summary_of_the_paper": "The paper focuses on the problem of incorporating symmetries into RL agents’ representations to improve data efficiency. It proposes a model-based representation learning method parameterized by a group action that is based on encoding states and state-action pairs as elements of the group's representation in latent space. The proposed method is evaluated against a number of baselines in the data-limited Atari setting.", "summary_of_the_review": "I recommend rejecting this paper due to a) the limited empirical benefit of incorporating equivariance into the proposed representation learning method and b) the lack of clarity in both the paper’s exposition and its claims about the proposed method’s properties. With additional work and a more informative evaluation setting, I think the paper will be in a much better position for consideration in future conferences.", "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": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "CNI1l-nvuup", "reviewer": "Reviewer_1BRN", "summary": "", "questions": "", "limitations": "", "rating": 8, "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 the paper to be very well written. I think the idea and method were clearly explained despite the complexity of the overall model, combined with all the moving parts, well done! \n\nThe author's perspective on equivariance in RL is interesting and novel (to the best of my knowledge). I was convinced by the motivations outlined and the consequently design choices. The proposed method makes sense, and is backed by interesting results.\n\nI think it would benefit the paper to make it clear early on that the choice of the loss depends on a choice of subgroup blocks that depends on prior knowledge the environment (e.g. 2d translation, rotation, etc.) as this didn't come clearly to me until Example 1 is presented. I appreciate that the appendix includes a discussion on the choice of group for Atari games, but it would be better to have more general discussion in the main text. \n\nI understand that authors wanted to show performance across many different environments to showcase the robustness of the performance, but I would have liked to see more in-depth experimental results in one or two tasks including learning curves, specially given that data-efficiency is the main desired property. \n\n\nminor comment:\nI believe the b_{\\eta}  above Eq 18 should be g_{\\eta}?", "summary_of_the_paper": "The paper studies the question of recovering representations that are equivariant under the agent's action and under the symmetries of the environment. The representation are symmetry-based, using the state-transition symmetry (to create a strong inductive bias robust to symmetries that are not present in the reward) in a latent transition model (to enable learning of the transition model in the latent space, and to allow for linear assumption of the group action through the representation). A decomposed structure is imposed on the representation of the symmetry group, such that it allows to represent the latent embedding space as a direct product of subgroups. The authors propose 3 losses that encode the aforementioned constraints of the latent state and state-action embedding, as well as the invariance constraint of the reward.", "summary_of_the_review": "The paper is well written and introduces a very interesting perspective on equivariant representations in RL. The method proposed makes sense to me, and the results seem promising. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "BHQYNuv3xa", "reviewer": "Reviewer_c4oF", "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": "Strengths:\n\n* Leveraging invariance and equivariance as inductive bias for representation learning in RL makes sense. The proposed way to encode the equivariance / invariance constraints in the latent representation space is novel and interesting.\n* The writing is in general clear and not too difficult to follow.\n\nWeaknesses:\n\n* One of my concerns about this paper is the performance comparison with SPR. From Table 1, the performance of (EqR, $L_R$) is very close to SPR in the median score. Since on Atari median score is often considered as a more reasonable performance measure than mean score, I would not be convinced that EqR outperforms SPR.\n* Besides, it seems (EqR, $L_R + L_{GET}$) performs worse (in terms of median score) than (EqR, $L_R$) in Table 1 and Figure 3(b). The novel part of this paper, $L_{GET}$, does not contribute much to the performance. Given that the method is also very similar to SPR in some implementation choices, I am not convinced that the proposed method is better.\n* To help readers better understand the math in Section 3.2 and Section 4, I suggest the authors give more examples during these 2 sections to provide a context (possibly with the pendulum example).\n\nOther comments:\n\n* In Equation 5, can the same $g$ act both on element in $\\mathbb{S}$ and on element in $\\mathbb{S} \\times \\mathbb{A}$ ? \n* In the first paragraph of Section 4, $\\mathcal{G}_T$ and $\\mathcal{G}_R$ are not defined.\n* I am not quite getting the sense of using $\\tau_\\theta(s_{t'})\\tau_\\theta(s_t)^{-1}$ as $\\tau_g$. Can the authors provide some explanations?\n* Typo: above Equation 17, duplicate \"using\".\n* Above Equation 18, $b_\\eta$ and $b$ are not consistent with Figure 2.", "summary_of_the_paper": "This paper presents a latent variable model for representation learning in RL, taking into consideration both equivariance to an agent’s action and symmetry transformations of the environment. The proposed method (EqR) outperforms several previous methods in data-efficient (100k) Atari benchmark.", "summary_of_the_review": "While I do find the idea of this paper novel and interesting, the empirical evidence does not convince me of its superiority over prior methods, in particular SPR. I will not recommend acceptance for its current form.", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes to learn a latent space representation such that some linear equivariance and symmetry constraints are respected in the latent space, with the goal to improve sample efficiency. One core idea is that the latent space is also the same as the space of linear transformation used in the constraints, which is shown to simplify some of the mathematical derivations. Experiments on the Atari 100K benchmark demonstrate a statistical improvement over the SPR baseline when using the SE(2) group of linear transformations as latent space.

Following the discussion period, most reviewers were in favor of acceptance. However, one reviewer remained unconvinced, and after carefully reading the paper, I actually share the same concerns, i.e., that it is unclear under which conditions the proposed approach actually works, and what makes it work. I believe that, as a research community, we should value understanding over moving the needle on benchmarks, especially when proposing such a complex method as this one (see Fig. 5).

More specifically:

1. The method is only evaluated on Atari games, showing some improvements when using SE(2), and arguing that there are corresponding symmetries in such games. There is however no analysis demonstrating (or even hinting at the fact) that the proposed technique is actually learning to take advantage of such symmetries (NB: I had a quick look at the animation added by the authors in the supplementary material, but I do not see if/how they help on this point). Even if analyzing representations on Atari may be tricky, I believe that given the motivation of this new algorithm, it *must* be evaluated on some toy example (e.g., the pendulum mentioned throughout the paper) to validate that it is learning what we want it to learn (although I also agree with the authors that experimenting on a more complex benchmark like Atari is equally important).

2. The idea of embedding states into the same space as transformations is interesting, and brings some advantages when writing down equations, as demonstrated by the authors. However, there is no justification besides mathematical convenience, and it doesn't seem intuitive to me at all that why this should be a good idea, considering that it ties the state representation to the mathematical representation of group transformations. For instance, what does the spcial group element $e$ mean for a state? And this coupling makes it difficult to interpret the effect of using a different group of transformations: for instance when moving from GL(2) to SE(2), is the observed benefit because we are using only specific transformations, or simply because we are reducing the dimensionality of the state embedding? (note that in Fig. 4(c) the MLP variant has similar performance to GL(2), and based on my understanding they use the same embedding dimensionality  ==> I believe it would be important to check what would happen with an MLP variant using the same dimensionality as SE(2))

3. The effect of the $L_{GET}$ loss is not convincing, as pointed out by several reviewers. I think it would have been an opportunity for the authors to investigate why, especially since it seems to work in some games and not others. But just focusing on "here are the 17/26 games where it works better" doesn't really bring added value here. Do these games have some specific properties that make them better candidates to take advantage of $L_{GET}$? This could have been a very interesting insight if that was the case, but as it is now, I am not sure what we can learn from that.

4. There are several implementation "details", some moving the final algorithm farther from its theoretical justification, that are not ablated, making it difficult to understand their impact (ex: using target networks, the choice of the value of M, using projections onto the unit sphere of some arbitrary dimensionality, how the $s'$ state is chosen in $L_{GET}$)

As a result, we have here an algorithm with some interesting theoretical background, but with a lot of moving components which -- when properly tweaked -- can lead to a statistically meaningful improvement on Atari 100K -- without really understanding why. I believe this is not quite enough for publication at ICLR, and I would encourage the authors to delve deeper into the understanding of their algorithm, which I hope will bring useful insights to the research community working on representation learning.
Decision: Reject