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: The proposed method combines 3 unsupervised representation learning objectives (latent dynamics modelling, goal-conditioned RL, and inverse dynamics modeling) in order to pre-train representations for RL. The experiments on the Atari 100K benchmark show that SGI improves performance over previous representation pretraining baselines (ATC, CPT), and is competitive with exploration-based approaches (APT, VISR) which require more pretraining data and interactions with the environment.  Summary of experimental results:  (1) Data-efficiency: Table1 shows that SGI achieves comparable performance as APT and CPT (unsupervised exploration methods) but using far fewer data and environment interactions (250M steps for APT, 4B steps for CPT, 6M for SGI).  (2) Varying the quality of the offline pre-training data (collected by a random policy, SGI-R; an exploratory policy, SGI-E; suboptimal DQN actions, SGI-W; and DQN agent at different checkpoints, SGI-M): SGI’s performance increases as data comes from better-performing or more-exploratory policies.  (3) Varying the size of the encoder used in SGI-M: SGI-M/S (small), SGI-M/L (large): Larger model size results in increased performance.  (4) Pretraining vs. size of the network (Section 5.2): Model size and amount of pretraining data should increase in parallel.  (5) Ablation study (Section 5.3, Table 3): Inverse dynamics modeling objective seems to be the most critical to performance. Performance increases as more objectives are used.  (6) How to finetune pre-trained representations properly: Reducing the finetuning LRs improves performance.