Data-Efficient Reinforcement Learning with Self-Predictive Representations
Authors: Max Schwarzer
Ankesh Anand
Rishab Goel
R Devon Hjelm
Aaron Courville
Philip Bachman
Venue: ICLR 2021 Spotlight (2021)
Keywords: Reinforcement Learning
Self-Supervised Learning
Representation Learning
Sample Efficiency

Abstract: While deep reinforcement learning excels at solving tasks where large amounts of data can be collected through virtually unlimited interaction with the environment, learning from limited interaction remains a key challenge. We posit that an agent can learn more efficiently if we augment reward maximization with self-supervised objectives based on structure in its visual input and sequential interaction with the environment.  Our method, Self-Predictive Representations (SPR), trains an agent to predict its own latent state representations multiple steps into the future. We compute target representations for future states using an encoder which is an exponential moving average of the agent’s parameters and we make predictions using a learned transition model.  On its own,  this future prediction objective outperforms prior methods for sample-efficient deep RL from pixels. We further improve performance by adding data augmentation to the future prediction loss, which forces the agent’s representations to be consistent across multiple views of an observation.  Our full self-supervised objective, which combines future prediction and data augmentation, achieves a median human-normalized score of 0.415 on Atari in a setting limited to 100k steps of environment interaction, which represents a 55% relative improvement over the previous state-of-the-art. Notably, even in this limited data regime, SPR exceeds expert human scores on 7 out of 26 games. We’ve made the code associated with this work available at https://github.com/mila-iqia/spr.

Published as a conference paper at ICLR 2021
DATA- REINFORCEMENT LEARNING WITH SELF-PREDICTIVE REPRESENTATIONS
Max Schwarzer"
Ankesh Anand"
Rishab Goel Mila
Mila, Université de Montréal
Mila, Université de Montréal
Microsoft Research
Devon Hjelm
Aaron Courville
Philip Bachman Microsoft Research
Microsoft Research
Mila, Université de Montréal
Mila, Université de Montréal
CIFAR Fellow
ABSTRACT
While deep reinforcement learning excels solving tasks where large amounts of data can be collected through virtually unlimited interaction with the environment, learning from limited interaction remains a key challenge. We posit that an agent can learn more efficiently if we augment reward maximization with self-supervised objectives based on structure in visual input and sequential interaction with the environment. Our method, Self-Predictio Predictide Representations (SPR), trains an agent to predict its own latent state representations multiple steps into the future. We compute targel representations for future states using encoder which is an exponential moving average of the agent's parameters and we make predictions using a learned transition model. On its own, this future prediction objective outperforms prior methods for sample-efficient deep RI. from pixels. We further improve performance by adding data augmentation the future prediction loss, which forces the agent's representations to be consistent across multiple views of observation. Our full self-supervised objective, which combines future prediction and data augmentation, achieves median human-normalized score of 0.415 on Atari which represents a 55% relative improvement over the previous state-of-the--e Notably, even in this limited data regime, SPR exceeds expert human scores on 7 out of 26 games. We've 've made the code associated with this work available at https: /////ttttuaaaa com mi la iiii/ sspr.
setting limited 100k steps of environment interaction,
INTRODUCTION
Deep Reinforcement Learning (deep RL, François-Lavet et al., 2018) has proven to be an indispens- able tool for training successful agents on difficult sequential pocissonn-mkinn problems (Bellemare et al., 2013; Tassa et al.. 2018). The success of deep RL is particularly noteworthy in highly complex, strategic games such StarCraft (Vinyals al., 2019) and DoTA2 (OpenAI al., 2019), where deep RL agents now surpass expert human performance in some scenarios.
Deep RL involves training agents based on large neural networks using large amounts of data (Sutton, 2019), trend evident across both model- based (Schrittwieser et al., 2020) and model-free (Badia et al., 2020) learning. The sample complexity of such state-of-the-are agents is often incredibly high: MuZero (Schrittwieser et al., 2020) and Agent-57 (Badia et al., 2020) use 10-50 years of experience per Atari game, and OpenAl Five (OpenAl et al., 2019) uses 45,000 years of experience to accomplish its remarkable performance. This is clearly impractical: unlike casily-simulated environments such as video games, collecting interaction data for many real-world tasks is costly, making improved data efficiency a prerequisite for successful use of deep RL in these settings (Dulac-Arnold et al., 2019).
Equul contribution; the order of first authors was determined by coin flip. ankesh anand) @mi 1a quebec
sschwarmm

---

Published as a conference paper at ICLR 2021
SPR
0.425
SPR
0./04
)ug.)
sug.)
0.268
0.357
CURL
0.175
CURL
0
DER
0.161
DER
0.443
0.00 0.05
0.10
0.15
0.20
0.25
0.30
0.35
0.40
0.0
0.1
0.2
0.3
0.4
0.5
0.6
0.7
Median Human Normalized Score
Mean Saaaaa--oooaaaiiaae Score
Figure 1: Median and Mean Human- -ormalizzd scores of different methods across 26 games in the Atari 100k benchmark (Kaiser et al., 2019), averaged over 10 random seeds for SPR, and 5 seeds for most other methods except CURL,, which uses 20. Each method is allowed access to only 100k environment steps or 400k frames per game. (*) indicates that the method uses data augmentation. SPR achieves state-of-art results on both mean and median human-normalized scores. Note that, even without data augmentation, SPR still outperforms all prior methods on both metrics.
Meanwhile, new self-supervised representation learning methods have significantly improved data efficiency when learning new vision and language tasks, particularly low data regimes or semi- supervised learning (Xie et al., 2019; Hénaff et al., 2019; Chen et al., 2020b). Self- Self-supervised methods improve data efficiency by leveraging a nearly limitless supply of training signal from tasks generated on-the-fly, based "views" drawn from the natural structure of the data (e.g., image patches, data augmentation or temporal proximity, see Doersch et al., 2015; Oord et al., 2018; Hjelm et al., 2019; Tian et al., 2019; Bachman et al., 2019; He al., 2020; Chen al., 2020a).
Motivated by successes in semi semi-supervised and self-supervised learning (Tarvainen & Valpola, 2017; Xie et al., 2019; Grill al., 2020), we train better state representations for RL by forcing representa- tions to be temporally predictive and consistent when subject to data augmentation. Specifically, we extend strong model-free agent by adding dynamics model which predicts future latent represen- tations provided by parameter- wise exponential moving average of the agent itself. We also add data augmentation to the future prediction task, which enforces consistency across different views of each observation. Contrary to some methods (Kaiser et al., 2019; Hafner et al., 2019), our dynamics model operates entirely in the latent space and does not rely on reconstructing raw states.
We evaluate method, which call Self-Predictive Representations (SPR), the 26 games in the Atari 100k benchmark (Kaiser et al., 2019), where agents are allowed only 100k steps environment interaction (producing 400k frames of input) per game, which roughly corresponds to two hours of real-time experience. Notably, the human experts in Mnih et al. (2015) and Van Hasselt et al. (2016) were given the same amount of time learn these games, budget of 100k steps permits reasonable comparison in terms data efficiency.
In experiments, augment modified version of Data-Efficient Rainbow (DER) (van Hasselt et al., 2019) with the SPR loss, and evaluate versions of SPR with and without data augmentation. We find that each version superior controlled baselines. When coupled with data augmentation, SPR achieves a median score of 0.415, which a state-of-the-art result on this benchmark, outperforming prior methods by a sigmificant margin. Notably, SPR also outperforms human expert scores on out of 26 games while using roughly the same amount of in in-game experience.
METHOD
We consider reinforcement learning (RL) in the standard Markov Decision Process (MDP) setting where an agent interacts with its environment in episodes, each consisting of sequences of observa- tions, actions and rewards. We use $t, a and r; to denote the state, the action taken by the agent and the reward received at timestep t. We seek to train an agent whose expected cumulative reward in each episode is maximized. To do this, combine strong model-free RI. algorithm, Rainbow (Hessel
2

---

Published as a conference paper at ICLR 2021
fo online encoder
learning head
Q G-Learning Loss
21
0
de
8i aug.
conv. transifion model
exponential moving avg.
9o online projection
2
prediction
+ gt
Cosine Similarity Loss
m
target encoder
Zi+b
target projection
WL+k
fm
9m
S14-1: + oug.
Figure 2: An illustration of the full SPR method. Representations from the online encoder are used in the reinforcement learning task and for prediction of future representations from the target encoder via the transition model. The target encoder and projection head are defined as an exponential moving average of their online counterparts and only the Athh step of future prediction, but in practice we compute the loss over all steps from 1 to K. Note: our implementation for this paper includes go in the Q-learning head.
updated via gradient descent. For brevity, illustrate
et al., 2018), with Self-Predictive Representations as auxiliary loss to improve sample efficiency, We now describe our overall approach detail.
2.1 DEEP Q-LEARNING
We focus the Atari Learning Environment (Bellemare al., 2013), challenging setting where the agent takes discrete actions while receiving purely visual, pixel-based observations. A prominent method for solving Atari, Deep Q Networks (Mnih al., 2015), trains a neural network Qo to approximate the agent's current Q-function (policy evaluation) while updating the agent's policy greedily with respect to this Q-function (policy improvement). This involves minimizing the error between predictions from Qo and target value estimated by Qe. earlier version of the network:
cDON = =(Qo(st)) (r. + max Q5 Qc(s4+1,a)))
(1)
Various improvements have been made over the original DQN: Distributional RL (Bellemare et al., 2017) models the full distribution of future reward rather than just the mean. Dueling DQN (Wang et al.. 2016) decouples the value of state from the advantage of taking a given action in that state, Double DQN (Van Hasselt al., 2016) modifies the Q-learning update to avoid overestimation due to the max operation, among many others. Rainbow (Hessel et al., 201 8) consolidates these improvements into a single combined algorithm and has been adapted to work well data-limited regimes (van Hasselt et al., 2019).
2.2 SELF-PREDICTIV REPRESENTATION
For our auxiliary loss, we start with the intuition that encouraging state representations to be predictive of future states given future actions should improve the data efficiency of RL algorithms. Let (SEt*K, Oc:++) k) denote a sequence of K + previously experienced states and actions sampled from a replay buffer, where K is the maximum number of steps into the future which we want to predict. Our method has four main components which we describe below:
Online and Target networks: We use an online encoder fo to transform observed states s, into representations 2, & (st). We use these representations in an objective that encourages them be predictive of future observations up some fixed temporal offset K, given sequence of K actions
3

---

Published as a conference paper at ICLR 2021
to perform. We augment each observation 8; independently when using data augmentation. Rather than predicting representations produced by the online encoder, we follow prior work (Tarvainen Valpola, 2017; Grill al., 2020) by computing target representations for future states using target encoder fm. whose parameters are an exponential moving average (EMA) of the online encoder parameters, Denoting the parameters of fo as for those of fm as 0m, and the EMA coefficient as E [0, 1), the update rule for 0m is:
On + r0m + +1--)00.
(2)
The target encoder is not updated via gradient descent. The special case = 0, fn = 8, is noteworthy, as it performs well when regularization is already provided by data augmentation.
Transition Model: For the prediction objective, we generate sequence of K predictions ê 1:411:44-K of future state representations &++1**K using an action-conditionel Iransition model h. We compute 2++1:+*K iteratively: 4 h(&c+R, (14x), starting from & _ 2, _ fo(s). We compute 2141:t+K K by applying the target encoder fm to the observed future states st-1:t+k Zl+a & fm (S++2). The transition model and prediction loss operate in the latent space, thus avoiding pixel-based reconstruction objectives. We describe the architecture of h in section 2.3.
We use online and target projection heads Jo and 9m (Chen et al., 2020a) to project online and target representations to smaller latent space, and apply an additional prediction head (Grill et al., 2020) to the online projections predict the target projections:
Projection Heads:
det - (()0(2)+))), V2s+a E &441+KK DE+k gm V2z+k 2441+
(3)
The target projection head parameters are given by an EMA of the online projection head parame- ters, using the same update as the online and target encoders.
Prediction Loss: We compute the future prediction loss for SPR by summing over cosine similari- ties! between the predicted and observed representations at timesteps t + for S S K:
CSPR (Sut+K Oc:+k)
Deth
(
(1
(4)
110e+x 110c+1l122
where Út-k and D are computed from (Se:t+ K; 0.:++k) k>) as we just described.
We call our method Self-Predictive Representations (SPR), following the predictive nature of the objective and the of exponential moving average target network similar (Tarvainen & Valpola, 2017: He et al., 2020). During training, we combine the SPR loss with the Q-learning loss for Rainbow. The SPR loss affects for Por 9 and h. The Q-learning loss affects fo and the Q-learning head, which contains additional layers specific to Rainbow. Denoting the Q-learning loss from Rainbow as CRL, our full optimization objective is: L Lioual CRL + XLSPR
Unlike some other proposed methods for representation learning in reinforcement learning (Srinivas al.. 2020), SPR can be used with or without data augmentation, including in contexts where data augmentation is unavailable or counterproductive. Morcover, compared to related work on contrastive representation learning, SPR does not use negative samples, which may require careful design of contrastive tasks, large batch sizes (Chen et al., 2020a), or the use of a buffer to emulate large batch sizes (He et al., 2020)
2.3 TRANSITION MODEL ARCHITECTURE
For the transition model h, we apply a convolutional network directly to the 64 x x spatial output of the convolutional encoder fo. The network comprises two 64-channel convolutional layers with 3 x 3 filters, with batch normalization (loffe & Szegedy, 2015) after the first convolution and ReLU nonlinearities after each convolution. We append a one-hot vector representing the action taken to each location in the input to the first convolutional layer, similar to Schrittwieser et al. (2020). We maximum prediction depth of K 5, and truncate calculation of the SPR loss episode boundaries to avoid encoding environment reset dynamics into the model.
Cosine similarity is linearly related the "normalized 1.2" loss used in BYOL (Grill al., 2020).
4

---

Published as a conference paper at ICLR 2021
Algorithm 1: Self-Predictive Representations
Denote parameters of online encoder fo and projection 9o as 0, Denote parameters of target encoder fm and projection Iin 6m
Denote parameters of transition model h, predictor and Q-learning head as Φ Denote the maximum prediction depth as K, batch size as N initialize replay buffer B while Training do collect experience as.a, s') with (0., Φ) and add to buffer B sample a minibatch of sequences of (8, a, r, r') for i in r0nge(0, N) do
augmentation then + augment(s*) s)); augment(s")
end zú + (s)) p+0
1/ online representations
for
k in (I, .... K) do 2 f- h(2k-11 0
// latent states via transition model
b + fm (ske
// target representations // projections
ii - (2h0) Th + Im (#f)
sit + (
(
[IBE
// SPR loss at step
end
1 +- XP + RL loss(s' a a (rrtrrts" ;0)) 11 Add RL loss for batch with 0.
end
1 N カ== -i= 12 O,r + optimize 1),1) Om f 70. + (1 -
average loss over minibatch // update online parameters
(m 7m
// update target parameters
end
2.4 DATA AUGMENTATION
When using augmentation, (2021), consisting of small random shifts and color jitter. We normalize activations to lie in [0, 1] at the output of the convolutional encoder and transition model, in Schritwwieser al. (2020). We use Kornia (Riba et al., 2020) for efficient GPU-based data augmentations.
the same set of image augmentations in DrQ from Yarats et al.
When not using augmentation, find that SPR performs better when dropout (Srivastava al., 2014) with probability 0.5 is applied at each layer in the online and target encoders. This is consistent with Laine & Aila (2017); Tarvainen & Valpola (2017), who find that adding noise inside the network important when not using image-specifie augmentation, as proposed by Bachman et al. (2014).
2.5 IMPLEMENTATION DETAILS
For our Atari experiments, we largely follow van Hasselt et al. (2019) for DQN hyperparameters, with four exceptions. We follow DrQ (Yarats et al., 2021) by: using the 3-layer convolutional encoder from Mnih et al. (2015), using 10-step returns instead of 20-step returns for Q-learning, and not using separate DQN target network when using augmentation. We also perform two gradient steps per environment step instead of one. We show results for this configuration with and without augmentation Table and confirm that these changes are not themselves responsible for our performance. We reuse the first layer of the DQN MLP head the SPR projection head go. When using dueling DON (Wang et al., 2016), Io concatenates the outputs the first layers of the value and advantage heads. When these layers noisy (Fortunato et al., 2018), 111,, 9o does not use the noisy parameters. Finally, we parameterize the predictor q as a linear layer. We use 7 = 0.99 when augmentation is disabled and when enabled. For Lioual CR ALSPR
use À 2.
5

---

Published as a conference paper at ICLR 2021
Hyperparameters were tuned over subset of games (following Mnih al., 2015; Machado et al., 2018). We list the complete hyperparameters in Table
Our implementation uses rlpyt (Stooke & Abbeel, 2019) and PyToreh (Paszke et al., 2019). We find that SPR modestly increases the time required for training, which we discuss in more detail in Appendix D.
3
RELATED WORK
3.1 DATA-EFFICIENS RL:
number of works have sought to improve sample efficiency in deep RL. SiMPLe (Kaiser et al., 2019) learns pixel-level transition model for Atari to generate simulated training data, achieving strong results on several games in the 100k frame setting, at the cost of requiring several weeks for training. However, van Hasselt al. (2019) and Kielak (2020) introduce variants of Rainbow (Hessel et al., 2018) tuned for sample efficiency, Data-Efficient Rainbow (DER) and OTRainbow, which achieve comparable or superior performance with far less computation.
In the context of continuous control, several works propose to leverage a latent-space model trained on reconstruction loss improve sample efficiency (Hafner al., 2019; Lee et al., 2019; Hafner et al., 2020). Most recently. DrQ (Yarats et al., 2021) and RAD (Laskin et al., 2020) have found that applying modest image augmentation can substantially improve sample efficiency in reinforcement learning, yielding better results than prior model- found to improve generalization of reinforcement learning methods (Combes al., 2018; Laskin et al., 2020) in multi-task and transfer settings. We show that data augmentation can be more effectively leveraged in reinforcement learning by forcing representations to be consistent between different augmented views of an observation while also predicting future latent states.
methods. Data augmentation has also been
3.2 REPRESENTATION LEARNING IN RL:
Representation learning has a long history of use in RL - see Lesort et al. (2018). For example, CURL (Srinivas et al., 2020) proposed combination of image augmentation and contrastive loss to perform representation learning for RL. However, follow-up results from RAD (Laskin et al., 2020) suggest that most of the benefits of CURL come from image augmentation, not its contrastive loss.
CPC (Oord et al., 2018), CPCIAction (Guo et al., 2018), ST-DIM (Anand et al., 2019) and DRIML (Mazoure et al., 2020) propose to optimize various temporal contrastive losses in rein- forcement learning environments. We perform an ablation comparing such temporal contrastive losses to our method in Section 5. Kipf er al. (2019) propose to learn object-oriented contrastive representations by training structured transition model based on graph neural network.
SPR bears some resemblance to DeepMDP (Gelada al., 2019), which trains a transition model with an unnormalized L2 loss to predict representations of future states, along with reward prediction objective. However, DeepMDP uses its online encoder to generate prediction targets rather than employing a target encoder, and is thus prone representational collapse (sec. C.5 in Gelada et al. (2019)). To mitigate this issue, DeepMDP relies on an additional observation reconstruction objective. In contrast, our model is self-supervised, trained entirely in the latent space, and uses normalized loss. Our ablations (sec. 5) demonstrate that using a target encoder has a large impact on our method.
SPR is also similar PBL (Guo et al., 2020), which directly predicts representations of future states. However, PBL uses two separate target networks trained via gradient descent, whereas SPR uses single target encoder, updated without backpropagation. Moreover, PBL studies multi- task generalization in the asymptotic limits of data, whereas SPR is concerned with single-task performance in low data regimes, using 0.01% as much data as PBL. Unlike PBL, SPR additionally enforces consistency across augmentations, which empirically provides large boost in performance.
4 RESULTS
We test SPR on the sample-efficient Atari setting introduced by Kaiser et al. (2019) and van Hasselt al. (2019). In this task, only 100,000 environment steps of training data are available equivalent to

---

Published as a conference paper at ICLR 2021
400,000 frames, just under two hours compared to the typical standard of 50,000,000 environment steps, or roughly 39 days of experience. When used without data augmentation, SPR demonstrates scores comparable to the previous best result from Yarats al. (2021). When combined with data augmentation, SPR achieves a median human-normalized score of 0.415, which a new state-of-the- art result this task. SPR achieves super-human performance setting: Boxing, Krull, Kangaroo, Road Runner, James Bond and Crazy Climber, compared to maximum of two for any previous methods, and achieves scores higher than DrQ (the previous state-of-the-an method) on 23 out of 26 games. See Table 1 for aggregate metrics and Figure 3 for visualization of results. full list of scores is presented in Table 4 in the appendix. For consistency with previous works, we report human and random scores from Wang et al. (2016).
seven games in this data-limited
Table Performance of different methods on the 26 Atari games considered by Kaiser et (2019) after 100k environment steps. Results are recorded at the end of training and averaged over 10
random seeds for SPR, 20 for CURL, and for other methods. SPR outperforms prior methods on all aggregate metrics, and exceeds expert human performance on 7 out of 26 games.
Metric
Random
Human
SumPLe
DER
OTRainbow
CURL
DrQ
SPR (no Aug)
SPR
Mcan Human-Norm's
0.000
1.000
0.443
0,285
0.264
0.381
0.357
0.463
0.704
Median Human Norm'd d'
0.000
0.161
0.204
0.175
0.268
0.307
Mean DQN@ 50M Norm'd
0.000
23.382
0,232
0,239
0.197
0.325
0.171
0,336
0.510
Median DQN@ 50M-Norm'd
0.000
0.994
0.142
0.103
0.142
0.131
0.225
Games Superhuman
N/A
2
2
3.0
Numan-Normalizz Score
2.0
1.5
0.5
0.0
SimPLe
OTRainbow
CuRL
SPR (no aug) SPR
Algorithm
Figure 3: boxplot of the distribution of human-normalized scores across the 26 Atari games under consideration, after 100k environment steps. The whiskers represent the interquartile range of norman-normalized scores over the 26 games. Scores for each game are recorded at the end of training and averaged over 10 random seeds for SPR, 20 for CURL, and 5 for other methods.
4.1 EVALUATION
We evaluate the performance different methods by computing the average episodic return the end of training. We normalize scores with respect to expert human scores account for different scales of scores in each game, as done in previous works. The human-normalized score of an agent on a game is calculated as
agent score random score and aggregated across the 26 games by mean or median.
human score random score
We find that human scores on some games are so high that differences between methods are washed out by normalization, making it hard for these games to influence aggregate metrics. Moreover, we find that the median score is typically only influenced by a handful of games. Both these factors compound together to make the median overall performance. To address this, we also report NON-normalized scores, defined analogously to scores and calculated using scores from DQN agents (Mnih al., 2015) trained over million steps, and report both mean and median those metrics all results and ablations, and plot the distribution of scores over all the games in Figure 3.
score an unreliable metric for judging
7

---

Published as a conference paper at ICLR 2021
Table Scores on the 26 Atari games under consideration for ablated variants of SPR. All variants listed here use data augmentation.
Variant
Score median
SQN@55M-NNrmalize Score
mean
mean
median
SPR
0.704
0,415
0.510
0.361
1 1-sepp SPR
0.570
0.301
0.337
0.346
Non-temporal SPR
0.507
0.271
0.326
0.295
Quadratic SPR
0.047
0.040
-0.016
0.031
SPR without projections
0.437
0.171
0.247
0.174
Rainbow (controlled, w/ aug.)
0.480
0.346
0.284
0.278
Additionally, we note that the standard evaluation protocol of evaluating over only 500,000 frames per game is problematic, as the quantity we are trying to measure is expected return over episodes. Because episodes may last up to up to 108,000 frames, this method may collect as few four complete episodes. As variance of results is already a concern in deep RL (see Henderson et al., 2018), we recommend evaluating over 100 episodes irrespective of their length. Moreover, to address findings from Henderson et al. (2018) 8) that comparisons based on small numbers of random seeds are unreliable, we average our results over ten random seeds, twice as many as most previous works.
ANALYSIS
The target encoder We find that using a separate target encoder is vital in all cases. A variant of SPR in which target representations are generated by the online encoder without a stopgradient (as done by e.g., Gelada et al., 2019) exhibits catastrophically reduced performance, with median human--ormalizee score of 0.278 with augmentation versus 0.415 for SPR. However, there is more flexibility in the EMA constant used for the target encoder. When using augmentation, value of = 0 performs best, while without augmentation we use 7 = 0.99. The success of interesting, since the related method BYOL reports very poor representation learning performance in this case. We hypothesize that optimizing reinforcement learning objective in parallel with the SPR loss explains this difference, as provides an additional gradient which discourages representational collapse. Full results for these experiments are presented in Appendix C.
1.2
Mean
Dynamics modeling is key A key dis- tinction between SPR other recent ap- proaches leveraging representation learning for reinforcement lcarning, such as CURL (Srinivas et al., 2020) and DRIML (Mazoure et al., 2020), is our use of an explicit multi- step dynamics model. To illustrate the im- pact of dynamics modeling, we test SPR with variety of prediction depths K. Two of these ablations, one with no dynamics modeling and one that models only single step of dynamics, are presented in Table (as Non-temporal SPR and /-step SPR), and all are visualized in Figure 4. We find that extended dynamics modeling consistently improves performance up to roughly K 5. Moving beyond this continues to improve performance
SPR FRr-Normllzzed Score
Median
1.1
1.0
0.9
0.8
0
1
2
5
6
7
8
K
Figure 4: Performance of SPR with various prediction depths. Results averaged across ten seeds per game, for all 26 games. To equalize the importance of games, calculate an SPR-normalized score analogously to human-normalized scores, and show its mean and median across all 26 games. All other hyperparameters are identical to those used for SPR with augmentation.
subset of games, at the cost increased computation. Note that the non-temporal ablation test is similar to
8

---

Published as a conference paper at ICLR 2021
using BYOL (Grill et al., 2020) an auxiliary task, with particular architecture choices made for the projection layer and predictor.
Comparison with contrastive losses Though many recent works in representation learning employ contrastive learning, we find that SPR consistently outperforms both temporal and non-temporal variants of contrastive losses (see Table 6, appendix), including CURL (Srinivas et al., 2020).
Using a quadratic loss causes collapse SPR's of cosine similarity objective (or normalized 1.2 loss) sets in contrast to some previous works, such as DeepMDP (Gelada et al., 2019), which have learned latent dynamics models by minimizing an un-normalized L2 loss over predictions of future latents. To examine the importance of this objective, we test a variant of SPR that minimizes un-normalized L.2 loss (Quadratic SPR in Table 2), and find that it performs only slightly better than random. This is consistent with results from Gelada et al. (2019), who find that Deep.MDP's representations are prone to collapse, and use an auxiliary reconstruction objective to prevent this.
Projections are critical Another distinguishing feature of SPR is its use of projection and predic- tion networks. We test variant of SPR that uses neither, instead computing the SPR loss directly over the 64 x 7 x 7 convolutional feature map used by the transition model (SPR withoul projections in Table 2). We find that this variant has inferior performance, and suggest two possible explanations. First, the convolutional network represents only small fraction of the capacity of SPR's 's's network, containing only some 80,000 parameters out of a total of three to four million. Employing the first layer of the DQN head as projection thus allows the SPR objective to affect far more of the network, while in this variant its impact is limited. Second, the effects of SPR in forcing invariance to augmentation may be undesirable at this level; as the convolutional feature map is the product of only three layers, may be challenging to learn features that are simultaneously rich and invariant.
FUTURE WORK
Recent work in both visual (Chen et al., 2020b) and language representation learning (Brown et al., 2020) has suggested that self-supervised models trained on large datasets perform exceedingly well on downstream problems with limited data, often outperforming methods trained using only tpec-specific data. Future works could similarly exploit large corpora of unlabelled data, perhaps from multiple MDPs or raw videos, to further improve the performance of RL methods in low-data regimes. As the SPR objective is unsupervised, it could be directly applied in such settings.
Another interesting direction is to use the transition model learned by SPR for planning. MuZero (Schrittwieser al., 2020) has demonstrated that planning with model supervised via reward and value prediction can work extremely well given sufficient (massive) amounts of data. It remains unclear whether such models can work well low-data regimes, and whether augmenting such models with self-supervised objectives such as SPR can improve their data efficiency.
would also be interesting examine whether self-supervised methods like SPR can improve generalization to unseen tasks or changes in environment, similar to how unsupervised pretraining on ImageNet can generalize to other datasets (He et al., 2020; Grill et al., 2020).
7 CONCLUSION
In this paper we introduced Self-Predictive Representations (SPR), a self-supervised representation learning algorithm designed to improve the data efficiency of deep reinforcement learning agents. SPR learns representations that are both temporally predictive and consistent across different views of environment observations, by directly predicting representations of future states produced by a target encoder. SPR achieves state-of-the-art performance on the 100k steps Atari benchmark, demonstrating significant improvements over prior work. Our experiments show that SPR is highly robust, and is able to outperform the previous state of the art when either data augmentation or temporal prediction is disabled. We identify important directions for future work, and hope continued research at the intersection of sule--sue learning and reinforcement learning leads to algorithms which rival the efficiency and robustness of humans.

---

Published as a conference paper at ICLR 2021
ACKNOWLEDGEMENTS
We are grateful for the collaborative research environment provided by Mila and Microsoft Research. We would also like to acknowledge Hitachi for providing funding support for this project. We thank Nitarshan Rajkumar and Evan Racah for providing feeback on an earlier draft; Denis Yarats and Aravind Srinivas for answering questions about DrQ and CURL; Michal Valko, Sherjil Ozair and the B team for long discussions about BYOL, and Phong Nguyen for helpful discussions. Finally, we thank Compute Canada and Microsoft Research for providing computational resources used in this project.
10

---

Published as a conference paper at ICLR 2021
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning atari. In NeurIPS, 2019. Cited on page 6.
Philip Bachman, Ouais al Sharif, and Doina Precup. Learning with pseudo-ensembles. In NeurIPS, 2014. Cited on page 5.
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In NeurIPS, 2019, Cited on page 2.
Adrià Puigdomenech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In ICML, 2020. Cited on page 1.
Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47, 2013. Cited on pages 1 and 3.
Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforcement learning. In ICML, 2017. Cited on page 3.
Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et Language models are few-shot learners. In NeurlPS, 2020. Cited on page 9.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. ICML, 2020a. Cited on pages 2, 4, and 16.
Ting Chen, Simon Kornblith. Kevin Swersky, Mohammad Norouzi, and Geoffrey Hinton, Big self-supervised models are strong semi cemi-supervised learners. In NeurlPS, 2020b. Cited on pages 2 and 9.
Remi Tachet des Combes, Philip Bachman, and Harm van Seijen. Learning invariances for policy generalization. arXiv preprint urXiv: htti::v800002500 2018. Cited on page 6.
Carl Doersch, Abhinav Gupta, and Alexei A Efros. Unsupervised visual representation learning by context prediction. In ICCV, 2015. Cited on page
Gabriel Aanalc Daniel Mankowitz, and Todd Hester. Challenges of real-world reinforcement learning. In /CML., 2019. Cited on page 1.
Meire Fortunato, Mohammad Gheshlaghi Azar, Bilal Piot, Jacob Menick, Matteo Hessel, Ian Osband, Alex Graves, Volodymyr Mnih, Remi Munos, Demis Hassabis, Olivier Pietquin, Charles Blundell, and Shane Legg. Noisy networks for exploration. In ICLR, 2018. Cited on page 5.
Vincent François-Lavet, Peter Henderson, Riashat Islam, Mare G Bellemare, and Joelle Pineau. An introduction to deep reinforcement learning. arXiv preprint arXiv:1811.12560 2018. Cited on page 1.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. In ICML, 2019. Cited on pages 6, 8, and 9.
Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre H Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, al. Bootstrap your own latent: A new approach to self- suff-vupe learning. In NeurIPS, 2020. Cited on pages 2, 4, 9, and 17.
Daniel Guo, Bernardo Avila Pires, Bilal Piot, Jean-bastien bastien Grill, Florent Altché, Rémi Munos, and Mo- hammad Gheshlaghi Azar. Bootstrap latent-predictive representations for multitask reinforcement learning. ICML, 2020. Cited on page 6.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A Pires, and Rémi Munos. Neural predictive belief representations. ICMI., 2018. Cited pages 6 and 16.
11

---

Published as a conference paper at ICLR 2021
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In JCML, 2019. Cited on pages and 6.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. ICLR, 2020. Cited on page
Kaiming Hc, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020. Cited pages 2, 4, and 9.
Olivier J Hénaff, Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, SM Eslami, and Aaron van den Oord. Data-efficient image recognition with contrastive predictive coding, arXiv preprint arXiv:
2019. Cited on page 2.
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In AAAI, 2018. Cited on page 8.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Gheshlaghi Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In AAAI, 2018. Cited on pages 2, 3, and 6.
R Devon Hjelm, Alex Fedorov, Samuel Aavoie-Marchildon. Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. In ICLR, 2019. Cited on page 2.
Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICMI,, 2015. Cited on page 4.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Mitos, Blažej Osiński, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, et al. Model based reinforcement learning for atari. In ICLR, 2019. Cited on pages 2, 6, 7, 15, and 18.
Kacper Piotr Kielak. Do recent advancements in model-based deep reinforcement learning really improve data efficiency?, 2020. URL https: /ttp::s///ppnrevieww net. forum?id Bke9ulHFwB. Cited on pages and 15.
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. In ICLR, 2019. Cited on page 6.
Samuli Laine and Timo Aila. Temporal ensembling for semi-supervised learning. In ICLR, 2017. Cited on page
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Reinforcement learning with augmented data. In NewrIPS, 2020. Cited on page 6.
Alex X Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model. arXiv preprint arXiv:1907.0099 2019. Cited on page
Timothée Lesort, Natalia Díaz- Fodz-Rodrgg Sran-Franois Goudou, and David Filliat. State representa- tion learning for control: An overview. Neural Networks, 108, 2018. Cited page 6.
Marlos C Machado, Marc G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61:523-562, 2018. Cited page 6.
Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. In NeurlPS, 2020. Cited on pages 6 and 8.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning. Nature, 518(7540), 2015. Cited on pages 2, 3, 5, 6, and 7.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint / 2018. Cited pages 2, 6, and 16.
12

---

Published as a conference paper at ICLR 2021
OpenAI, Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyshaw Dębiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, llya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06600 2019. Cited on page 1.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer. James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep learning library. In NeurIPS, 2019, Cited on page 6.
Edgar Riba. Dmytro Mishkin, Daniel Ponsa, Ethan Rublee, and Gary Bradski. Kornia: an open source differentiable computer vision library for pytorch. In The IEEE Winter Conference on Applications of Computer Vision, 2020. Cited on page 5.
Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 2020. Cited on pages 1, 4, 5, and 9.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. In ICML, 2020. Cited on pages 4, 6, 8, 9, and 16.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, llya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research, 15(1), 2014. Cited on page 5.
Adam Stooke and Pieter Abbeel. rlpyt: A research code base for deep reinforcement learning in pytorch. arXiv preprint arXiv: 01099,01500,, 2019. Cited on page 6.
Richard Sutton. The bitter lesson. Incomplete Ideas (blog). 2019. Ancompleteideas net htt.a//c.e..c..eeeeeeeeaaaaaaaae
http: http://www.
Cited on page 1.
Antti Tarvainen and Harri Valpola. Mean teachers are better role models: Weight-averaged consistency targets improve supe-super deep learning results. NeurlPS, 2017. Cited on pages 2, 4, and 17.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint 2018. Cited page 1.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive multiview coding, arXiv preprint /rX.v:/906.05849, 2019. Cited on page 2.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. In AAAI, 2016. Cited on pages 2 and 3.
Hado P van Hasselt, Matteo Hessel, and John Aslanides. When use parametric models in reinforcement learning? In NeurIPS, 2019. Cited on pages 2, 3, 5, 6, and 16.
Oriol Vinyals, Igor Babuschkin, Wojciech M Czarnecki, Michaël Mathieu, Andrew Dudzik, Junyoung Chung, David Choi, Richard Powell, Timo Ewalds, Petko Georgiev, al. Grandmaster level in starcraft ii using multi-agent reinforcement learning. Nature, 2019. Cited on page 1.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Pro- ceedings of Machine Learning Research, pp. 9929-9939. PMLR, 13-18 Jul 2020. URL http:/ http://prroceedings www.ppress////mmmmmmmmmmmmm. / vvlllll99 9/ www///knnk.../// html. Cited on page 16.
Ziyu Wang, Tom Schaul, Matteo Hessel, Hado Hasselt, Mare Lanctot, and Nando Freitas. Ducling network architectures for deep reinforcement learning. ICML, 2016. Cited on pages 3, 5, and 7.
Qizhe Xie, Zihang Dai, Eduard Hovy. Minh- TaahTThagg Luong, and Quoc Le. Unsupervised data augmentation for consistency training. arXiv preprint arXiv: /904. /2848, 2019. Cited page 2.
13

---

Published as a conference paper at ICLR 2021
Denis Yarats, Ilya Kostrikov, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations. 2021. URL https:/ wwwpr/eeeeeeeeeee/e//p/e. http://wwww..tttttmmdddeddd.a.e 6sT Cited pages 5, 6, and
A ATARI DETAILS
provide full set of hyperparameters used both the augmentation and no-augmentation cases in Table 3, including new hyperparameters for SPR.
Table 3: Hyperparameters for SPR on Atari, with and without aug- mentation.
Parameter
Setting (for both variations)
Gray-scaling
True
Observation down-sampling
84x84
Frames stacked
Action repetitions
Reward clipping
[-1,
Terminal on loss of life
True
Max frumes per episode
108K
Update
Distributional
Dueling
Thee
Support Q-distribation
51
Discount factor
0.99
Minibatch size
32
Optimizer
Adam
Optimizer: learning rale
0.0001
Optimizer: 81
Oplimizer: 82
0.999
Optimizer:
Max gradient norm
10
Priority exponent
0.5
Priority correction Exploration
0.4 +
Noisy nets
Noisy nets parameter
Training steps
100K
Evaluation trajectories
Min replay size for sampling
2000
Replay period every
step
Updates per step
Multi-step return length
network: channels
32, 64, 64
Q network: filter size
x 8, 4 4, 3
network: stride
2,
network: hidden units
256
Nom linearity
ReLL
Target network: update period
À (SPR loss coefficient
(Prediction Depth)
Parameter
With Augmentation
Without Augmentation
Data Augmcntation
Random shifts (=1 pixcls) Intensity(scale=0..05
Noac
Dropout
(EMA coefficient)
0.99
A.1 FULL RESULTS
We provide full results across all 26 games for the methods considered, including SPR with and without augmentation, in Table 4. Methods are ordered in rough order of their date release or publication.
14

---

Published as a conference paper at ICLR 2021
Table Mean episodic returns on the 26 Atari games considered by Kaiser et al. (2019) after 100k environment steps. The results are recorded at the end of training and averaged over 10 random seeds. SPR outperforms prior methods on all aggregate metrics, and exceeds expert human performance on 7 out of 26 games while using similar amount of experience.
Game
Random
Human
SimPL.
DER
OTRainbow
CURL
DrQ
SPR (no
SPR
Alien
227.8
7127.7
616.9
739.9
824.7
558.2
771.2
847.2
801.5
Amidar
5.8
1719.5
88.0
188.6
82.8
142.1
102.8
142.7
176.3
Assault
600.6
571.0
Asterix
210.0
8503.3
1128.3
470.8
628.5
734.5
603.5
820.2
977.8
Bank Heist
753.1
34.2
182.1
131.6
168.9
425.6
380.9
BattleZone
2360.0
37187.5
5184.4
10124.6
400006
14870.0
12954.0
10738.0
166510
Boxing
0.1
12.1
2.5
1.2
6.0
12.7
36.8
Chomperoommand
811.0
7387.8
1246.9
861.8
1033.3
1058,5
780.3
667.1
974.8
Crazy Climber
42923.6
Demon Attack
152.1
1971.0
208.1
508.0
71188
817.6
1113.4
370.1
$45.2
Freeway
29.6
20.3
27,9
25.0
26.7
9.8
16.1
24.4
Frostbite
4334.7
254.7
231.6
1181.3
331.1
1657.4
1821.5
Gopher
257.6
2412.5
771.0
349,5
778.0
669.3
636.3
774.5
715.2
Jameshond
29.0
302.8
125.3
301.6
112.3
471.0
236.0
367.2
365 4
Kanguroo
$2.0
303.00
323.1
779.3
605.4
872.5
940.6
1359.5
3276.4
Krull
1598.0
2663.5
4539.9
2851.5
3277.9
4229.6
4018.1
3123.1
3688. 9
Kung Fu Master
258.5
22736.3
17257.2
14346.1
5722.2
143077,8
9111.0
15469.7
13192.7
Pong
-20.7
14.6
12.8
-19.3
1.3
-16.5
-8.5
166.0
-5.9
Eye
218.4
-13.6
124.0
Qhert
163.9
13455.0
1288.8
1152.9
509.3
1042.4
854-4
606.6
669
Road Runser
11.5
7845.0
5640.6
9600.0
2696.7
$661,0
8895,1
10511.0
14220.5
Senquest
68.4
42054.7
683,3
354.1
286.9
384.5
301.2
$80.8
Up Down
$33.4
11693.2
3350.3
2877.4
2847.6
2955.2
3180.8
6604.6
28138.5
Mcan Hunan-Norn'd
0.000
L,000
0.443
0.285
0.264
0,381
0.357
0.463
0,704
Median Human-Norue d
0.000
1,000
0. 144
0.161
0.204
0.175
0.268
0,307
0.415
Norm'd
Median DQN 30M-Norm'd
0.000
0.994
0.118
0.142
0.103
0.142
0.131
0.225
0.361
# Superhoman
2
2
Table 5: Scores on the 26 Atari games under consideration for our controlled Rainbow implementation with and without augmentation, compared to previous methods. The high mean DQN-normalizz score of our DQN without augmentation is due to an atypically high score on Private Eye, a hard exploration game on which the original DQN achieves a low score.
Variant
Human-Normalized Score
DQN@ D0NN-MMM-oraalized Score median
mean
median
mean
Rainbow (controlled, no aug)
0.240
0.204
0.374
0.149
OTRainbow
0.264
0.204
0.197
0.103
DER
0.285
0.161
0.239
0.142
Rainbow (controlled, w/ aug)
0.480
0.346
0.284
0.278
DrQ
0.357
0.268
0.171
0.131
A.2 CONTROLLED BASELINES
To ensure that the minor hyper-parameter changes we make to the DER baseline are not solely responsible for improved performance, perform controlled experiments using the same hyper- parameters and same random seeds for baselines. We find that our controlled Rainbow implementation without augmentation is slightly stronger than Rainbow (Kielak, 2020), while with augmentation enabled our results are somewhat stronger than DrQ.² None of these methods, however, close to the performance of SPR.
Rainbow but comparable Overtrained
This is perhaps surprising, given that the model used by DrQ omits many of the components of Rainhow.
15

---

Published as a conference paper at ICLR 2021
Table 6: Scores on the 26 Atari games under consideration for various contrastive alternatives SPR
implemented in our codebase. All variants listed here use data augmentation.
Variant
Human-Normalized Score
DQN@ 50MMooomonmll Score
mean
median
mean
median
SPR
0.704
0.415
0.510
0.361
Rainbow (controlled)
0,480
0.346
0.284
0.278
Non-temporal contrastive
0.379
0.200
0.268
0.179
1step contrastive
0,473
0.231
0.280
0.213
5-step contrastive
0.506
0.172
0.239
0. 0.142
Uniformity loss
0,422
0.176
0.271
0.144
B COMPARISON WITH A CONTRASTIVE LOSS
To compare SPR with alternative methods drawn from contrastive learning, we examine several variants of contrastive losses based on InfoNCE (Oord al., 2018):
A contrastive loss based solely different views of the same state, similar to CURL (Srinivas al., 2020).
A temporal contrastive loss with both augmentation and where targets are drawn one step in the future, equivalent to single-step CPC (Oord et al., 2018).
A temporal contrastive loss with explicit dynamics model, similar to CPCIAction (Guo et al., 2018). Predictions are made up to five steps in the future, and encodings of every state except St+-k used negative samples for St+-k-
A soft contrastive approach inspired by Wang & Isola (2020), who propose to decouple the repulsive and attractive effects of contrastive learning into two separate losses, one of which is similar to the SPR objective and encourages representations to be invariant to augmentation or noise, and one of which encourages representations to be uniformly distributed on the unit hypersphere. We optimize this uniformity objective joinely with the SPR loss, which takes the role the "invariance" objective proposed by (Wang Isola, 2020), We use = 2 in the uniformity loss, and give it weight equal to that given to the SPR loss, based on hyperparameters used by Wang & Isola (2020).
To create as fair comparison possible, use the same augmentation (random shifts and intensity) and the same Rainbow hyperparameters as in SPR with augmentation. As in SPR, we calculate contrastive losses using the output of the first layer of the Q-head MLP, with a bilinear elassifier (as in Oord et al., 2018). Following Chen et al. (2020a), we use annealed cosine similarities with temperature of 0.1 in the contrastive loss. We present results in Table 6.
Although all of these variants outperform the previous contrastive result on this task, CURL, none of them substantially improve performance over the controlled Rainbow they use baseline. We consider these results broadly consistent with those of CURL, which observes relatively small performance boost over their baseline, Data- Efff-Efffceent Rainbow (van Hasselt et al., 2019).
C THE ROLE OF THE TARGET ENCODER IN SPR
We consider several variants of SPR with the target network modified, and present aggregate metrics for these experiments in Table 7. We first evaluate a variant of SPR in which target representations are drawn from the online encoder and gradients allowed to propagate into the online encoder through predictable. We find
them, effectively allowing the encoder to learn to make its representations that this leads to drastic reductions in performance both with and without augmentation, which we attribute representational collapse.
16

---

Published as a conference paper at ICLR 2021
Table Scores on the 26 Atari games under consideration for variants of SPR with different target encoder schemes, without augmentation.
Variant
Score median
DQN @ 50M-Normalized Score
Without Augmentation
mean
mean
median
SPR
0.463
0.307
0.336
0.225 0.233
No Stopgradient
0.375
0.208
0.301
With Augmentation
SPR
0.704
0.415
0.510
0.361
No Stopgradient
0.515
0.278
0.344
0.231
12-
1.02-
11-
0.98
0.8
0.94
0.309 0.908 0.565
0.50
0.975
0.054
0.9
0.785
0.536
0.0
0.909
0.008
0.995
0.59
0.978 0.944
0.9
0.785 0.536
0.0
Figure 5: Performance on a subset of 10 Atari games for different values of the EMA parameter with augmentation (left) and without (right). Scores are averaged across 10 seeds per game for each value of r. Self-normalized score is calculated separately for the augmentation and no-augmentation cases.
To illustrate the influence of the EMA constant 5, we evaluate 7 at 9 values logarithmically inter- polated between 0.999 and 0 on subset of 10 Atari games. 4 We use 10 seeds per game, and evaluate SPR both with and without augmentation; parameters other than T are identical those listed in Table 3. To equalize the importance of games in this analysis, we normalize by the av- erage score across all tested values of for each game to calculate a self-normalized score, as scoresns agent
uverage score random score
We test SPR both with and without augmentation, and calculate the celf-normalized score separately between these cases. Results shown in Figure 5. With augmentation, we observe clear peak in performance at r = 0, equivalent to a target encoder with no EMA-based smoothing. Without augmentation, however, the story is less clear, and the method appears less sensitive (note y-axis scales). We use 0.99 in this case, based on its reasonable performance and consistency with prior work (e.g.. Grill et al., 2020). Overall, however, we note that SPR does appear overly sensitive unlike purely unsupervised methods such as B BYO;; in no case does SPR fail to train.
We hypothesize that the difference between the augmentation and no-augmentation cases is partially due to augmentation rendering the stabilizing effect of using an EMA target network (e.g., as observed by Grill al., 2020; Tarvainen & Valpola, 2017) redundant. Prior work has already noted that using an EMA target network can slow down learning early in training (Tarvainen & Valpola, 2017); in our context, where limited number of environment samples are taken in parallel with optimization, this may "waste" environment samples by collecting them with an inferior policy. To resolve this, Tarvainen Valpola (2017) proposed to increase 7 over course training, slowing down changes to the target network later in training. It is possible that doing here could allow SPR to achieve the best of both worlds, but it would require tuning an additional hyperparameter, the schedule by which is increased, and we thus leave this topic for future work.
'T E (0.999, 0.9976, 0.9944, 0.9867, 0.9684, 0.925, 0.8222, 0.5783, 0}. Pongg Breakout, Up N Down, Kangaroo, Bank Heist, Assault, Boxing, BattleZone, Frostbite and Crazy Climber
17

---

Published as a conference paper at ICLR 2021
D WALL CLOCK TIMES
We report wall-clock runtimes for a selection of methods in Table 8. SPR with augmentation for 100K steps on Atari takes around 4 and half to finish complete training and evaluation run single game. We find that using data augmentation adds an overhead, and SPR without augmentation can run in just hours.
SPR's wall-clock run-time compares very favorably to previous works such as SimPLe (Kaiser et al., 2019), which requires roughly three weeks train on a GPU comparable to those used for SPR.
Table Wall-clock runtimes for various algorithms for complete training and evaluation run on single Atari game using P100 GPU. Rainbow (controlled) is roughly comparable to DrQ, although its runtime will differ due different DQN hyperparameters. Runtime for SimPLe is taken from its v3 version on Arxiv, although the latest version doesn't It mention runtime. All runtimes are approximate, as exact running times vary from game to game.
Model
Runtime in hours (100k env steps)
SPR
4.6
Rainbow (controlled)
2.1
SPR (No aug)
3.0
Rainbow (controlled, aug)
1.4
SimPLe
500
18


Reviews:
[{"review_id": "rvr3cXLRm9X", "reviewer": "AnonReviewer2", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes a method to improve data-efficiency of DQN (specially, the Rainbow DQN) by training the reinforcement learning agent to jointly minimize the DQN loss and a loss for multi-step predictions on its own latent states.\n\nPros\n+ The proposed approach achieves signficant performance improvement as compared to existing algorithms in the low data regime on Atari games.\n+ The paper is well-written and generally easy to follow. Fig. 2 provides a nice illustration on the algorithm.\n\nCons\nThe approach provides an interesting combination of several heuristics, but it isn't always clear how these work and whether these are necessary components. \n- Both the target encoder network and the target projection head uses the EMA trick, and an additional prediction head q is used. There is some discussion on the EMA trick in the experiment section for the target encoder, but there is little discussion on the the target projection head and the prediction head.\n- Why is the prediction loss chosen to be negative cosine similarity? Does using quadratic loss work?\n\nMinor comments\n- Alg. 1: a loop over experiences in the minibatch seems to be missing, and RL loss(s, a, r, s'; \\theta_{o}) should be a function of the mini-batch not just a single experience (s, a, r, s'). Is s' perturbed as well?\n- Specifying the network parameters in notations will improve clarity (though many papers don't). In this paper, sometimes the network parameters are included in a notation, but sometimes not (e.g. L^{RL} and $L^{RL}_{\\theta}).\n- Strictly speaking, cosine similarity = 1 - 0.5 * normalized l2 loss (not really proportional)\n- Is the algorithm's performance sensitive to the choice of K?\n- The paper states that \"using a separate target encoder is vital in all cases\", this is confusing as the paper states later that with data augmentation, this is not needed (\\tau=0 works best).", "extra_scores": {}, "extra_text": {}}, {"review_id": "r_vVCVXUnde", "reviewer": "AnonReviewer3", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes Self-Predictive Representations (SPR), a self-supervised representation learning algorithm designed to improve the data efficiency of deep reinforcement learning agents. The main idea is to maintain a target encoder network and a target projection network (the moving average of the corresponding online networks) which can be used as a similarity prediction loss for transitions in the latent space. This loss essentially replaces the reconstruction loss typically used in similar latent space models. \n\nThe paper is very well written and easy to read and understand. Figure 2 clearly demonstrates the proposed method and the text does a good job at filling in the details. The experiments are relatively comprehensive (using the 100K sample threshold on Atari games proposed by SimPLe) and the authors compared the performance of the proposed methods with majority of key works on rl from pixels. The appendix also provides more detailed results and the implementation details. The authors also included an anonymous link to a repo which includes the code which is always a plus. The implementation is based on Rainbow as mentioned in the paper.\n\nThe main idea of he paper is very intriguing: that the learned state representation can be improved by self forward prediction. The idea is generic enough to be used in similar problems to improve the sota models.  However, I wish the Analysis section was bigger than half a page as I have hard time understanding *where* this improvement is coming from. There are interesting  ablation studies and insights in Section 5 but the authors could provide more. For example, the role of input augmentation, prediction horizon (dynamic modeling in the text), the magic of target network as well as the importance of tau, could be disentangle and explored more. \n\nThe paper can also be improved further by including a wall-clock time analysis. The method is clearly outperforming Rainbow but it comes at the cost extra parameters and computation. The authors listed \"not relying on reconstructing raw states\" as a benefit of the proposed method but I'm curious how does the proposed method compares to such methods, both the ones that unroll in the pixel space (e.g. SimPLe) and the ones that only use it as an auxiliary loss during training (e.g. Dreamer). By comparison I mean in terms of performance and wall clock time. Regardless, Dreamer seems to be a an interesting base model to compare to. ", "extra_scores": {}, "extra_text": {}}, {"review_id": "dgjodqg5wGm", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose learning self-supervised representations that are consistent across future time steps.  To do this, the authors augment a modified Data Efficient Rainbow with the self-supervised architecture from BYOL; however, instead of applying it to two augmentations of the same view, they match a sequence of states to each other.  The online sequence is generated as follows: the current state at time t is passed through the online encoder to generate z_t, from which representations {z_t+1,...,z_t+k} are generated via an action-conditioned transition model.  These k states are then matched with the target encoding of the observations {o_t+1,...,o_t+k} via cosine similarity.  They experimented with applying augmentations, as well as dropout of the features of both online and target networks (without augmentation).  The reward is also predicted from the state z_t, to ensure that the representation learned is still useful for the task.  The authors evaluate SPR on sample-efficient Atari, which is limited to only 100k environment steps (400k frames).\n\nThe paper admirably integrates a powerful self-supervised representation learning technique with the framework of reinforcement learning, by enabling the contrasting of multiple consecutive states via the use of a multi-step dynamics model.  However, inherent to the process of adapting these techniques is the need to explore design decisions and assumptions made, and their effects.\n\nHere are a few questions to consider, that the reader would potentially be interested in:\n1) Why was the original state not also matched against a target encoding (with a different augmentation) - essentially 0-step BYOL?\n2) To improve the quality of the dynamically generated states as well as the transition model itself, would it not be better to learn from the multi-step reward signal as well?   This would ensure consistency between the predicted states and the encoded states.  The gradient can be stopped before it flows through the online encoder.\n3) In BYOL, the authors match one augmentation passed through the online network to another augmentation passed through the target network.  By symmetry, however, the second augmentation passed through the online network should also match the first augmentation passed through the target.  Thus, BYOL is computed as a sum of these two cosine similarity losses.  Would it not also make sense for SPR to apply this symmetry?  The representation from the past should match that of the future, but is it not desirable for future representations to match the past for consistency (within a reasonable frame length)?\n4) The authors of BYOL report sensitivity to the transformation set used and the batch size, despite not relying on negatives.  Do these considerations transfer to SPR?\n5) Representations learned from contrastive self-supervised techniques such as BYOL and SimCLR have been shown to be good for a variety of applications, such as linear evaluation, semi-supervised tasks, and transfer learning.  Do these benefits also apply when adapted to RL via SPR - for example, transferring across tasks?\n\nAlthough there are other benchmarks that could be used to evaluate SPR (such as the DeepMind Control Suite) and present a stronger case, the existing experimental results are already compelling.  Ultimately, the idea is presented clearly and the writing is solid.  However, a more careful analysis of the implications of the design decisions made and how the pros and cons of these self-supervised representation learning approaches transfer when integrated into RL would separate this paper from an incremental work.", "extra_scores": {}, "extra_text": {}}, {"review_id": "F83WVxCHxNy", "reviewer": "AnonReviewer4", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper introduces a method to improve sample-efficiency in reinforcement learning, called Self-Predictive Representations. The authors use predictions in the latent space by using a learned transition model. To further improve their method, they add data-augmentation to enforce consistent the representation to be consistent over multiple views of an observation. Their method is tested on the 100k steps Atari environment and compared to other state-of-the-art methods. The results show the proposed method to significantly outperform previously suggested methods.\n\nIn its current form, the paper is marginally below the threshold in my opinion.\n\nThe reasons for my decision are as follows:\n-\tThe main point of contention is the choice to use different numbers for random seeding. While I understand the motivation stated by the authors to make results comparable to the original baselines/state of the art, it results in several problems. The error bars in Figure 3 are not meaningful as they are based on different numbers of seeds for different methods. Furthermore, I am surprised that the authors did cite Henderson et al., 2018, but do not address the problems that arise from using only a small number of random seeds (as shown in Henderson et al., 2018, Figure 5).\n-\tThe paper is not consistent in notation. The authors introduce MDPs without states but observations, which is simply not the correct definition. They move on and replace observations with states in section 2.2, without every introducing or defining it (states). Lastly, Fig.2 uses X_t, which is also never introduced or defined.\n-\tIn 2.5, the authors communicate their choice for different parameters as based on ‘early experiments’. Without further clarification, this statement is not meaningful as it is highly ambiguous.\n\nDespite the aforementioned issues, the paper is mostly well-written and the presented idea is valid and important. As the main criticism is the choice of random seeds in the experiments, I would accept the paper if this is accordingly addressed. \n\n------------------------------------------------------------------\n\nThe authors sufficiently addressed all questions in their rebuttal and I will therefore increase my score.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors propose self predictive representations (predicting the agents own future latents of a forward model with data augmentation) as a means of improving the data efficiency of agents. The reviewers found the paper to be compelling (after the authors made adjustments) and pointed out that the method is likely generic and might be widely applicable. Experimental improvements in the work are significant, and the method is well explored.
Decision: Accept (Spotlight)