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