Decoupling Representation Learning from Reinforcement Learning
Authors: Adam Stooke
Kimin Lee
Pieter Abbeel
Michael Laskin
Venue:  (2021)
Keywords: reinforcement learning
representation learning
unsupervised learning

Abstract: In an effort to overcome limitations of reward-driven feature learning in deep reinforcement learning (RL) from images, we propose decoupling representation learning from policy learning.  To this end, we introduce a new unsupervised learning (UL) task, called Augmented Temporal Contrast (ATC), which trains a convolutional encoder to associate pairs of observations separated by a short time difference, under image augmentations and using a contrastive loss.  In online RL experiments, we show that training the encoder exclusively using ATC matches or outperforms end-to-end RL in most environments.  Additionally, we benchmark several leading UL algorithms by pre-training encoders on expert demonstrations and using them, with weights frozen, in RL agents; we find that agents using ATC-trained encoders outperform all others.  We also train multi-task encoders on data from multiple environments and show generalization to different downstream RL tasks.  Finally, we ablate components of ATC, and introduce a new data augmentation to enable replay of (compressed) latent images from pre-trained encoders when RL requires augmentation.  Our experiments span visually diverse RL benchmarks in DeepMind Control, DeepMind Lab, and Atari, and our complete code is available at \url{hidden url}.

Under review as a conference paper at ICLR 2021
DECOUPLING REPRESENTATION LEARNING FROM REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
In an effort to overcome limitations of reward-driven feature learning in deep re- inforcement learning (RL) from images, we propose decoupling representation learning from policy learning. To this end, we introduce new unsupervised learning (UL) task, called Augmented Temporal Contrast (ATC), which trains a convolutional encoder associate pairs of observations separated by short time difference, under image augmentations and using a contrastive loss. In online RI. experiments, we show that training the encoder exclusively using ATC matches or outperforms end-to-end RL in most environments. Additionally, we benchmark several leading UL algorithms by pre-training encoders on expert demonstrations and using them, with weights frozen, in RL agents; we find that agents using ATC-trained encoders outperform all others. We also train multi-task encoders on data from multiple environments and show generalization to different down- stream RL tasks. Finally, we ablate components ATC, and introduce a new data augmentation enable replay of (compressed) latent images from pre- preitriined en- coders when RL requires augmentation. Our experiments span visually diverse RL benchmarks in DeepMind Control, DeepMind Lab, and Atari, and our com- plete code is available at hidduur
INTRODUCTION
Ever since the first fully-Icarned approach succeeded at playing Atari games from screen images (Mnih al., 2015), standard practice in deep reinforcement learning (RL) has been to learn visual features and a control policy jointly, end-to-end. Several such deep RL. algorithms have matured (Hessel al. 2018 Schulman et al.] [017]] Mnih et al. 2016). HHaamooja et al.] 2018)) and have been successfully applied to domains ranging from real-world (Levine et al. 20166] Kalashnikov cet al. 2018) and simulated robotics (Lee et al. 2019 Laskin et al. 2020a Hatner et al. (020) to sophisticated video games (Berner al. 2019 Jaderberg et al. 20199) and even high-fidelity driving simulators (Dosovitskiy et al. 2017). While the simplicity of end-to-end methods is appealing, relying the reward function to learn visual features can be severely limiting. For example, leaves features difficult to acquire under sparse rewards, and it can narrow their utility to a single task. Although our intent broader than to focus on either sparse-reward or multi-task settings, they arise naturally in our studies. We investigate how to learn visual representations which are agnostic rewards, without degrading the control policy.
A number of recent works have significantly improved RL performance by introducing auxiliary losses, which are unsupervised tasks that provide feature-leariing signal to the convolution neural network (CNN) encoder, additionally to the RL loss (Jaderberg et al. 2017 van den Oord et al.) 2018 Laskin et al. 2020b Guo et al., 2020 Schwarzer et al.] 2020). Meanwhile,
the field of computer vision, recent efforts in unsupervised and self-supervised learning (Chen et al.) 2020 Grill et al. 2020 He et al. 2019) have demonstrated that powerful feature extractors can be learned without labels, as evidenced by their usefulness for downstream tasks such as ImageNet classifica- tion. Together, these advances suggest that visual features for RL could possibly be learned entirely without rewards, which would grant greater flexibility to improve overall learning performance. To our knowledge, however, no single unsupervised learning (UL) task has been shown adequate for this purpose in general vision- based environments.

---

Under review as a conference paper at ICLR 2021
In this paper, we demonstrate the first decoupling of representation learning from reinforcement learning that performs as well as or better than end-to-end RL. We update the encoder weights using only UL and train control policy independently. stands in contrast to previous state-of-the-ar methods, which have trained the UL and RI objectives jointly, or Laskin al. (2020b), which observed diminished performance with decoupled encoders.
the (compressed) latent images. This capability
Our main enabling contribution a new unsupervised task tailored to reinforcement learning, which call Augmented Temporal Contrast (ATC). ATC requires model to associate observations from nearby time steps within the same trajectory (Anand et al. 2019). Observations are encoded via a convolutional neural network (shared with the RL agent) into small latent space, where the InfoNCE loss is applied (van den Oord et al., 2018). Within each randomly sampled training batch, the positive observation, Ot+k, for every anchor, Ot, serves as negative for all other anchors. For regularization, observations undergo stochastic data augmentation (Laskin et al., 2020b) prior to encoding, namely random shift (Kostrikov et al., 2020), and momentum encoder (He et al. [2020] Laskin et al.) 2020b) is used to process the positives. A learned predictor layer further processes the anchor code (Grill et al. 2020 Chen et al. 2020) prior to contrasting, In summary, our algorithm is novel combination of elements that enables generic learning of the structure of observations and transitions in MDPs without requiring rewards or actions as input.
We include extensive experimental studies establishing the effectiveness of our algorithm in vi- sually diverse range of common RI. environments: DeepMind Control Suite (DMControl: Tassa jet 2018), DeepMind Lab (DMLab: [Beattie et al. al.[ 2016), and Atari (Bellemare et al.) (013). Our experiments span discrete and continuous control, 2D and 3D visuals, and both on-policy and off policy RL algorithms. Complete code for all of our experiments is available at hiddenurl empirical contributions are summarized as follows:
Our
Online RL with UL: We find that the convolutional encoder trained solely with the unsupervised ATC objective can fully replace the end-to-end RL encoder without degrading policy performance. ATC achieves nearly equal or greater performance in all DMControl and DMLab environments tested and in 5 of the 8 Atari games tested. In the other 3 Atari games, using ATC as an auxiliary loss or for weight initialization still brings improvements over end-to-end RL.
Encoder Pre-Training Benchmarks: We pre-train the convolutional encoder to convergence on ex- pert demonstrations, and evaluate it by training an RL agent using the encoder with weights frozen. We find that ATC matches or outperforms all prior UL algorithms as tested across all domains, demonstrating that ATC is a state-of-the-an UL algorithm for RL.
Multi-Task Task Encoders: An encoder is trained on demonstrations from multiple environments, and is evaluated, with weights frozen, in separate downstream RL agents. A single encoder trained on four DMControl environments generalizes successfully, performing equal better than end-to-end RL in four held-out environments. Similar attempts to generalize across eight diverse Atari games result in mixed performance, confirming some limited feature sharing among games.
Ablations and Encoder Analysis: Components of ATC are ablated, showing their individual effects. Additionally, data augmentation is shown to be necessary in DMControl during RL even when using a frozen encoder. We introduce a new augmentation, subpixel random shift, which matches performance while augmenting the latent images, unlocking computation and memory benefits.
2 RELATED WORK
Several recent works have used unsuperviiee//ssfssspprvieedd representation learning methods to improve performance in RL. The UNREAL agent (Jaderberg al. 2017) introduced unsupervised auxiliary tasks to deep RL, including the Pixel Control task, a Q-learning method requiring predic- tions of screen changes in discrete control environments, which has become a standard in DMLab (Hessel et al. 2019). CPC (van den Oord et al. 2018) applied contrastive losses over multiple time steps as an auxiliary task for the convolutional and recurrent layers of RL agents, and it has been extended with future action-conditioning (Guo et al. 2018). Recently, PBL (Guo et al. 2020) sur- passed these methods with an auxiliary loss of forward and backward predictions in the recurrent latent space using partial agent histories. Where the trend is of increasing sophistication in auxiliary recurrent architectures, our algorithm is markedly simpler, requiring only observations, and yet proves sufficient in partially observed settings (POMDPs).
2

---

Under review as a conference paper at ICLR 2021
2019) introduced various temporal, contrastive losses, including ones that operate on "local" features from an intermediate layer within the encoder, without data augmenta- tion. CURL (Laskin et al. 2020b) introduced augmented, contrastive auxiliary task similar to ours, including a momentum encoder but without temporal contrast. [Mzoure et al. (020)) provided extensive analysis pertaining to InfoNCE losses functions of successive time steps in MDPs, including local features in their auxiliary loss (DRIML) similar to ST-DIM, and finally conducted experiments using global temporal contrast of augmented observations in the Procgen (Cobbe al. 2019) environment. Most recently, MPR (Schwarzer et al. (0200) combined data augmentation with multi-step. convolutional forward modeling and a similarity loss to improve DQN agents in the Atari 100k benchmark. [Hfnne et al. (2019) 2020): [ee et al. (2019) proposed to leverage world-modeling in latent-space for continuous control. small number model-free methods have attempted to decouple encoder training from the RL loss as ablations, but have met reduced performance relative to end-to-end RL {Laskin et al. 2020b Lee et al. 2020). None have previously been shown effee-
ST-DIM (Anand et al.
tive in diverse collection of RL environments as ours (Bellemare et al. 2013 Tassa et al. 2018
Beattie et al. 2016).
Finn et al. ( Ha & Schmidhuber (2018) example works which pretrained encoder features
in advance using image reconstruction losses such as the VAE (Kingma & Welling (013)) Devin al. (2018); Kipf al. (2019) pretrained object-centric representations, the latter learning for- ward model by way of contrastive losses; Yan et al.] (2020) introduced a similar technique to learn encoders supporting manipulation of deformable objects by traditional control methods. MERLIN (Wayne et al. 2018 trained a convolutional encoder and sophisticated memory module online, de- tached from the RL agent, which learned read-only accesses memory. It used reconstruction and losses and achieved high performance in DMLab- extreme partial observability. Our loss function may benefit those settings, as it outperforms similar reconstruction losses in our experiments. Decoupling unsupervised pretraining from downstream tasks is common in computer vision (Hénaff al., 2019 He et al. 2019 Chen et al. 2020) and has favorable properties of providing task agnostic features which can be used for training smaller task- specific networks, yielding significant gains in computational efficiency over end-to-end methods.
one-step satent-predictio
environments with
3 AUGMENTED TEMPORAL CONTRAST
Our unsupervised learning task, Augmented Temporal Contrast (ATC), requires a model to associate an obser- vation, Os, with one from a specified, near-future time step. OL+R- Within each training batch, we apply stochas- tic data augmentation to the observations (Laskin et al. 2020b), namely random shift (Kostrikov et al. 2020), which is simple to implement and provides highly effec- tive regularization in most cases. The augmented obser- vations are encoded into a small latent space where a con- trastive loss is applied. This task encourages the learned encoder to extract meaningful elements of the structure of the MDP from observations.
CONTRAST W
P
2ith
2 ZV
MLP hr
Momentum
Encnder
LINEAR g
LINEAR g3
CONV fo AUG(O)
CONV ff
Our architecture for ATC consists of four learned com- ponents (i) convolutional encoder, fer which pro- cesses the anchor observation, or, into the latent image 2, (oA)))(),) (ii) linear global compressor, go to produce a small latent code vector e, = (z))), (iii) a residual predictor MLP, hig which acts as an implicit forward model to advance the code pt hy (ct) + ct and (iv) a contrastive transformation matrix, W. To pro- cess the positive observation, Or+ into the target code we use a momentum encoder (He et al. 2019) parameterized as slowly moving av- erage of the weights from the learned encoder and com-
AUG (K) AUG AUG(O))...AAGIIO AUG(O
Figure 1: Augmented Temporal Contrast augmented observations are processed through a learned encoder fo, compressor, g and residual predictor his and are associated through contrastive loss with a positive example from h time steps later, processed through a momentum encoder.
Cr+kk = 9% [fa(AU
pressor layer:
θ+ (1 -0) + r0
--
(1-7)ō -) TΦ.
3

---

Under review as a conference paper at ICLR 2021
The complete architecture is shown in Figure the RL agent.
1 The convolutional encoder, fo, alone is shared with
We employ the InfoNCE loss (Gutmann & Hyvärinen 2010 van den Oord et al. 2018) using log- its computed bilinearly, as 2t our implementation, every anchor in the training batch utilizes the positives corresponding to all other anchors its negative examples. Denoting an observation indexed from dataset as 0,, and its positive as 024, the logits can be written as ligt p,W Cj+; ; our loss function in practice is:
LATC = CATC= -Eo log expl,
(2)
Eo exp (xp
EXPERIMENTS
4.1 EVALUATION ENVIRONMENTS AND ALGORIHTMM
We evaluate ATC on three standard, visually diverse RL benchmarks the DeepMind control suite (DMControl; Tassa al. 2018), Atari games in the Arcade Learning Environment (Bellemare al.) 2013). and Deep Lab (DMLab; Beattie et al. 2016). Atari requires discrete control in arcade- style games. DMControl is comprised of continuous control robotic locomotion and manipulation tasks. In contrast, DMLab requries the RL agent to reason in more visually complex 3D maze environments with partial observability.
We use ATC to enhance both on-policy and off-policy RI. algorithms. For DMControl, we use RAD- SAC (Laskin al., [2200a] Haarnoja al.) (2018) with the augmentation of Kostrikov et (020)) which randomly shifts the image in each coordinate (by up to pixels), replicating edge pixel values as necessary to restore the original image size. difference from prior work is that we use more downsampling in our convolutional network, by using strides (2, 2, 2, 1) instead of 1, to reduce the convolution output image by 25x For both Atari and DMLab, we use PPO (Shhulman et al. 2017). In Atari, we use feed-forward agents, sticky actions, and no end-of-life boundaries for RL episodes. In DMLab we used recurrent, L.TTM agents receiving only single time-step image input, the four-layer convolution encoder from Jaderberg et al. (2019). and tuned the entropy bonus for each level. In the online setting, the ATC is trained using small replay buffer recent experiences.
We include all our own baselines for fair comparison and provide complete settings in an appendix. Unless otherwise noted, each curve represents minimum of 3 random seeds. The bold lines show the average, and the lightly shaded area around each curve represents the maximum extent of the best and worst seeds each checkpoint.
4.2 ONLINE RI. WITH ATC
DMControl In the online setting, we found ATC be capable of training the encoder by itself (i.e., with encoder fully detached from any RL gradient update), achieving essentially equal or better scores versus end-to-end RL in all six environments we tested, Figure In CARTPOLE- SWINGUP SPARSE, where rewards only received once the pole reaches vertical, ATC training enabled the agent to master the task significantly faster. The encoder is trained with one update for every RL update to the policy, using the same batch size, except in CHEETAI--RUN which required twice the ATC updates.
DMLab We experimented with two kinds of levels in DMLab: EXPLORE- COW.ORREGGALLLLLATOOON which requires repeatedly navigating maze whose layout randomized every episode, and LASERTAG_ THREE_REE..O PRoReE,OPPONETTS which requires fast reflexes to pursue and tag enemies at dis- tance. We found ATC capable of training fully detached encoders while achieving equal or better performance than end-to-end RI. Results are shown in Figure Both environments exhibit spar-
For our input image size 84 x 84, the convolution output image is 7 x 7 rather than 35 x 35. Perfor- mance remains largely unchanged, except for a small decrease in the HAL.-CHEETAH environment, but the experiments run significantly faster and use less GPU memory.
4

---

Under review as a conference paper at ICLR 2021
Ball in Cup: Catch
Cartpol:: Swingup
Cartpale:: Swingup Sparse
Chectah: Run
Walker: Walk
Hopper: Hop
800
800
Encoder Traning
800
900
800
150
500
ATC
400
600
400
400
500
400
103
200
200
200
Eny Steps x301
X107
x30
x100
X301
1000
Figure 2: Online encoder training by ATC, fully detached from RL training, performs as well as end-to-end RL in DMControl, and better in pparse-reward environments (environment steps shown, see appendix for action repeats). Each curve is 10 random seeds.
sity which greater the "large" version than the "small" version, which our algorithm addresses, discussed next.
In EXPLORE, the goal object rarely seen, especially carly on, making its appearance difficult to learn. We therefore introduced prioritized sampling for ATC with priorities corresponding to empirical absolute returns: p x 1 + Rats where Rabs = Ei=o 7'/r to train more frequently on more informative scenes D Whereas uniform-ATC performs slightly below RL, uniform-ATC outperforms RL and nearly matches using ATC (uniform)
auxiliary task. By considering the encoder as a stand-alone feature extractor separate from the policy, no importance sampling correction is required.
In LASERTAG, enemies are often seen, but the reward of tagging one is rarely achieved by the random agent. ATC learns the relevant features anyway, boosting performance while the RL-only agent remains at zero average score. We found that increasing the rate of UL training to do twice as many update further improved the score to match the ATC-auxiliary agent, showing flexibility to address the representation-leaniie bottleneck when opponents are dispersed.
Explore Large
Explore Small
Lasertag Large
Laseriag Small
Encoder Tralting
60
15
Enceder Tranning (Lavertag)
3.5
ATC-pri
40
100
ATC2x
RL+ATE
20
0.5
RL-ATC
0.0
0.0
12.5
Agent X10
x10F
x107
k10F
Figure 3: Online encoder training by ATC, fully detached from the RL agent, performs as well or better than end-to-end RI. in DMLab (1 agent step = 4 environment steps, the standard action repeat). Prioritized ATC replay (EXPlOrE) increased ATC training (Lasertag) addresses sparsities nearly match performance of RL. with ATC as an auxiliary loss (RL+ATC). Each curve 3 random seeds.
Atari We tested diverse subset of eight Atari games, shown in Figure We found detached- encoder training to work as well as end-to-end RI. in five games, but performance suffered in BREAKOUT and SPACE INVADERS in particular. Using ATC as an auxiliary task, however, im- proves performance in these games and others. We found it helpful to anneal the amount of UL training over the course RL in Atari (details in an appendix). Notably, we found several games, including SPACE INVADERS, to benefit from using ATC only to initialize encoder weights, done using an initial 100k transitions gathered with a uniform random policy. Some of our remaining experiments provide more insights into the challenges of this domain.
4.3 ENCODER PRE-TRAINING BENCHMARKS
To benchmark the effectiveness of different UL algorithms for RL, we propose a new evaluation methodology that is similar to how UL pre-training techniques are measured in computer vision (see
In EXPL.ORF_ SEPL.ORF_GOAL_. LOCATIONS,,, the only reward is +10, carned when reaching the goal object. Since the ATC batch size was 512 but the RI. batch size was 1024, performing twice many UL. updates still only consumed the same amount of encoder training data RT. We did not fine-tune for batch size.
5

---

Under review as a conference paper at ICLR 2021
1000
gravitar
6000
ms_ m1s_pacman
2000
200
300
750
4000
Encoder Training
1000
100
500
2000
RL+ATC AL:AIn-init
0.0
12.5
0
0.0
12.5
25
250
12.5
25.0
0.0
12.5
25
0
RL
x10
X10
x10
K10"
pong
seaquest
space invaders
20
15
2000
Score
1000
1000
12.5
25.0
12.5
0.0
Agent Steps
x10
Figure 4:
Online encoder training using ATC, fully detached from the RL agent, works well in 5 of games tested (1 agent step =4 environment steps, the standard action repeat). 6 of 8 games benefit significantly from using ATC as an auxiliary loss or for weight initialization. Each curve is 8 random seeds.
e.g. Chen et al. 22020) Grill et al. 2020)): (i) collect data set composed of expert demonstrations from each environment; (ii) pre-train the CNN encoder with that data offline using UL: (iii) evaluate by using RI. to learn a control policy while keeping the encoder weights frozen. This procedure isolates the asymptotic performance of each UI. algorithm for RI.. For convenience, drew expert demonstrations from partially-tranned RL agents, and every UL algorithm trained on the same data for each environment. Our RL agents used the same post-encoder architectures as in the online experiments. Further details about pre-training by each algorithm are provided in an appendix.
DMControl We compare ATC against two competing algorithms: Augmented Contrast (AC), from CURL (Laskin al. (0200), which uses the same observation for the anchor and the positive, and a VAE (Kingma & Welling 2013), for which we found better performance by introducing time delay to the target observation (VAE-T). We found ATC to match or outperform the other Further, ATC is the only to
algorithms, in all four test environments, as shown in Figure 5 match or outperform the reference end-to-end RL across cases.
Ball-in--ap: Catch
Cartpolle: Swingup
Cheetah: Run
Walkr: Walk
1000
800
600
Encoder
800
800
Training
Reum
600
600
400
600
ATC tours)
None (RL)
400
AC (CURL)
200
200
200
None (rand)
Ery Stops
X10
X10
x10
x10*
Figure 5: RL in DMControl, using encoders pre-trained on expert demonstrations using UL, with weights frozen across all domains, ATC outperforms prior methods and the end-to and-to-end RL reference. Each curve is mininum of random seeds.
DMLab We compare against both Pixel Control (Jadeberg et al. 2017 Hessel al.] 2019) and CPC (van den Oord et al. 2018). which have been shown to bring strong benefits in DMLab. While all algorithms perform similarly well in EXPLORE, ATC performs significantly better in LASERTAG, Figure 6 Our algorithm is simpler than Pixel Con- trol and CPC in the sense that it uses neither actions, deconvolution, recurrence.
Lasertag
Enoodar Traning
135
15
k00
Contrcl Nore (HL)
52
125 Stees
25.0
0.2
22.5
Figure 6: RL in DMLab, using pre-trained en- coders with weights frozen-in LASERTAG espe- cially, ATC outperforms leading prior UL. algo- rithms.
Atari We compare against Pixel Control, VAE-T, and a basic inverse model which predicts actions between pairs of observations. We also compare

---

Under review as a conference paper at ICLR 2021
against Opatio-Tempore Deep InfoMax (ST-DIM), which uses temporal contrastive losses with "lo- cal" features from an intermediate convolution layer to ensure attention to the whole screen; it was shown to produce detailed game- state knowledge when applied to individual frames (Anand et al. 2019). Of the four games shown in Figure 7 ATC the only UL algorithm to match the end-to-end RL reference in GRAVITAR and BREAKOUT, and it performs best in SPACE INVADERS.
breakout
gravitar
qbert
space invaders
250
800
1500
Pre-Treining
Encoder
200
ATC (ours)
150
600
10000
PixelControl
VAE-T
100
400
5000
500
Inverse
0
None
0.0
12.5
25.0
0.0
12.5
25,0
0.0
12.5
25,0
0.0
12.5
25.0
None (rand)
Agent Stepsx10*
X10
x10
X10
Figure 7: RL in Atari, using pre-trained encoders with weights frozen ATC outperforms several leading, prior UL. algorithms and exceeds the end-to-end RL reference in of the 4 games tested.
4.4 MULTI-TASK ENCODERS
In the offline setting, we conducted initial explorations into the capability of ATC to learn multi-task encoders, simply by pre-training demonstrations from multiple environments. We evaluate the encoder by using it, with frozen weights, in separate RL agents learning each downstream task.
DMControl Figure g shows results in DMControl, where we pretrained using only the four environments in the top row. Although the encoder was never trained on the HOPPER, PENDULUM, nor FINGER domains, the multi-task encoder supports efficient RL in them. PENDULUM-SWING and SETTPOLESSWINGPP... SWINGUP-SPARS: SPARSE stand out as challenging environments which benefited from cross-domain and cross-tast pre-training, respectively. The pretraining was remarkably efficient, requiring only 20,000 updates to the encoder.
Ball-in-Cup: Catch
Cartpole: Swingup
Cheetah: Run
Walker: Walk
1000
1000
500
500
500
2
2.5
5.0
X10
x10
X10
X10
Cartpole: Swingup Sparse
Hopper: Hop
Pendulum: Swingup
Finger: Spin
750
Weww'
1000
Encoder
Resum
500
100
500
500
Pre-Training
4-Task
250
50
Mhall
None (RL)
0
0
0.0
0.5
1.0
0
2
4
Env Staps
110
x10
x10F
Figure& 8: Separate RL agents using single encoder with weights frozen after pre-training expert demon- strations from the four top environments. The encoder generalizes four environments, bottom row, where sparse reward tasks especially benefit from transfer. Each curve is minimum random seeds.
Atari Atari proved more challenging domain for learning multi-task encoders. Learning all eight games together in Figure [T] in the appendix, resulted in diminished performance relative to single-game pretraining in three of the eight. The decrease was partially alleviated by widening the encoder with twice as many filters per layer, indicating that representation capacity is limiting factor. To test generalization, we conducted a seven-get seven-game pre-training experiment where we test the encoder on the held-out game. Most games suffered diminished performance (although still perform significantly higher than a frozen random encoder), confirming the limited extent to which visual features transfer across these games.

---

Under review as a conference paper at ICLR 2021
breakout
250
Encodor meemtwiiigg Con: T-16
200
150
Con: T-6
100
Con: T-]
Sm. T-16
inguts
Random
RI-trained
UL (with shift)
UL (without shift)
Tw1
50rps
X10*
Figure 10: An example scene from BREAKOUT, where low- performance UL encoder (without shift) focuses on the paddle. Introducing random shift and sequence data makes the high- performance UL encoder (full ATC) focus the ball, as does the encoder from a fully-trained, end-to-end RL agent.
Figure 9: BREAKOUT benefits from contrasting against negatives from neighboring time steps.
4.5 ABLATIONS AND ENCODER ANALYSIS
Random Shift in ATC offline experiments, we discovered random shift augmentations to be helpful in all domains. To knowledge, this is the first application of random shift to 3D visual environments as in LMLab. Atari, we found performance in shift, but reducing the probability of applying random shift to each observation from 1.0 to 0.1 alleviated the effect while still bringing benefits in other games, so we used this setting in our main experiments. Results are shown in Figure 12 in appendix.
to suffer from random
Random Shift in RL In DMControl, found the best results when using random shift during RL, even when training with a frozen encoder. This is evidence that the augmentation regularizes only the representation but also the policy, which first processes the latent image into 50- dimensional vector. To unlock computation and memory benefits of replaying only the latent images for the RL agent, we attempted apply data augmentation the smallest possible random shifts to be too extreme. Instead, we introduce a new augmentation, subpixel random shift, which linearly interpolates among neighboring pixels. As shown in Figure [ in the appendix, this augmentation restores performance when applied to the latent images, allowing a pre-trained encoder to entirely bypassed during policy training updates.
the latent image. But we found
Temporal Contrast on Sequences In BREAKOUT alone, we discovered that composing the UL training batch of trajectory segments, rather than individual transitions, gave significant benefit. Treating all elements of the training batch independently provides "hard" negatives, since the en- coder must distinguish between neighboring time steps. This setting had no effect in the other Atari games tested, and we found equal better performance using individual transitions in DMControl and DMLab. Figure further shows that using similarity loss (Grill et al. 2020) does not capture the benefit.
Encoder Analysis We analyzed the learned encoders in BREAKOUT to further study this abla- tion effect. Similar to ZZagorykoo & Komodakis (2016), we compute spatial attention maps by mean-pooling the absolute values of the activations along the channel dimension and follow with 2-dimensional spatial softmax. Figure 10 shows the attention of four different encoders the displayed scene. The poorly performing OL encoder heavily utilizes the paddle to distinguish the observation. The UL encoder trained with random shift and sequence data, however, focuses near the ball, as does the fully-trained RL encoder. (The random encoder mostly highlights the bricks, which less relevant for control.) In appendix, we include other example encoder analyses from Atari and DMLaa which show ATC-trained encoders attending only to key objects on the game screen, while RL-trained encoders additionally attend to potentially distracting features such as game score.
5 CONCLUSION
Reward-free representation learning from images provides flexibility and insights for improving deep RL agents. We have shown broad range of cases where our new unsupervised learning algorithm can fully replace RL for training convolutional encoders while maintaining improving
8

---

Under review as a conference paper at ICLR 2021
online performance. In a small number of environments- few of the Atari dames-including the RL loss for encoder training still surpasses our UL-only method, leaving opportunities for further improvements in UL for RL.
Our preliminary efforts to use actions as inputs (into the predictor MLP) or as prediction outputs (inverse loss) with ATC did not immediately yield strong improvements. We experimented only with random shift, but other augmentations may be useful, as well. In multi-task encoder training, technique avoids any need for sophisticated reward-balancing (Hessel et al. 2019), but more advanced training methods may still help when the required features are in conflict, as in Atari, or they otherwise impact our loss function unequally. On the theoretical side, may be helpful to analyze the effects of domain shift on the policy when detached representation is learned online.
One obvious application of our offline methodology would be in the batch RL setting, where the agent learns from a fixed data set. Our offline experiments showed that relatively small number of transitions are sufficient to learn rich representations by UL, and the lower limit could be further explored. Overall, we hope that our algorithm and experiments spur further developments leveraging unsupervised learning reinforcement learning.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Mare-Alexandre Côté, and R Devon Hjelm. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, 2019.
Charles Beattie, Joel Z. Leibo, Denis Teplyashin, Tom Ward, Mareus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, Julian Schrittwieser, Keith Ander- son, Sarah York, Max Cant, Adam Cain, Adrian Bolton, Stephen Gaffney, Helen King, Demis Hassabis, Shane Legg, and Stig Petersen. Deepmind lab. arXiv preprint arXiv: 1612.03801, 2016.
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: 253-279, 2013.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, al. Dota with large scale deep reinforcement learning. arXiv preprint arXiv: 19/2.0668, 2019.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. arXiv:2002.05500 2020.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural genera- tion benchmark reinforcement learning. arXiv preprint 1rXiv:191..01188, 2019.
Coline Devin, Pieter Abbeel, Trevor Darrell, and Sergey Levine. Deep object-centric representa- tions for generalizable robot learning. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 7111-7118 18. IEEE, 2018
Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. Carla: An open urban driving simulator. arXiv preprint 1rxiv:/711.0339 2017.
C. Finn, Xin Yu Tan, Yan Duan, T. Darrell, S. Levine, and P. Abbeel. Deep spatial autoencoders for visuomotor learning. In 2016 IEEE International Conference on Robotics and Automation (ICRA), pp. 512-519, 2016.
Jean-Bi -aan-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec. Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent: new approach to self-supervised learning. arXiv preprint WrXiv:2006.07733, 2020.
Daniel Guo. Bernardo Avila Pires, Bilal Piot, Jean-bastien Grill, Florent Altché, Rémi Munos, and Mohammad Gheshlaghi Azar. Bootstrap latent-predictive representations for multitask reinforce- ment learning. arXiv preprint arXiv:2004. /4646, 2020.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A Pires, and Rémi Munos. Neural predictive belief representations. arXiv preprint arXiv:1811.06400 2018.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In International Conference on Artificial Intelligence and Statistics, 2010.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803. 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning. 2018
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, 2019.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
10

---

Under review as a conference paper at ICLR 2021
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. arXiv preprint arXiv:1911.05722 2019.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In IEEE/CVF Conference Computer Vision and Pattern Recognition, 2020.
Olivier 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: 09272, 2019.
Matteo Hessel, Joseph Modayil, Hado van Hasselt, Tom Schaul, Georg Ostrovski. Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. AAAI Conference on Artificial Intelligence, 2018.
Matteo Hessel, Hubert Soyer, Lasse Espeholt, Wojeiech Czarnecki, Simon Schmitt, and Hado van Hasselt. Multi-task deep reinforcement learning with popart. In AAAI Conference on Artificial Intelligence. 2019.
Max Jaderberg, Volodymyr Mnih, Wojciech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. In International Conference on Learning Representations, 2017.
Max Jaderberg, Wojeiech M Czarnekki, lain Dunning, Luke Marris, Guy Lever, Antonio Garcia Castaneda, Charles Beattie, Neil C Rabinowitz, Ari S Morcos, Avraham Ruderman, et al. Human- level performance in 3d multiplayer games with population- reinforcement learning. Sci- ence, 364(6443))8899-88866 2019.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, al. Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation. arXiv preprint MXXiv:880.112299 2018.
Diederik Kingma and Max Welling. arXiv:1312.6114 14, 2013.
Auto-encoding variational bayes.
arXiv preprint
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. arXiv preprini arXiv:1911.122
2019.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. arXiv preprint arXiv:2004.14999, 2020a.
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Curl: Contrastive unsupervised representa- tions for reinforcement learning. In International Conference on Machine Learning, 2020b.
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:09907000900 000000 2019.
Kuang-Huee Lee, Ian Fischer, Anthony Liu, Yijie Guo, Honglak Lee, John Canny, and Sergio Guadarrama. Predictive information accelerates learning in rl. Advances in Neural Information Processing Systems, 33, 2020.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. The Journal of Machine Learning Research, 17(1):333-1373, 2016.
Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, and R Devon Hjelm. Deep reinforcement and infomax learning. arXiv preprint arXiv:20066072/7. 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540::529-555 2015.
11

---

Under review as a conference paper at ICLR 2021
Volodymyr Mnih, Adria Puigdomencch Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforeement learning. In International Conference Machine Learning, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
Max Schwarzer, Ankesh Anand, Rishab Goel, R Devon Hjelm, Aaron Courville, and Philip Bach- man. Data-efficient reinforcement learning with momentum predictive representations. arXiv preprint arXiv:2007.05599, 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint
2018.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding. arXiv preprint arXiv:1807.053: 0073444777778 2018.
Greg Wayne, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack Rae, Piotr Mirowski, Joel Leibo, Adam Santoro, al. Unsupervised predictive memory goal-directed agent. arXiv preprint 1803: 2018.
Wilson Yan, Ashwin Vangipuram, Pieter Abbeel, and Lerrel Pinto. Learning predictive represen- tations for deformable objects using contrastive estimation. arXiv preprint UrX1v:2003.05436, 2020.
Sergey Zagoruyko and Nikos Komodakis. Paying more attention attention: Improving the perfor- mance of convolutional neural networks via attention transfer. arXiv preprint arXiv:16/2.0:3 2016.
12

---

Under review as a conference paper at ICLR 2021
A
APPENDIX
A.1 ALGORITHMS
Algorithm 1
Online RL with decoupled ATC encoder (steps distinct from end-to-end RL. in blue)
Require: OATC: 0x
& ATC model parameters (encoder ʃe thru contrast W), policy parameters
S 0
replay buffer of observations >> initialize momentum encoder (conv and linear only)
2: BATC f AATC
3: repeat
4: Sample environment and policy, through encoder:
5: for to m do
a minibatch
6:
7( (-)fa(s); a)), M(s,a,s)) (s}
7:
b store observations (delete oldest if full)
8:
8
9:
end for
10:
Update policy by given RL formula:
off-policy
11:
for to n do Ox - o + RL a, 0x)
D given number RL. updates per minihatch stop gradient into encoder
13: end for
Update encoder (and contrastive model) by ATC:
15:
for to p do
8, 8,84
b sample observations: anchors and positives
17:
OATC OATC - AATC VoreNoare Vo CATO (8, (s,84)
> ATC gradient update
OATC + (1-7)0AAC + update momentum encoder (conv and linear only)
19:
end for
20:
until converged
21:
return Encoder fo and policy T
A.2
ADDITIONAL FIGURES
allen
breakout
Trostbite
gravitar
800
2000
250
600
Pre-Training
7-garre
1000
100
200
200
400
B-garre wide
8-game
12.5
12.5
1-game
x10
x10
100
x10
Nona (rand)
pong
abert
seaquest
space invaders
15000
2000
10000
3000
1000
500
-20
12.5
25.0
12.
25.0
0.0
12.5
25.0
12.5
25.0
Agent Stepsx
x10
x10
x10'
Figure 11: RL using multi-task encoders (all with weights frozen) for eight Atari games gives mixed perfor- mance, partially improved by increased network capacity (8-game-wide). Training on 7 games and testing on the held-out one yields diminished but non-zero performance, showing some limited feature transfer between games.
In subpixel random shift, pixels linearly weighted average of the four nearest pixels to randomly chosen coordinate location. We used uniformly random horizontal and vertical shifts, and tested maximum displacements in (+) {0.1, 0.25, 0.5, 0.75, 1.0) pixels (with "edge" mode padding +1). We found 0.5 to work well in all tested domains, restoring the performance of raw image augmentation but climinating convolutions entirely from the RL training updates.
13

---

Under review as a conference paper at ICLR 2021
breakout
gravitar
invaders
dmlab lasertag
700
Random Shift Probability
200
600
800
15
1.0
500
10
0.0
100
300
200
0.0
12.5
25.0
12.5
25,0
0.0
12.5
25.0
12.5
25.0
Agent Steps x10"
x10*
K10"
x10"
Figure 12: Random shift augmentation helps in some Atari games and hurts in others, but applying with probability 0.1 is performant middle ground. DMLab henefits from random shift. (Offline pre-training.)
Ball-in-Cup: Catch
Cartpole: Swingup
Cheetah:: Run
Walker: Walk
800
500
800
Random Shift Latent 2.5.5 (sutptoet)
600
Dbservation 14
11 (ptee) None
400
200
200
200
0
0.0
2.5
5.0
Trv Steps
X10F
K10
X10
H107
Figure 13: Even after pre-training encoders for DMControl using random shift, RL requires augmentation our subpixel augmentation acts the (compressed) latent image, permitting its use in the replay buffer.
Stacked inputs
Random
RL-trained
UL (with shift) UL (without shift)
Figure 14: Attention map in BREAKOUT which shows the RL-trained encoder focusing game score, whereas UL ATC encoder focuses properly on the paddle and ball.
Stacked inputs
Random
Pixel Control
RL-trained
ATC (ours)
Figure 15: Attention map in LASERTAG. UL encoder with pixel control focuses on the score, while UL encoder with the proposed ATC focuses properly on the coin similar RL-trained encoder.
Stacked inputs
Random
Pixel Control
RL-trained
ATC (ours)
Attention map in the LASERTAG which shows that UI. encoders focus properly the enemy similar to RL-trained encoder.
Figure 16:
14

---

Under review as a conference paper at ICLR 2021
A.3 RL SETTINGS
Table 1: DMControl, RAD-SAC Hyperparameters.
HYPERPARAMETER
VALUE
OBSERVATION RENDERING
(84, 84), RGB
RANDOM SHIFT PAD
#4
REPLAY BUFFER SIZE
le5
INITIAL STEPS
led
STACKED FRAMES
ACTION REPEAT
(FINGER, WALKER)
(CARTPOLE)
(REST)
OPTIMIZER
ADAM
(81, 82) (fo, Te. Q+)
.99))
(81, 82) (a)
.999)
LEARNING RATE (fo, To, Qa)
2e-4 (CHEETAH)
1e-3 (REST)
LEARNING RATE (a)
1e-4
BATCH SIZE
512 (CHEETAH, PENDULUM)
256 (REST)
Q FUNCTION EMA 7
0.01
CRITIC TARGET UPDATE FREQ
CONVOLUTION FILTERS
32,
CONVOLUTION STRIDES [2, 1]
CONVOLUTION FILTER SIZE
ENCODER EMA 7
0.05
LATENT DIMENSION
50
HIDDEN UNITS (MLP)
[1024, 1024
DISCOUNT Y
.99
INITIAL TEMPERATURE
0.1
Table 2: Atari, PPO Hyperparameters.
HYPERPARAMETER
VALUE
OBSERVATION RENDERING
(84, 84), GREY
STACKED FRAMES
4
ACTION REPEAT
4
OPTIMIZER
ADAM
LEARNING RATE
2.5e-1
PARALLEL ENVIRONMENTS
16
SAMPLING INTERVAL
128
LIKELIHOOD RATIO CLIP,
0.1
PPO EPOCHS
4
PPO MINIBATCHES
4
CONVOLUTION FILTERS
32, 64, 64]
CONVOLUTION SIZES
8,4,3
CONVOLUTION STRIDES
4, 1]
HIDDEN UNITS (MLP)
[12]
DISCOUNT .99
GENERALIZED ADVANTAGE ESTIMATION À
0.95
LEARNING RATE ANNEALING
LINEAR
ENTROPY BONUS COEFFICIENT
0.01
EPISODIC LIVES
FALSE
REPEAT ACTION PROBABILITY
0.25
REWARD CLIPPING
+1
VALUE LOSS CORFFICIENT
15

---

Under review as a conference paper at ICLR 2021
Table 3: DMLab, PPO Hyperparameters.
HYPERPARAMETER
VALUE
OBSERVATION RENDERING
(72, 96), RGB
STACKED FRAMES
1
ACTION REPEAT
OPTIMIZER
ADAM
LEARNING RATE
2.5e-4
PARALLEL ENVIRONMENTS
16
SAMPLING INTERVAL
128
LIKELIHOOD RATIO CLIP,
PPO EPOCHS
PPO MINIBATCHES
2
CONVOLUTION FILTERS
[32, 64, 64, 64
CONVOLUTION FILTER SIZES
18, 3,3)
CONVOLUTION STRIDES
4, 2, 1,
HIDDEN UNITS (LSTM)
[56]]
SKIP CONNECTIONS
CONV 3, LSTM
DISCOUNT
.99
GENERALIZED ADVANTAGE ESTIMATION
0.97
LEARNING RATE ANNEALING
NONE
ENTROPY BONUS COEFFICIENT
0.01 (EXPLORE)
0.0003 (LASERTAG)
VALUE COEFFICIENT
0.5
16

---

Under review as a conference paper at ICLR 2021
A.4 ONLINE ATC SETTINGS
Table 4: Common ATC Hyperparameters.
HYPERPARAMETER
VALUE
RANDOM SHIFT PAD
t4
LEARNING RATE
1e-3
LEARNING RATE ANNEALING
COSINE
TARGET UPDATE INTERVAL
1
TARGET UPDATE
0.01
PREDICTOR HIDDEN SIZES, he
[12]
REPLAY BUFFER SIZE
le5
Table 5: DMControl ATC Hyperparameters.
HYPERPARAMETER
VALUE
RANDOM SHIFT PROBABILITY
BATCH SIZE
AS RL (INDIVIDUAL OBSERVATIONS)
TEMPORAL SHIFT, &
MIN AGENT STEPS UL MIN AGENT STEPS TO RL UL UPDATE SCHEDULE
le4
le4
AS RL
(2x CHEETAH)
LATENT SIZE
128
Table Atari ATC Hyperparameters.
HYPERFARAMETER
VALUE
RANDOM SHIFT PROBABILITY
0.1
BATCH SIZE
512 (32 TRAJECTORIES TIME STEPS)
TEMPORAL SHIFT, & 3
MIN AGENT STEPS TO UL
5e4
MIN AGENT STEPS TO RL
1e5
UL UPDATE SCHEDULE
ANNEALED QUADRATICALLY FROM PER SAMPLER ITERATION
(1el ONCE AT 1e5 STEPS FOR WEIGHT INITIALIZATION)
LATENT SIZE
256
Table 7: DMLab ATC Hyperparameters.
HYPERPARAMETER
VALUE
RANDOM SHIFT PROBABILITY
1
BATCH SIZE
512 (INDIVIDUAL OBSERVATIONS)
TEMPORAL SHIFT, 3
MIN AGENT STEPS TO
MIN AGENT STEPS TO RL
1e5
UL UPDATE SCHEDULE
2 PER SAMPLER ITERATION
LATENT SIZE
256
17

---

Under review as a conference paper at ICLR 2021
A.5 OFFLINE PRE-TRAINING DETAILS
We conducted coarse hyperparameter sweeps to tune each competing UL algorithm. In all cases, the best setting is the one shown in comparisons.
When our VAFs include a time difference between input and reconstruction observations, we include one hidden layer with action additionally input between the encoder and decoder. We tried both 1.0 and 0.1 KL-divergence weight in the VAE loss, and found 0.1 to perform better in both DMControl and Atari.
DMControl For the VAE, we experimented with ( and time step difference between input and reconstruction target observations and training for either le4 or 5e4 updates. The best settings were 1-step temporal, and 5e4 updates, with batch size 128. ATC used 1-step temporal, 5e4 updates (although this can be significantly decreased), and batch size 256 (including CHEETAH). The pre- training data set consisted of the first 5e4 transitions from RAD-SAC agent learning each task, including 5c3 random actions. Within this span, CARTPOLE and BALLJIN _CUP learned completely, but WALKER and CHEETAH reached average returns of 514 and 630, respectively (collected without the compressive convolution).
DMLab For Pixel Control, we used the settings from Hessel et al. (0199 (see the appendix therein), except we used only empirical returns, computed offline (without bootstrapping). For CPC, we tried training batch shapes, batch x time (64, 8), (32, 16), (16, 32), and found the setting with rollouts of length 16 to be best. We contrasted all elements of the batch against each other, rather than only forward constrasts. In all cases we also used 16 steps to warmup the LSTM. For all al- gorithms tried learning rates 3e-4 and 1e-3 and both 5e4 and 1.5e5 updates. For ATC and CPC, the lower learning rate and higher number of updates helped in LASERTAG especially. The pretraining data 125e3 samples from partially trained RL agents receiving average returns of 127 and 6 in LOPPORE_GGOALLLLCAATTIII SMALL and LASERTAG. THRREE_OPPONEEEE SMMALL, respectively.
Atari For the VAE, we experimented with O, 1, and 3 time step difference between input and reconstruction target, and found to work best. For ST-DIM we experimented with 1, 3, and time steps differences, and batch sizes from 64 to 256, learning rates 1e-3 and 5e-4. Likewise, 3-step delay worked best. For the inverse model, we tried and 3-step predictions, with I-step working better overall, and found random shift augmentation to help. For pixel control, used the settings in Jaderberg et al.] (0017), again with full empirical returns. We ran each algorithm for up to 1e5 updates, although final ATC results used 5e4 updates. We ran each RL agent with and without observation normalization on the latent image and observed no difference in performance. Pretraining data was 125e3 samples sourced from the replay buffer of DQN agents trained for 15e6 steps with epsilon- epsilon-greedy e 0.1. Evaluation scores were:
Table 8: Atari Pre-Training Data Source Agents.
GAME
EVALUATION SCORE
ALIEN
1,800
BREAKOUT
279
FROSTBITE
1, 1,400
GRAVITAR
390
PONG
QBERT
8, 800
SEAQUEST
11,000
SPACE INVADERS
1,200
18


Reviews:
[{"review_id": "e6miFpBSNCB", "reviewer": "AnonReviewer1", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary of the paper:\nThe paper aims to define an unsupervised pretraining architecture that can be used to pretrain representations for a reinforcement learning agent. The proposed solution is called \"Augmented Temporal Contrast\" (ATC). It consists of an encoder-compressor-predictor architecture for an (augmented) observation that is trained in a small latent space by minimizing the InfoNCE loss between that prediction, and the encoding of some future state.\nThe representation used (the pretrained encoder) is then frozen, and a policy network is trained consuming these representations. The authors then evaluate their approach on several different tasks: a control domain (DMcontrol), Mazenagivation (DMLab) and Atari. They show that their approach - without finetuning the representations - achieves comparable results to end-to-end reinforcement learning on most of the different domains. They also show that adding their defined loss as a regularizer always helps learning. They finally show that their representations can generalize out-of-domain by running several multi-task experiments, while only having pretrained on one domain.\n\nCommentary on the goal of the paper:\n\nThe goal of the paper is extremely important: separating learning representations from learning policies would enable better transfer, possible more sample efficiency and lower variance in outcome. \n\nStrengths:\n- The authors propose a well-designed solution that combines existing approaches in a well thought-out way. \n- The paper is extremely well written\n- The paper has an extensive empirical section. \n- Results are generally very good.\n\nWeaknesses:\n- The goal of the paper is a bit vague. As I said above, I agree and understand the desire for the separation of representation learning and reinforcement learning. However, the paper would have been stronger if it had concentrated on a single benefit of this separation, and evaluated their approach on that. While they say that their goal is to investigate \"how to learn representations which are agnostic to rewards\", this is too general as well. (One example would be to say that the decoupling makes for better generalizes to new MDPs - but this is not the focus of the analysis, just an aspect. The paper is, unfortunately, less convincing for it.)\n- Along a similar line of thought, the results, while strong, are not as convincing as they could be, because the paper does not focus on the benefits of reward-agnostic representations learned by ATC. The results that consider the generalization advantage (the Multi-Task learning experiments) are weaker, without the paper offering an analysis as to why.\n\nIn total, I would argue that this is a well-written paper with interesting analysis that could be a lot stronger by narrowing the scope of the contained argument. I argue for rejection, because I can see an updated version of this paper to be a great paper.", "extra_scores": {}, "extra_text": {}}, {"review_id": "dYg8nU0sOFm", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "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 an unsupervised task called Augmented Temporal Contrast which associates pairs of observations separated in time using a contrastive loss. The paper uses the task in several training regimes (in online RL, pretraining and multi-task RL).\n\nPros:\n- Well written and structured paper.\n- Interesting, general and simple to implement task.\n- Evaluated in several training regimes and on several environments.\n- Improves sample efficiency on most of the environments and setups, and improves over prior methods.\n- The attention maps in the paper and appendix are great and although they may be hand picked(?) examples, it highlights the issue with many approaches that can't model a goal rarely seen.\n\nCons:\n- You write you use multiple seeds but I don't see anywhere details on this? Consider adding it to the tables in appendix.\n- In Figure 3 one agent step is 4x environment frames? I suggest to make it clear in plot or in caption. In Figure 2, environment frames is used.\n\nComments/questions:\n- Wrt. 4.1, to what extent is the \"small\" replay for DMLab necessary over just using observations in batches/unrolls? Looking at Table 3 I can't see how large the replay is? 10k as in SAC or smaller?\n\nUpdate: Not all the experiments are particular thorough and the novelty less than expected.", "extra_scores": {}, "extra_text": {}}, {"review_id": "AYBhxVXFT3", "reviewer": "AnonReviewer2", "summary": "", "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": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes an auxiliary task for learning better representations for reinforcement learning. The idea is interesting and is a very active area of research at the moment.\n\nMy main concern is that the paper is mostly an empirical evaluation, as the novel algorithm is mostly an extra contrastive loss. Further, it seems their method is almost exclusively meant for pixel-based environments (see point 1 below), so the authors should make this point more explicit. Given this empirical emphasis, I feel the authors could have performed a deeper exploration into understanding _why_ their proposed algorithm performs the way it does in the different environments considered. In particular, some very specific design decisions were made in evaluation (see, for instance, points 2, 5, 9, 11, 12, 13, 14, 15, 16 below). \nThe clarity of exposition could also use some improvement, as I detail below.\n\nMain questions/concerns:\n1. In Section 3, the authors write \"This task encourages the learned encoder to extract meaningful elements of the structure of the MDP from observations.\". This assumes some type of continuity in pixel space, relative to MDP dynamics, which is in general not true.\n2. In Figure 1, why doesn't the momentum encoder go through a residual predictor?\n3. It seems $\\theta$, $\\phi$, $\\bar{\\theta}$, and $\\bar{\\phi}$ are all updated independently, is this the case? What is the actual training regimen? Are the ATC and regular RL networks trained concurrently?\n4. It would help if you include a proper algorithm in the paper.\n5. Above equation (2), the authors say \"In our implementation, the positives from all other elements...\". It's not clear what \"positives from  all other elements\" means.\n6. In section 4.1 the authors say \"multiple seeds were run\", please specify how many.\n7. In section 4.2 the authors say they are capable of \"training the encoder online, fully detached from the RL agent\", but how is it fully detached if they share the conv layers?\n8. It's not clear what the difference between ATC and UL training is. In some experiments the authors use ATC, in others UL. Are they the same thing? For example, in Figure 3, which ones are ATC? Also in Figure 14 vs Figures 15 and 16?\n9. In Figure 3, why does one environment compare with pri and the other with 2x, but not both in both environments?\n10. In the **Atari** subsection on the comment of detached training, the authors point out subpar performance on Breakout and SpaceInvaders. On SpaceInvaders it's possible the screen changes could cause issues, but what do the authors think cause the subpar performance in Breakout?\n11. In point (iii) of section 4.3, what's the network architecture used for training the RL part?\n12. In section 4.3 the authrs say they \"drew expert demonstrations from partially-trained RL agents\". Were these all drawn from the same checkpoint?\n13. In the **DMControl** section, the VAE is trying to reproduce a frame $T$ steps in the future? What is the value of $T$ used? Did you try different values?\n14. Similar questin for the **Atari** subsection. Also for this section, does your VAE try to predict individual frames or stacked frames (as frame stacking is common in Atari experiments)?\n15. In Figure 8 top, are these after pre-training the encoder? If that is the case, regular RL would have used fewer frames in comparison, no? Where would RL be if left to train for longer?\n16. In section 4.5 please clarify what \"random shift augmentations\" are.\n17. In the **Encoder analysis** subsection, what do you mean by \"attention\"?\n18. In Figures 14, 15, and 16 it's not at all clear what we're supposed to be looking for, nor how they show that ATC/UL is focusing on the score/enemy and the others are not.\n\nMinor comments:\n1. At the bottom of page 2, the term \"POMDP\" has not been introduced yet.", "extra_scores": {}, "extra_text": {}}, {"review_id": "67WeylJ5OSx", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "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": "**Summary:**\nThis paper presents a new unsupervised learning method for learning latent representations for visual RL control domains. The method, Augmented Temporal Contrast (ATC), can be used alone to learn a representation to be combined with an RL algorithm, or as an auxiliary task in an end-to-end system. ATC matches or outperforms comparable end-to-end systems in several environments. The paper provides an extensive experimental study to support its claims.\n\n**Strengths:**\nThe paper is clearly written, and all of the main points are well articulated. ATC appears sufficiently novel, and is applicable to a wide variety of domains, and can be deployed in various configurations (e.g. auxiliary task, unsupervised pre-training, etc…). Included is a thorough experimental study that effectively demonstrates the performance of the method.\n\n**Weaknesses:**\nAlthough the experimental study seem thorough. I could not find the actual number of independent runs (seeds) for each domain listed anywhere. This information should be included so that the reader can better evaluate the variance of each method, and make more confident conclusions. \n\n**Recommendation:**\nOverall I vote to accept. The method presented in the paper is not revolutionary, but it appears to be novel and significant enough to be of interest to deep RL practitioners.\n\n**Questions:**\nHow many independent runs (seeds) were used in each of the domains? Can this information be included in the main text?\n\n**After Author Response and Discussion:**\nThanks to the authors for their responses. After reading the other reviews and the author responses, I am lowering my score to 5. I think that the number of independent runs used (especially on the smaller domains), and the way the results are presented with the min-max extent makes me less convinced of the results than I was in the initial review. Adding many more independent runs (seeds), especially on the smaller domains, would improve my confidence a lot. Overall I think the paper is of interest to the community, but the experiments and their analysis could be improved.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces ATC, which is a contrastive learning on observations separated in time, to learn representations that do not need to take rewards into consideration. These learned representations allow, for the first time, a real disentanglement between representation learning and control, as the agent can simply load such a representation, “freeze it”, and still recover performance of end-to-end deep reinforcement learning agents.

Overall, all reviewers agree this is a promising direction. Nevertheless, there has been extensive discussion (with the authors and privately) about the significance of the reported results due to the small number of seeds. On one hand, there’s the argument that there is a wide range of experiments and that should compensate for a small number of seeds in individual experiments. On the other hand, there are experiments with as little as two seeds (e.g., DMControl multi-env) and this can be seen at most as anecdotal evidence. There’s also the argument that we, as a community, should be striving for more reliable and meaningful experiments in reinforcement learning. Moreover, there have been concerns about how “variance” is being reported (max and min performance) and, although the authors replied to that, an alternative plotting was never shown.

Importantly, at this point it is not clear how many seeds were used in each experiment (Figures 6, 7, 9, 11, 12, 13 do not report the number of seeds used). It is said that each curve represents a minimum of 3 random seeds, but that is very informal and not that useful. Exactly stating the number of seeds would be the right thing to do, not to mention that in the rebuttal it is said that 8-game pretraining for Atari multi-env uses 2 seeds, contradicting the original claim. Also, sometimes, different methods, in the same experiment, are  “averaged” across different numbers of seeds (“DMLab offline -- ATC is 4 seeds, PC and CPC are 2 seeds each”). This is particularly problematic because of the small number of seeds and potentially high variance. Reporting the max over 4 numbers drawn from a Gaussian distribution is very likely to lead to a larger number than when reporting the max over 2 numbers drawn from the same Gaussian distribution. 

I do acknowledge the effort to increase the number of seeds during the rebuttal phase, but it is hard to accept a paper with unknown results. We have very little evidence to believe that going from 2 seeds to 5 seeds is not going to change the results. The reviewers couldn’t agree on the variance of this process as well. Some say the variance of PPO is low between runs when using the same hyper parameters while others mention papers (e.g., Deep RL that matters) that show how much variance one can have across these methods. Thus, I cannot accept this paper conditioned on more seeds being added to the final version because we don’t know what the results will look like. Since this paper is mostly an empirical study, it should have thorough experiments and a careful analysis of the results, but the small number of seeds prevents that in my opinion. Thus, as difficult as it is given the promising direction of the paper, I’m recommending its rejection. I strongly encourage the authors to increase the number of runs in their experiments and to use a more standard measure of variability (e.g., standard error, standard deviation) when reporting their results. This will then be a very strong submission for a future conference.

Decision: Reject