Learning Temporally-Consistent Representations for Data-Efficient Reinforcement Learning

1 INTRODUCTION . Learning sample-efficient continuous-control directly from pixels is a challenging task for reinforcement learning ( RL ) agents . Despite recent advances , modern approaches require hundreds of thousands , if not millions , of agent-environment interactions to attain acceptable performance ( Kalashnikov et al. , 2018 ; Jaderberg et al. , 2019 ) . In real-world settings , this magnitude of data collection may take a prohibitive amount of time . Therefore , advances in learning efficiency for continuouscontrol in pixel-based state spaces would significantly accelerate the success of real-world RL applications such as industrial robotics and autonomous vehicles ( Dulac-Arnold et al. , 2021 ) . The standard RL loop ties policy optimization and representation learning together with a single supervisory signal . This dual-objective , single-signal learning scheme results in a representation learning bottleneck that ultimately produces agents that understand their environment slowly throughout training ( Shelhamer et al. , 2017 ) , thereby reducing performance in the early stages of learning . Approaches that have seen the most success in alleviating the bottleneck rely on methods that explicitly encourage the distillation of useful information from the state space . Recently , data augmentation ( Chen et al. , 2020a ; b ; Hénaff et al. , 2020 ) has been studied in the context of RL as a way to improve sample efficiency ( Yarats et al. , 2021b ; Laskin et al. , 2020b ) . Other approaches to improve sample efficiency in RL add auxiliary tasks to the learning process . For example , image reconstruction ( Yarats et al. , 2020 ) , latent variable models ( Lee et al. , 2020a ; Gelada et al. , 2019 ) , variational objectives ( Lee et al. , 2020b ) , and contrastive learning ( Laskin et al. , 2020a ; Stooke et al. , 2021 ) . While the aforementioned approaches improve performance over base RL algorithms , they do not fully leverage the long-term temporal connection between actions , future states , and rewards . For example , some approaches do not interact with the temporal axis of the RL problem ’ s underlying Markov decision process ( MDP ) ( Yarats et al. , 2020 ; 2021b ; Laskin et al. , 2020b ) , or are only concerned with next-step dynamics ( Laskin et al. , 2020a ; Stooke et al. , 2021 ; Gelada et al. , 2019 ) . As we will show with our work , ignoring the long-term temporal axis of the MDP withholds information that can help agents to understand consequences of their environment interactions more quickly . Approaches that do attempt to capture longer-term information either require a significant amount of random exploration and pre-training ( Lee et al. , 2020a ; b ) , which counteracts the goal of minimizing agent-environment interactions , or condition their predictions of the future on the full-dimensional state space ( Lee et al. , 2020b ) , thereby not encouraging latent representations themselves to be predictive of latent representations multiple steps into the future . In this paper , we revisit the problem of sample-efficient continuous-control in image-based state spaces . We introduce k-Step Latent ( KSL ) 1 , a representation learning module that aids RL agents with continuous control on pixels . In contrast to aforementioned prior methods , KSL ’ s auxiliary task directly exploits the RL problem ’ s underlying MDP , requires no pre-training , and encourages latent representations to be predictive of latent representations that are multiple steps into the future . KSL agents learn to create temporally-coherent representations by projecting states into the latent space such that latent representations , when conditioned upon a series of actions , are predictive of latent representations that are multiple steps into the future . KSL ’ s auxiliary task propels agents to new state-of-the-art results on both the data efficiency ( 100k steps ) and asymptotic performance ( 500k steps ) checkpoints in the popular PlaNet benchmark suite ( Hafner et al. , 2019 ) . This paper makes the following contributions . First , KSL augments the learning process of the popular Soft Actor-Critic ( Haarnoja et al. , 2018a ) algorithm to produce new state-of-the-art results in the PlaNet benchmark suite . KSL significantly outperforms current methods such as DrQ ( Yarats et al. , 2021b ) and RAD ( Laskin et al. , 2020b ) . Across all six tasks in PlaNet , KSL produces an average improvement over DrQ of 46.9 % and 14.5 % at the 100k and 500k steps mark , respectively , and an average improvement over RAD of 58.1 % and 21.6 % at the 100k and 500k steps mark , respectively . Second , we analyze the learned encoders and latent projections that are produced by several methods . Our analyses show that KSL creates representations that are more strongly tied to reward , are less sensitive to perturbation in the state space , and move more smoothly over time than other baseline methods . Also , our analyses show that KSL creates encoders that generalize to unseen tasks better than other baseline methods . 2 BACKGROUND . We consider the usual formulation of RL agents that act within a Markov decision process ( MDP ) , defined by the tuple ( S , A , P , R , γ ) , where S is the state space , A is the action space , P ( st+1|st , at ) is the transition probability distribution , R : S × A → R is a reward function that maps states and actions to a scalar feedback signal , and γ ∈ [ 0 , 1 ) is a discount factor . The agent ’ s goal is to learn an action-selection policy π ( a|s ) that maximizes the sum of discounted rewards over the time horizon T of the given task : arg maxπ Ea∼π , s∼P [ ∑T t=1 γ tR ( st , at ) ] . Specifically , we examine learning agents in state spaces comprised of images o ∈ O. Image-based states are high-dimensional and may give an incomplete description of the environment . For example , incomplete information may arise when the camera ’ s point-of-view shows occluded objects . As such , it is common to stack a series of n consecutive images to help alleviate the problems associated with incomplete information : st = { ot , ... , ot−n+1 } ( Mnih et al. , 2015 ) . 2.1 SOFT ACTOR-CRITIC . Soft Actor-Critic ( SAC ) ( Haarnoja et al. , 2018a ; b ) is a current state-of-the-art off-policy , model-free RL algorithm for continuous control . SAC learns a state-action value-function criticQθ , a stochastic actor πφ , and a temperature α that weighs between maximizing reward and entropy in an augmented RL objective : Est , at∼π [ ∑ tR ( st , at ) + αH ( π ( ·|st ) ) ] . SAC ’ s critic is learned by minimizing the squared Bellman error over trajectories τ = ( st , at , rt , st+1 ) sampled from a replay memory D : Lcritic ( θ ) = Eτ∼D [ ( Qθ ( st , at ) − ( rt + γy ) ) 2 ] , ( 1 ) 1Code available at : https : //github.com/anon-researcher-repo/ksl where the temporal-difference target , y , is computed by sampling an action from the current policy : y = Ea′∼π [ Qθ̄ ( st+1 , a′ ) − α log πφ ( a′|st+1 ) ] . ( 2 ) The target critic Qθ̄ is updated as an exponential moving average ( EMA ) : θ̄ ← ζθ̄+ ( 1− ζ ) θ where ζ controls the “ speed ” at which θ̄ tracks θ ( e.g. , He et al . ( 2020 ) ) . SAC ’ s actor parameterizes a multivariate GaussianN ( µ , σ ) where µ is a vector of means and σ is the diagonal of the covariance matrix . The actor is updated via minimizing : Lactor ( φ ) = −Ea∼π , τ∼D [ Qθ ( st , a ) − α log πφ ( a|st ) ] , ( 3 ) and α is simply learned against a target entropy value . 3 k-STEP LATENT 3.1 MOTIVATION . We design KSL to address the representation learning bottleneck of base RL algorithms by providing an auxiliary task that encourages representational consistency of states that are nearby in time . States that are nearby in time are likely to share high levels of mutual information ( Mobahi et al. , 2009 ) , which implies that the states ’ representations should as well . KSL contains a recurrent mechanism that encourages representations of states , when conditioned upon actions , to be predictive of states that are multiple steps into the future . Intuitively , we hypothesize that agents who understand the long-term consequences of their actions would perform better than agents who do not . 3.2 TRAINING DETAILS . A visual depiction of KSL can be seen in Figure 1 and pseudocode can be found in Appendix B . At each learning step , we uniformly sample M trajectories of length k , τ = { ( s1 , a1 , ... , ak−1 , sk ) i } Mi=1 , from the replay memory , where s1 can be any state in the episode . For each state in each trajectory , we apply a random translation augmentation ( Laskin et al. , 2020b ) . Using augmentations within KSL encourages representations to be invariant to perturbations over the temporal axis . KSL contains an online and momentum pathway , which are highlighted in Figure 1 as orange and green , respectively . At step j = 1 of trajectory2 i , sj is fed through an online encoder zo = Φo ( sj ) and sj+1 is fed through a momentum encoder zm = Φm ( sj+1 ) . Next , zo is concatenated with aj and fed through a learned transition model z′o = T ( [ zo|aj ] ) . This transition model enables KSL to learn temporal projections from latent representations of state sj and applied action aj to a representation of the following state sj+1 in the form of z′o . Then , z ′ o and zm are fed through their respective projection modules qo = Ψo ( z′o ) , ym = Ψm ( zm ) . Such modules after an encoder have been shown to improve the learned representations of the encoder ( Chen et al. , 2020a ) . Finally , qo is fed through a prediction head yo = P ( qo ) . 2Note that representations through both pathways are computed for every step and trajectory . We omit batch notation for steps and trajectories in the paragraph for clarity . KSL ’ s loss for step j is computed as the distance between the ` 2 normalized versions of the final two vectors from both paths : ‖ ( yo/‖yo‖ ) − ( ym/‖ym‖ ) ‖22 . This operation is performed in a loop over all k steps per each trajectory in τ , with one notable difference at each step after j = 1 . Instead of encoding every s along the online path , we recursively feed z′o from the previous step j − 1 into the learned transition model T . This recurrent operation is shown with the red arrow in Figure 1 . Therefore , Φo is only used to compute a direct representation of the first state within each trajectory . The minimization of KSL ’ s loss requires Φo to create latent projections such that multi-step consequences can be inferred from any arbitrary starting state using the transition model T . We formulate the loss function as an average over the computed distance for every state and trajectory in a given batch : LKSL = 1 Mk M∑ i=1 k∑ j=1 ∣∣∣∣∣ ∣∣∣∣∣ y ( i , j ) o‖y ( i , j ) o ‖ − y ( i , j ) m ‖y ( i , j ) m ‖ ∣∣∣∣∣ ∣∣∣∣∣ 2 2 , ( 4 ) with y ( i , j ) p denoting the final vector of pathway p for trajectory i at step j . We implement a stop gradient to disallow backpropagation from Equation 4 through the momentum pathway . Instead , we update the momentum encoder and momentum projection module as an EMA of their online counterparts . We use the state encoding of the online encoder Φo as input into the critic Qθ and actor πφ . While we allow the gradient of Qθ to update weights within Φo , we implement a stop gradient that disallows signal from the actor ’ s computation graph ( Equation 2 and Equation 3 ) to propagate into Φo . Gradient flow is depicted in Figure 1 as purple , dotted lines . Φm is an EMA of Φo and Qθ̄ is an EMA Qθ , so it is natural for Φm to transform states for input into Qθ̄ . See Appendix F for an expanded description of how KSL fits within SAC . We also provide ablations across architecture choices and values of k in Appendix E and Appendix G. We note that KSL is not restricted to SAC and can be applied alongside any RL algorithm for representation learning in tasks with continuous or discrete action spaces .

Summary: This paper proposes a representation-learning method (k-step latent, or KSL) that uses a self-supervised auxiliary loss between recurrently-predicted action-conditioned representations of the state space and non-recurrently predicted target representations, in the style of BYOL. The method is trained using two separate optimizers on different parts of the models, to avoid interference in the statistics maintained by the optimizers. Results at 100k and 500k steps on 6 tasks from the DM control suite (from pixels) compared to methods using alternative self-supervised auxiliary losses show that the proposed method improves data efficiency. Analysis of the learned latent representations shows that those from KSL produce more robust encoders and are more consistent with the underlying MDP.