CoBERL: Contrastive BERT for Reinforcement Learning

Introduction

Developing sample efficient reinforcement learning (RL) agents that only rely on raw high dimensional inputs is challenging. Specifically, it is difficult as it often requires us to simultaneously train several neural networks based on sparse environment feedback and strong correlation between consecutive observations. This problem is particularly severe when the networks are large and densely connected, like in the case of the transformer (Vaswani et al., 2017) due to noisy gradients often found in RL problems. Outside of the RL domain, transformers have proven to be very expressive (Brown et al., 2020) and, as such, they are of particular interest for complex domains like RL.

In this paper, we propose to tackle this shortcoming by taking inspiration from Bidirectional Encoder Representations for Transformers (BERT; Devlin et al., 2019), and its successes on difficult sequence prediction and reasoning tasks. Specifically we propose a novel agent, named Contrastive BERT for RL (CoBERL), that combines a new contrastive representation learning objective with architectural improvements that effectively combine LSTMs with transformers.

For representation learning we take inspiration from previous work showing that contrastive objectives improve the performance of agents (Fortunato et al., 2019; Srinivas et al., 2020b; Kostrikov et al., 2020; Mitrovic et al., 2021). Specifically, we combine the paradigm of masked prediction from BERT (Devlin et al., 2019) with the contrastive approach of Representation Learning via Invariant Causal Mechanisms (ReLIC; Mitrovic et al., 2021). Extending the BERT masked prediction to RL is not trivial; unlike in language tasks, there are no discrete targets in RL. To circumvent this issue, we extend ReLIC to the time domain and use it as a proxy supervision signal for the masked prediction. Such a signal aims to learn self-attention-consistent representations that contain the appropriate information for the agent to effectively incorporate previously observed knowledge in the transformer weights. Critically, this objective can be applied to different RL domains as it does not require any data augmentations and thus circumvents the need for much domain knowledge. This is another advantage compared to the original ReLIC and which require to hand-design augmentations.

In terms of architecture, we base CoBERL on Gated Transformer-XL (GTrXL; Parisotto et al., 2020) and Long-Short Term Memories (LSTMs; Hochreiter & Schmidhuber, 1997). GTrXL is an adaptation of a transformer architecture specific for RL domains. We combine GTrXL with LSTMs using a gate trained via an RL loss. This allows the agent to learn to exploit the representations offered by the transformer only when an environment requires it, and avoid the extra complexity when not needed.

We extensively test our proposed agent across a widely varied set of environments and tasks ranging from 2D platform games to 3D first-person and third-person view tasks. Specifically, we test it in the control domain using DeepMind Control Suite (Tassa et al., 2018) and probe its memory abilities using DMLab-30 (Beattie et al., 2016). We also test our agent on all 57 Atari games (Bellemare et al., 2013). Our main contributions are:

A novel contrastive representation learning objective that combines the masked prediction from BERT with a generalization of ReLIC to the time domain; with this we learn self-attention consistent representations and extend BERT-like training to RL using contrastive objectives without the need of hand-engineered augmentations.

An improved architecture that, using a gate, allows us to flexibly integrate knowledge from both the transformer and the LSTM.

Improved data efficiency (and in some cases also overall performance) on a varied set of environments. Also, we show that individually both our contrastive loss and the architecture improvements play a role in improving performance.

Method

To tackle the problem of data efficiency in deep reinforcement learning, we propose two modifications to the status quo. First, we introduce a novel representation learning objective aimed at learning better representations by enforcing self-attention consistency in the prediction of masked inputs. Second, we propose an architectural improvement to combine the strength of LSTMs and transformers.

It has been empirically observed that performing RL directly from high dimensional observations (e.g. raw pixels) is sample inefficient (Lake et al., 2017); however, there is evidence that learning from low-dimensional state based features is significantly more efficient (e.g. Tassa et al., 2018). Thus, if state information could be effectively extracted from raw observations it may then be possible to learn from these as fast as from states. However, unlike in the supervised or unsupervised setting, learning representations in RL is complicated by non-stationarity in the training distribution and strong correlation between observations at adjacent timesteps. Furthermore, given the often sparse reward signal coming from the environment, learning representations in RL has to be achieved with little to no supervision. Approaches to address these issues can be broadly classified into two classes. The first class uses auxiliary self-supervised losses to accelerate the learning speed in model-free RL algorithms (Schmidhuber, 1990; Jaderberg et al., 2016; Oord et al., 2018; Srinivas et al., 2020a). The second class learns a world model and uses this to collect imagined rollouts, which then act as extra data to train the RL algorithm reducing the samples required from the environment (Sutton, 1990; Ha & Schmidhuber, 2018; Kaiser et al., 2019; Schrittwieser et al., 2020). CoBERL is part of the first set of methods, as it uses a self-supervised loss to improve data efficiency. In particular, we take inspiration from the striking progresses made in recent years in both masked language modelling (BERT, Devlin et al., 2019) and contrastive learning (Oord et al., 2018; Chen et al., 2020).

From BERT we borrow the combination of bidirectional processing in transformers (rather than left-to-right or right-to-left, as is common with RNN-based models such as LSTMs) with a masked prediction setup. With this combination the model is forced to focus on the context provided by the surrounding timesteps to solve its objective (Voita et al., 2019). Thus, when the model is asked to reconstruct a particular masked frame it does so by attending to all the relevant frames in the trajectory. This is of particular relevance in RL, where state aliasing is a source of uncertainty in value estimation, and so we believe that attending to other frames in the sequence could be a way to provide extra evidence to reduce this uncertainty.

However, unlike in BERT where the input is a discrete vocabulary for language learning and targets are available, in RL inputs consist of images, rewards and actions that do not form a finite or discrete set and targets are not available. Thus, we must construct proxy targets and the corresponding proxy tasks to solve. For this we use contrastive learning, and we derive our contrastive loss from ReLIC (Mitrovic et al., 2021). ReLIC creates a series of augmentations from the original data and then enforces invariant prediction of proxy targets across augmentations through an invariance regularizer, yielding improved generalization guarantees. Compared to ReLIC, CoBERL does not use data augmentations. Instead we rely on the sequential nature of our input data to create the necessary groupings of similar and dissimilar points needed for contrastive learning. Not having a need for augmentations is critical for RL as each domain would require handcrafting different augmentations. Augmentations also make other methods less data efficient. Finally, we do not use an additional encoder network as in ReLIC, thus making CoBERL fully end-to-end.

We now set out to explain the details of the auxiliary loss. In a batch of sampled sequences, before feeding embeddings into the transformer stack, $15\%$ of the embeddings are replaced with a fixed token denoting masking. Then, let the set $\mathcal{T}$ represent indices in the sequence that have been randomly masked and let $t\in\mathcal{T}$. For the $i$-th training sequence in the batch, for each index $t\in\mathcal{T}$, let $x_{t}^{i}$ be the output of the GTrXL and $y_{t}^{i}$ the corresponding input to the GTrXL from the encoder (see Fig 1B). Let $\phi(\cdot,\cdot)$ be the inner product defined on the space of critic embeddings, i.e. $\phi(x,y)=g(x)^{T}g(y)$, where $g$ is a critic function. The critic separates the embeddings used for the contrastive proxy task and the downstream RL task. Details of the critic function are in App. B. This separation is needed since the proxy and downstream tasks are related but not identical, and as such the appropriate representations will likely not be the same. As a side benefit, the critic can be used to reduce the dimensionality of the dot-product. To learn the embedding $x_{t}^{i}$ at mask locations $t$, we use $y_{t}^{i}$ as a positive example and the sets $\{y^{b}_{t}\}_{b=0,b\neq i}^{B}$ and $\{x^{b}_{t}\}_{b=0,b\neq i}^{B}$ as the negative examples with $B$ the number of sequences in the minibatch. We model $q^{t}_{x}$ as

with $q^{t}_{x}$ denoting $q^{t}_{x}\left(x_{t}^{i}|\{y^{b}_{t}\}_{b=0}^{B},\{x^{b}_{t}\}_{b=0,b\neq i}^{B}\right)$; $q^{t}_{y}$ is computed analogously (see Fig 1C). In order to enforce self-attention consistency in the learned representations, we explicitly regularize the similarity between the pairs of transformer embeddings and inputs through Kullback-Leibler regularization from ReLIC. Specifically, we look at the similarity between appropriate embeddings and inputs, and within the sets of embeddings and inputs separately. To this end, we define:

with $p^{t}_{x}$ and $s^{t}_{x}$ shorthand for $p^{t}_{x}(x_{t}^{i}|\{y^{b}_{t}\}_{b=0}^{B})$ and $s^{t}_{x}(x_{t}^{i}|\{x^{b}_{t}\}_{b=0}^{B})$, respectively; $p^{t}_{y}(y_{t}^{i}|\{x^{b}_{t}\}_{b=0}^{B})$ and $s^{t}_{y}(y_{t}^{i}|\{y^{b}_{t}\}_{b=0}^{B})$ defined analogously (see Fig 1D). All together, the final objective takes the form:

with $\textrm{sg}(\cdot)$ indicating a stop-gradient. As in our RL objective, we use the full batch of sequences that are sampled from the replay buffer to optimize this contrastive objective. Finally, we optimize a weighted sum of the RL objective and $\mathcal{L}(X,Y)$ (see App. C for the details on the weighting).

2 Architecture of CoBERL.

While transformers have proven very effective at connecting long-range data dependencies in natural language processing (Vaswani et al., 2017; Brown et al., 2020; Devlin et al., 2019) and computer vision (Carion et al., 2020; Dosovitskiy et al., 2021), in the RL setting they are difficult to train and are prone to overfitting (Parisotto et al., 2020). In contrast, LSTMs have long been demonstrated to be useful in RL. Although less able to capture long range dependencies due to their sequential nature, LSTMs capture recent dependencies effectively. We propose a simple but powerful architectural change: we add an LSTM layer on top of the GTrXL with an extra gated residual connection between the LSTM and GTrXL, modulated by the input to the GTrXL (see Fig 1A). Finally we also have a skip connection from the transformer input to the LSTM output.

More concretely, let $Y_{t}$ be the output of the encoder network at time $t$, then the additional module can be defined by the following equations (see Fig 1A, $\textrm{Gate}()$, below, has the same form as other gates internal to GTrXL):

The architecture of CoBERL is based on the idea that LSTMs and Transformer can be complementary. In particular, the LSTM is known for having a short contextual memory. However, by putting a transformer before the LSTM, the embeddings provided to the LSTM have already benefited from the ability of the Transformer to process long contextual dependencies, thus helping the LSTM in this respect. The second benefit works in a complementary direction. Transformers suffer from a quadratic computational complexity with respect to the sequence length. Our idea is that by letting an LSTM process some of this contextual information we can reduce the memory size of the Transformer up to the point where we see a loss in performance. This hypothesis come from studies showing that adding recurrent networks in architectures has the ability to closely emulate the behavior of non-recurrent but deeper models, but it does so with far fewer parameters (Schwarzschild et al., 2021). Having fewer parameters is an aspect of particular interest in RL, where gradients are noisy and hence training larger model is more complicated than supervised learning (McCandlish et al., 2018). Both hypotheses have been empirically confirmed by our ablations in section 4.2.

The learnt gate (eq. 4 and Appendix K), was done to give CoBERL the ability to initially skip the output of the transformer stack until the weights of this are warmed-up. Once the transformer starts to output useful information the agent will be able to tune the weights in the gate to learn from them. We hypothesize that this would also help in terms of data efficiency as at the beginning of training the agent could only use the LSTM to get off the ground and start collecting relevant data to train the Transformer. Finally, the skip connection on the output of the LSTM was taken from Kapturowski et al. (2018).

Since the architecture is agnostic to the choice of RL regime we evaluate it in both on-policy and off-policy settings. For on-policy, we use V-MPO (Song et al., 2019), and for off-policy we use R2D2 (Kapturowski et al., 2018).

Recurrent Replay Distributed DQN (R2D2; Kapturowski et al., 2018) demonstrates how replay and the RL learning objective can be adapted to work well for agents with recurrent architectures. Given its competitive performance on Atari-57, we implement our CoBERL architecture in the context of Recurrent Replay Distributed DQN (Kapturowski et al., 2018). We effectively replace the LSTM with our gated transformer and LSTM combination and add the contrastive representation learning loss. With R2D2 we thus leverage the benefits of distributed experience collection, storing the recurrent agent state in the replay buffer, and "burning in" a portion of the unrolled network with replayed sequences during training.

Given V-MPO’s strong performance on DMLab-30, in particular in conjunction with the GTrXL architecture (Parisotto et al., 2020) which is a key component of CoBERL, we use V-MPO and DMLab-30 to demonstrate CoBERL’s use with on-policy algorithms. V-MPO is an on-policy adaptation of Maximum a Posteriori Policy Optimization (MPO) (Abdolmaleki et al., 2018). To avoid high variance often found in policy gradient methods, V-MPO uses a target distribution for policy updates, subject to a sample-based KL constraint, and gradients are calculated to partially move the parameters towards the target, again subject to a KL constraint. Unlike MPO, V-MPO uses a learned state-value function $V(s)$ instead of a state-action value function.

Related Work

The transformer architecture (Vaswani et al., 2017) has recently emerged as one of the best performing approaches in language modelling (Dai et al., 2019; Brown et al., 2020) and question answering (Dehghani et al., 2018; Yang et al., 2019). More recently it has also been successfully applied to computer vision (Dosovitskiy et al., 2021). Given the similarities of sequential data processing in language modelling and reinforcement learning, transformers have also been successfully applied to the RL domain, where as motivation for GTrXL, Parisotto et al. (2020) noted that extra gating was helpful to train transformers for RL due to the high variance of the gradients in RL relative to that of (un)supervised learning problems. In this work, we build upon GTrXL and demonstrate that, perhaps for RL: attention is not all you need, and by combining GTrXL in the right way with an LSTM, superior performance is attained. We reason that this demonstrates the advantage of both forms of memory representation: the all-to-all attention of transformers combined with the sequential processing of LSTMs. In doing so, we demonstrate that care should be taken in how LSTMs and transformers are combined and show a simple gating is most effective in our experiments. Also, unlike GTrXL, we show that using an unsupervised representation learning loss that enforces self-attention consistency is an effective way to enhance data efficiency when using transformers in RL.

Recently contrastive learning (Hadsell et al., 2006; Gutmann & Hyvärinen, 2010; Oord et al., 2018) has emerged as a very performant paradigm for unsupervised representation learning, in some cases even surpassing supervised learning (Chen et al., 2020; Caron et al., 2020; Mitrovic et al., 2021). These methods have also been leveraged in an RL setting with the hope of improving performance. Apart from MRA (Fortunato et al., 2019) mentioned above, one of the early examples of this is CURL (Srinivas et al., 2020a) which combines Q-Learning with a separate encoder used for representation learning with the InfoNCE loss from CPC (Oord et al., 2018). More recent examples use contrastive learning for predicting future latent states (Schwarzer et al., 2020; Mazoure et al., 2020), defining a policy similarity embeddings (Agarwal et al., 2021) and learning abstract representations of state-action pairs (Liu et al., 2021). The closest work to our is M-CURL (Zhu et al., 2020). Like our work, it combines mask prediction, transformers, and contrastive learning, but there are a few key differences. First, unlike M-CURL who use a separate policy network, CoBERL computes Q-values based on the output of the transformer. Second, CoBERL combines the transformer architecture with a learnt gate (eq. 5 in App. K to produce the input for the Q-network, while M-CURL uses the transformer as an additional embedding network (critic) for the computation of the contrastive loss. Third, while CoBERL uses an extension of ReLIC (Mitrovic et al., 2021) to the time domain and operates on the inputs and outputs of the transformer, M-CURL uses CPC (Oord et al., 2018) with a momentum encoder as in (Srinivas et al., 2020a) and compares encodings from the transformer with the separate momentum encoder.

Experiments

We provide empirical evidence to show that CoBERL i) improves data efficiency across a wide range of environments and tasks, and ii) needs all its components to maximise its performance. In our experiments, we demonstrate performance on Atari57 (Bellemare et al., 2013), the DeepMind Control Suite (Tassa et al., 2018), and the DMLab-30 (Beattie et al., 2016). Recently, Dreamer V2 (Hafner et al., 2020b) has emerged as a strong model-based agent across Atari57 and DeepMind Control Suite; we therefore include it as a reference point for performance on these domains.

For all experiments, we report scores at the end of training. All results are averaged over five seeds and reported with standard error. For a more comprehensive description of the evaluation done, as well as detailed information on the distributed setup used see App. A.1. Ablations are run on 7 Atari games chosen to match the ones in the original DQN publication (Mnih et al., 2013), and on all the 30 DMLab games. The hyper-parameters of all the baselines are tuned individually to maximise performance (see App. C.5 for the detailed procedure). We use a ResNet as the encoder for CoBERL. We use Peng’s $Q(\lambda)$ as our loss (Peng & Williams, 1994). To ensure that this is comparable to R2D2, we also run an R2D2 baseline with this loss. For thorough description of the architectures for all environments see App. B.

For DMLab-30 we use V-MPO (Song et al., 2019) to directly compare CoBERL with (Parisotto et al., 2020) and also demonstrate how CoBERL may be applied to both on and off-policy learning. The experiments were run using a Podracer setup (Hessel et al., 2021), details of which may be found in App. A.2. CoBERL is trained for 10 billion steps on all 30 DMLab-30 games at the same time, to mirror the exact multi-task setup presented in (Parisotto et al., 2020). Compared to (Parisotto et al., 2020) we have two differences. Firstly, all the networks run without pixel control loss (Jaderberg et al., 2016) so as not to confound our contrastive loss with the pixel control loss. Secondly all the models used a fixed set of hyperparameters with 3 random seeds, whereas in (Parisotto et al., 2020) the results were averaged across hyperparameters.

To test the generality of our approach, we analyze the performance of our model on a wide range of environments. We show results on the Arcade Learning Environment (Bellemare et al., 2013), DeepMind Lab (Beattie et al., 2016), as well as the DeepMind Control Suite (Tassa et al., 2018). To help with comparisons, in Atari-57 and DeepMind Control we introduce an additional baseline, which we name R2D2-GTrXL. In this variant of R2D2 the LSTM is replaced by GTrXL. R2D2-GTrXL has no unsupervised learning. This way we are able to observe how GTrXL is affected by the change to an off-policy agent (R2D2), from its original V-MPO implementation in Parisotto et al. (2020). We also perform an additional ablation analysis by removing the contrastive loss from CoBERL (see Sec. 4.2.1). This baseline demonstrates the importance of contrastive learning in these domains, and we show that the combination of an LSTM and transformer is superior to either alone.

As commonly done in literature (Mnih et al., 2015; Hessel et al., 2018; Machado et al., 2018; Hafner et al., 2020b), we measure performance on all 57 Atari games after running for $200$ million frames. As detailed in App. C, we use the standard Atari frame pre-processing to obtain the $84$x$84$ gray-scaled frames that are used as input to our agent. We do not use frame stacking.

Table 1 shows the results of all the agents where published results are available. CoBERL shows the most games above average human performance and significantly higher overall mean performance. Interestingly, the performance of R2D2-GTrXL shows that the addition of GTrXL is not sufficient to obtain the improvement in performance that CoBERL exhibits (in 4.2 we will demonstrate that both the contrastive loss and LSTM contribute to this improvement). R2D2-GTrXL also exhibits slightly better median than CoBERL, showing that R2D2-GTrXL is indeed a powerful variant on Atari. Additionally, we observe that the difference in performance in CoBERL is higher when examining the lower percentiles. This suggests that CoBERL causes an improvement in data efficiency. To confirm this tendency we calculated the area under the curve (AUC) for the learning curves presented in Appendix F and we perform a t-test analysis on it (please see App.I for details on how AUC was calculated). Figure 2 present this analysis, and it confirms a significant difference with respect to the model used, t(114)=3.438, p=.008, with CoBERL(M=7380.23, SD=942.29) being better than GTrXL (M=5192.76, SD=942.01), thus showing that CoBERL is more data efficient.

We also perform experiments on the DeepMind Control Suite (Tassa et al., 2018). While the action space in this domain is typically treated as continuous, we discretize the action space in our experiments to apply the same architecture as in Atari and DMLab-30. For more details on the number of actions for each task see App. C.4. We do not use pre-processing on the environment frames. Finally, CoBERL is trained only from pixels without state information.

We include six tasks popular in current literature: ball_in_cup catch, cartpole swingup, cheetah run, finger spin, reacher easy, and walker walk. Most previous work on these specific tasks has emphasized data efficiency as most are trivial to solve even with the baseline—D4PG-Pixels—in the original dataset paper (Tassa et al., 2018). We thus include 6 other tasks that are difficult to solve with D4PG-Pixels and are relatively less explored: acrobot swingup, cartpole swingup_sparse, fish swim, fish upright, pendulum swingup, and swimmer swimmer6. In Table 2 we show results on CoBERL, R2D2-gTRXL, R2D2, CURL (Srinivas et al., 2020a), Dreamer (Hafner et al., 2020a), Soft Actor Critic (Haarnoja et al., 2018) on pixels as demonstrated in (Srinivas et al., 2020a), and D4PG-Pixels (Tassa et al., 2018). CURL, DREAMER, and Pixel SAC are for reference only as they represent the state the art for low-data experiments (500K environment steps). These three are not perfectly comparable baselines; however, D4PG-Pixels is run on a comparable scale with 100 million environment steps. Because CoBERL relies on large scale distributed experience, we have a much larger number of available environment steps per gradient update. We run for 100M environment steps as with D4PG-Pixels, and we compute performance for our approaches by taking the evaluation performance of the final 10% of steps. Across the majority of tasks, CoBERL outperforms D4PG-Pixels. The increase in performance is especially apparent for the more difficult tasks. For most of the easier tasks, the performance difference between the CoBERL, R2D2-GTrXL, and R2D2 is negligible. For ball_in_cup catch, cartpole swingup, finger spin and reacher easy, even the original R2D2 agent performs on par with the D4PG-Pixels baseline. On more difficult tasks such as fish swim, and swimmer swimmer6, there is a very large, appreciable difference between CoBERL, R2D2, and R2D2-GTrXL. The combination of the LSTM and transformer specifically makes a large difference here especially compared to D4PG-Pixels. Interestingly, this architecture is also very important for situations where the R2D2-based approaches underperform. For cheetah run and walker walk, CoBERL dramatically narrows the performance gap between R2D2 and state of the art (in App. F we report learning curves for each game). Figure 3 shows that also in this domain, we see a significant AUC improvement for CoBERL when compared to R2D2-gTrXL, t(24)=1.609, p=.03.

To test CoBERL in a challenging 3 dimensional environment we run it in DmLab-30 (Beattie et al., 2016). The agent was trained at the same time on all the 30 tasks, following the setup of GTrXl (Parisotto et al., 2020), which we use as our baseline. In Figure 4A we show the final results on the DMLab-30 domain. If we look at all the 30 games, CoBERL reaches a substantially higher score than GTrXL (CoBERL=115.47% $\pm$ 4.21%, GTrXL=101.54% $\pm$ 0.50%, t(60)=4.37, p=1.14e-5, Figure 4A). In Figure 4B we analysed data efficiency by computing the AUC and associated average statistics. In DMLab the difference between models is even more significant than Atari, t(60)=6.097, p=9.39e-09, with CoBERL (M=1193.28, SD=383.29) having better average AUC than GTrXL (M=764.09, SD=224.51), hence showing that our methods scales well to more complex domains. (see Appendix F for learning curves)

2 Ablations

In Sec. 2, we explained contributions that are essential to CoBERL. We now disentangle the added benefit of these two separate contributions. Moreover, we run a set of ablations to understand the role of model size on the results. Ablations are run on 7 Atari games chosen to match the ones in the original DQN publication (Mnih et al., 2013), and on all the 30 DMLab games.

In Table 3 we show that our contrastive loss contributes to a significant gain in performance, both in Atari and DMLab-30, when compared to CoBERL without it. Also, in challenging environments like DmLab-30, CoBERL without extra loss is still superior to the relative baseline. The only case where we do not see and advantage of using the auxiliary loss is if we consider the median score on the reduced ablation set of Atari games. However in the case of the DmLab-30, where we consider a larger set of levels (7 vs. 30), there is a clear benefit of the auxiliary loss.

Moreover, Table 4 reports a comparison between our loss, SimCLR (Chen et al., 2020) and CURL (Srinivas et al., 2020a). Although simpler than both SimCLR - which in its original implementation requires handcrafted augmentations - and CURL - which requires an additional network - our contrastive method shows improved performance. These experiments where run only on Atari to reduce computational costs while still being sufficient for the analysis. We also ran one extra ablation study where we did not use masking of the input to make our setup as close as possible to (Mitrovic et al., 2021). As seen on Table 4, the column CoBERL w/o masking shows that removing the masking technique derived from the language literature makes the results substantially worse.

2.2 Impact of architectural changes

Table 5 shows the effects of removing the LSTM from CoBERL (column “w/o LSTM”), as well as removing the gate and its associated skip connection (column “w/o Gate”). In both cases CoBERL performs substantially worse showing that both components are needed. Finally, we also experimented with substituting the learned gate with either a sum or a concatenation. The results, presented in Appendix D, show that in most occasions these alternatives decrease performance, but not as substantially as removing the LSTM, gate or skip connections. Our hypothesis is that the learned gate gives more flexibility in complex environments, we leave it open for future work to explore this.

2.3 Impact of number of parameters

Table 6 compares the models in terms of the number of parameters. For Atari, the number of parameters added by CoBERL over the R2D2(GTrXL) baseline is very limited; however, CoBERL still produces a significant gain in performance. We also tried to move the LSTM before the transformer module (column “CoBERL with LSTM before”). In this case the representations for the contrastive loss were taken from before the LSTM. Interestingly, this setting performs worse, despite having the same number of parameters as CoBERL. This goes in the direction of our hypothesis that having the LSTM after the Transformer could allow the former to exploit the extra context provided by the latter. For DMLab-30, it is worth noting that CoBERL has a memory size of $256$, whereas GTrXL has a memory of size $512$ resulting in substantially fewer parameters. Nevertheless, the discrepancies between models are even more pronounced, even though the number of parameters is either exactly the same (“CoBERL with LSTM before”) or higher (GTrXL). This ablation is of particular interest as it shows that the results are driven by the particular architectural choice rather than the added parameters. Also it helps supporting our idea that the LSTM, by processing a certain amount of contextual information, allow for a reduced memory size on the Transformer, without a loss in performance and with a reduce computation complexity (given fewer parameters).

Conclusions

We proposed a novel RL agent, Contrastive BERT for RL (CoBERL), which introduces a new contrastive representation learning loss that enables the agent to efficiently learn consistent representations. This, paired with an improved architecture, resulted in better data efficiency on a varied set of environments and tasks. Moreover, through an extensive set of ablation experiments we confirmed that all CoBERL components are necessary to achieve the performance of the final agent. Critically, CoBERL is fully end-to-end as it does not require an extra encoder network (vs CURL or ReLIC), or data augmentations (vs SIMCLR and ReLIC) and it has fewer parameters then gTrXL. To conclude, we have shown that our auxiliary loss and architecture provide an effective and general means to efficiently train large attentional models in RL. (For extra conclusions and limitations see App.J)

Optional reproducibility statement

To help with reproducibility we have included extensive details for

We also report the pseudo-code for the algorithm and the auxiliary loss in Appendix H

References

Appendix A Setup details

Following R2D2, the distributed system consists of several parts: actors, a replay buffer, a learner, and an evaluator. Additionally, we introduce a centralized batched inference process to make more efficient use of actor resources.

Actors: We use $512$ processes to interact with independent copies of the environment, called actors. They send the following information to a central batch inference process:

$r_{t-1}$: the reward at the previous time, initialized with $r_{-1}=0$.

$a_{t-1}$: the action at the previous time, $a_{-1}$ is initialized to .

$h_{t-1}$: recurrent state at the previous time, is initialized with $h_{-1}=0$.

They block until they receive $Q(x_{t},a;\theta)$. The $l$-th actor picks $a_{t}$ using an $\epsilon_{l}$-greedy policy. As R2D2, the value of $\epsilon_{l}$ is computed following:

where $\epsilon=0.4$ and $\alpha=7$. After that is computed, the actors send the experienced transition information to the replay buffer.

Batch inference process: This central batch inference process receives the inputs mentioned above from all actors. This process has the same architecture as the learner with weights that are fetched from the learner every $0.5$ seconds. The process blocks until a sufficient amount of actors have sent inputs, forming a batch. We use a batch size of $64$ in our experiments. After a batch is formed, the neural network of the agent is run to compute $Q(x_{t},a,\theta)$ for the whole batch, and these values are sent to their corresponding actors.

Replay buffer: it stores fixed-length sequences of transitions $T=(\omega_{s})_{s=t}^{t+L-1}$ along with their priorities $p_{T}$, where $L$ is the trace length we use. A transition is of the form $\omega_{s}=(r_{s-1},a_{s-1},h_{s-1},x_{s},a_{s},h_{s},r_{s},x_{s+1})$. Concretely, this consists of the following elements:

$a_{s-1}$: action done by the agent at the previous time.

$h_{s-1}$: recurrent state (in our case hidden state of the LSTM) at the previous time.

$x_{s}$: observation provided by the environment at the current time.

$a_{s}$: action done by the agent at the current time.

$h_{s}$: recurrent state (in our case hidden state of the LSTM) at the current time.

$x_{s+1}$: observation provided by the environment at the next time.

The sequences never cross episode boundaries and they are stored into the buffer in an overlapping fashion, by an amount which we call the replay period. Finally, concerning the priorities, we followed the same prioritization scheme proposed by Kapturowski et al. (2018) using a mixture of max and mean of the TD-errors in the sequence with priority exponent $\eta=0.9$.

Evaluator: the evaluator shares the same network architecture as the learner, with weights that are fetched from the learner every episode. Unlike the actors, the experience produced by the evaluator is not sent to the replay buffer. The evaluator acts in the same way as the actors, except that all the computation is done within the single CPU process instead of delegating inference to the batch inference process. At the end of $5$ episodes the results of those $5$ episodes are average and reported. In this paper we report the average performance provided by such reports over the last $5\%$ frames (for example, on Atari this is the average of all the performance reports obtained when the total frames consumed by actors is between 190M and 200M frames).

Learner: The learner contains two identical networks called the online and target networks with different weights $\theta$ and $\theta^{\prime}$ respectively (Mnih et al., 2015). The target network’s weights $\theta^{\prime}$ are updated to $\theta$ every $400$ optimization steps. $\theta$ is updated by executing the following sequence of instructions:

First, the learner samples a batch of size $64$ (batch size) of fixed-length sequences of transitions from the replay buffer, with each transition being of length $L$: $T_{i}=(\omega^{i}_{s})_{s=t}^{t+L-1}$.

Then, a forward pass is done on the online network and the target with inputs $(x^{i}_{s},r^{i}_{s-1},a^{i}_{s-1},h^{i}_{s-1})_{s=t}^{t+H}$ in order to obtain the state-action values $\{(Q(x^{i}_{s},a;\theta),Q(x^{i}_{s},a;\theta^{\prime})\}$.

With $\{(Q(x^{i}_{s},a;\theta),Q(x^{i}_{s},a;\theta^{\prime})\}$, the $Q(\lambda)$ loss is computed.

The online network is used again to compute the auxiliary contrastive loss.

Both losses are summed (with by weighting the auxiliary loss by $0.1$ as described in C), and optimized with an Adam optimizer.

Finally, the priorities are computed for the sampled sequence of transitions and updated in the replay buffer.

A.2 V-MPO distributed setup

For on-policy training, we used a Podracer setup similar to (Hessel et al., 2021) for fast usage of experience from actors by learners.

TPU learning and acting: As in the Sebulba setup of (Hessel et al., 2021), acting and learning network computations were co-located on a set of TPU chips, split into a ratio of 3 cores used for learning for every 1 core used for inference. This ratio then scales with the total number of chips used.

Environment execution: Due to the size of the recurrent states used by CoBERL and stored on the host CPU, it was not possible to execute the environments locally. To proceed we used 64 remote environment servers which serve only to step multiple copies of the environment. 1024 concurrent episodes were processed to balance frames per second, latency between acting and learning, and memory usage of the agent states on the host CPUs.

A.3 Computation used

We train the agent with a single TPU v2-based learner, performing approximately $5$ network updates per second (each update on a mini-batch of $64$ sequences of length $80$ for Atari and $120$ for Control). We use $512$ actors, using $4$ actors per CPU core, with each one performing $\sim 64$ environment steps per second on Atari. Finally for the batch inference process a TPU v2, which allows all actors to achieve the speed we have described. In particular, we used 8 TPU cores for learning and 2 for inference.

We train the agent with 4 hosts each with 8 TPU v2 cores. Each of the 8 cores per host was split into 6 for learning and 2 for inference. We separately used 64 remote CPU environment servers to step 1024 concurrent environment episodes using the actions returned from inference. The learner updates were made up of a mini-batch of 120 sequences, each of length 95 frames. This setup enabled 4.6 network updates per second, or 53.4k frames per second.

A.4 Complexity analysis

As stated, the agent consists of layers of convolutions, transformer layers, and linear layers. Therefore the complexity is $max\{O(n^{2}\cdot d),O(k\cdot n\cdot d^{2})\}$, where $k$ is the kernel size in the case of convolutions, $n$ is the size of trajectories, and $d$ is the size of hidden layers.

Appendix B Architecture description

As shown in Fig. 1, observations $O_{i}$ are encoded using an encoder. In this work, the encoder we have used is a ResNet-47 encoder. Those 47 layers are divided in $4$ groups which have the following characteristics:

An initial stride-2 convolution with filter size $3$x$3$ ($1\cdot 4$ layers).

Number of residual bottleneck blocks (in order): $(2,4,6,2)$. Each block has 3 convolutional layers with ReLU activations, with filter sizes $1$x$1$, $3$x$3$, and $1$x$1$ respectively ($(2+4+6+2)\cdot 3$ layers).

Number of channels for the last convolution in each block: $(64,128,256,512)$.

Number of channels for the non-last convolutions in each block: $(16,32,64,128)$.

Group norm is applied after each group, with a group size of $8$.

After this observation encoding step, a final 2-layer MLP with ReLU activations of sizes $(512,448)$ is applied. The previous reward and one-hot encoded action are concatenated and projected with a linear layer into a $64$-dimensional vector. This $64$-dimensional vector is concatenated with the $448$-dimensional encoded input to have a final $512$-dimensional output.

B.2 Transformer

As described in Section 2, the output of the encoder is fed to a Gated Transformer XL. For Atari and Control, the transformer has the following characteristics:

For DmLab the transformer has the following characteristics:

Memory size: $256$ for CoBERL and $512$ for gTrXL.

the GTrXL baseline is identical, but with a Memory size of $512$.

B.3 LSTM and Value head

For both R2D2 and V-MPO the outputs of the transformer and encoder are passed through a GRU transform to obtain a $512$-dimensional vector. After that, an LSTM with $512$ hidden units is applied. The the value function is estimated differently depending on the RL algorithm used.

Following the LSTM, a Linear layer of size $512$ is used, followed by a ReLU activation. Finally, to compute the Q values from that $512$ vector a dueling head is used, as in Kapturowski et al. (2018), a dueling head is is used, which requires a linear projection to the number of actions of the task, and another projection to a unidimensional vector.

Following the LSTM, a 2 layer MLP with size $512$ and $30$ (i.e. the number of levels in DMLab) is used. In the MLP we use ReLU activation. As we are interested in the multi-task setting where a single agent learns a large number of tasks with differing reward scales, we used PopArt (van Hasselt et al., 2016) for the value function estimation (see Table. LABEL:tab:hyperparameters_VMPO for details).

B.4 Critic Function

For DmLab-30 (V-MPO), we used a 2 layer MLP with hidden sizes 512 and 128. For Atari and Control Suite (R2D2) we used a single layer of size 512.

Appendix C Hyperparameters

For the experiments in Atari57 and the DeepMind Control suite, CoBERL uses the R2D2 distributed setup. We use $512$ actors for all our experiments. We do not constrain the amount of replay done for each experience trajectory that actors deposit in the buffer. However, we have found empirical replay frequency per data point to be close among all our experiments (with an expected value of $1.5$ samples per data point). We use a separate evaluator process that shares weights with our learner in order to measure the performance of our agents. We report scores at the end of training. The hyperparameters and architecture we choose for these two domains are the same with two exceptions: i) we use a shorter trace length for Atari ($80$ instead of $120$) as the environment does not require a long context to inform decisions, and ii) we use a squashing function on Atari and the Control Suite to transform our $Q$ values (as done in (Kapturowski et al., 2018)) since reward structures vary highly in magnitude between tasks.

We use the commonly used input pre-processing on Atari and DMLab frames, shown on Tab. 8. One difference with the original work of Mnih et al. (2015), is that we do not use frame stacking, as we rely on our memory systems to be able to integrate information from the past, as done in Kapturowski et al. (2018). ALE is publicly available at https://github.com/mgbellemare/Arcade-Learning-Environment.

C.2 Control Suite pre-processing

As mentioned in 4, we use no pre-processing on the frames received from the control environment.

C.3 DmLab pre-processing

C.4 Control environment discretization

As mentioned, we discretize the space assigning two possibilities (1 and -1) to each dimension and taking the Cartesian product of all dimensions, which results in $2^{n}$ possible actions. For the cartpole tasks, we take a diagonal approach, utilizing each unit vector in the action space and then dividing each unit vector into 5 possibilities with the non-zero coordinate ranging from -1 to 1. The amount of actions this results in is outlined on Tab. 9.

C.5 Hyperparameters Used

We list all hyperparameters used here for completeness.

We started by optimizing the hyperparameters of GTrXL highlighted in bold in Tab LABEL:tab:hyperparameters_R2D2 by doing a sweep over 10 Atari games: Seaquest, Qbert, Frostbite, Ms Pacman, Space Invaders, Gravitar, Solaris, Hero, Venture, Montezuma Revenge. Following this, the hyperparameters were kept fixed throughout all the experiments.

Table LABEL:tab:hyperparameters_R2D2 reports all the hyperparameters of the R2D2 experiments, both the fixed ones and the ones with the ranges over which we did the sweep. The fixed hyper-parameters were taken from Kapturowski et al. (2018). We then choose a set of hyper-paramters (both for the architecture and the algorithm) to sweep over to maximise the performance of gTrXL in this off-policy setting, given that there was not prior literature on this. The hyper-paramters over which we did the sweep for the algorithm were chose in accordance to Kapturowski et al. (2018) and the related sweep. And for the architecture hyper-parameters we based our choice on Parisotto et al. (2020). We believe that in this way we ensure to have a properly tuned baseline that enforces a fair comparison. Table LABEL:tab:hyperparameters_R2D2_chosen reports the chosen hyperparameters that we found to optimize the performance of GTrXL on the 10 Atari games. We then moved to CoBERL. CoBERL, in comparison to the baseline GTrXL has two extra hyperparamters: ‘Contrastive loss weight’ and ‘Contrastive loss mask rate’. The former was tuned, whereas the latter we kept equal to 0.15 as done in . Consequently, we re-ran the same procedure as before, but we fixed all the previous hyperparameters optimized for GTrXL (see Tab. LABEL:tab:hyperparameters_R2D2_chosen and we perform a grid search over the same 10 Atari games to find the value of ‘Contrastive loss weight’ that maximized performance. The values over which we did the search are 0.01, 0.1 and 1. We ended-up picking 1, although the difference between 0.1 and 1 was minimal. Table 12 reports the 2 extra parameters used for CoBERL.

For DmLAB we optimized the hyperparameters of GTrXL on all the 30 games. Table LABEL:tab:hyperparameters_VMPO reports both the fixed ones and the ones with the ranges over which. The fixed hyper-parameters were taken directly from Parisotto et al. (2020) and we sweep over ‘’Epsilon Alpha”, “Target Update Period” and ’Memory size’ to make sure we maximised performance of this baseline. Table LABEL:tab:hyperparameters_VMPO_chosen reports the chosen hyperparameters that we found to optimize the performance of GTrXL on DmLAB. We then moved to CoBERL. Again, by keeping fixed all the previous hyperparameters optimized for GTrXL we perform a grid search over all the 30 games to find the value of ‘Contrastive loss weight’ that maximized performance. The values over which we did the search were 0.1 and 1. We did not find any significant difference between the two values, so we left it equal to 1 such that the two losses would have the same effect. Table 15 reports the 2 extra parameters used for CoBERL and the reduced memory size, in accordance with our hypothesis that the LSTM on top of Transformer would help reducing the size of the memory especially in last.

Appendix D Additional ablations

Table 16 shows the results of several gating mechanisms that we have investigated. As we can observe the GRU gate is a clear improvement especially on DMLab, only being harmful in median on the reduced ablation set of Atari games.

Appendix E Game scores

Appendix F Learning Curves

F.2 DMControl Learning curves

F.3 DMLab Learning curves

Appendix G Licenses

The The Arcade Learning Environment Bellemare et al. (2013) is released as free, open-source software under the terms of the GNU General Public License. The latest version of the source code is publicly available at: http://arcadelearningenvironment.org

DeepMind Control Suite Tassa et al. (2018) is released as free, open-source software under the terms of Apache-2.0 License. The latest version of the source code is publicly available at: https://github.com/deepmind/dm_control/blob/master/dm_control/suite/README.md

DmLab Beattie et al. (2016) is released as free, open-source software under the terms of Apache-2.0 License. The latest version of the source code is publicly available at: https://github.com/deepmind/lab/tree/master/game_scripts/levels/contributed/dmlab30

Appendix H Pseudo-code

Pseudo-code for the auxiliary loss calculation

Appendix I Area under the curve

For all the levels we calculated the AUC by integrating composite Simpson’s rule with a delta(x) of 5 steps. We use the intergate package from scipy (Virtanen et al., 2020).

Appendix J Limitations and Future Work

A limitation of our method is that it relies on single time step information to compute its auxiliary objective. Such objective could naturally be adapted to operate on temporally-extended patches, and/or action-conditioned inputs. Also, as done in the original BERT (Devlin et al., 2019), it could be possible to add a CLS token at the beginning of each sequence sent to the Transformer and then train the CLS token with RL gradients. In this way it would be possible to directly use the embeddings of the CLS token as a sequence summary and hence provide more context to the policy estimation network. We regard those ideas as promising future research avenues.

Appendix K Extra information about the gate employed in equation 4

The gate we use in equation 4 is derived directly from the one used in GTrXL (Parisotto et al., 2020). We report here the details for better clarity.

The gate is the defined in the following way:

$W_{z},U_{z},W_{g},U_{g},W_{r}andU_{r}$ are set of linear weights and $b_{g}$ is a bias.

In our case x is the output of the transformer network and y is the output of the encoder network.

Appendix L Extra analysis

We also attempted higher masking rate, but as shown in figure 8 they seem to perform worse, probably because high level of masking are reducing to much the number of frame present in the sequence, hence removing the information need to successfully perform the auxiliary task.

Even when calulcate at 100% human score, the AUC shows a significant advantage of CoBERL over gTrXL t(60)=2.616, p=0.0011. CoBERL M=545.18, SD=282.58; gTrXL M=374.49, SD=282.58.