Pretrained Encoders are All You Need
Abstract
Data-efﬁciency and generalization are key challenges in deep learning and deep reinforcement learning as many models are trained on largescale, domain-speciﬁc, and expensive-to-label datasets.
Self-supervised models trained on largescale uncurated datasets have shown successful transfer to diverse settings.
We investigate using pretrained image representations and spatiotemporal attention for state representation learning in Atari.
We also explore ﬁne-tuning pretrained representations with self-supervised techniques, i.e., contrastive predictive coding, spatiotemporal contrastive learning, and augmentations.
Our results show that pretrained representations are at par with state-of-the-art self-supervised methods trained on domain-speciﬁc data.
Pretrained representations, thus, yield data and compute-efﬁcient state representations.
https: //github.com/PAL-ML/
PEARL_v
1. Introduction
Data-efﬁciency and generalization are key challenges in deep learning (DL) and deep reinforcement learning (RL), especially for real-world deployment and scalability.
Selfsupervised learning (SSL) has been used in computer vision and natural language processing (Chen et al., 2020c;
Radford et al., 2021; Henaff, 2020;
He et al., 2020; Devlin et al., 2019; Radford et al.) to learn from large-scale unlabeled/uncurated data and do a few-shot transfer to labeled data.
Deep RL has also leveraged self-supervised learning for data-efﬁciency (Srinivas et al., 2020; Laskin et al., 2020).
RL in pixel space is also sample-inefﬁcient, and state representations can help sample-efﬁcient, robust and generalizable RL Lake et al., Kaiser et al., 2019; Tassa et al., 2018).
Previous work investigated self-supervised state representation learning in RL, but involved training on large-scale domain-speciﬁc data (Anand et al., 2019).
We investigate how pre-trained models can be leveraged for state representation learning in Atari.
In particular, we focus on using pretrained image representation learning and attention models.
We also investigate self-supervised ﬁnetuning of pretrained representations using state-of-the-art (SotA) self-supervised methods.
Our results show that our pretrained representations, particularly ‘zoomed in’ representations, perform as well as the SotA self-supervised state representation learning models trained on large domain-speciﬁc datasets.
Adding pretrained temporal attention does not particularly improve performance, possibly because the pretrained image representations already leverage attention.
Adding self-supervised ﬁne-tuning on domain-speciﬁc also does not particularly improve performance, especially in a data-constrained setting.
We make 3 key contributions: i. a new methodology for using pretrained image representations for sample-efﬁcient and generalizable state representation learning in RL; ii.
evaluations of using pretrained temporal attention models with static image representation for temporal data; iii.
evaluating self-supervised ﬁne-tuning of pretrained representations using SotA SSL techniques on domain-speciﬁc data.
2. Related Work
Our work lies at the intersection of 5 key areas:
SelfSupervised Learning (SSL), Self-Supervised RL, Attention,
Domain generalization, and Pretraining.
Unlike previous work in Self-Supervised RL, Domain generalization, and
Pretraining, we focus on state representation learning in RL, not on the RL.
Unlike previous work in state representation learning for RL (Anand et al., 2019), we leverage pretrained models, not trained on domain-speciﬁc data, demonstrating data-efﬁcient and generalizable learning.
Finally, like previous work, we leverage spatial-temporal attention, particularly optical ﬂow (Yuezhang et al., 2018), and also use SotA
SSL, but only for ﬁne-tuning pretrained representations.
Our results show that good pretrained embeddings perform competitively, without necessarily needing augmentations from spatio-temporal attention or self-supervised ﬁne-tuning.
Self-Supervised Learning (SSL): SSL has been used in natural language processing and computer vision (Devlin et al., 2018; Henaff, 2020;
He et al., 2020; Chen et al., 2020a) using both contrastive ((Oord et al., 2018; Bachman et al., 2019;
He et al., 2020; Chen et al., 2020b; Radford et al., 2021) and purely predictive methods (Grill et al., 2020).
SSL on large-scale uncurated datasets has shown promising few-shot transfer to diverse labeled data (Chen et al., 2020c;
Radford et al., 2021).
We use pretrained selfsupervised models for sample-efﬁcient state representation learning in RL.
Also, we use self-supervised ﬁne-tuning of pretrained representations on domain-speciﬁc data.
Self-Supervised RL: SSL has been used in model-based (Kaiser et al., 2019;
Ha & Schmidhuber, 2018; Schrittwieser et al., 2020; Hafner et al., 2019) and model-free RL (Jaderberg et al., 2016; Shelhamer et al., 2016; Oord et al., 2018).
Learning is either reconstruction-based (Jaderberg et al., 2016;
Higgins et al., 2017; Yarats et al., 2019) or constrastive (Sermanet et al., 2018; Warde-Farley et al., 2018;
Anand et al., 2019; Oord et al., 2018; Srinivas et al., 2020), and the predictions are in pixel space (Jaderberg et al., 2016) or latent space (Oord et al., 2018; Guo et al., 2020).
Image augmentations (Srinivas et al., 2020; Laskin et al., 2020), spatiotemporal structures (Sermanet et al., 2018; Aytar et al., 2018;
Oord et al., 2018; Anand et al., 2019), and scene or object representations (Burgess et al., 2019; Zhu et al., 2018; Greff et al., 2019; van Steenkiste et al., 2019) have also been used.
Representation learning has also been decoupled from RL (Anand et al., 2019; Stooke et al., 2020).
We focus on state representation learning in RL, instead of the RL itself, and unlike previous work, leverage pretrained self-supervised representations not trained on domain-speciﬁc data.
Attention: Attention has been used in RL (Zhang et al., 2018; Gregor et al., 2018; Manchin et al., 2019; Salter et al., 2019), especially to add robustness and interpretability (Sorokin et al., 2015; Mott et al., 2019).
Optical ﬂow-based attention has also been used in RL (Yuezhang et al., 2018).
We explore spatio-temporal attention, including optical ﬂow, to improve image representations for temporal data.
Domain generalization: Transfer learning uses representations from one domain to generalize to different domains (Rusu et al., 2016; Oquab et al., 2014; Donahue et al., 2013).
Domain adaptation adapts from one domain to another domain using data from the target domain (Bousmalis et al., 2016; Ganin et al., 2016; Wulfmeier et al., 2017), whereas domain randomisation covers a distribution of environments during training to generalize (Sadeghi & Levine, 2016; Andrychowicz et al., 2020; Viereck et al., 2017; Held et al., 2017;
Tobin et al., 2017; Peng et al., 2018).
We use pretrained models, not trained on Atari-speciﬁc data, for sample-efﬁcient state representation learning in RL.
Pretraining:
Unsupervised pretraining has been leveraged in RL, e.g., by maximizing diversity of states (Liu & Abbeel, 2021) or skills (Eysenbach et al., 2018; Hansen et al., 2019;
Sharma et al., 2019).
Intrinsic rewards for exploration (Pathak et al., 2016; Sekar et al., 2020), predicting state dynamics (Anderson et al., 2015), and reward-free representations (Schwarzer et al.) have also been used.
Transferable skills can also be learned (Campos et al., 2021).
We use pretrained models for state representation learning in RL.
3. Approach State representation learning in RL: Deep reinforcement
learning has leveraged the expressive power of deep learning to create end-to-end models for RL in high dimensional spaces.
RL in high dimensional spaces, e.g., pixels spaces, however, is sample-inefﬁcient (Lake et al., 2016; Kaiser et al., 2019) and learning policies from state representations could be more sample-efﬁcient, robust, and generalizable (Tassa et al., 2018;
Liu et al., 2021; Eslami et al., 2018).
Thus, we decouple state representation learning from RL (Stooke et al., 2020) and aimed to investigate state representation learning in RL (Anand et al., 2019).
Pretrained models with self-supervised ﬁne-tuning: Pretrained models have been leveraged for few-shot learning (Chen et al., 2020c; Henaff, 2020;
Radford et al., 2021), but not for state representation learning in RL.
State representation learning in RL has leveraged SSL (Anand et al., 2019), but involved domain-speciﬁc and data-intensive training.
We aimed to leverage pretrained models for sample-efﬁcient state representation learning in RL.
We decided on 3 key explorations: i. Pretrained image representations , including ‘zooming in’ via grid-based patches; ii.
Spatio-temporal attention , e.g., using optical ﬂow; iii.
Self-supervised ﬁnetuning on domain-speciﬁc data via constrastive losses.
Our framework, PEARL (Pretrained Encoder and Attention for Representation Learning), has 3 components (Figure 1).
3.1. Pretrained Image Representations
Recent work compared image representations from different supervised, self-supervised, and weakly supervised models, and found that CLIP (Radford et al., 2021), a weakly su- pervised image representation learning model, produced the best representations for few-shot learning (Khan et al., 2021).
Thus, we selected CLIP as our image encoder.
In addition to image representations for full-image, we also considered ‘zooming in’ via grid-based patches, i.e., equallysized non-overlapping patches covering the whole image.
3.2. Spatio-temporal Attention
For spatial attention, we considered a SotA supervised object detection model, i.e., EfﬁcientDet (Tan et al., 2019), and a SotA SSL model, i.e., Dino (Caron et al., 2021).
For temporal attention, we decided to compare optical ﬂow using RAFT, a SotA pretrained model (Teed & Deng, 2020), with image difference using structural similarity (Wang et al., 2004).
We used image difference particularly because objects appear/disappear in video animations and optical ﬂow may not be able to capture sudden ﬂowless changes.
Finally, we compare two types of attention: i. mask-based attention, which highlights the relevant regions in the full image; ii.
patch-based attention, which crops out the relevant patches and calculates embeddings for each by zooming in.
3.3. Self-supervised Fine-tuning
We considered 3 SotA methods for self-supervised ﬁnetuning of pretrained representations using domain-speciﬁc data: i. Image augmentations (color jitter, random crop, and gaussian blur) (Chen et al., 2020b); ii.
Spatio-temporal constrastive learning (ST-DIM)
(Anand et al., 2019); iii.
Contrastive predictive learning (CPC)
(Oord et al., 2018).
4. Experiments
We outline our experiments and ﬁndings below.
Similar to the SotA work (Anand et al., 2019), our evaluations use a linear probe (Alain & Bengio, 2018) with a { 70, 10, 20 } % { train, validation, and test } -split with 50k data points and early stopping.
We share the F1 score for each Atari game.
4.1. Pretrained Image Representations
Setup: We tried 3 nxn-grid patch sizes (1x1, 2x2, and 4x4).
Each patch generated a 512-size embedding using CLIP and we concatenated the embeddings.
We tried 6 different conﬁgurations: 1. 1x1, i.e., full image (FI); 2. 2x2; 3.
1x1+2x2; 4. 4x4; 5.
1x1+4x4; 6. 1x
+2x2+4x4.
Figurecompares the results for our 6 conﬁgurations to the SotA results.
SotA refers to the top performance of all the models in (Anand et al., 2019) – ST-DIM was top in all, exceptgames, CPC was top in 3, and Pixel-Pred was top in 4.
Results:
On average, using at least 5 embeddings from 1 full image and 4 2x2 patches, CLIP’s pretrained embeddings per- form better than SotA.
Even our worst conﬁguration (1x1) is on average better than a VAE trained on domain-speciﬁc data (Anand et al., 2019).
With 21 patches (1x+2x+4x), our model performs better than the SotA in all, exceptgames, and has an average of 5% better performance than
SotA. Overall, the performance improves with increasing number of patches.
Thus, zooming in gives performance improvements at the cost of bigger embedding sizes.
4.2. Spatial-Temporal Attention
Setup: We compared Optical Flow Mask (FM) with Image
Difference Mask (DM), and also, each mask combined with full image, i.e., FM+ (FM+FI) and DM+ (DM+FI).
We evaluated patch-wise attention using full image combined with 4 patches selected from 4x4 and 2x2 patches, weighted by Optical Flow (FP5) and Image Difference (DP5).
Each image ‘combination’ refers to a concatenation of the image embeddings.
We compared each of our 3 settings with their equivalent size vanilla embeddings, i.e., embeddings using just the pretrained encoder (Section 4.1): i. DM and FM with FI (1x512 embedding); ii.
DM+
and FM+ with FI+FI (2x512 embedding); iii.
DP5 and FP5 with FI+2x2 (5x embedding).
We did not evaluate EfﬁcientDet and Dino as the attention masks did not look promising (Figure 3).
Results: Our results (Figure 4) show that optical ﬂow masks (FM) alone are on average better than image difference (DM) masks.
However, image difference masks combined with the full image (DM+) are better than optical ﬂow masks combined with the full image (FM+).
Also, both FM+ and
DM+ are both slightly better than the two full image embeddings concatenated (FI+), and the difference is biggest in games like Pong and Montezuma’s Revenge.
Finally, full image + 4 non-grid patches from image difference (DP5) are slightly better than full image + 4 non-grid patches from optical ﬂow (FP5), but they are both worse than the full image + 4 2x2 grid patches.
Thus, n grid-based patches may be better than n non-grid patches, which may zoom in more but not cover the full image.
Also, temporal masks combined with full image are slightly better than two copies of the full image.
However, the difference is not signiﬁcant and is game-dependent, possibly because our encoder model,
CLIP, already incorporates ‘good enough’ attention (Fig 3).
4.3. Self-supervised Fine-tuning
We evaluated if state-of-the-art self-supervised methods could be used to ﬁne-tune pretrained embeddings using domain-speciﬁc data.
We evaluated three self-supervised methods: i. Augmentations (random crop, color jitter, and gaussian blur) with an MLP layer and pairwise contrastive loss (Chen et al., 2020c); ii.
ST-DIM (Anand et al., 2019) with separate temporal (T-DIM), spatial (S-DIM), and spatio-temporal contrast (ST-DIM); iii. CPC, replacing
CPC’s original encoder with our CLIP encoder and then adding a linear layer followed by a Gated Recurrent Unit.
We compare each of the evaluations with their equivalent embedding-size baselines from Section 4.1 – FI for Blur, Jit,
Crop, CPC, and T-DIM-1x (T1x); 2x2 for T-DIM 2x2 (T2x) and S-DIM (s2x); 1x1+2x2 for ST-DIM 1x1+2x2 (ST1,2x).
Our results (Figure 5) show that on average, ﬁne-tuning with CPC leads to the most improvement (3%), followed by T-DIM-1x1 and gaussian blur.
Color jitter, on average, makes no change and the rest lead to a 2-4% drop in performance compared to the equivalent vanilla cases (Section
Maximum improvements are using gaussian blur in
Breakout (15.5%), CPC in Montezuma’s revenge (10.8%), and T-DIM in Breakout (6.3%), but most improvements are game and method-dependent.
Thus, overall, there are no signiﬁcant improvements using self-supervised ﬁne-tuning, possibly due to data or model constraints and especially since we are not ﬁne-tuning the pretrained CLIP model.
5. Conclusion
Data-efﬁciency and generalization are key challenges in deep learning and deep RL.
Most deep RL models are trained from scratch on domain-speciﬁc data, but recent research shows that self-supervised models trained on largescale uncurated data have promising few-shot transfer.
We investigated the use of pretrained models for state representation learning in RL.
Our results show that pretrained self-supervised models, not trained on domain-speciﬁc data, give competitive performance compared to SotA selfsupervised models, trained on large-scale domain-speciﬁc data.
Thus, pretrained models enable data-efﬁcient and generalizable state representation learning for RL.
Moreover, our framework, PEARL (Pretrained Encoder and
Attention for Representation Learning), investigates not only using pretrained image representations but also pretrained spatio-temporal attention.
Our pretrained image representation model also uses attention and our results show that attention helps state representation learning.
Finally, even though self-supervised ﬁne-tuning on domainspeciﬁc data did not signiﬁcantly improve pretrained representations, it could be because we froze our pretrained model and had a few trainable layers with limited data.
State representations are key in RL.
We believe that ‘Pretrained Encoders are All You Need’ (PEAYN) for dataefﬁcient and generalizable state representation learning in
RL, and will hopefully be a stepping stone to data-efﬁcient, generalizable, and interpretable RL.
No PEAYN, no gain :)
. Caption A. Atari Game Abbreviations
Table 1 shows the list of all games and their respective abbreviations used in our paper.