Pretraining Representations for Data-Efficient Reinforcement Learning
Authors: Max Schwarzer
Nitarshan Rajkumar
Michael Noukhovitch
Ankesh Anand
Laurent Charlin
R Devon Hjelm
Philip Bachman
Aaron Courville
Venue: NeurIPS 2021 Poster (2021)
Keywords: Pretraining
Sample Efficiency
Reinforcement Learning
Self-Supervised Learning

Abstract: Data efficiency is a key challenge for deep reinforcement learning. We address this problem by using unlabeled data to pretrain an encoder which is then finetuned on a small amount of task-specific data. To encourage learning representations which capture diverse aspects of the underlying MDP, we employ a combination of latent dynamics modelling and unsupervised goal-conditioned RL. When limited to 100k steps of interaction on Atari games (equivalent to two hours of human experience), our approach significantly surpasses prior work combining offline representation pretraining with task-specific finetuning, and compares favourably with other pretraining methods that require orders of magnitude more data. Our approach shows particular promise when combined with larger models as well as more diverse, task-aligned observational data -- approaching human-level performance and data-efficiency on Atari in our best setting.

Pretraining Representations for Data-Efficient Reinforcement Learning
Max Schwarzer*1., Nitarshan Rajkumar" 1,2, Michael Noukhovitch 1.2 Ankesh Anand 1.2 Laurent Charlin 13,55 Devon Hjelm' 1.4 , Philip Bachman', Aaron Courville
Mila, Université de Montréal, HEEC Montréal, *Microsoft Research, $CIFAR
Abstract
Data efficiency is key challenge for deep reinforcement learning. We address this problem by using unlabeled data to pretrain an encoder which is then finetuned on small amount of task-specific data. To encourage learning representations which capture diverse aspects of the underlying MDP, we employ a combination of latent dynamics modelling and unsupervised goal-conditioned RL. When limited 100k steps of interaction on Atari games (equivalent to two hours of human experience), our approach significantly surpasses prior work combining offline representation pretraining with task-specific finetuning, and compares favourably with other pretraining methods that require orders of magnitude more data. Our approach shows particular promise when combined with larger models as well more diverse, task-aligned observational data approaching human-level performance and data-efficiency on Atari in best setting. We provide code associated with this work at https: //github. com/mila--a iqia/SGI.
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 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 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 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 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 train agents explore their environment in 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 interaction 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).
agents cannot be run faster than real-time, may require significant human
[mmxx. schwarzer,, nitarshan) @mila.quehec
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Unsupervised Pre-Training
1. Self- Predictve Representations
3. Goal conditionddd RL
Shared Feature Extractor
coline
Coenn Siilrrty Low
sompled gool
8op4 ic
taraal
compuale dolann apee
d(24,9)
orlia enoockr
2. Inverse Modelling
compule distance Is gocl
d( d( .)
comn iranaiice hoe model Risk
moweee Pip.
Inwersn Modial Loss
Reward dl (f),.) 18)
Se+1
pesiction
G-Learning Loss
monnder
Task Task-Specifi RL
Ppokk oug.
leaminng head
Learming
Zt
SPR Loss
Figure 1: A schematic diagram showing our two stage pretrain-ther therin--ttennineuuee method. All unsuper- vised training losses and task-specific RL use the shared torso the left.
To this end, we propose use a combination 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 reinforcement learning,
combination of latent dynamics modeling, unsupervised and inverse dynamics modeling- with the intuition that collection of such objectives capture more information about the dynamical and temporal aspects of the environment of interest than any single objective. We refer to method 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 basclines such ATC (Stooke al., 2021). Our results are competitive with exploration- based 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 the following sections (with a brief overview Appendix A).
2 Representation Learning Objectives
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 (c.g., Gelada et al., 2019; Guo et al., 2018; Stooke et al., 2021; Schwarzer al., 2021) in which 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 al., 2021; Degris& & Modayil, 2012), we expect this to outperform monolithic representation learning methods such temporal contrastive learning (c.g., Stooke al., 2021). In deciding which

---

tasks to use, we consider questions an adequate representation should be able to answer about its environment, including:
If take action a in state 8, what state s' am likely to see next?
If transitioned from state state s', what action did I take?
What action a would take in state 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 fully- unsupervised 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 RI. The second question corresponds to inverse dynamics modeling, which has also been commonly used in the past (Lessrrt et al., 2018). The third question corresponds to self-supervised goal-conditioned reinforcement learning which has the advantage of being structurally similar 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 representa- tion space provided by 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 al.. 2021) is representation learning algorithm developed for data-efficient reinforcement learning. SPR learns a latent-space transition model, directly predicting representations of future states without reconstruction negative samples. As in its base algorithm, Rainbow (Hessel et al., 2018), SPR learns a convolutional encoder, denoted as for which produces representations of states 2₁ fo())). SPR then uses dynamics model h to recursively estimate the representations of future states, ē &( a+++), beginning from ε - 2t- These representations projected lower-dimensional space by a projection function po to produce dc+* _ (o(
Simultancously. SPR uses target encoder fm to produce target representations ē _ fm ($++ which are further projected by a target projection function pm to produce Jr4. - Pm (Zc+&). SPR then maximizes the cosine similarity between these predictions and targets, using a learned linear prediction function g translate from á to y/:
K g((c+k) Deth
SPR C9
(Sta+K K2 (044k) =
(1)
[1g(0++))llz 110e+all2
k=1
The parameters of these target modules On are defined as exponential moving average of the parameters 0, so and po: On- = 70m + (1-78 -)...
2.2 Goal-Conditioned Reinforcement Learning
Inspired by works such as Dabney et al. (2021) that show that modeling many different value functions a useful representation learning objective, we propose to augment SPR with an unsupervised goal- conditioned reinforcement learning objective. We define goals to be normalized vectors of the same size the output of the agent's convolutional encoder (3,136- or 4, 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 (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 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 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 (I (Lesort et al., 2018), in which the model is trained to predict a from s; and s1+1- Because this is a classification task (in discrete control) or
3

---

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(@t- & +++ for each & E .0, using a two-layer ML.P trained by cross-entropy.
K-1),
Related Work
Data-Efticiency 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, 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 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 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 uully-unsupervieed interaction with the environment, during which the agent 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) 8), 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 wussivee-scaee 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 al. (2018) for this in Section 4.2.
Visual Representation Learning Computer vision has seen a series of dramatic advances in self- supervised 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 small quantity of labeled data, in 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 auxiliary task. Gelada et al. (2019) propose simple next-step prediction task, coupled with reward prediction, but found 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 temporal contrastive objective for representation learning, and show improvement in continual RL settings. Concurrently, SPR (Schwarzer et al., 2021) proposed 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 temporal- contrastive methods to pretrain representations. Anand et al. (2019) show that representations from encoders trained with ST-DIM contain 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 cata--effcieeey. Concurrent our work. FERM (Zhan et al., 2020) propose contrastive pretraining from human demonstrations in 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 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.
4 Experimental Details
In our experiments, We seek to address two main challenges for the deployment of RL agents in the real world (Dulac-Arnold et al., 2020): (I) training the RL agent with a limited budget of interactions in the real environment, and (2) leveraging existing interaction data of arbitrary quality.
4.1 Environment and Evaluation
To address the first challenge, we focus our experimentation on the Atari 100k benchmark introduced by Kaiser et al. (2019), in which agents are allowed only 100k steps of interaction with their environment. 2 This is roughly equivalent to the two hours human testers were given to learn these games by Mnih et al. (2015), providing baseline of human -ample-efficiency.
Atari is also an ideal setting due to its complex observational spaces and diverse tasks, with 26 different games included in the Atari 100k benchmark. These factors have led to Atari's s extensive use for representation learning and exploratory pretraining (Anand et al., 2019; Stooke al., 2021; Campos et al., 2021), and specifically Atari 100k for data-efficient RL (e.g., Kaiser et al., 2019; Kostrikov et al., 2021; Schwarzer et al., 2021), including finetuning after exploratory pretraining (e.g.. Hansen et al., 2020; Liu & Abbeel, 2021), 1). providing strong baselines compare to.
Our evaluation metric for an agent on a game is noraa-normmlized score (HNS), defined as agent score-randorg score
human_sco scwoesssoee-eee wow.ssrrroooooooyoooo rrandoooommmmmomo score' We calculate this per game by averaging scores over 100 evaluation trajectories the end of training, and across 10 random seeds for training. We report both mean (Mn) and median (Mdn) HNS over the 26 Atari-100K games, as well as on how many games a method achieves super-human performance (>H) and greater than random performance (>0). Following the guidelines of Agarwal et al. (2021) we also report interquartile mean HNS (IQM) and quantify uncertainty via bootstrapping in Appendix B.
4.2 Pretraining Data
The second challenge pertains to pretraining data. Although some prior work on offline represen- tational pretraining has focused on expert-quality data (Stooke al., 2021), we expect real-world pretraining data to be of greatly varying quality. We thus construct four different pretraining datasets to approximate different data quality scenarios.
(R)andom To assess performance near the lower limit of data quality, we use a random policy to gather a dataset of 6M transitions for each game. To encourage the agent to venture further from the starting state, we execute each action for a random number of steps sampled from a Geometric( } distribution.
(E)xploratory To emulate slightly better data that covers a larger range of the state space, we use an exploratory policy. Specifically, we employ the IDF (inverse dynamics) variant of the algorithm proposed by Burda et al. (2018). We log the first 6M steps from an agent trained in each game. This algorithm achieves setter-than-randmm performance on only 70% of tasks, setting far below the performance of more modern unsupervised exploration methods.
To create higher-quality datasets, we follow Stooke et al. (2021) and use experience gathered during the training of standard DQN agents (Mnih et al., 2015). We opt to use the publicly-availabe DQN Replay dataset (Agarwal et al., 2020), which contains data from training for 50M steps (across all 57 games, with five different random seeds). Although we might prefer to use data from recent
Note that sticky actions disabled under this benchmark.
5

---

unsupervised exploration methods such as APT (Liu & Abbeel, 2021), VISR (Hansen et al., 2020), or CPT (Campos et al.. 2021), none of these works provide code or datasets, making this impractical. We address using data collected from on-task agents with behavioural cloning baseline in Section 5, with surprising findings relative to prior work.
(W)eak We first generate a weak dataset by selecting the first IM steps for each of the 5 available runs in the DQN Replay dataset. This data is generated with an e-greedy policy with high, gradually decaying e, leading to substantial action diversity and many suboptimal exploratory actions. Although the behavioral policies used to generate this agent are not especially competent (see Table 1), they have above-random performance on almost all games, suggesting that that this dataset includes more task-relevant transitions.
(M)ixed Finally, for a realistic best-case scenario, we create dataset of both medium and low- quality data. To simulate real-world collection of data from different policies, we concatenate multiple checkpoints evenly spread throughout training of a DQN. We believe this is also a reasonable approximation for data from a modern unsupervised exploration method such CPT (Campos et al.. 2021); as shown in Table 1, the agent for this dataset has performance in between CPT and VISR, with median closer to CPT and mean closer to VISR. This data is also lower quality than the expert data originally used in the method most similar to ours, ATC (Stooke et al., 2021). We create dataset of 3M steps and larger dataset of 6M steps; all M experiments use the 3M step dataset unless otherwise noted.
We compare the agents used for our datasets to those for unsupervised explo- ration pretraining baselines in Table 1. We estimate the performance of the Weak and Mixed agents as the average of the corresponding logged evaluations in the Dopamine (Castro et al., 2018) baselines. Even our largest dataset is quite small compared to the amounts of data gath- ered by unsupervised exploration methods (see the "Data" column in Table I); this is intentional, as we believe that unsuper- vised interactional data may be expensive in real world applications. We show the performance of the non-random data col- lection policies in Table 2 (note that fully- random policy has a score of by defini- tion).
Table I: Performance of agents used in pretraining data collection compared to external baselines on 26 Atari games (Kaiser et al., 2019)
Method
Mdn
Mean
>
>0
Data
Exploratory Pretraining Baselines
VISR@0
0.056
0.817
5
19
250M
APT@O
0.038
0.476
2
18
250M
C
0.809
4.945
12
25
4B
Offine Datasets
Exploratory Weak2 Mixed
0.039
0.042
0
18
6M
0.028
0.056
0
23
5M
0.618
1.266
10
26
3M
Calculated from ICLR 2021 OpenReview submis-
sion; unreported in arXiv version.
2 Upper-bound estimate from averaging evaluation performance of corresponding agents in Dopamine.
4.3 Training Details
We optimize our three representation learning objectives jointly during unsupervised pretraining, summing their losses. During finetuning, we optimize only the reinforcement learning and forward dynamics losses, following Schwarzer et al. (2021) (see Section 5.5), and lower the learning rates for the pretrained encoder and dynamics model by two orders of magnitude (see Section 5.4).
We consider the standard three-layer convolutional encoder introduced by Mnih et al. (2015), ResNet inspired by Espeholt et al. (2018), as well as an enlarged ResNet of the same design. In other respects, our implementation matches that of SPR and is based on its nublielyrreleased code. Full implementation and hyperparameter details are provided in Appendix C. We refer to agents by the model architecture and pretraining dataset type used: SGI-R is pretrained on Random, SGI-E on Exploratory, SGI-W on Weak, and SGI-M on Mixed. To investigate scaling, we vary the size of the encoder used in SGI-M: the standard Mnih et al. (2015) encoder is SGI-M/S (for small), our standard ResNet is simply SGI-M and using larger ResNet is SGI-M/L (for large)'. For SGI-M/L we also use the 6M step dataset described earlier. All ablations are conducted in comparison to SGI-M unless
1 OOur data-collection agents are weaker than those used by ATC on seven of the eight games they consider.
'See Appendix C for details on these networks

---

otherwise noted. Finally, agents without pretraining are denoted SGI-None; SGI-None/S would be roughly equivalent SPR (Schwarzer et al., 2021).
For baselines, we compare to no-pretraining Atari 100k methods (Kaiser et al., 2019; van Hasselt al., 2019; Kostrikov et al., 2021; Schwarzer al., 2021). For our models trained on Random and Exploratory data we compare against previous pretraining-vi--explraate Atari 100k (Liu & Abbeel, 2021; Hansen et al., 2020; Campos et al., 2021). In the higher quality data regime, we compare to recent work on data-agnostic unsupervised pretraining, ATC (Stooke et al., 2021), as well as behavioural cloning (BC).
approaches applied to
Table 2: HNS on Atari100k for SGI and baselines.
Method
Mdn
Mn
>
>0
Data
5 Results and Discussion
No Pretraining (Finetuning Only)
We find that SGI performs competitively on the Atari- benchmark; presenting aggregate results in Table 2, and full per-game data in Ap- pendix Our best setting, SGI-M/L, achieves median HNS of 0.753, approaching human- level samplc-efficiency and outperforming all comparable methods except the recently pro- posed CPT (Campos et al., 2021). With less data and smaller model, SGI-M achieves median HNS of 0.679, significantly outperforming the prior method ATC on the same data (ATC-M). Meanwhile, SGI-E achieves a median HNS of 0.456, matching exceeding other exploratory methods such as APT (Liu & Abbeel, 2021) and VISR (Hansen et al., 2020), well as ATC-E.
SimPLe
0.144
0.443
2
0
DER
0.161
0.285
2
26
0
DrQ
0.268
0.357
2
24
0
SPR
0.415
0.704
7
26
0
SGI-None
0.343
0.565
3
26
0
Exploratory Pretraining + Finetuming
Method
Mdn
Mn
>H
>00
Data
VISR
0.095
1,281
7
21
250M
APT
0.475
0.666
7
250M
CPT@02
0.809
4.945
12
25
4000M
ATC-R² ATC-E²
0.191
0.472
4
26
6M
0.237
0.462
3
26
6M
SGI-R SGI-E
0.326
0.888
5
26
6M
0.456
0.838
6
6M
Offline-data Pretraining Finetuning
Pretraining data efficiency SGI achieves strong performance with only limited pretrain- ing data; our largest dataset contains 6M transi- tions, or roughly 4.5 days of experience. This compares favourably to recent works on unsu- pervised exploration such APT CPT. which require far larger amounts of data and environ- ment interaction (250M steps or 193 days for APT, 4B steps or 8.45 years for CPT). We expect SGI would perform even better if used these large-data settingss, as we find that it scales ro- bustly with data (see Section 5.2).
Method
Mdn
Mn
>H
>0
Data
ATC-Wa
0,219
0.587
4
3M
ATC-M
0.204
0.780
5
26
3M
BC-M@0
0.139
0.227
0
23
3M
BC-M
0.548
0.858
8
26
3M
SGI-W
0.589
1.144
8
5M
SGI-M/S
0.423
0.914
8
3M
SGI-M
0.679
1.149
9
26
3M
SGI-M/L
0.753
1.598
9
26
6M
APT claims 0.6955, but we calculate 0.666 from their reported per-game scores.
CPT @0 does not do any finetuning.
Our implementation (see Appendix C.6)
Behavioural cloning is a strong baseline Al- though ATC pretrains with expert data, they did not investigate behavioral cloning as a base- line pretraining objective. We do so on our Mixed dataset, the only one to be generated by policies with significantly above-random perfor- mance. Behavioral cloning without finetuning (BC-M@0) performs poorly, perhaps due to the varying behavioural quality in the dataset. But when finetuned, BC-M yields very respectable performance, surpassing ATC-M but not SGI- M. All fine-tuning settings for BC-M match SGI- M.
2.0
Fire-tuning score (Whtuned)
1.5-
-0.5
Sci Trend
Random
Exploratory
Weak
-1.5
Mixed
-2.0
Score (Whitaned)
Figure 2: SGI finetuning performance vs. pretrain- ing data score for all combinations of game and dataset. Data score is estimated as clipped return per episode, trend calculated via kernel regression. Values whitened per-game for clarity.
5.1 Data quality matters
In principle, SGI can be used with any offline datasct but we demonstrate that it scales with
7

---

0.8
Large Reillet
Medium Reshit
0.7
Mnth et (2015)
similak
1044 HNS
0.5
0.6
0.4
346
S+1
0.
0.2
0.1
25089 75000 125000 150000 175000
Pretraining steps)
Figure 3: Finetuning performance of SGI for different CNN sizes and amounts of pretrained data from the Mixed dataset. We plot IQM HNS with confidence intervals (see Appendix B).
Figure 4: Average cosine similarity between rep- resentations over pretraining, averaged across the 26 Atari 100k games. indicates representations are identical, 0 perfect dissimilarity.
the quality of its data. Near the lower bound of data quality where all actions are selected randomly, SGI-R still provides some benefit over an atherwise-identical rundomly-initiaizze
agent (SGI-None) 16 out of 26 games, with similar median but 57% higher mean HNS. With better data from an exploratory policy, SGI-E improves on 16/26 games, gets 33% higher median HNS, and surpasses APT (Liu & Abbeel, 2021) which used 40 times more pretraining data. With similarly weak data but possibly more task-specific transitions, SGI-W gets 72% higher median HNS compared SGI-None and with realistic data from mixture of policies, SGI-M improves to 98%.
Importantly, the pattern observe is very different from what would be expected for imitation learning. In particular, SGI-W's strong performance shows that expert data is not required. To characterize this, we plot the average clipped reward experienced per episode for each of our pretraining datasets in Figure 2. Normalizing across tasks, we find a strong positive correlation between task reward engagement (p < 0.0001) and finetuning performance. Moreover, find diminishing returns to further task engagement.
5.2 Pretraining unlocks the value of larger networks
The three-layer network introduced by Mnih al. (2015) has become a fixture of deep reinforcement learning, and has been employed by previous works examining pretraining in this field (e.g. Liu Abbeel, 2021; Stooke et al., 2021). However, we find that representational pretraining with this network (SGI-M/S) provides only minor benefits compared to training from scratch. In contrast, larger networks struggle without pretraining but shine when pretrained as shown in Figure
This finding is consistent with recent work in unsupervised representation learning for classifi- cation, which has observed that unsupervised pretraining benefits disproportionately from larger networks (Chen et al., 2020a). In particular, our results suggest that model size should increase in par- allel with amount of pretraining data, matching recent work on scaling in language modeling (Kaplan et al., 2020; Hernandez et al., 2021). SGI thus provides simple way use unlabeled data to extend the benefits of large networks, already well-known in the large-data regime (e.g., Schrittwieser et al., 2021; Espeholt et al., 2018), to data-efficient RI.
5.3 Combining SGI's objectives improves performance
We test all possible combinations of our three SSL. objectives, denoted by combinations of the letters S. G and I to indicate which objectives they employ. Results Table 3 show that performance monotonically increases as more objectives are used, with inverse dynamics modeling combined with either of the other objectives performing respectably well. This illustrates the importance of using multiple objectives to obtain representational coverage of the MDP.
Unclipped rewards available for the offline DQN dataset.

---

We note that including inverse modeling appears to be critical, and hypothesize that this is related to represen- tational collapse. To measure this, we plot the average cosine similarity between representations Y/ of differ- ent states for several pretraining methods in Figure 4, using our ResNet encoder on the Mixed dataset. We observe that S. G and S+G all show some degree repre- sentational collapse, while configurations that include inverse dynamics modeling avoid representational col- lapse, as does ATC, whose contrastive loss implicitly optimizes for representational diversity (Wang& & Isola, 2020). Intriguingly, we observe that increased represen- tational diversity does not necessarily improve perfor- mance. For example, SGI outperforms ATC, G+I and I In finetuning but has less diverse pretraining representations. We also observe that adding SPR (S)
Table 3: HNS on Atari 100K for pre- training ablations of SGI.
Pretraining
Mdn
Mean
>> H
None
0.343
0.565
S
0,009
-0.054
G
0.060
0.181
0.411
0.943
S+G
0.029
0.098
0
G+I
0.512
1.004
S+I
0.629
0.978
SGI-M
0.679
1.149
consistently pulls representations towards collapse (compare S+I and I, S+G and G, and SGI and G+I); how this relates to performance is question for future work.
5.4 Naively finetuning ruins pretrained representations
We find that properly finetuning pretrained representa- tions is critical, as results in Table 4 show. Although allowing pretrained weights change freely during fine- tuning is better than initializing from scratch (Naive vs No Pretrain), freezing the pretrained encoder (Frozen) leads to better performance than either. SGI's approach of reducing finetuning learning rates for pretrained pa- rameters leads to superior performance (Reduced LRs, equivalent SGI-M).
Table 4: HNS on Atari 100K for fine- tuning schemes for SGI.
Method
Mdn
Mean
H
No pretrain
0.343
0,565
3
Naive
0.429
0.845
8
Frozen
0.499
0.971
8
Reduced LRs
0.679
1.149
9
We thus hypothesize that representations learned by SGI are being disrupted by gradients early in finetuning, in phenomenon analogous catastrophic forgetting (Zenke et al., 2017; Hadsell et al., 2020). As representations may not generalize between different value functions across training (Dabney et al., 2021), allowing the encoder strongly adapt early in training could make worse at modeling later value functions, compared to the neutral initialization from SGI. We also note that there is long history in computer vision of employing specialized finetuning hyperparameters (Li et al., 2020; Chu et al., 2016) when transferring tasks.
5.5 Not all SSL objectives are beneficial during finetuning
Although SGI uses S during finetuning, we experiment with a variant that optimizes only the standard DQN objective, roughly equivalent to using DrQ (Kostrikov al., 2021) with DQN hyperparameters set match SGI. We find that finetuning with has a large impact with or without pretraining (compare None and Only entries in Table 5.). Although, SGI without S manages to achieve roughly the same mean with S, removing S harms performance on performance on 19 out of 26 games and reduces median normalized score by roughly 33%. We also find no benefit to using all of SGI's constituent objectives during finetuning (All Losses in Table 5) compared to using S alone, although the gap between them is not statistically significant for metrics other than median (see Figure 5d).
Table 5: HNS on Atari 100K for finetun- ing ablations of SGI.
Finetune SSL
Mdn
Mean
>
Without SGI pretraining
None S Only
0.161
0.315
2
score as SGI
0.343
0.565
With SGI pretraining
None
0.452
1.114
8
SGI
0.397
1.011
S Only
0.679
1.149
9

---

6 Conclusion
We present SGI, a fully self-supervised (reward-free) approach representation learning for re- inforcement learning, which uses a combination of pretraining objectives to encourage the agent to learn multiple aspects of environment dynamics. We demonstrate that SGI enables significant improvements on the Atari 100k data-efficiency benchmark, especially in comparison to unsupervised exploration approaches which require orders of magnitude more pretraining data. Investigating the various components of SGI, we find that it scales robustly with higher-quality pretraining data and larger models, that all three of the self-supervised objectives contribute to the success of this approach, and that careful reduction of fine-tuning learning rates is critical to optimal performance.
7 Acknowledgements
We would like to thank Mila and Compute Canada for computational resources. We would like to thank Rishabh Agarwal for useful discussions. Aaron Courville would like to acknowledge financial support from Microsoft Research and Hitachi. We would also like to thank Wancong (Kevin) Zhang for his helpful comments to the codebase GitHub.
References
Agarwal, R., Schuurmans, D., and Norouzi, M. An optimistic perspective offline reinforcement learning. ICML., 2020. 5
Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A., and Bellemare, M. G. Deep reinforcement learning at the edge of the statistical precipice. NeurIPS, 2021. 5, 15, 19
Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M.-A., and Hjelm, R. D. Unsupervised state representation learning in atari. NeurIPS, 2019, 5,5
Andrychowicz, M., Wolski, F., Ray, A., Schneider, J., Fong. R., Welinder, P., McGrew, B., Tobin, J., Abbeel, O. P., and Zaremba, W. Hindsight experience replay. In Advances in neural information processing systems, pp. 5048-5058, 2017.
Ba. J. L., Kiros, J. R., and Hinton, G. E. Layer normalization. arXiv preprint arXiv:1607.06450. 2016. 18
Bachman, P., Hjelm, R. D., and Buchwalter, W. Learning representations by maximizing mutual information across views. In NeurlPS, 2019.
Bellemare, M. G., Dabney, W., and Munos, R. A distributional perspective on reinforcement learning. ICML, 2017. 15
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A.,
Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. arXiv preprint
14X1v:2005.14165, 2020. 1
Burda. Y., Edwards, H.. Pathak, D., Storkey, A., Darrell, T. and Efros, A. A. Large-scale study of curiosity-driven learning, 2018. 5,
Campos, V., Sprechmann, P., Hansen, S. S., Barreto, A., Blundell, C., Vitvitskyi, A., Kapturowski, S., and Badia, A. Coverage principle for discovering transferable behavior in reinforcement learning, 2021. URL https:// /ttpt:/e///ppeevviiee
Framework for Deep Reinforcement Learning. 2018. URL IdddddIhhhJJJJdddd 1, 4, 5, 6, 10. 6
Castro, P. S., Moitra, S., Gelada, C., Kumar, S., and Bellemare, M. G. Dopamine: A Research
Chen, T., Korbblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual representations. /CML, 2020a. 1, 4, 8
Chen, T., Kornblith, S., Swersky, K., Norouzi, M., and Hinton, G. Big self-supervised models are strong semi-supervised Iearners. In NeurIPS, 2020b. 4
10

---

Chu. B., Madhavan, V., Beijbom, O., Hoffman, J., and Darrell, T. Best Practices for Fine-Tuning Visual Classifiers to New Domains, pp. 435-442. 11 2016. ISBN 978-3-319-49408-1 doi: 10.10779988-33319-9900
Dabney, W., Barreto, A., Rowland, M., Dadashi, R., Quan, J., Bellemare, M. G., and Silver, D. The ialue-impoveemen path: Towards better representations for reinforcement learning. 2021. 3, 9
Degris, T. and Modayil, J. Scaling-up knowledge for a cognizant robot. AAAI Spring Symposium on Designing Intelligent Robots: Reintegrating AI., 2012. 2
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. In ACL, 2019, 1
Dubey, R., Agrawal, P., Pathak, D., Griffiths, T., and Efros, A. Investigating human priors for playing video games. In International Conference on Machine Learning, 2018.
Dulac-Amold, G., Levine, N., Mankowitz, J., Li, J., Paduraru, C., Gowal, S., and Hester, T. An empirical investigation of the challenges of real-works reinforcement learning. 2020. 5
Espeholt, L., Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dunning, I., Legg, S., and Kavukeuoglu, K. IMPALA: scalable distributed deep- with importance weighted actor-learner architectures. CoRR, abs/1802.01561, 2018. URL http:/ http:://axxvvvioo org/ abs/ 1802. 01561. 6, 8, 18
Eysenbach, B., Gupta, A., Ibarz, J., and Levine, S. Diversity is all you need: Learning skills without reward function. International Conference on Learning Representations, 2018. 4
Gelada, C., Kumar, S., Buckman, J., Nachum, O., and Bellemare, M. G. Deepmdp: Learning continuous latent space models for representation learning, ICML, 2019. 2, 4
Grill, J.-B., Strub, F., Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., Doersch, C., Pires, B. A., Guo, 7. D., Azar, M. G., et al. Bootstrap your own latent: A new approach to self-supervised learning. In NeurlPS, 2020. 1,
Guo, D., Pires. A., Piot, B., Grill, J.-b., Altché, F., Munos, R., and Azar, M. G. Bootstrap latent- preent-rediitive representations for multitask reinforcement learning. In ICML, 2020. 4
Guo. Z. D., Azar, M. G., Piot, B., Pires, B. A., and Munos, R. Neural predictive belief representations. ICML, 2018. 2
Hadsell, R., Rao, D., Rusu, A. A., and Pascanu, R. Embracing change: Continual learning in deep neural networks. Trends in Cognitive Sciences, 2020. doi: http:///bbb.org//01116f.....i.i 04. URL http:/ http:://www.... www.://wienneeeee...... com/ wwt./ccc/cca/eeceeeeeeeeeeeeeeee sciescc/eeeeceecceee http://cicccerriiiaaee http://iwccoacccccc/ap/a///aaaaa
Hansen, S., Dabney, W., Barreto, A., Warde-Farley, D., de Wiele, V., and Mnih, Fast task inference with variational intrinsic successor features. In ICLR, 2020. 1, 2, 4, 6,
Harris, C. R., Millman, K. J., van der Walt, S. J., Gommers, R., Virtanen, P., Cournapeau, D., Wieser, E., Taylor, J., Berg, S., Smith, N. J., Kern, R., Picus, M., Hoyer, S., van Kerkwijk, M. H., Brett, M., Haldane, A., del R' 10, J. F., Wiebe, M., Peterson, P., G' Grard-Marchaan P., Sheppard, K., Reddy, T., Weckesser, W., Abbasi, H., Gohlke, C., and Oliphant, E. Array programming with NumPy. Nature, 585(7825)35773622 September 2020. doi: 10.1038/441586-02--2666
URL
https://doi: /ttp::://doo.org1o...o.o 1038/ 0000000///4/88/ s015ss444566-01 020- 2649- 15
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. In CVPR, 2016. 18
He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020.
Hénaff, J., Srinivas, A., De Fauw, J., Razavi, A., Doersch, C., Eslami, S., and Oord, A. d. Data- efficient image recognition with contrastive predictive coding. arXiv preprint arXiv: 1905, 09272, 2019. 1, 4

---

Hemannez, D., Kaplan, J., Henighan, T., and McCandlish, S. Scaling laws for transfer, 2021. 8
Hessel, M., Modayil, J., van Hasselt, H., Schaul, T., Ostrovski, G., Dabney, W., Horgan, D., Piot, B., Azar, M. G., and Silver, D. Rainbow: Combining improvements in deep reinforcement learning. In AAAI, 2018. 3, 15
Hessel, M., Danihelka, I., Viola, F., Guez, A., Schmitt, S., Sifre, L., Weber, T., Silver, D., and van Has- selt, H. Muesli: Combining improvements in policy optimization. arXiv preprint ar arXiv:2104.06159 2021. 2
Hjelm, R. D., Fedorov, A., Lavoie-Marchildon S., Grewal, K., Bachman, P., Trischler, A., and Bengio, Y. Learning deep representations by mutual information estimation and maximization. ICLR, 2019. 4
loffe, S. and Szegedy, C. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, 2015. 18
Kaiser, L., Babaeizadeh, M., Mitos, P., Osiński, B., Campbell, R. H., Czechowski, K., Erhan, D.,
Finn, C., Kozakowski, P., Levine, S., et al. Model based reinforcement learning for atari. In ICLR, 2019. 4, 5, 6, 7, 21, 22, 23, 24
Kaplan, J., McCandlish, S., Henighan, T., Brown, T. B., Chess, B., Child, R., Gray, S., Radford, A., Wu, J., and Amodei, D. Scaling laws for neural language models. arXiv preprini arXiv:2001.0836/
2020. 8
Kielak, K. P. Do recent advancements in model-based deep reinforcement learning really improve data efficiency?, 2020. URL https:://// hatp:://pppeeeeeeeee net peet/forum?? wtrrfformmmmmmeeeer ddi.kdeuk---eeed 4
Kostrikov, I., Yarats, D., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In ICLR, 2021. 4, 5, 7, 18
Lake, B. M., Ullman, D., Tenenbaum, J. B., and Gershman, S. J. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017. 1
Lesort, T., Diaz-Rodríguez, N., Goudou, J.-F, and Filliat, D. State representation learning for control: overview. Neural Networks, 108, 2018. 3.4
Li, H., Chaudhari, P., Yang. H., Lam, M., Ravichandran, A., Bhotika, R., and Soatto, Rethinking the hyperparameters for fine-tuning. In International Conference on Learning Representations, 2020. URL https: / /ttpp//eeeeeeeeeeeformnmmaaaaaae www.ee/oooromdd/1t
9
Liu, H. and Abbeel, P. Unsupervised active pre-training for reinforcement learning, 2021. URL 1, 2, 4, 5, 6, 7, 8, 21
https://openreviev
Mazoure, B., Combes, R. T. d., Doan, T., Bachman, P., and Hjelm, R. D. Deep reinforcement and infomax learning. In NeurIPS, 2020. 4
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. Veness, J., Bellemare, M. G., Graves, A., Ried- miller, M., M, Fidjeland, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. Nature, 518(7540). 2015. 5, 6, 8, 15, 18
Ng, A. Harada, D., and Russell, S. Policy invariance under reward transformations: Theory and application to reward shaping. In ICML, 1999. 18
Oord, A. v. d., Li, Y., and Vinyals, O. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807 1807.03748 2018. 4
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J.. Chanan, G., Killeen, T., Lin. Z., Gimelshein,
N., Antiga, L., et al. Pytorch: An imperative style, high- high-performance deep learning library. In NeuriPS, 2019. 15
Pathak, D., Agrawal, P., Efros, A., and Darrell, T. Curiosity-driven exploration by self-supervised prediction. In ICML. PMLR, 2017. 4
12

---

Perez, E., Strub, F., de Vries, H., Dumoulin, V., and Courville, A. C. Film: Visual reasoning with general conditioning layer. AAAI, 2018. 18
Sandler, M., Howard, A., Zhu, M., Zhmoginov, A., and Chen, L.-C. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 45 4510-4520, 2018. 18
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L.., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 2021.
Schwarzer, M., Anand, A., Goel, R., Hjelm, R. D., Courville, A., and Bachman, P. Data-efficient reinforcement learning with self-predictive representations. In ICLR, 2021. 2,3, 4, 5, 6, 7, 15, 18, 21, 22
Sekar, R., Rybkin, O., Daniilidis, K., Abbeel, P., Hafner, D., and Pathak, D. Planning to explore via self-supervised world models. In ICMI,, 2020.
Sharma, A., Gu, S., Levine, S., Kumar, V., and Hausman, K. Dynamics- unsupervised discovery of skills. In International Conference on Learning Representations, 2019.
Sohn, K., Berthelot, D., Li, C.-L., Zhang, Z., Carlini, N., Cubuk, E. D., Kurakin, A., Zhang, H., and Raffel, Fixmatch: Simplifying eemi-supereieed learning with consistency and confidence. arXiv preprint arXiv:2001.07683 2020. 4
Stooke, A. and Abbeel, P. rlpyt: A research code base for deep reinforcement learning in pytorch. arXiv preprint arXiv: 1909.01500, 2019. 15
Stooke, A., Lee, K., Abbeel, and Laskin, M. Decoupling representation learning from reinforce- ment learning, 2021. 1, 2, 4, 5, 6, 7, 8, 14,
Sutton, R. S. and Barto, A. G. Reinforcement Learning: An Introduction. The MIT Press, second edition, 2018. URL http: http://incomeeeeeeeeeeee deeplleeeeeeeeseee.e.....eeseeee wwt/book/ttaiaeeee teeels//heebaee book-2nn html. 15
Tan, M. and Le, Q. Efficientnet: Rethinking model scaling for convolutional neural networks. In International Conference on Machine Learning, pp. 6105-61 14, 2019. 18
Tarvainen, A. and Valpola, H. Mean teachers are better role models: Weight-averaged consistency targets improve semi cumi-supervised deep learning results. In NeurIPS, 2017. 4
Tsividis, P., Pouncy, T., Xu, J. L., Tenenbaum, J. B., and Gershman, S. J. Human learning in atari. In AAAI Spring Symposia, 2017.
Van Hasselt, H., Guez, A., and Silver, D. Deep reinforcement learning with double q-learning. In AAAI, 2016. 15
van Hasselt, H. P., Hessel, M., and Aslanides, J. When to use parametric models in reinforcement learning? In NeurlPS, 2019. 4,
Wang. Т. and Isola, P. Understanding contrastive representation learning through alignment and uniformity on the hypersphere, 2020.
Wang, Z., Schaul, T. Hessel, M., Hasselt, H., Lanctot, M., and Freitas, N. Dueling network architectures for deep reinforcement learning. In ICML, 2016. 15
Yang, M. and Nachum, O. Representation matters: Offline pretraining for sequential decision making. arXiv preprint ar 1rxiv:2102.05885 2021. 5
Yarats, D., Fergus, R., Lazaric, A., and Pinto, L. Reinforcement learning with prototypical represen-
tations. arXiv preprint arXiv:2102.11277 2021. 4
Zenke, F., Poole, B., and Ganguli, S. Continual learning through synaptic intelligence. In ICMI., 2017. 9
13

---

Zhan, A., Zhao, P., Pinto, L., Abbeel, P., and Laskin, M. A framework for efficient robotic manipula- tion, 2020.
Zhang, C., Vinyals, O., Munos, R., and Bengio, S. A study on overfitting deep reinforcement learning. arXiv preprint 1rXiv:1804.06899, 2018. 1
14


Reviews:
[{"review_id": "yBoGTzDQj_5", "reviewer": "Reviewer_KfDr", "summary": "This paper proposes a combination of three existing self-supervised objectives (self predictive representations, goal conditioned rl and inverse dynamics modeling) to improve the performance of RL agents, after pre-training on unsupervised data and then fine tuning on a small supervised (i.e. task specific) dataset. The authors analyzed the significance of each one of these objectives as well as the importance of the number of parameters and the impact of pre-training data quality.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "==== Writing\n\nThe paper is clearly written and is easy to read and understand.  The paper does not introduce any new model or objective but rather analysis a combination of already existing objectives for pre-training and the effect of each one on the performance of the agent on the downstream task. Due to the high number of possible combinations of self-supervised objectives, I really enjoyed the notations used by the paper to make it accessible and easy to read and understand. However, these notations and abbreviations can become overwhelming and I suggest the authors to include an additional summary table which describes them all in one glance.\n\n==== Methodology\n\nIn short, the paper is an empirical study of the importance of the pre-training objectives on the performance of the rl agent. And such studies are hard to do by nature because the number of possible combinations can grow exponentially, making experimentations harder. The authors, did a good job keeping the number of experiments at a (relatively) low number by fixing many of variables e.g. the rl algorithm itself which is fixed to Q-learning and the environment which is fixed at Atari. \n\nAt the same time, this is a major shortcoming of the paper. For example, one may argue that the gains in Table 2 vs SimPLe with no pre-training can be described by using a fundamentally more efficient rl method (Q-learning vs PPO) just because the paper does not include any study to demonstrate the same performance gains for PPO. Or ATC vs SGI comparison can be \"explained away\" by insignificance of augmentation for Atari. Or one may argue that SGI is observing better performances because learning a representation of Atari games is easy and SGI will perform the same when the environment is hard to compress (at least naively) such as the The Distracting Control Suite. To summarize, I believe the paper suffers from overfitting on its fixed assets e.g. environment (Atari) and rl method (Q), and in lack of any experiment to cover other cases, this is hard to support or oppose. Ofcourse, it is impossible to run all possible variations but a selective few can substantially support the main claims of the paper. At the same time, the claims can be rewritten in a less generalize form to limit it on the scope of the experiments.\n\nAnother major issue of the paper is how it represents the representation learning (no pun intended). The paper claims \"stronger\" representations however this claim is only support by demonstrating stronger performance on a single downstream task (this is another reason why Atari is not a good environment for this study vs multi-task environments). I encourage the authors to use other representation learning metrics such as Mutual I(z;x) or Linear Regression of reward from z or the number of latent active units.\n\nOverall, I find a disparity between the scope of the claims and the provided empirical evidence which can be improved by toning down the claims and adding more empirical evidence.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Minimal description of the limitation and and social impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "nuMJBJRXCTt", "reviewer": "Reviewer_Tq2R", "summary": "Focusing on the performance of an RL agent given a finite number of samples from ATARI in the regime where it is able to pretrain on a offline dataset, the paper shows that the final performance of the agent is improved by the addition of several additional auxiliary losses during pretraining, along with adjustments to the finetuning step.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Broadly speaking I believe this paper makes a strong contribution by showing that the proposed techniques improve the finetuned performance of the model, bringing it to substantially higher performance relative to other models in similar regimes. \n\nThe evaluation of the models is detailed, the ablations and variations that are considered are reasonable and executed well, and the presentation is good. I believe that work like this which clearly sets out the benefit (and weaknesses) of pretraining approaches is worthy of acceptance. For that reason I am giving a recommendation of 7.\n\nHowever I have one significant comment, and a set of more minor ones.\n\n1) A substantial part of the improvement in mean/median scores on Atari, as shown in Table 4, comes from the particular finetuning strategy employed (specifically, the reduced LR setting). Without that trick, the performance of the network drops substantially, roughly equivalent to removing all but the Inverse Modeling head from the finetuning tasks (comparing to Table 3), or reverting back to the behavioral cloning baseline (comparing to Table 2). This has substantial impact, and while it is discussed in Section 5.4, I believe it should be give a substantially more central role to the paper, including asking some further questions to probe exactly the cause of the effect (for instance, is performance recovered with a frozen encoder if you make the Q-Learning network deeper?).\n\nAdditionally, I have a set of more minor comments:\n\n* [Line 73] You're claiming these questions are policy-indepenent, but isn't the inverse modelling question (line 70) policy-dependent for non-trivial environments? Similarly for the question in line 71, it seems like that is policy-dependent as well.\n\n* [Line 111] It seems worthwhile to me to cite some earlier examples of inverse dynamics modeling here.\n\n* [Line 113] re: \"it is naturally not prone to representational collapse\". That seems like a broad claim which depends a lot on the situation, can you cite or explain?\n\n* [Figure 3] This plot strikes me as begging to be extended both further to the right and in model size. Especially because you bring up the comparison to [Kaplan, 2020], it would be interesting to see these points (especially if they were extended) plotted as log-logs to try to see whether there are visible patterns in the scaling properties.\n\n* [Line 389] You refer to an All Losses row in Table 5, but no such row exists.\n\n* [Line 391] nit: Can you make clear that Figure 5d is in the appendix?\n\n\nFinally, I would like to note that while I am familiar with the subject space, I am not fully up to date on recent developments in this area and am not a strong judgement of novelty past the references given in the paper (relative to which I think this submission is sufficiently novel). I will take a look at the other reviews to ensure that I didn't miss anything in my survey.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the societal impacts relevant to the paper are tangential and I believe the authors sufficiently describe them.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "POjamxjGkR", "reviewer": "Reviewer_5yeQ", "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.\n\nSummary of experimental results:\n\n(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).\n\n(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.\n\n(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.\n\n(4) Pretraining vs. size of the network (Section 5.2): Model size and amount of pretraining data should increase in parallel.\n\n(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.\n\n(6) How to finetune pre-trained representations properly: Reducing the finetuning LRs improves performance.\n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality & Significance: The largest contribution of the paper are the extensive experiments and empirical analyses (Section 4), which I believe are valuable for understanding how pre-trained representations can help RL with learning efficiency. (See “Summary” above for a list of empirical findings reported in the paper.)\n\nThe proposed method itself (SGI) is less novel -- it is a combination of three unsupervised representation learning objectives that operate in latent space provided by a shared encoder. However, the ablation analysis (Section 5.3) is helpful in understanding how the combination of the three objectives contributes to SGI’s performance.\n\nQuality: The experiments are thoughtfully designed and well-explained (evaluation metrics, datasets, empirical analyses and discussion of results).\n\nClarity: The paper is very well-written; the problem motivation, experimental design, related works, evaluation metrics, etc. are all clearly explained.\n\nMinor nit suggestions that did not affect my score:\n- L71: Missing question mark “?”\n- L293: “settingss” -> “settings”\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "LMjbLXXQkiw", "reviewer": "Reviewer_2aRe", "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. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper introduces a rather simple yet effective approach for learning good representations in an unsupervised way which can be helpful for downstream RL tasks. This problem is very important as it can lead to more sample efficient algorithms and thus enable more real-world applications of RL where interactions with an environment are expensive. The paper is clear overall, the experiments are well designed, and the approach is thoroughly evaluated and compared with many relevant baselines and ablations. The results also look quite promising and the method clearly improves performance and sample efficiency relative to prior art. \n\nI particularly liked the analysis of the algorithm's dependence on the quality of the dataset and found the insights very valuable even if not necessarily surprising. The paper also contains many ablations thus resulting in a better understanding of how each component contributes to the final results. \n\nHowever, I think there are a few additional baseline variants / ablations which could further improve the paper and our understanding of how the various algorithms compare. Would it be possible to add experiments with BC, VISR, APT, CPT, ATC, SGI using the all of the four different datasets considered, namely R, E, W, and M? Some of these are already included in Table 2 but completing this space would be useful to gain a full picture. For Table 5, could you add results with G only and I only?\n\nFrom Table 4, it looks like Frozen is comparable with BC-M so it seems that reducing the learning rate during finetuning is quite important to maximize performance. Hence, it would be interesting to see the effect of reducing the learning rate on the other baselines which perform finetuning on downstream tasks. Is the right conclusion that reducing the learning rate during finetuning plays a more important role than the objective function used for pretraining the encoder? I think further analyzing this question would be of interest for the readers. \n\nI think it would also be valuable to add experiments in a different domain in order to better understand whether the conclusions drawn from the Atari experiments generalize to other domains as well. A few suggestions are evaluating the methods on continuous control tasks from MuJoCo or DMC. \n\nGiven the above reasons, I tend to recommend acceptance of this paper, but encourage the authors to take into account the feedback to further improve this paper. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "None", "review_text": "", "extra_scores": {"ethical_concerns": "None", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Taking as evaluation setting Atari 100k, the paper demonstrates the effectiveness of pretraining an agent torso/encoder on a collection of auxiliary self supervised losses, greatly improving data efficiency and final performance of the agent when training on a small of amount of experience on the downstream tasks (in conjunction with careful tuning of fine-tuning learning step).

Whilst the proposed unsupervised losses are not novel (a combination of latent dynamics modelling, goal-conditioned RL, and inverse dynamics modeling), the most impactful contribution of the submission is the thorough evaluation of the method, its extensive, carefully designed experiments, ablations and empirical analyses, as well as the solid choice of baselines. Although limited to Atari 100k, these experiments are well described and discussed, forming a very solid foundation for eventual follow-up work investigating if these results can generalize to other domains or training regimes,

I believe the paper will be of great interest to the broader community, as it sheds light on viable training regimes for agents in real-world application, where access to real environment experience is scarce.
Decision: Accept (Poster)