C O BERL: C ONTRASTIVE BERT FOR R EINFORCEMENT L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
Many reinforcement learning (RL) agents require a large amount of experience to solve tasks.
We propose Contrastive BERT for RL (C O BERL), an agent that combines a new contrastive loss and a hybrid LSTM-transformer architecture to tackle the challenge of improving data efﬁciency.
C O BERL enables efﬁcient and robust learning from pixels across a wide variety of domains.
We use bidirectional masked prediction in combination with a generalization of a recent contrastive method to learn better representations for RL, without the need of hand engineered data augmentations.
We ﬁnd that C O BERL consistently improves data efﬁciency across the full Atari suite, a set of control tasks and a challenging 3D environment, and often it also increases ﬁnal score performance.
I NTRODUCTION
Developing sample efﬁcient reinforcement learning (RL) agents that only rely on raw high dimensional inputs is challenging.
Speciﬁcally, it is difﬁcult 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 difﬁcult sequence prediction and reasoning tasks.
Speciﬁcally we propose a novel agent, named Contrastive BERT for
RL (C O BERL), 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).
Speciﬁcally, we combine the paradigm of masked prediction from BERT (Devlin et al., 2019) with the contrastive approach of Representation Learning via Invariant Causal
Mechanisms (R E LIC; 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 R E LIC 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 R E LIC and which require to hand-design augmentations.
In terms of architecture, we base C O BERL 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 speciﬁc 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.
ENCODER
TRANSFORMER
GATE
LSTM
V
Y t
X t
RL Loss
ENCODER
TRANSFORMER
GATE
LSTM
V
ENCODER
TRANSFORMER
GATE
LSTM
V
Y
X X
Y
ENCODER
TRANSFORMER
GATE
LSTM
V
X N
Y N ...
RL Loss
RL Loss RL Loss
CONTRASTIVE LOSS
CROSSENTROPY CROSSENTROPY
A
B
C
D
We extensively test our proposed agent across a widely varied set of environments and tasks ranging from 2D platform games to 3D ﬁrst-person and third-person view tasks.
Speciﬁcally 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 R E LIC to the time domain; with this we learn selfattention 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 ﬂexibly integrate knowledge from both the transformer and the LSTM.
• Improved data efﬁciency (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.
M ETHOD
To tackle the problem of data efﬁciency in deep reinforcement learning, we propose two modiﬁcations 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.
2. R EPRESENTATION L EARNING IN REINFORCEMENT LEARNING
It has been empirically observed that performing RL directly from high dimensional observations (e.g. raw pixels) is sample inefﬁcient (Lake et al., 2017); however, there is evidence that learning from low-dimensional state based features is signiﬁcantly more efﬁcient (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 classiﬁed into two classes.
The ﬁrst 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).
C O BERL is part of the ﬁrst set of methods, as it uses a self-supervised loss to improve data efﬁciency.
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 ﬁnite 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
R E LIC (Mitrovic et al., 2021).
R E LIC 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 R E LIC, C O BERL 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 efﬁcient.
Finally, we do not use an additional encoder network as in R E LIC, thus making C O BERL 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 ﬁxed token denoting masking.
Then, let the set T represent indices in the sequence that have been randomly masked and let t ∈ T .
For the i -th training sequence in the batch, for each index t ∈ T , let x it be the output of the GTrXL and
y it the corresponding input to the GTrXL from the encoder (see Fig 1B).
Let φ ( · , · ) be the inner product deﬁned on the space of critic embeddings, i.e. φ ( 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 beneﬁt, the critic can be used to reduce the dimensionality of the dot-product.
To learn the embedding x it at mask locations t , we use y it as a positive example and the sets { y bt } Bb =0 ,b (cid:54)
=
i and { x bt } Bb =0 ,b (cid:54) =
i as the negative examples with B the number of sequences in the minibatch.
We model q tx as
[Equation]
with q tx denoting q tx (cid:16) x it |{ y bt } Bb =0 , { x bt } Bb =0 ,b (cid:54) =
i (cid:17) ; q ty 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 R E LIC.
Speciﬁcally, we look at the similarity between appropriate embeddings and inputs, and within the sets of embeddings and inputs separately.
To this end, we deﬁne:
[Equation]
with p tx and s tx shorthand for p tx ( x it |{ y bt } Bb =0 ) and s tx ( x it |{ x bt } Bb =0 ) , respectively; p ty ( y it |{ x bt } Bb =0 ) and s ty ( y it |{ y bt } Bb =0 ) deﬁned analogously (see Fig 1D).
All together, the ﬁnal objective takes the form:
L ( X, Y )
= − (cid:88)
[Equation]
with sg ( · ) 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 L ( X, Y ) (see App.
C for the details on the weighting).
2. A RCHITECTURE OF C O BERL.
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 difﬁcult to train and are prone to overﬁtting (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 deﬁned by the following equations (see Fig 1A, Gate () , below, has the same form as other gates internal to GTrXL):
[Equation]
The architecture of C O BERL 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 beneﬁted from the ability of the Transformer to process long contextual dependencies, thus helping the LSTM in this respect.
The second beneﬁt 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 conﬁrmed by our ablations in section 4.2.
The learnt gate (eq. 4 and Appendix K), was done to give C O BERL 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 efﬁciency 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.
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 R2D (Kapturowski et al., 2018).
R2D2 Agent 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 beneﬁts 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.
V-MPO Agent 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.
R ELATED W ORK
Transformers in RL.
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 efﬁciency when using transformers in RL.
Contrastive Learning.
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), deﬁning 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,
C O BERL computes Q-values based on the output of the transformer.
Second, C O BERL 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 C O BERL uses an extension of R E LIC (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.
E XPERIMENTS
We provide empirical evidence to show that C O BERL i) improves data efﬁciency 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 V
(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 ﬁve 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 C O BERL.
We use Peng’s Q ( λ ) 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 C O BERL with (Parisotto et al., 2020) and also demonstrate how C O BERL 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.
C O BERL 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 ﬁxed set of hyperparameters with 3 random seeds, whereas in (Parisotto et al., 2020) the results were averaged across hyperparameters.
4. T ESTING C O BERL ACROSS SEVERAL ENVIRONMENTS
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.
We also perform an additional ablation analysis by removing the contrastive loss from C O BERL (see
Sec. 4.2.1).
This baseline demonstrates the importance of contrastive learning in these domains
[Table]
and we show that the combination of an LSTM and transformer is superior to either alone
[Table]
Coberl gTrXL
[Table]
Atari 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 xgray-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.
C O BERL shows the most games above average human performance and signiﬁcantly higher overall mean performance.
Interestingly, the performance of R2D2-GTrXL shows that the addition of GTrXL is not sufﬁcient to obtain the improvement in performance that C O BERL 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 C O BERL, showing that R2D2-GTrXL is indeed a powerful variant on Atari.
Additionally, we observe that the difference in performance in C O BERL is higher when examining the lower percentiles.
This suggests that C O BERL causes an improvement in data efﬁciency.
To conﬁrm 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 conﬁrms a signiﬁcant difference with respect to the model used, t(114)=3.438, p=.008, with C O BERL(M=7380.23, SD=942.29) being better than
GTrXL (M=5192.76, SD=942.01), thus showing that C O BERL is more data efﬁcient.
Control Suite 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, C O BERL is trained only from pixels without state information.
[Table]
[Figure]
A r e a U n d e r t h e
C u r v e
( A
U C )
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 speciﬁc tasks has emphasized data efﬁciency 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 difﬁcult 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 C O BERL, 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
C O BERL 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 ﬁnal 10% of steps.
Across the majority of tasks, C O BERL outperforms D4PG-Pixels.
The increase in performance is especially apparent for the more difﬁcult tasks.
For most of the easier tasks, the performance difference between the C O BERL, R2D2-GTrXL, and R2D2 is negligible.
For ball_in_cup catch , cartpole swingup , finger spin and reacher easy , even the original R2D agent performs on par with the D4PG-Pixels baseline.
On more difﬁcult tasks such as fish swim , and swimmer swimmer6 , there is a very large, appreciable difference between C O BERL, R2D2, and R2D2-GTrXL.
The combination of the LSTM and transformer speciﬁcally 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 , C O BERL 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 signiﬁcant AUC improvement for C O BERL when compared to R2D2-gTrXL, t(24)=1.609, p=.03.
DMLab-30.
To test C O BERL in a challenging 3 dimensional environment we run it in DmLab30 (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 ﬁnal results on the DMLab-30 domain.
If we look at all the 30 games, C O BERL reaches a
A
B A
B
H u m a n
N o r m a l
i s e d S c
o r e
Coberl gTrXL Coberl gTrXL
Coberl gTrXL
A r e a U n d e r t h e
C u r v e
( A
U C ) substantially higher score than GTrXL (C O BERL=115.47% ± 4.21%, GTrXL=101.54% ± 0.50%, t(60)=4.37, p=1.14e-5, Figure 4A).
In Figure 4B we analysed data efﬁciency by computing the AUC and associated average statistics.
In DMLab the difference between models is even more signiﬁcant than Atari, t(60)=6.097, p=9.39e-09, with C O BERL
(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)
4. A BLATIONS
In Sec. 2, we explained contributions that are essential to C O BERL.
We now disentangle the added beneﬁt 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.
I MPACT OF AUXILIARY LOSSES
In Table 3 we show that our contrastive loss contributes to a signiﬁcant gain in performance, both in
Atari and DMLab-30, when compared to C O BERL without it.
Also, in challenging environments like DmLab-30, C O BERL 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 beneﬁt 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 sufﬁcient 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.
[Table]
[Table]
4.2. I MPACT OF ARCHITECTURAL CHANGES
Table 5 shows the effects of removing the LSTM from C O BERL (column “w/o LSTM”), as well as removing the gate and its associated skip connection (column “w/o Gate”).
In both cases C O BERL 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 ﬂexibility in complex environments, we leave it open for future work to explore this.
[Table]
I MPACT OF NUMBER OF PARAMETERS
Table 6 compares the models in terms of the number of parameters.
For Atari, the number of parameters added by C O BERL over the R2D2(GTrXL) baseline is very limited; however, C O BERL still produces a signiﬁcant gain in performance.
We also tried to move the LSTM before the transformer module (column “C O BERL 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 C O BERL.
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 C O BERL 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 (“C O BERL 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)
[Table]
C ONCLUSIONS
We proposed a novel RL agent, Contrastive BERT for RL (C O BERL), which introduces a new contrastive representation learning loss that enables the agent to efﬁciently learn consistent representations.
This, paired with an improved architecture, resulted in better data efﬁciency on a varied set of environments and tasks.
Moreover, through an extensive set of ablation experiments we conﬁrmed that all C O BERL components are necessary to achieve the performance of the ﬁnal agent.
Critically,
C O BERL is fully end-to-end as it does not require an extra encoder network (vs CURL or R E LIC), or data augmentations (vs SIMCLR and R E LIC) 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 efﬁciently train large attentional models in RL.
(For extra conclusions and limitations see App.
J)
O PTIONAL REPRODUCIBILITY STATEMENT
To help with reproducibility we have included extensive details for
• Setup details in Appendix A • Architecture description in Appendix B • Hyper-parameters chosen in Appendix C
We also report the pseudo-code for the algorithm and the auxiliary loss in Appendix H
A S ETUP DETAILS A. R2D2 D ISTRIBUTED SYSTEM SETUP
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 efﬁcient 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:
• x t : the observation at time t .
• 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 0 .
• h t − 1 : recurrent state at the previous time, is initialized with h − 1 = 0 .
They block until they receive Q ( x t , a ; θ ) .
The l -th actor picks a t using an (cid:15) l -greedy policy.
As R2D2, the value of (cid:15) l is computed following:
[Equation]
where (cid:15) =
4 and α = 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 sufﬁcient 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, θ ) for the whole batch, and these values are sent to their corresponding actors.
Replay buffer: it stores ﬁxed-length sequences of transitions T = ( ω s ) t
+ L − 1 s = t along with their priorities p T , where L is the trace length we use.
A transition is of the form ω 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:
• r s − 1 : reward at the previous time.
• 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.
• r s : reward 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 η = 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 θ and θ (cid:48) respectively (Mnih et al., 2015).
The target network’s weights θ (cid:48) are updated to θ every 400 optimization steps.
θ is updated by executing the following sequence of instructions:
• First, the learner samples a batch of size 64 (batch size) of ﬁxed-length sequences of transitions from the replay buffer, with each transition being of length L : T i = ( ω is ) t + L − 1 s = t .
Then, a forward pass is done on the online network and the target with inputs ( x is , r is − 1 , a is − 1 , h is − 1 ) t
+ H s =
t
in order to obtain the state-action values
{ ( Q ( x is , a ; θ ) , Q ( x is , a ; θ (cid:48) ) } .
• With { ( Q ( x is , a ; θ ) , Q ( x is , a ; θ (cid:48) ) } , the Q ( λ ) 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. 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 C O BERL 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. C OMPUTATION USED
R2D2 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 andfor Control).
We use 512 actors, using 4 actors per CPU core, with each one performing ∼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.
V-MPO 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. C OMPLEXITY ANALYSIS
As stated, the agent consists of layers of convolutions, transformer layers, and linear layers.
Therefore the complexity is max { O ( n 2 · d ) , O ( k · n · 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.
B A RCHITECTURE DESCRIPTION B. E NCODER
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 ﬁlter size 3 x 3 ( 1 · 4 layers).
• Number of residual bottleneck blocks (in order): (2 , 4 , 6 , 2) .
Each block has 3 convolutional layers with ReLU activations, with ﬁlter sizes 1 x 1 , 3 x 3 , and 1 x 1 respectively ( (2+4+6+2) ·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 ﬁnal 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 ﬁnal 512 -dimensional output.
B. T RANSFORMER
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:
• Number of layers: 8 . • Memory size: 64 . • Hidden dimensions: 512 .
• Number of heads: 8 .
• Attention size: 64 .
• Output size: 512 .
• Activation function: GeLU.
For DmLab the transformer has the following characteristics:
• Number of layers: 12 . • Memory size: 256 for C O BERL and 512 for gTrXL.
• Hidden dimensions: 128 .
• Number of heads: 4 .
• Attention size: 64 .
• Output size: 512 .
• Activation function: ReLU.
the GTrXL baseline is identical, but with a Memory size of 512 .
B. LSTM AND V ALUE 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.
R2D2 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.
V-MPO 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.
13 for details).
B. C RITIC F UNCTION
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.
C H YPERPARAMETERS
For the experiments in Atari57 and the DeepMind Control suite, C O BERL 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 .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.
C. A TARI AND DML AB PRE - PROCESSING
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.
ALE is publicly available at https://github.com/mgbellemare/
Arcade-Learning-Environment .
[Table]
C. C ONTROL S UITE PRE - PROCESSING
As mentioned in 4, we use no pre-processing on the frames received from the control environment.
C.
[Table]
PROCESSING
[Table]
C. C ONTROL ENVIRONMENT DISCRETIZATION
As mentioned
[Table]
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.
[Table]
C. H YPERPARAMETERS U SED
We list all hyperparameters used here for completeness.
We started by optimizing the hyperparameters of GTrXL highlighted in bold in Tab 10 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 ﬁxed throughout all the experiments.
Table 10 reports all the hyperparameters of the R2D2 experiments, both the ﬁxed ones and the ones with the ranges over which we did the sweep.
The ﬁxed hyper-parameters were taken from
Kapturowski et al.
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.
We believe that in this way we ensure to have a properly tuned baseline that enforces a fair comparison.
Tablereports the chosen hyperparameters that we found to optimize the performance of GTrXL on the
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 [11].
Consequently, we re-ran the same procedure as before, but we ﬁxed all the previous hyperparameters optimized for GTrXL (see Tab.and
we perform a grid search over the same 10 Atari games to ﬁnd 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 C O BERL.
For DmLAB we optimized the hyperparameters of GTrXL on all the 30 games.
Table 13 reports both the ﬁxed ones and the ones with the ranges over which.
The ﬁxed hyper-parameters were taken directly from Parisotto et al.
and we sweep over ‘’Epsilon Alpha”, “Target Update Period” and ’Memory size’ to make sure we maximised performance of this baseline.
Table 14 reports the chosen hyperparameters that we found to optimize the performance of GTrXL on DmLAB.
We then moved to CoBERL.
Again, by keeping ﬁxed all the previous hyperparameters optimized for GTrXL we perform a grid search over all the 30 games to ﬁnd 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 ﬁnd any signiﬁcant 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.
[Table]
=
[Table] [Table]
[Table] [Table]
[Table] [Table]
D A DDITIONAL 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
[Table] [Table]
E G AME SCORES
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
enduro
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
ﬁshing
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
frostbite
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
gopher
[Table]
[Table] [Table]
[Table] [Table]
[Table]
gravitar
[Table] [Table]
[Table] [Table]
[Table] [Table]
hero
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
ice
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table]
jamesbond
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
kangaroo
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
97 ±
[Table]
krull
[Table]
24 ± 19282.
[Table]
15 ± 3103.
3675.08 ± 5575.
2512.84 ± 1288.
space invaders 37497.11 ± 9973.
19113.28 ± 8512.
4704.35 ± 2763.
star gunner 103870.95 ± 23124.
78559.08 ± 42110.
88169.92 ± 13783.
surround 9.37 ± 0.
9.40 ± 0.
8.60 ± 1.
tennis 14.08 ± 11.
9.20 ± 11.
-0.04 ± 0.
time pilot 36085.82 ± 6934.
14250.32 ± 1651.
21051.93 ± 4156.
tutankham 33.93 ± 28.
21.66 ± 13.
79.51 ± 40.
up n down 407873.64 ± 96615.
237945.96 ± 101020.
438679.20 ± 38137.
venture 1754.74 ± 229.
1632.54 ± 247.
1624.82 ± 100.
video pinball 248629.59 ± 231446.
48305.17 ± 42432.
339391.16 ± 194812.
wizard of wor 19454.60 ± 7061.
14127.73 ± 9153.
18270.16 ± 12166.
yars revenge 385904.98 ± 264624.
217128.25 ± 28248.
248086.47 ± 162482.
zaxxon 20411.27 ± 12535.
25063.44 ± 10385.
31679.84 ± 10447.
DmLab Levels coberl gtrxl rooms collect good objects test 9 .
74 ± 0 .9 .
69 ± 0 .rooms exploit deferred effects test 56 .
57 ± 3 .54 .
87 ± 0 .rooms select nonmatching object 64 .
11 ± 10 .58 .
48 ± 4 .rooms watermaze 56 .
61 ± 1 .51 .
72 ± 3 .rooms keys doors puzzle 33 .
99 ± 5 .35 .
53 ± 8 .language select described object 611 .
02 ± 1 .605 .
59 ± 14 .language select located object 606 .
05 ± 3 .584 .
23 ± 3 .language execute random task 232 .
67 ± 28 .198 .
45 ± 15 .language answer quantitative question 327 .
26 ± 4 .297 .
57 ± 9 .lasertag one opponent large 15 .
47 ± 3 .11 .
65 ± 3 .lasertag three opponents large 22 . 99 ± 3 .28 .
46 ± 4 .lasertag one opponent small 32 .
64 ± 3 .28 .
62 ± 3 .lasertag
three opponents small 42 .
96 ± 3 .34 .
54 ± 0 .natlab
ﬁxed large map 46 .
03 ± 8 .43 .
46 ± 14 .natlab varying map regrowth 31 .
07 ± 6 .27 .
62 ± 9 .natlab varying map randomized 48 .
97 ± 8 .32 .
79 ± 11 .platforms hard 58 .
10 ± 7 .35 .
67 ± 23 .platforms random 85 .
15 ± 1 .69 .
01 ± 1 .psychlab continuous recognition 58 .
91 ± 2 .54 .
96 ± 3 .psychlab arbitrary visuomotor mapping 51 .
56 ± 0 .59 .
96 ± 0 .psychlab sequential comparison 33 .
56 ± 0 .35 . 89 ± 0 .psychlab visual search 77 .
96 ± 1 .76 .
92 ± 0 .explore object locations small 83 .
07 ± 5 .71 .
85 ± 6 .explore object locations large 62 .
57 ± 5 .60 .
61 ± 2 .explore obstructed goals small 2 .
85 ± 5 .245 .
40 ± 10 .explore obstructed goals large 109 .
61 ± 1 .79 .
48 ± 6 .explore goal locations small 369 .
20 ± 3 .340 .
19 ± 2 .explore goal locations large 139 .
74 ± 9 .123 .
20 ± 14 .explore object rewards few 71 .
47 ± 14 .69 .
02 ± 0 .explore object rewards many 104 .
62 ± 2 .94 .
97 ± 1 .
F
L EARNING C
URVES
F.
A TARI L EARNING CURVES
[Figure]
DMC ONTROL L
[Figure]
E p
i s o d e
R e t u r n
DML AB L EARNING CURVES
[Figure]
G L ICENSES
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.
License.
The latest version of the source code is publicly available at: https://github.com/ deepmind/lab/tree/master/game_scripts/levels/contributed/dmlab
H P SEUDO - CODE
Algorithm 1 Pseudo-code for CoBERL training_iterations ←initialise_weights(VisualEncoder, TransformerStack, Gate, ValueNetwork) while training_iterations ≤ max_training_iterations do sampled_batch ← sample_experience() encoded_images ← VisualEncoder(sampled_batch) encoded_actions ← OneHotActionEncoder(sampled_batch)
combined_inputs ← concat(encoded_images, previous_rewards, encoded_actions) transformer_inputs, transformer_targets, mask ← create_masks_targets(combined_inputs)
output_transformer_contrastive ← TransformerStack(transformer_inputs, causal_mask
=False) contrastive_loss ← compute_aux_loss(output_transformer_contrastive, transformer_targets, mask)
output_transformer_RL ← TransformerStack(transformer_inputs, causal_mask=
True) gated_output ← Gate(combined_inputs, output_transformer_RL)
lstm_output ← LSTM(gated_output) combined_inputs_for_value_net
← concat(lstm_output, combined_inputs) value_estimation ← ValueNetwork(combined_inputs_for_value_net) rl_loss ← compute_rl_loss(value_estimation, extra_args) total_loss ← rl_loss + contrastive_loss end while auxiliary
=
) 9# Compute labels index labels_idx = arange(
]) labels
= logits
logits_
= matmul(input2, jnp.transpose(input2)) logits_12 =
, jnp.transpose(input2))
logits_
= matmul(input2, jnp.transpose(input1))
# Calculate invariance penalty.
inv_penalty =
kl_with_logits( # [B * T] stop_gradient(logits_11), logits_22, axis=-1) inv_penalty += kl_with_logits( stop_gradient(logits_12), logits_22, axis=-1) inv_penalty += kl_with_logits( stop_gradient(logits_21), logits_11, axis=-1) inv_penalty += kl_with_logits( stop_gradient(logits_12), logits_21, axis=-1) inv_penalty = inv_penalty / 4.
# [B * T] 32logits_11 = logits_11 - mask
* 1e  logits_22 = logits_22 - mask * 1e 35logits_1211 = concatenate([logits_12, logits_11], axis=-1)
logits_2122 = concatenate([logits_21, logits_22], axis=-1) 38loss_12 = -sum(labels * log_softmax(logits_1211), axis=-1) loss_21 = -sum(labels * log_softmax(logits_2122), axis=-1) 41loss = reshape(loss_12 + loss_21, [batch_size, seq_dim])
* mask_ext inv_penalty =
reshape(inv_penalty, [batch_size, seq_dim])
* mask_ext 44loss = mean(loss + inv_penalty) 46return
loss
I A REA UNDER THE CURVE
For all the levels we calculated the AUC by integrating composite Simpson’s rule with a delta(x) ofsteps.
We use the intergate package from scipy (Virtanen et al., 2020).
J L IMITATIONS AND F UTURE W ORK
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.
K E XTRA INFORMATION ABOUT THE GATE EMPLOYED IN EQUATION
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 deﬁned in the following way:
[Equation] [Equation]
where:
[Equation]
and:
[Equation]
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.
L E XTRA ANALYSIS
We also attempted higher masking rate, but as shown in ﬁgure L 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 signiﬁcant advantage of C O BERL over gTrXL t(60)=2.616, p=0.0011.
C O BERL M=545.18, SD=282.58; gTrXL M=374.49, SD=282.58.