EqR: Equivariant Representations for Data-Efficient Reinforcement Learning

1 INTRODUCTION . 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 . 2 RELATED 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 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 . 3 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 ∈ 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 ) . 3.2 MDP HOMOMORPHISM AND SYMMETRIC MDPS . 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 .

Summary: 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.