Pretraining Representations for Data-Efficient Reinforcement Learning

1 Introduction . Deep reinforcement learning ( RL ) methods often focus on training networks tabula rasa from random initializations without using any prior knowledge about the environment . In contrast , humans rely a great deal on visual and dynamics priors about the world to perform decision making efficiently ( Dubey et al. , 2018 ; Lake et al. , 2017 ) . Thus , it is not surprising that RL algorithms which learn tabula rasa suffer from severe overfitting ( Zhang et al. , 2018 ) and poor sample efficiency compared to humans ( Tsividis et al. , 2017 ) . Self-supervised learning ( SSL ) provides a promising approach to learning useful priors from past data or experiences . SSL methods leverage unlabelled data to learn strong representations , which can be used to bootstrap learning on downstream tasks . Pretraining with self-supervised learning has been shown to be quite successful in vision ( Hénaff et al. , 2019 ; Grill et al. , 2020 ; Chen et al. , 2020a ) and language ( Devlin et al. , 2019 ; Brown et al. , 2020 ) settings . Pretraining can also be used in an RL context to learn priors over representations or dynamics . One approach to pretraining for RL is to train agents to explore their environment in an unsupervised fashion , forcing them to learn useful skills and representations ( Hansen et al. , 2020 ; Liu & Abbeel , 2021 ; Campos et al. , 2021 ) . Unfortunately , current unsupervised exploration methods require months or years of real-time experience , which may be impractical for real-world systems with limits and costs to interaction — agents can not be run faster than real-time , may require significant human oversight for safety , and can be expensive to run in parallel . It is thus important to develop pretraining methods that work with practical quantities of data , and ideally that can be applied offline to fixed datasets collected from prior experiments or expert demonstrations ( as in Stooke et al. , 2021 ) . ∗ { max.schwarzer , nitarshan } @ mila.quebec 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . To this end , we propose to use a combination of self-supervised objectives for representation learning on offline data , requiring orders of magnitude less pretraining data than existing methods , while approaching human-level data-efficiency when finetuned on downstream tasks . We summarize our work below : RL-aligned representation learning objectives : We propose to pretrain representations using a combination of latent dynamics modeling , unsupervised goal-conditioned reinforcement learning , and inverse dynamics modeling – with the intuition that a collection of such objectives can capture more information about the dynamical and temporal aspects of the environment of interest than any single objective . We refer to our method as SGI ( SPR , Goal-conditioned RL and Inverse modeling ) . Significant advances for data-efficiency on Atari : SGI ’ s combination of objectives performs better than any in isolation and significantly improves performance over previous representation pretraining baselines such as ATC ( Stooke et al. , 2021 ) . Our results are competitive with explorationbased approaches such as APT or VISR ( Liu & Abbeel , 2021 ; Hansen et al. , 2020 ) , which require two to three orders of magnitude more pretraining data and the ability to interact with the environment during training , while SGI can learn with a small offline dataset of exploration data . Scaling with data quality and model size : SGI ’ s performance also scales with data quality and quantity , increasing as data comes from better-performing or more-exploratory policies . Additionally , we find that SGI ’ s performance scales robustly with model size ; while larger models are unstable or bring limited benefits in standard RL , SGI pretraining allows their finetuning performance to significantly exceed that of smaller networks . We assume familiarity with RL in the following sections ( with a brief overview in Appendix A ) . 2 Representation Learning Objectives . A wide range of SSL objectives have been proposed for RL which leverage various aspects of the structure available in agent interactions . For example , the temporal dynamics of an environment can be exploited to create a forward prediction task ( e.g. , Gelada et al. , 2019 ; Guo et al. , 2018 ; Stooke et al. , 2021 ; Schwarzer et al. , 2021 ) in which an agent is trained to predict its immediate future observations , perhaps conditioned on a sequence of actions to perform . Structure in RL goes far beyond forward dynamics , however . We propose to combine multiple representation learning objectives , covering different agent-centric and temporal aspects of the MDP . Based on the intuition that knowledge of an environment is best represented in multiple ways ( Hessel et al. , 2021 ; Degris & Modayil , 2012 ) , we expect this to outperform monolithic representation learning methods such as temporal contrastive learning ( e.g. , Stooke et al. , 2021 ) . In deciding which tasks to use , we consider questions an adequate representation should be able to answer about its environment , including : • If I take action a in state s , what state s′ am I likely to see next ? • If I transitioned from state s to state s′ , what action a did I take ? • What action a would I take in state s so that I reach another state s′ as soon as possible Note that none of these questions are tied to task reward , allowing them to be answered in a fullyunsupervised fashion . Additionally , these are questions about the environment , and not any specific policy , allowing them to be used in offline pretraining with arbitrary behavioral policies . In general , the first question may be answered by forward dynamics modeling , which as mentioned above is well-established in RL . The second question corresponds to inverse dynamics modeling , which has also been commonly used in the past ( Lesort et al. , 2018 ) . The third question corresponds to self-supervised goal-conditioned reinforcement learning which has the advantage of being structurally similar to the downstream target task , as both require learning to control the environment . To facilitate their joint use , we formulate these objectives so that they operate in the latent representation space provided by a shared encoder . We provide an overview of these components in Figure 1 and describe them in greater detail below ; we also provide detailed pseudocode in Appendix D . 2.1 Self-Predictive Representations . SPR ( Schwarzer et al. , 2021 ) is a representation learning algorithm developed for data-efficient reinforcement learning . SPR learns a latent-space transition model , directly predicting representations of future states without reconstruction or negative samples . As in its base algorithm , Rainbow ( Hessel et al. , 2018 ) , SPR learns a convolutional encoder , denoted as fo , which produces representations of states as zt = fo ( st ) . SPR then uses a dynamics model h to recursively estimate the representations of future states , as ẑt+k+1 = h ( ẑt+k , at+k ) , beginning from ẑt , zt . These representations are projected to a lower-dimensional space by a projection function po to produce ŷt+k , po ( ẑt+k ) . Simultaneously , SPR uses a target encoder fm to produce target representations z̃t+k , fm ( st+k ) , which are further projected by a target projection function pm to produce ỹt+k , pm ( z̃t+k ) . SPR then maximizes the cosine similarity between these predictions and targets , using a learned linear prediction function q to translate from ŷ to ỹ : LSPRθ ( st : t+K , at : t+K ) = − K∑ k=1 q ( ŷt+k ) · ỹt+k ||q ( ŷt+k ) ||2 · ||ỹt+k||2 . ( 1 ) The parameters of these target modules θm are defined as an exponential moving average of the parameters θo of fo and po : θm = τθm + ( 1− τ ) θo . 2.2 Goal-Conditioned Reinforcement Learning . Inspired by works such as Dabney et al . ( 2021 ) that show that modeling many different value functions is a useful representation learning objective , we propose to augment SPR with an unsupervised goalconditioned reinforcement learning objective . We define goals g to be normalized vectors of the same size as the output of the agent ’ s convolutional encoder ( 3,136- or 4,704-dimensional vectors , for the architectures we consider ) . We use these goals to annotate transitions with synthetic rewards , and train a modified version of Rainbow ( Hessel et al. , 2018 ) to estimate Q ( st , a , g ) , the expected return from taking action a in state st to reach goal g if optimal actions are taken in subsequent states . We select goals using a scheme inspired by hindsight experience replay ( Andrychowicz et al. , 2017 ) , seeking to generate goal vectors that are both semantically meaningful and highly diverse . As in hindsight experience replay , we begin by sampling a state from another trajectory or the future of the current trajectory . However , we take the additional step of applying stochastic noise to encourage goals to lie somewhat off of the current representation manifold . We provide details in Appendix C.2 . 2.3 Inverse Dynamics Modeling . We propose to use an inverse dynamics modeling task ( Lesort et al. , 2018 ) , in which the model is trained to predict at from st and st+1 . Because this is a classification task ( in discrete control ) or regression task ( continuous control ) , it is naturally not prone to representational collapse , which may complement and stabilize our other objectives . We directly integrate inverse modeling into the rollout structure of SPR , modeling p ( at+k|ŷt+k , ỹt+k+1 ) for each k ∈ ( 0 , . . . , K−1 ) , using a two-layer MLP trained by cross-entropy . 3 Related Work . Data-Efficiency In order to address data efficiency in RL , Kaiser et al . ( 2019 ) introduced the Atari 100k benchmark , in which agents are limited to 100,000 steps of environment interaction , and proposed SimPLe , a model-based algorithm that substantially outperformed previous model-free methods . However , van Hasselt et al . ( 2019 ) and Kielak ( 2020 ) found that simply modifying the hyperparameters of existing model-free algorithms allowed them to exceed SimPLe ’ s performance . Later , DrQ ( Kostrikov et al. , 2021 ) found that adding mild image augmentation to model-free methods dramatically enhanced their sample efficiency , while SPR ( Schwarzer et al. , 2021 ) combined data augmentation with an auxiliary self-supervised learning objective . SGI employs SPR as one of its objectives in offline pretraining , leading to significant improvements in data-efficiency . Exploratory pretraining A number of recent works have sought to improve reinforcement learning via the addition of an unsupervised pretraining stage prior to finetuning on the target task . One common approach has been to allow the agent a period of fully-unsupervised interaction with the environment , during which the agent is trained to maximize a surrogate exploration-based task such as the diversity of the states it encounters , as in APT ( Liu & Abbeel , 2021 ) or ProtoRL ( Yarats et al. , 2021 ) , or to learn a set of skills associated with different paths through the environment , as in DIAYN ( Eysenbach et al. , 2018 ) , VISR ( Hansen et al. , 2020 ) , and DADS ( Sharma et al. , 2019 ) . Others have proposed to use self-supervised objectives to generate intrinsic rewards encouraging agents to visit new states ; e.g . Pathak et al . ( 2017 ) and Burda et al . ( 2018 ) use the loss of an inverse dynamics model like that used in SGI , while Sekar et al . ( 2020 ) uses the disagreement between an ensemble of latent-space dynamics models . Finally , Campos et al . ( 2021 ) report strong results based on massive-scale unsupervised pretraining . Many of these methods are used to pretrain agents that are later adapted to specific reinforcement learning tasks . However , SGI differs in that it can be used offline and is agnostic to how data is collected . As such , if no pre-existing offline data is available , one of the methods above can be used to generate a dataset for SGI ; we use Burda et al . ( 2018 ) for this in Section 4.2 . Visual Representation Learning Computer vision has seen a series of dramatic advances in selfsupervised representation learning , including contrastive methods ( Oord et al. , 2018 ; Hjelm et al. , 2019 ; Bachman et al. , 2019 ; He et al. , 2020 ; Chen et al. , 2020a ) as well as purely predictive ones ( Grill et al. , 2020 ) . Variants of these approaches have also been shown to improve performance when coupled with a small quantity of labeled data , in a semi-supervised setting ( Chen et al. , 2020b ; Hénaff et al. , 2019 ) , and several self-supervised methods have been designed specifically for this case ( for example , Sohn et al. , 2020 ; Tarvainen & Valpola , 2017 ) . These advances have spurred similar growth in methods aimed specifically at improving performance in RL . We refer the reader to Lesort et al . ( 2018 ) for a review of earlier methods , including inverse dynamics modeling which is used in SGI . Recent research has focused on leveraging latent-space dynamics modeling as an auxiliary task . Gelada et al . ( 2019 ) propose a simple next-step prediction task , coupled with reward prediction , but found it is prone to latent space collapse and requires an auxiliary reconstruction loss for experiments on Atari . Guo et al . ( 2020 ) use a pair of networks for both forward and backward prediction , and show improved performance in extremely large-data multi-task settings . Mazoure et al . ( 2020 ) use a temporal contrastive objective for representation learning , and show improvement in continual RL settings . Concurrently , SPR ( Schwarzer et al. , 2021 ) proposed a multi-step latent prediction task with similarities to BYOL ( Grill et al. , 2020 ) . Two works similar to ours , Anand et al . ( 2019 ) and Stooke et al . ( 2021 ) , propose reward-free temporalcontrastive methods to pretrain representations . Anand et al . ( 2019 ) show that representations from encoders trained with ST-DIM contain a great deal of information about environment states , but they do not examine whether or not representations learned via their method are , in fact , useful for reinforcement learning . However , Stooke et al . ( 2021 ) employ a similar algorithm and find only relatively minor improvements in performance compared to standard baselines in the large-data regime ; our controlled comparisons show that SGI ’ s representations are far better for data-efficiency . Concurrent to our work , FERM ( Zhan et al. , 2020 ) propose contrastive pretraining from human demonstrations in a robotics setting . As FERM is quite similar to ATC , we are optimistic that our improvements over ATC in Atari 100k would translate to FERM ’ s setting . Finally , Yang & Nachum ( 2021 ) propose a family of algorithms for learning representations for DRL from state observations ; we believe that these might be promising in DRL from pixels , but they have not yet been adapted to that setting .

Summary: This paper proposes a new approach for pretraining an encoder using unlabeled data which is then used for a downstream RL task, leading to large improvements in sample efficiency. The approach is based on latent dynamics modeling and unsupervised goal-conditioned RL. The approach shows significant gains in performance and sample efficiency on Atari 100k and scales well with model size and better data.