D ECOUPLING R EPRESENTATION L EARNING FROM R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
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 ﬁnd 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 hiddenurl .
I NTRODUCTION
Ever since the ﬁrst fully-learned approach succeeded at playing Atari games from screen images (Mnih et 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 et al., 2018; Schulman et al., 2017; Mnih et al., 2016; Haarnoja et al., 2018) and have been successfully applied to domains ranging from real-world (Levine et al., 2016;
Kalashnikov et al., 2018) and simulated robotics (Lee et al., 2019;
Laskin et al., 2020a; Hafner et al., 2020) to sophisticated video games (Berner et al., 2019; Jaderberg et al., 2019), and even high-ﬁdelity driving simulators (Dosovitskiy et al., 2017).
While the simplicity of end-to-end methods is appealing, relying on the reward function to learn visual features can be severely limiting.
For example, it leaves features difﬁcult to acquire under sparse rewards, and it can narrow their utility to a single task.
Although our intent is 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 to rewards, without degrading the control policy.
A number of recent works have signiﬁcantly improved RL performance by introducing auxiliary losses, which are unsupervised tasks that provide feature-learning 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, in the ﬁeld 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 classiﬁcation.
Together, these advances suggest that visual features for RL could possibly be learned entirely without rewards, which would grant greater ﬂexibility 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.
In this paper, we demonstrate the ﬁrst 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 a control policy independently, on the (compressed) latent images.
This capability stands in contrast to previous state-of-the-art methods, which have trained the UL and RL objectives jointly, or Laskin et al.
(2020b), which observed diminished performance with decoupled encoders.
Our main enabling contribution is a new unsupervised task tailored to reinforcement learning, which we call Augmented Temporal Contrast (ATC).
ATC requires a 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 a small latent space, where the
InfoNCE loss is applied (van den Oord et al., 2018).
Within each randomly sampled training batch, the positive observation, o t + k , for every anchor, o t , 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 a 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 a 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 a visually diverse range of common RL environments: DeepMind Control Suite (DMControl; Tassa et al. 2018), DeepMind Lab (DMLab; Beattie et al. 2016), and Atari (Bellemare et al., 2013).
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 .
Our empirical contributions are summarized as follows:
Online RL with UL :
We ﬁnd 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 expert demonstrations, and evaluate it by training an RL agent using the encoder with weights frozen.
We ﬁnd that ATC matches or outperforms all prior UL algorithms as tested across all domains, demonstrating that ATC is a state-of-the-art UL algorithm for RL.
Multi-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 or better than end-to-end RL in four held-out environments.
Similar attempts to generalize across eight diverse Atari games result in mixed performance, conﬁrming 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 beneﬁts.
R ELATED W ORK
Several recent works have used unsupervised/self-supervised representation learning methods to improve performance in RL.
The UNREAL agent (Jaderberg et al., 2017) introduced unsupervised auxiliary tasks to deep RL, including the Pixel Control task, a Q-learning method requiring predictions 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) surpassed 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 it proves sufﬁcient in partially observed settings (POMDPs).
ST-DIM (Anand et al., 2019) introduced various temporal, contrastive losses, including ones that operate on “local” features from an intermediate layer within the encoder, without data augmentation.
CURL (Laskin et al., 2020b) introduced an augmented, contrastive auxiliary task similar to ours, including a momentum encoder but without temporal contrast.
Mazoure et al.
(2020) provided extensive analysis pertaining to InfoNCE losses on functions of successive time steps in MDPs, including local features in their auxiliary loss (DRIML) similar to ST-DIM, and ﬁnally conducted experiments using global temporal contrast of augmented observations in the Procgen (Cobbe et al., 2019) environment.
Most recently, MPR (Schwarzer et al., 2020) combined data augmentation with multi-step, convolutional forward modeling and a similarity loss to improve DQN agents in the Atari 100k benchmark.
Hafner et al.
(2019; 2020); Lee et al.
(2019) proposed to leverage world-modeling in a latent-space for continuous control.
A small number of 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 effective in as diverse a collection of RL environments as ours (Bellemare et al., 2013; Tassa et al., 2018;
Beattie et al., 2016).
Finn et al.
; Ha & Schmidhuber (2018) are example works which pretrained encoder features in advance using image reconstruction losses such as the VAE (Kingma & Welling, 2013).
Devin et al.
; Kipf et al.
(2019) pretrained object-centric representations, the latter learning a forward 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, detached from the RL agent, which learned read-only accesses to memory.
It used reconstruction and one-step latent-prediction losses and achieved high performance in DMLab-like environments with extreme partial observability.
Our loss function may beneﬁt those settings, as it outperforms similar reconstruction losses in our experiments.
Decoupling unsupervised pretraining from downstream tasks is common in computer vision (H ´ enaff et 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 taskspeciﬁc networks, yielding signiﬁcant gains in computational efﬁciency over end-to-end methods.
A UGMENTED T EMPORAL C ONTRAST
Our unsupervised learning task, Augmented Temporal
Contrast (ATC), requires a model to associate an observation, o t , with one from a speciﬁed, near-future time step, o t + k .
Within each training batch, we apply stochastic data augmentation to the observations (Laskin et al., 2020b), namely random shift (Kostrikov et al., 2020), which is simple to implement and provides highly effective regularization in most cases.
The augmented observations are encoded into a small latent space where a contrastive loss is applied.
This task encourages the learned encoder to extract meaningful elements of the structure of the MDP from observations.
Our architecture for ATC consists of four learned components - (i) a convolutional encoder , f θ , which processes the anchor observation, o t , into the latent image z t = f θ ( AUG ( o t )) , (ii) a linear global compressor , g φ to produce a small latent code vector c t =
g φ ( z t ) , (iii) a residual predictor MLP, h ψ , which acts as an implicit forward model to advance the code p t = h ψ ( c t ) + c t , and (iv) a contrastive transformation matrix, W .
To process the positive observation, o t + k into the target code ¯ c t + k =
g ¯ φ ( f ¯ θ ( AUG ( o t + k )) , we use a momentum encoder (He et al., 2019) parameterized as a slowly moving average of the weights from the learned encoder and compressor layer:
[Equation]
The complete architecture is shown in Figure 1.
The convolutional encoder, f θ , alone is shared with the RL agent.
We employ the InfoNCE loss (Gutmann & Hyv¨arinen, 2010; van den Oord et al., 2018) using logits computed bilinearly, as l = p t W ¯ c t + k .
In our implementation, every anchor in the training batch utilizes the positives corresponding to all other anchors as its negative examples.
Denoting an observation indexed from dataset O as o i , and its positive as o i + , the logits can be written as l i,j + = p
i W ¯ c j + ; our loss function in practice is:
[Equation]
E XPERIMENTS 4. E VALUATION E NVIRONMENTS AND A LGORITHMS
We evaluate ATC on three standard, visually diverse RL benchmarks - the DeepMind control suite (DMControl; Tassa et al. 2018), Atari games in the Arcade Learning Environment (Bellemare et al., 2013), and DeepMind Lab (DMLab; Beattie et al. 2016).
Atari requires discrete control in arcadestyle 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 RL algorithms.
For DMControl, we use RADSAC (Laskin et al., 2020a; Haarnoja et al., 2018) with the augmentation of Kostrikov et al.
(2020), which randomly shifts the image in each coordinate (by up to 4 pixels), replicating edge pixel values as necessary to restore the original image size.
A difference from prior work is that we use more downsampling in our convolutional network, by using strides (2 , 2 , 2 , 1) instead of (2 , 1 , 1 , 1) to reduce the convolution output image by 25x.
For both Atari and DMLab, we use PPO (Schulman 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, LSTM agents receiving only a single time-step image input, the four-layer convolution encoder from Jaderberg et al.
, and we tuned the entropy bonus for each level.
In the online setting, the ATC loss is trained using small replay buffer of recent experiences.
We include all our own baselines for fair comparison and provide complete settings in an appendix.
Unless otherwise noted, each curve represents a 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 at each checkpoint.
4. O NLINE RL WITH ATC
DMControl
In the online setting, we found ATC to 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 2.
In C ARTPOLE -S WINGUPS PARSE , where rewards are only received once the pole reaches vertical, ATC training enabled the agent to master the task signiﬁcantly faster.
The encoder is trained with one update for every RL update to the policy, using the same batch size, except in C HEETAH -R UN , which required twice the
ATC updates.
DMLab We experimented with two kinds of levels in DMLab: E XPLORE G OAL L OCATIONS , which requires repeatedly navigating a maze whose layout is randomized every episode, and
L ASERTAG T HREE O PPONENTS , which requires fast reﬂexes to pursue and tag enemies at a distance.
We found ATC capable of training fully detached encoders while achieving equal or better performance than end-to-end RL.
Results are shown in Figure 3.
Both environments exhibit spar-
[Figure]
sity which is greater in the “large” version than the “small” version, which our algorithm addresses, discussed next.
In E XPLORE , the goal object is rarely seen, especially early on, making its appearance difﬁcult to learn.
We therefore introduced prioritized sampling for ATC , with priorities corresponding to empirical absolute returns: p ∝ 1 + R abs , where R abs = (cid:80) nt =
γ t | r t | , to train more frequently on more informative scenes.
2 Whereas uniform-ATC performs slightly below RL, uniform-ATC outperforms RL and nearly matches using ATC (uniform) as an auxiliary task.
By considering the encoder as a stand-alone feature extractor separate from the policy, no importance sampling correction is required.
In L ASERTAG , 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 updates 3 further improved the score to match the ATC-auxiliary agent, showing ﬂexibility to address the representation-learning bottleneck when opponents are dispersed
[Figure]
Atari We tested a diverse subset of eight Atari games, shown in Figure 4.
We found detachedencoder training to work as well as end-to-end RL in ﬁve games
[Figure]
but performance suffered in
B REAKOUT and S PACE I NVADERS in particular.
Using ATC as an auxiliary task, however, improves performance in these games and others.
We found it helpful to anneal the amount of UL training over the course of RL in Atari (details in an appendix).
Notably, we found several games, including S PACE I NVADERS , to beneﬁt 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
[Figure]
4. E NCODER P RE -T RAINING B ENCHMARKS
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
XPLORE G OAL L OCATIONS
Since the ATC batch size was 512 but the RL batch size was 1024, performing twice as many UL updates still only consumed the same amount of encoder training data as RL.
We did not ﬁne-tune for batch size.
[Figure]
e.g. Chen et al.
(2020); Grill et al.
(2020)): (i) collect a data set composed of expert demonstrations from each environment; (ii) pre-train the CNN encoder with that data ofﬂine using UL; (iii) evaluate by using RL to learn a control policy while keeping the encoder weights frozen.
This procedure isolates the asymptotic performance of each UL algorithm for RL.
For convenience, we drew expert demonstrations from partially-trained RL agents, and every UL algorithm trained on the same data set 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 et al., 2020b), 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 a time delay to the target observation (VAE-T).
We found ATC to match or outperform the other algorithms, in all four test environments, as shown in Figure 5.
Further, ATC is the only one to match or outperform the reference end-to-end RL across all cases
[Figure] [Figure]
DMLab We compare against both Pixel Control (Jaderberg et al., 2017; Hessel et al., 2019) and
CPC (van den Oord et al., 2018), which have been shown to bring strong beneﬁts in DMLab.
While all algorithms perform similarly well in E XPLORE ,
ATC performs signiﬁcantly better in L ASERTAG ,
Figure 6.
Our algorithm is simpler than Pixel Control and CPC in the sense that it uses neither actions, deconvolution, nor recurrence.
Atari We compare against Pixel Control, VAE-T, and a basic inverse model which predicts actions between pairs of observations.
We also compare against Spatio-Temporal Deep InfoMax (ST-DIM), which uses temporal contrastive losses with “local” 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 is the only UL algorithm to match the end-to-end
RL reference in G RAVITAR and B REAKOUT , and it performs best in S PACE I NVADERS .
[Figure]
4. M ULTI -T ASK E NCODERS
In the ofﬂine setting, we conducted initial explorations into the capability of ATC to learn multi-task encoders, simply by pre-training on demonstrations from multiple environments.
We evaluate the encoder by using it, with frozen weights, in separate RL agents learning each downstream task.
DMControl Figure 8 shows our results in DMControl, where we pretrained using only the four environments in the top row.
Although the encoder was never trained on the H OPPER , P ENDULUM , nor F INGER domains, the multi-task encoder supports efﬁcient RL in them.
P ENDULUM -S WINGUP and C ARTPOLE -S WINGUP -S PARSE stand out as challenging environments which beneﬁted from cross-domain and cross-task pre-training, respectively.
The pretraining was remarkably efﬁcient, requiring only 20,000 updates to the encoder
[Figure]
Atari Atari proved a more challenging domain for learning multi-task encoders.
Learning all eight games together in Figure 11, 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 ﬁlters per layer, indicating that representation capacity is a limiting factor.
To test generalization [Figure] we conducted a seven-game pre-training experiment where we test the encoder on the held-out game.
Most games suffered diminished performance (although still perform signiﬁcantly higher than a frozen random encoder), conﬁrming the limited extent to which visual features transfer across these games
[Figure] [Figure]
[Figure]
4. A BLATIONS AND E NCODER A NALYSIS
Random Shift in ATC
In ofﬂine experiments, we discovered random shift augmentations to be helpful in all domains.
To our knowledge, this is the ﬁrst application of random shift to 3D visual environments as in DMLab.
In Atari we found performance in G RAVITAR to suffer from random shift but reducing the probability of applying random shift to each observation from 1.0 to 0.
alleviated the effect while still bringing beneﬁts in other games, so we used this setting in our main experiments.
Results are shown in Figure 12 in an appendix
Random Shift in RL In DMControl, we found the best results when using random shift during
RL, even when training with a frozen encoder.
This is evidence that the augmentation regularizes not only the representation but also the policy, which ﬁrst processes the latent image into a 50dimensional vector.
To unlock computation and memory beneﬁts of replaying only the latent images for the RL agent we attempted to apply data augmentation to the latent image.
But we found 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 Figurein the appendix, this augmentation restores performance when applied to the latent images, allowing a pre-trained encoder to be entirely bypassed during policy training updates.
Temporal Contrast on Sequences In B REAKOUT alone, we discovered that composing the UL training batch of trajectory segments, rather than individual transitions, gave a signiﬁcant beneﬁt.
Treating all elements of the training batch independently provides “hard” negatives, since the encoder must distinguish between neighboring time steps.
This setting had no effect in the other Atari games tested, and we found equal or better performance using individual transitions in DMControl and DMLab.
Figure 9 further shows that using a similarity loss (Grill et al., 2020) does not capture the beneﬁt.
Encoder Analysis We analyzed the learned encoders in B REAKOUT to further study this ablation effect.
Similar to Zagoruyko & Komodakis (2016), we compute spatial attention maps by mean-pooling the absolute values of the activations along the channel dimension and follow with a 2-dimensional spatial softmax.
Figure 10 shows the attention of four different encoders on the displayed scene.
The poorly performing UL 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 are less relevant for control.)
In an appendix, we include other example encoder analyses from Atari and DMLab 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.
C ONCLUSION
Reward-free representation learning from images provides ﬂexibility and insights for improving deep RL agents.
We have shown a broad range of cases where our new unsupervised learning algorithm can fully replace RL for training convolutional encoders while maintaining or improving online performance.
In a small number of environments–a few of the Atari games–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, our 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 conﬂict, as in Atari, or if they otherwise impact our loss function unequally.
On the theoretical side, it may be helpful to analyze the effects of domain shift on the policy when a detached representation is learned online.
One obvious application of our ofﬂine methodology would be in the batch RL setting, where the agent learns from a ﬁxed data set.
Our ofﬂine experiments showed that a relatively small number of transitions are sufﬁcient 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 for reinforcement learning.
and policy π φ A. A DDITIONAL F IGURES
[Figure]
In subpixel random shift
[Figure]
new pixels are a linearly weighted average of the four nearest pixels to a randomly chosen coordinate location.
We used uniformly random horizontal and vertical shifts
[Figure]
and tested maximum displacements in ( ± ) { 0 . 1 , 0 . 25 , 0 . 5 , 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 eliminating convolutions entirely from the RL training updates
[Figure] [Figure]
[Figure] [Figure]
[Figure]
RL-trained encoder
[Figure]
A. RL S ETTINGS
[Table] [Table]
[Table] [Table]
[Table]
A. O NLINE ATC S ETTINGS
[Table]
[Table] [Table]
[Table]
A. O FFLINE P RE -T RAINING D ETAILS
We conducted coarse hyperparameter sweeps to tune each competing UL algorithm.
In all cases, the best setting is the one shown in our comparisons.
When our VAEs 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. 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 0 and 1 time step difference between input and reconstruction target observations and training for either 1e4 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 signiﬁcantly decreased), and batch size 256 (including C HEETAH ).
The pretraining data set consisted of the ﬁrst 5e4 transitions from a RAD-SAC agent learning each task, including 5e3 random actions.
Within this span, C ARTPOLE and B ALL IN C UP learned completely, but W ALKER and C HEETAH 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.
(see the appendix therein), except we used only empirical returns, computed ofﬂine (without bootstrapping).
For CPC, we tried training batch shapes, batch × time in (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 algorithms we 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 L ASERTAG especially.
The pretraining data was 125e3 samples from partially trained RL agents receiving average returns of 127 and 6 in E XPLORE G OAL L OCATIONS S MALL and L ASERTAG T HREE
O PPONENTS S MALL , respectively.
Atari For the VAE, we experimented with 0, 1, and 3 time step difference between input and reconstruction target, and found 3 to work best.
For ST-DIM we experimented with 1, 3, andtime 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 1- and 3-step predictions, with 1-step working better overall, and found random shift augmentation to help.
For pixel control, we used the settings in Jaderberg et al.
(2017), again with full empirical returns.
We ran each algorithm for up to 1e5 updates, although ﬁnal 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 15e steps with epsilon-greedy (cid:15)
= 0 .
Evaluation scores were:
[Table]