Year: 2021

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 nonlinear 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 other comparable methods on statistically reliable evaluation metrics.

heading: 1 INTRODUCTION
text: The recent success of deep reinforcement learning (François-Lavet et al., 2018) in applications to games such as Atari (Mnih et al., 2015), Go (Silver et 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 a 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 RL 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 is to place a central role on transformations of the data, where invariance and equivariance to a set of transformations imposes strong conditions on the learned representations. This viewpoint is particularly appealing in RL, where the agent is in control of some of these transformations through its own actions. Fig. 1 illustrates this concept using the example of a 2D pendulum. Moreover,
transformations naturally lead to a 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-supervised learning approaches that rely on a (predefined) set of transformations (Chen et al., 2020; Zbontar et al., 2021), and also within the context of RL (Yarats et al., 2021; Laskin et al., 2020b), further highlights the importance of transformations in data-efficient representation learning.
Motivated by these observations, this work develops a broader perspective on the notion of equivariant representation learning within RL. In particular, we integrate equivariance under the agent’s action and equivariance under the symmetries of the environment into a single latent variable model that is equivariant to an a 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 representations equivariant to continuous transformations of the environment resulting from agent’s action. We benchmark our approach, which we call Equivariant representations for RL (EqR), on the 26 games in the Atari 100K benchmark (Kaiser et al., 2019), where we 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 can be expressed using matrix Lie groups, including autonomous driving, navigation, and robotics.

heading: 2 RELATED WORK
text: The use of transformations, be it in data-augmentation or self-supervision, has become a 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. A 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 dataefficiency. 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 et al., 2020a; Gelada et al., 2019).
Other works in RL that are relevant to our objective are those that attempt to increase data-efficiency using a learned model of the environment. While some methods such as SimPLe (Kaiser et al., 2019), learn this transition model at the pixel level, the majority of methods use a 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 et al. (2020a) and DrQ (Yarats et al., 2021). Using a reconstruction-based method is 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-based 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-art results in sample-efficient representation learning in RL.

heading: 3 BACKGROUND

heading: 3.1 GROUPS AND THEIR REPRESENTATIONS
text: 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 ∈ G has a unique inverse, such that their composition gives the identity g−1g = e. Any subset G′ ≤ G that is closed under binary operation of the groups forms a subgroup. A group G can act on a set X by transforming its elements x ∈ X through a
bijection. We use α ∶ G×X↦ X to denote the group action, and for brevity replace α(g, x) with g ⋅x moving forward. The action captures some of the structure of G due to two constraints – the identity element acts trivially e ⋅ x = x; and composition of actions is equal to action of the composition, i.e., (gg′) ⋅ x = g ⋅ (g′ ⋅ x),∀g, g′ ∈ G. X is then called a G-set. Any G-action partitions X into orbits xG = {g ⋅ x ∣ g ∈ G}, and we denote the set of orbits under G-action using X/G. A G-action is transitive iff its action results in a single orbit.
Parameterizing Lie Groups In this work, we assume G is (any 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 is linear on real vector spaces. We use ρ(G) to denote a linear representation of G, and ρg ∶ RD → RD for the action (a.k.a. the representation) of g ∈ G. Two other greek letters τ and κ are also used for this purpose. Many such Lie groups are identifiable by their Lie algebra g = Lie(ρ(G)).1 This connection enables a simple parameterization of ρ(G) using a set of linear bases for their Lie algebra – that is ρg = exp(∑i βg,i E(i)), where exp(Y) = ∑∞j=0 Y j j! is the matrix exponential. We refer to this parameterization later in Section 4. Such linear representations in the form of invertible matrices can be used for both continuous transformations (e.g., 3D rotation) and finite groups (×90○ rotations).

heading: 3.2 MDP HOMOMORPHISM AND SYMMETRIC MDPS
text: We define an MDP as the 4-tupleM = ⟨S,A,R, T ⟩ where S and A are respectively the set of states and actions, R ∶ S × A → R is the reward function, and T ∶ S × A × S → R≥0 is the state transition function.2 For two MDPsM = ⟨S,A,R, T ⟩ and M̄ = ⟨S̄, Ā, R̄, T̄ ⟩, MDP homomorphism can be defined as a tupleH = ⟨hS, hA⟩ where hS ∶ S→ S̄ is the state mapping and hA ∶ S×A→ Ā 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) ∀s, a ∈ S ×A (1) 2) Equivariance of the deterministic transition model under the agent’s action:
T̄ (hS(s), hA(s, a)) = hS(T (s, a)) ∀s, a ∈ S ×A (2) A probabilistic variation of the above equation for a stochastic MDP (Bloem-Reddy & Teh, 2020) is:
T̄ (hS(s′) ∣ hS(s), hA(s, a)) = ∑ s′′∈[s′]h T (s′′ ∣ s, a) ∀s, s′ ∈ S, a ∈ A, (3)
where [s′]hS = hS−1(hS(s′)) is the equivalence class of s′ under hS. 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 forM.
Symmetric MDPs The automorphism group GM = 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) = R(g ⋅ ⟨s, a⟩)) ∀g ∈ GM, s ∈ S, a ∈ A (4) T (s′ ∣ s, a) = T (g ⋅ s′ ∣ g ⋅ ⟨s, a⟩) and g ⋅ T (s, a) = T (g ⋅ ⟨s, a⟩) ∀g ∈ GM, s, s′ ∈ S, a ∈ A (5)
We refer to a reward function R that satisfies Eq. (4) as a GM-invariant reward function and a deterministic transition function T that satisfies Eq. (5) as a GM-equivariant transition function. Note that this is a distinct notion from invariance and equivariance under agent’s action in the context of MDP homomorphism. Here, the action refers to the action of a 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 a 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) = Q(g ⋅ ⟨s, a⟩) and π(a, s) = π(g ⋅ ⟨a, s⟩) ∀g ∈ GM s, a ∈ S ×A. (6) 1This relation is bijective for ”simply connected” Lie groups. 2We ignore the discount factor for brevity.
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 M̄ = ⟨S̄, Ā, R̄, T̄ ⟩ is defined by S̄ = S/GM, Ā = A/GM, R̄(⟨s, a⟩GM) = R(s, a) and T̄ (s′GM ∣ ⟨s, a⟩GM) = T (s′ ∣ s, a). This results in symmetry-based model minimization of symmetric MDPs.

heading: 4 DESIDERATA FOR SYMMETRY-BASED REPRESENTATION IN RL
text: Separating Transition and Reward Symmetries One important choice is between using the symmetry group of the MDP (GM) versus the symmetry group of state-transitions (GT ), where GT 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 GM ≤ GT , i.e, the symmetries of a transition model contains the symmetries of the MDP. In fact it is easy to see that GM = GT ∩ GR, where GR 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 GT 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 a 2D map often has the symmetry of the Euclidean group, while the reward (e.g., arriving 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 model-free/model-based RL If the objective is to carry out modelfree RL, Eq. (6) motivates the need to learn action-value functions, or the policies that are invariant to symmetries of the MDP (GM). For a deterministic policy, the invariance of Eq. (6) becomes an equivariance constraint: g ⋅ π(s) = π(g ⋅ s). 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 in the agent’s action space are known permutations. However, if our objective is just to learn a symmetry-based model of the environment (i.e., transition and reward functions), Eq. (5) suggests that we need to learn a GT -equivariant transition function. Symmetries in a 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 irrelevant to the RL agent. Using state and action embeddings enables learning of the transition model in the latent space. Indeed the constraint on the model and the embedding is that of the MDP homomorphism 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-action pairs is linear through ρ(G) despite having non-linear transformations in the observation space.
From the fact that symmetries of states GS ≤ G is a subgroup of the state-action or transition symmetry, it follows that ρg ∈ ρ(G) can be divided into two parts: 1) τg ∈ τ(GS) the group representation acting on the state embedding, and; 2) κsg ∈ κ(G), the group representation for state-dependent action embedding.3
3This is because ρ(G) can be seen as a representation that is induced by the representation τ(GS) of its subgroup ρ = IndGGS τ .
At this point we can combine the requirement for an 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 our symmetric latent variable model ∀s, a ∈ S ×A and g ∈ G (see Figure 2):
T̄ (hS(s), hA(s, a)) = hS(T (s, a)) (7) τgT̄ (hS(s), hA(s, a)) = T̄ (τghS(s), κsghA(s, a)) (8)
Matrix Embedding of States and Actions While the design choice of this subsection is not necessary, we see that it can significantly simplify the constraints on a symmetry based model. We propose to use group representations for our state, and state action embeddings hS ∶ S → τ(GS) and hA ∶ S×A→ κ(G). This choice assumes that a G action on state and state-action pairs is transitive, so that each state, and state-action pair can be mapped to a (or at least one) group member. To emphasize this in our notation, we use κ(s) instead of hS(s) and similarly use τ(s, a) instead of hA(s, a) for state, and state-dependent 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:
τ(g ⋅ s) = τgτ(s) and κ(g ⋅ ⟨s, a⟩) = κsgκ(s, a) ∀s, a ∈ S ×A, g ∈ G. (9) This means that the symmetries of the state-action pairs are preserved and now take a linear form in the latent space. Note that while the embeddings are automatically equivariant, they may be equivariant 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̄ (τ(s), κ(s, a)) = κ(s, a)τ(s) (10) which simply transforms the state-embedding τ(s) through the linear group action of statedependent action encoding κ(s, a). Using this transition model, the action equivariance constraint of Eq. (7), and G-equivariance constraint of Eq. (8) simplifies to: ∀s, a ∈ S×A given a state transition triplet {s, a, s′}
τ(s′) = κ(s, a)τ(s) (11) τgκ(s, a)τ(s) = κsgκ(s, a)τgτ(s) (12)
In practice our model seeks to satisfy these two constraints via the direct minimization of appropriate loss functions, as will be discussed in Section 5.
Decomposition of the Latent Space The decomposition G = G1× . . .×GK into a direct product of subgroups can disentangle the factors of variation in the dataset (Higgins et al., 2018). 4 This gives us a way to represent the latent embedding space as a direct product of K subgroups of G, where each factor varies independently by actions of a subgroup of G on the latent embedding. Intuitively such a symmetry-based disentanglement provides an effective inductive bias particularly when there is modularity so that temporally coherent changes in the environment are due to the change of a (sparse) subset of factors. We impose this decomposed structure in the form of a direct sum for state representation and the state-dependent action representation – that is τ(s) = ⊕k τk(s) and κ(s, a) = ⊕k κk(s, a), where k ∈ {1, . . . ,K} and g = (g1, . . . , gK). Moreover the representation of the symmetry group G acting on the state embedding and the state-dependent action embedding is decomposed as τg = ⊕k τgk and κsg = ⊕k κsgk , where gk ∈ Gk for ∈ {1, . . . ,K} and g ∈ G. Combining this block structure with the Lie parameterization of Section 3.1 we get
τθ(s) =⊕ k exp(∑ i βi,k,θ(s)E(i)) κφ(τθ(s), a) =⊕ k exp(∑ i αi,k,φ(τθ(s), a)E(i)) (13)
where we use any standard neural network to implement the αφ and βθ functions above 5. As we can backpropagate through this function, the network parameters θ, φ can be learned end to end. The
4As noted by Caselles-Dupré et al. (2019), simply having a product structure in the latent space does not guarantee disentanglement, and further constraints are required. In this work, we do not impose any additional constraints for disentanglement.
5We denote both, the neural networks which maps to group representations, and network parameters by lowercase greek alphabets.
choice of the subgroup depends on the symmetries of the RL environment, and it only affects the set of bases {E(i)}i 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 (SE(2)) that comprise translation and rotations of Euclidean space (E). 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 a group that “contains” the group of interest as a subgroup. For example, if our stateactions only have 90○ rotational symmetry, we may use a more general group for the representation (e.g., SE(2)). The embedding function can define a homomorphism into the relevant subgroup.

heading: 5 LOSS FUNCTIONS
text: We consider a standard RL setup where the agent interacts with its environments in episodes and we have access to ({st, at, rt, st+1})t=1,..,T where st is the state, at is the action taken by the agent, rt is the reward received and st+1 is the observed next state at timestep t. Below we describe three loss functions that encode the equivariance/invariance 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 a loss function ` such as square loss6 that penalizes the difference between two sides:
LAET (θ, φ) = ` (τθ(st+1), κφ(st, at)τθ(st)) . (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 ofLAET encourages these symmetry transformations to capture state transitions resulting from the agent’s action.
Group Equivariant Transition Loss - Eq. (12) For this we need a st′ in addition to ⟨st, at⟩, where t′ can be any state (at different time step in the same or a different episode.) We find the group transformation that maps s to s′ the latent space using τg = τθ(st′)τθ(st)−1. Using this we can rewrite Equation Eq. (12) as
τθ(st′)τθ(st)−1 ´¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¹¶
τg
κφ(st, at)τθ(st) = κsgκφ(st, at) τθ(st′) ´¹¹¹¹¹¹¹¸¹¹¹¹¹¹¹¶ τgτθ(st) . (15)
Since the state-dependent action encoding κφ(st, at) for the pair ⟨st, at⟩ is also produced by a neural network, the only missing part in the equation above is κsg , the state-dependent action transformation. We use a neural network ρω ∶ τg ↦ κsg to infer it from state transformation τg . 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 τθ(s), κφ(s, a) are x-y coordinates plus angular velocity and torque respectively. If we rotate the pendulum using a rotation matrix τg , 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. ρω parameterizes this dependence.
A loss function ` could then measure the difference between left and the right hand side in the equation above
LGET (θ, φ,ω) = ` (τθ(st′)τθ(st)−1κφ(st, at)τθ(st), ρω(τθ(st′)τθ(st)−1)κφ(st, a)τθ(st′)) . (16)
Action Invariant Reward Loss - Eq. (1) While LAET and LGET enforce the equivariance of the latent transition model to an agent’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 rψ ∶ τθ(s)↦ R. We measure the difference between the predicted reward and the actual reward at time step t + 1:
LR(ψ, θ, φ) = (rψ(κφ(τθ(st), at)τθ(st)) − rt+1)2 . (17) 6In practice we use the normalized square loss `(Y ,Y ′) = ∥ Y∥Y ∥2 − Y ′ ∥Y ′∥2 ∥ 2
2

heading: 6 APPLICATION TO MODEL-FREE RL
text: Following previous success of using transition models for representation learning in model-free RL (Gelada et al., 2019; Schwarzer et al., 2021) we add the losses discussed above to the Temporal Difference (TD) error in Deep Q-learning. In practice, we need to make three modifications to our model/loss.
Target Network A trivial solution to both equivariance enforcing losses of Section 5 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 a target network to encode state st+1 and st′ in Eq. (16) in which the network parameters do not 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., τθ → τ ). Projection Head for Transition Losses Strict enforcement of symmetry constraints by our 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.

heading: 6.1 PUTTING IT ALL TOGETHER
text: ConsideringM consecutive state-actions {st∶t+M , at∶t+M} and x̂t = xt = τθ(st), we predict the state embeddings and the rewards of next M steps:
x̂t+m = κφ(x̂t+m−1, at+m−1)x̂t+m−1 and r̂t+k = rψ(x̂t+k) ∀m ∈ {1, . . . ,M} Note that we are using x̂ for M-step model prediction of the embedding to distinguish them from the latent embedding x, and the embedding produced by the target network x̄ = τ(st+m). The same applies to the M-step predicted reward r̂ and observed reward r. We then project these embedding using a projection head pζ to produce ẑt+m = pζ(x̂t+m) and z̄t+m = p(x̄t+m). Using this notation, our final expressions for LAET and LR are:
LAET = M
∑ m=1 ∥ ẑt+m∥ẑt+m∥2 − z̄t+m∥z̄t+m∥2
∥ 2
2
and LR = M
∑ m=1 (r̂t+m − rt+m)2 (18)
For LGET we need ⟨st, at, st+1⟩ and another state s′. From their embedding using the notation above we get τg = x̄t′xt−1, the linear transformation between them, and κsg = ρω(τg), the statedependent action transformation. Now for x̄t′ , we obtain the predicted next state from x̄t′ as x̄t′+1 = κsgκ(xt, at)x̄t′ = ρω(xt′x−1t )κ(xt, at)x̄t′ and from x̂t+1 as x̂t′+1 = τgx̂t+1. Before penalizing the difference between these embeddings, we project them to ŷt′+1 = bη(x̂t′+1) and ȳt′+1 = b(x̄t′+1) using projection head bη , to get the final expression for LGET :
LGET = ∥ ŷt′+1 ∥ŷt′+1∥2 − ȳt′+1∥ȳt′+1∥2
∥ 2
2
(19)
Q-learning We pass the representation xt to a Q-learning head qξ to learn policies based on the output of the Q-value estimator. The Q-value estimator is learnt by minimizing:
LDQN(ξ, θ) = (qξ(τθ(st), at) − (rt + γmax a qξ(τ(st+1), a)))2 (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 RL(Dabney et al., 2018), Dueling DQN (Wang et al., 2016), Double DQN (Van Hasselt et al., 2016). The total loss optimized by our model is:
L = LDQN + λ1LR + λ2LGET + λ3LAET (21)
where λ1, λ2 and λ3 are 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.

heading: 7 EXPERIMENTS
text: We test our method on a 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 sampleefficient Atari suite introduced by Kaiser et al. (2019), which consists of 26 games with only 100,000 environment steps of training data available. In our experiments, we use three types of simple connected Lie subgroup blocks including General Linear GL(2), Special Euclidean SE(2), and Translation T (2), see Appendix A for more details. Unless stated otherwise, our EqR model uses SE(2) subgroup blocks, withK = 12 blocks andM = 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 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 on top of SPR’s (Schwarzer et al., 2021), which is based on rlpyt (Stooke & Abbeel, 2019) and PyTorch (Paszke et al., 2019). We use the same underlying RL algorithm and hyperparameters used by SPR for fair comparision.
Evaluation Metrics We compute the average episodic return (the ‘game score’) at the end of training and normalize it with respect to human scores, as is standard practice. The human-normalized score (HNS) is given by agent score - random scorehuman score - random score . Since there is considerable variance across different runs, the mean and the median are not very reliable metrics. Instead, Agarwal et al. (2021) propose 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 a 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 variations of our model, EqR with LR and with LR + LGET , along with other comparable 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.
In order to better understand the effect of various modeling choices, loss functions and implementation details on the performance, we now consider different variations of EqR, with the same augmentation 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, we remove them and use an action
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-based inductive 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 transformations present in Atari games. However, the more restrictive T (2) slightly hurts the performance, while the more general GL(2) performs similarly to the MLP model. 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 LGET improves the performance slightly, while adding only LR improves the performance even further. We hypothesize that the reward loss is playing a role in both preventing representation collapse and preserving more information of the reward distribution 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 transformations 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 a total of 26 games, including this loss term leads to a statistically significant boost in performance. These 17 games are: ‘Alien’, ‘BankHeist’, ‘BattleZone’, ‘Boxing’, ‘ChopperCommand’, ‘CrazyClimber’, ‘DemonAttack’, ‘Freeway’, ‘Hero’, ‘Jamesbond’, ’MsPacman’, ‘Pong’, ‘PrivateEye’, ‘Qbert’, ‘RoadRunner’, ‘Seaquest’, ‘UpNDown’. The full list of game-wise scores for the ablation studies are presented in Tables 3 and 4 in Appendix B.1.

heading: 8 CONCLUSION
text: This paper considers three major symmetry-related constructs within a coherent framework. First, there is 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 is the group equivariant transition, which acts as an additional bias and ensures that the latent transition model itself is equivariant under symmetry transformations of state-action pairs.
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, as 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.

heading: 9 REPRODUCIBILITY STATEMENT
text: We have made every effort to ensure that our method is reproducible. Section 5 and Section 6 provide detailed descriptions of the loss functions and various implementation-related details. We present a 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.1 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.

heading: A SUBGROUP BLOCKS AND THEIR PARAMETERIZATION

heading: A.1 CHOICE OF GROUP
text: 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 action. But most of these games include symmetry transformations. For example, the screen often has multiple objects undergoing two dimensional(2-D) translations and rotations. In this case one can use blocks of the 2-D Special Euclidean Group SE(2). 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 a general method to parameterize each of these, based on the type of group.
GL(n) As the matrix representation of GL(n) is the set of invertible matrices which has a measure of 1 it is easy to parameterize it. We just generate n2 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 n parameters using a neural network and instead of using matrix multiplication use addition for the group action. Note that we can also use a matrix representation for T(n) but it is unnecessary and inefficient.
SE(n) Unlike GL(n) and T(n), parameterizing SE(n) is a bit tricky because it involves parameterizing SO(n). We use a homogeneous co-ordinate based representation of SE(n) =
{(R t 0 1 ) ,R ∈ SO(n) and t ∈ T (n)}. So we need n parameters for the t and another D = n(n−1) 2 parameters for SO(n) from the neural network. As explained in Section 4, we can use a Lie parameterization to get the elements of SO(n) by R = exp(∑Dd=1 βd E(d)) where E(d) denote D bases of the space of skew symmetric matrices and the βis are the parameters of the neural network. For example, in the case of SO(2) we can use the basis E(1) = ( 0 1−1 0) . Similarly, we can extend this to SO(n) by using a basis given by D n × n matrices E(ij) ∀{1 ≤ i < j ≤ n} whose only non-zero elements are E(ij)i,j = −1 and E (ij) j,i = 1.
Although Lie parameterization gives us a general recipe to output a representation of simple connected Lie groups like SO(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), we parameterize each rotation matrix in SO(n) using the product of rotations on D orthogonal planes in Rn: R = ∏ni=1∏1≤i<j≤nRij . Here Rij ∈ Rn×n is the rotation matrix in the i − j plane, and its non-zero elements besides the diagonal are the four values on the i, j rows and columns, which comprise the 2D rotation matrix that is Riji,j = [ cos(θi,j) sin(θi,j) − sin(θi,j) cos(θi,j)]. We have D parameters θi,j which we can obtain from a neural network.
The parameters in all the parameterization techniques mentioned here can be back-propagated. We summarize the number of parameters required from a neural network output, representation type and the associated group actions of different subgroups in Table 1.

heading: B ATARI DETAILS

heading: B.1 FULL RESULTS
text: We provide individual results on the 26 Atari games after 100K training steps. Our results are averaged 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 2 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 4 (a)).
• Table 4 compares EqR using SE(2) subgroup blocks with different loss terms included in the training objective (see Section 5 and Figure 4 (b)). The action equivariance transition loss, LAET is always included for EqR models.

heading: B.2 MODEL SCHEMATIC

heading: B.3 NETWORK ARCHITECTURE
text: We follow the baseline RL implementation of DrQ (Yarats et al., 2021) and SPR (Schwarzer et al., 2021) by using the 3-layer convolutional encoder from (Mnih et al., 2015) and then use a linear layer to 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 is a matrix for GL(n) and SE(n). We flatten it before we feed to other neural network like the Q-head qξ(⋅). For the action encoder κ(⋅) we use a simple 1 layer MLP with batchnorm, ReLU and a hidden size of 256. We concatenate the one-hot encodings of the actions with the state representations coming from τθ and pass it through the action encoder to get matrix representation of the group after parameterization.
For the reward predictor network rψ we use a 2-layered MLP with batchnorm, ReLU and a hidden size of 256.
For the Q-head qξ(⋅) we use 2-layered MLP as well. For the projection head pζ(⋅) we share the first layer of Q-head whereas for projection head bη(⋅) we use a single layer MLP.

heading: B.4 ALGORITHM
text: Algorithm 1: Equivariant Representations for RL Denote the parameters of online networks τθ, κφ, pζ , bη as Θo Denote the parameters of target networks τ , κ, p, b as Θc Denote the parameters of networks ρω , qξ as Φ Denote the dept of the prediction as M and batch size as N Initialize the replay buffer B while Training do
Collect {s, a, r, s′} using policy with (Θo,Φ) and add to the buffer B Sample a minibatch of M length sequences {s0∶M , a0∶M , r0∶M} ∼ B for i in range(0,N) do
if augmentation then si0∶M ← augment(si0∶M) end xi0 ← τθ(si0); // state representation x̂i0 ← xi0 li ← 0 for k in range(1,M + 1) do
x̂ik ← κφ(x̂ik−1, aik−1)x̂ik−1 ; // state transition by group action x̄ik ← τ(sik); // target state representation ẑik ← pζ(x̂ik), z̄ik ← pζ(x̄ik); // projections li ← li + λ2∥ ẑt+k∥ẑt+k∥2 − z̄t+k ∥z̄t+k∥2 ∥ 2 2; // compute LAET at step k r̂ik ← rψ(x̂ik); // predict rewards li ← li + λ1∥r̂ik − rik∥22; // compute LR at step k
end j ∼ {0, ..,N − 1}; // uniformly sample an index x̄j0 ← τ(s j 0); // encode the state for that index from the batch τ ig = x̄j0xi0 −1
; // find the group representation x̂j1 ← τ igx̂i1; // next state by group action x̄j1 ← ρω(τ ig)κ(xi0, ai0)x̄ j 0; // next state by action-embedding ŷi1 ← bη(x̂ j 1), ȳi1 ← b(x̄ j 1); // projections li ← li + λ3∥ ŷ i 1
∥ŷi1∥2 − ȳ
i 1
∥ȳi1∥2 ∥22; // compute LGET
li ← li +RLloss(x̂i0, ai0, ri0, x̄i1; qξ) end l ← 1
N ∑ N i=0 li; // average over minibatch Θo,Φ← optmize((Θo,Φ), l); // update online networks Θc ← Θo; // copy weights to target networks
end

heading: B.5 HYPERPARAMETERS
text: 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 et 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 - λ1, LAET - λ2 and LGET - λ3 are set to one whenever they are used in the model.

heading: Parameter Setting

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

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.

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.  In 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.  The 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.