E Q R: E QUIVARIANT R EPRESENTATIONS FOR D ATA -E FFICIENT R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 predeﬁned 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-efﬁcient 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.
I NTRODUCTION
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 inefﬁciency.
Notably, almost all success stories of RL rely on vast amounts of data or simulations with a huge computational overload.
More data-efﬁcient 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,
[Figure]
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 (predeﬁned) 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-efﬁcient 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 disentangled and equivariant to continuous transformations of the data.
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 obtain a median human normalized score of 0.418.
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.
R ELATED WORK
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 dataefﬁciency.
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-efﬁciency 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 efﬁciency 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 inefﬁcient because similar to pixel level models, one needs to learn potentially irrelevant details.
The fact that variations of model-free algorithms such as Data-Efﬁcient 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 conﬁrms 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 signiﬁcantly improve the sample efﬁciency, producing state-of-the-art results in sample-efﬁcient representation learning in RL.
B ACKGROUND 3. G ROUPS AND THEIR R EPRESENTATIONS
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 − 1 g = 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 x G
=
{ 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 ∶ R D
→ R D 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 identiﬁable by their
Lie algebra g =
Lie ( ρ (G)) .
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 =
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 ﬁnite groups ( × 90 ○ rotations).
3. MDP H OMOMORPHISM AND S YMMETRIC MDP S
We deﬁne an MDP as the 4-tuple M = ⟨ 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.
For two MDPs M = ⟨ S , A ,R,T ⟩ and ¯ M = ⟨ ¯ S , ¯ A , ¯ R, ¯ T ⟩ , MDP homomorphism can be deﬁned as a tuple H = ⟨ h S ,
h A ⟩ where h S ∶ S → ¯ S is the state mapping and h A ∶ S
× A → ¯ A is the state dependent action mapping.
These two mappings satisfy the following invariance and equivariance conditions:
1) Invariance of the reward:
[Equation]
2) Equivariance of the deterministic transition model under the agent’s action:
[Equation]
A probabilistic variation of the above equation for a stochastic MDP (Bloem-Reddy & Teh, 2020) is:
[Equation]
where [ s ′ ] h S = h S − 1 ( h S ( s ′ )) is the equivalence class of s ′ under h S .
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 G M = Aut (M) of an MDP identiﬁes the set of symmetry transformations of state-actions that preserve the reward and the transition dynamics:
[Equation]
⋅
[Equation]
s,a
[Equation]
) and g ⋅ T ( s,a )
[Equation]
g
[Equation]
We refer to a reward function R that satisﬁes Eq.
(4) as a G M -invariant reward function and a deterministic transition function T that satisﬁes Eq.
(5) as a G M -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 satisﬁes both Eq.
(4) and Eq.
(5), both the optimal action-value and optimal policy functions become invariant under G M action (Ravindran & Barto, 2001) – that is,
[Equation]
The connection of symmetric MDPs to MDP homomorphism is due to the fact that symmetries can be used to deﬁne a homomorphism H ∶ M ↦ ¯ M by collapsing the state-actions that form an orbit under G M .
Formally,the collapsed MDP ¯ M = ⟨ ¯ S , ¯ A , ¯ R, ¯ T ⟩ is deﬁned by ¯ S = S /G
M , ¯ A
= A /G M , ¯ R (⟨ s,a ⟩ G M )
=
R ( s,a ) and ¯ T ( s ′G M ∣ ⟨ s,a ⟩ G M ) =
T ( s ′ ∣ s,a ) .
This results in symmetry-based model minimization of symmetric MDPs.
D ESIDERATA FOR S YMMETRY -B ASED R EPRESENTATION IN RL Separating
Transition
and
Reward Symmetries One important choice is between using the symmetry group of the MDP ( G M ) versus the symmetry group of state-transitions ( G T ), where G T 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 G M ≤ G T , i.e, the symmetries of a transition model contains the symmetries of the MDP.
In fact it is easy to see that G M = G T ∩ G R , where G R 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 G T has two beneﬁts: 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 ( G M ).
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 efﬁciency 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 G T -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 beneﬁt 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.
T ¯ T ¯ T τ g κ sg τ g s ¯ s ¯ a a , , s ′
¯ s ′
State transition
T r a n s i t i o n S y mm e t r y
Group equivariant transition
Action equivariant transition
L a t e n t E m b e d d i n g ˜ s , ˜ a ˜ s ′ h 𝕊 h 𝕊 h s 𝔸
From the fact that symmetries of states G S ≤ G is a subgroup of the state-action or transition symmetry, it follows that ρ g ∈ ρ (G) can be divided into two parts: 1) τ g ∈ τ (G S )
the group representation acting on the state embedding, and; 2) κ sg ∈ κ (G) , the group representation for state-dependent action embedding.
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):
[Equation]
Matrix Embedding of States and Actions While the design choice of this subsection is not necessary, we see that it can signiﬁcantly simplify the constraints on a symmetry based model.
We propose to use group representations for our state, and state action embeddings h S ∶ S → τ (G S ) and h A ∶ 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 h S ( s ) and similarly use τ ( s,a ) instead of h A ( s,a ) for state, and state-dependent action embedding respectively.
This choice of embedding has several beneﬁts: First, the learned embeddings are automatically equivariant to symmetry transformations of the state, and state-actions:
[Equation]
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
[Equation]
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) simpliﬁes to: ∀ s,a ∈ S ×
A given a state transition triplet { s,a,s ′ }
[Equation]
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 = G 1 × ...
×G K into a direct product of subgroups can disentangle the factors of variation in the dataset (Higgins et al., 2018).
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 =
( g 1 ,...,g K ) .
Moreover the representation of the symmetry group G acting on the state embedding and the state-dependent action embedding is decomposed as τ g = ⊕
k τ g k and κ sg
= ⊕ k κ sg k , where g k ∈ G k for ∈ { 1 ,...,K } and g ∈ G .
Combining this block structure with the Lie parameterization of Section 3.1 we get
[Equation]
(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 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.
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 deﬁne a homomorphism into the relevant subgroup.
L OSS F UNCTIONS
We consider a standard RL setup where the agent interacts with its environments in episodes and we have access to ({ s t ,a t ,r t ,s t + 1 }) t = 1 ,..,T where s t is the state, a t is the action taken by the agent, r t is the reward received and s t + 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.
Given triplets ⟨ s t ,a t ,s t + 1 ⟩ from our dataset we simply apply a loss function (cid:96) such as square loss 6 that penalizes the difference between two sides:
[Equation]
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 L AET encourages these symmetry transformations to capture state transitions resulting from the agent’s action.
Group Equivariant Transition Loss - Eq.
For this we need a s t ′ in addition to ⟨ s t ,a t ⟩ , where t ′ can be any state (at different time step in the same or a different episode.)
We ﬁnd the group transformation that maps s to s ′ the latent space using τ g = τ θ ( s t ′ )
τ θ ( s t ) − 1 .
Using this we can rewrite Equation Eq.
(12) as
[Equation]
Since the state-dependent action encoding κ φ ( s t ,a t ) for the pair ⟨ s t ,a t ⟩ 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 reﬂection 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 reﬂection around the vertical axis, we expect that the effect of torque will be negated.
ρ ω parameterizes this dependence.
A loss function (cid:96) could then measure the difference between left and the right hand side in the equation above
[Equation]
Action Invariant Reward Loss - Eq.
(1) While L AET and L GET 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 :
[Equation]
[Equation]
A PPLICATION TO M ODEL - FREE RL
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 modiﬁcations 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 sufﬁcient specially in sparse reward settings.
Following Schwarzer et al.
(2021), we use a target network to encode state s t + 1 and s t ′ 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 (cid:96) 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. P UTTING IT A LL T OGETHER
Considering M consecutive state-actions { s t ∶ t + M ,a t ∶ t + M } and ˆ x t = x t = τ θ ( s t ) , we predict the state embeddings and the rewards of next M steps:
[Equation]
and
[Equation]
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 = τ ( s t + 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 ﬁnal expressions for L AET and L R are:
[Equation]
For L GET we need ⟨ s t ,a t ,s t + 1 ⟩ and another state s ′ .
From their embedding using the notation above we get τ g = ¯ x t ′ x t − 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 ′ +
= κ sg κ ( x t ,a t ) ¯ x t ′ =
ρ ω ( x t ′ x − 1 t ) κ ( x t ,a t ) ¯ x t ′ and from ˆ x t + 1 as ˆ x t ′ + 1 = τ g ˆ x 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 ﬁnal expression for L GET :
[Equation]
Q-learning We pass the representation x t 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:
[Equation] [Equation]
We use the data efﬁcient 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:
[Equation]
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.
E XPERIMENTS
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 sampleefﬁcient Atari suite introduced by Kaiser et al.
(2019), which consists of 26 games with only 100, 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, with K = 12 blocks and M = 5 steps during training.
L AET 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.
L GET , which makes transition model equivariant with respect to the symmetry transformation of state-actions, and L R 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
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
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 conﬁdence intervals (CI) with stratiﬁed 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 proﬁles 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.
Figureshows performance proﬁles for two variations of our model, EqR with L R and with L R + L GET , 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 signiﬁcant 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 L R .
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
[Figure]
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 L AET results in a considerable improvement over Rainbow with augmentation.
Adding L GET improves the performance slightly, while adding only L R 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 L GET and L R 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 beneﬁcial 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 signiﬁcant 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.
C ONCLUSION
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.
R EPRODUCIBILITY S TATEMENT
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.
A S UBGROUP BLOCKS AND THEIR PARAMETERIZATION A. C HOICE OF 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 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. I MPLEMENTING 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
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 n 2 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 inefﬁcient.
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 1 ) ,
R ∈ SO
( n ) and t ∈ T ( n )} .
So we need n parameters for the t and another D = n ( n − 1 )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 β i s are the parameters of the neural network.
For example, in the case of SO ( 2 ) we can use the basis E ( 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 R n :
R = ∏ ni
= 1 ∏ 1 ≤ i < j ≤ n R ij .
Here R ij ∈ R n × 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 R iji,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.
[Table]
B A TARI D ETAILS B. F ULL R ESULTS
We provide individual results on the 26 Atari games after 100K training steps.
Our results are averaged over 10 seeds
[Table]
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, L R , 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, L AET is always included for EqR models.
[Table]
Jamesbond
[Table] [Table]
scores are averaged over 10 seeds.
[Table]
[Figure]
B. N ETWORK A RCHITECTURE
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 deﬁnes our τ θ .
Note that the output of our encoder is a matrix for GL ( n ) and SE ( n ) .
We ﬂatten 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 ﬁrst layer of Q-head whereas for projection head b η (⋅) we use a single layer MLP.
B.
A LGORITHM
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 { s 0 ∶ M ,a 0 ∶ M ,r 0 ∶ M } ∼ B for i in range ( 0 ,N ) do if augmentation then s i
∶ M ← augment ( s i 0 ∶ M ) end x i
← τ θ ( s i 0 ) ; // state representation ˆ x i
← x il i ←for k in range ( 1 ,M + 1 ) do ˆ x ik ← κ φ ( ˆ x ik − 1 ,a ik − 1 ) ˆ x ik − 1 ; // state transition by
group action
¯ x ik ← τ ( s ik ) ; // target state representation ˆ z ik ← p ζ ( ˆ x ik ) , ¯ z ik ← p ζ ( ¯ x ik ) ; // projections l i ← l i + λ 2 ∥
ˆ z t + k
∥
ˆ z t + k ∥ 2 − ¯ z t + k
∥
¯ z t + k ∥ 2 ∥ 22 ; // compute L AET at step k ˆ r ik ← r ψ ( ˆ x ik ) ; // predict rewards l i ← l i + λ 1 ∥
ˆ r ik − r ik ∥ 22 ; // compute L R at step k end j ∼ { 0 ,..,N − 1 } ; // uniformly sample an index ¯ x j 0 ← τ ( s j 0 ) ; // encode the state for that index from the batch τ ig = ¯ x j 0 x
i 0 − 1 ; // find the group representation ˆ x j 1 ← τ ig ˆ x i 1 ; // next state by group action ¯ x j 1 ← ρ ω ( τ ig ) κ ( x i 0 ,a i 0 )
¯ x j 0 ; // next state by action-embedding ˆ y i 1 ← b η ( ˆ x j 1 ) , ¯ y i 1 ← b ( ¯ x j 1 ) ; // projections l i ← l i + λ 3 ∥
ˆ y i 1 ∥
ˆ y i 1 ∥ 2 − ¯ y i 1 ∥
¯ y i 1 ∥ 2 ∥ 22 ; // compute L GET l i ← l i + RLloss ( ˆ x i 0 ,a i 0 ,r i 0 , ¯ x i 1 ; q ξ ) end l ← 1 N ∑ Ni =
l i ; // average over minibatch Θ o , Φ ← optmize (( Θ o , Φ ) ,l ) ; // update online networks Θ c ← Θ o ; // copy weights to target networks end
B. H YPERPARAMETERS
In this section
[Table]
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 L R - λ 1 , L AET - λ 2 and L GET - λ 3 are set to one whenever they are used in the model.
[Table]