CoBERL: Contrastive BERT for Reinforcement Learning
Authors: Andrea Banino
Adria Puigdomenech Badia
Jacob C Walker
Tim Scholtes
Jovana Mitrovic
Charles Blundell
Venue: ICLR 2022 Spotlight (2022)
Keywords: Reinforcement Learning
Contrastive Learning
Representation Learning
Transformer
Deep Reinforcement Learning

Abstract: Many reinforcement learning (RL) agents require a large amount of experience to solve tasks. We propose Contrastive BERT for RL (COBERL), an agent that combines a new contrastive loss and a hybrid LSTM-transformer architecture to tackle the challenge of improving data efficiency. COBERL enables efficient 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 find that COBERL consistently improves data efficiency across the full Atari suite, a set of control tasks and a challenging 3D environment, and often it also increases final score performance.

Published as a conference paper at ICLR 2022
COBERL: CONTRASTIVE BERT FOR REINFORCE- MENT LEARNING
Andrea Banino I,
Adrià Puidomenech Badia !,+
Jacob Walker
Tim Scholtes I,+
Jovana Mitrovic
Charles Blundell
ABSTRACT
Many reinforcement learning (RI.) agents require a large amount experience to solve tasks. We propose Contrastive BERT for RI (CoBERL), an agent that combines a new contrastive loss and a hybrid tackle the challenge of improving data efficiency. COBERL enables efficient and robust learning from pixels across a wide variety of domains. We use bidirectional masked prediction in combination with a generalization of recent contrastive method to learn better representations for RL, without the need of hand engineered data augmentations. We find that COBERL consistently improves data efficiency across the full Atari suite, a set of control tasks and a challenging 3D environment, and often it also increases final score performance.
architecture to
INTRODUCTION
Developing sample efficient reinforcement learning (RL) agents that only rely on raw high dimen- sional 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, 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 novel agent, named Contrastive BERT for RL (CoBERL), that combines a new contrastive representation learning objective with architectural improvements that effectively combine L.STMs 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 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. representations that contain the appropriate
Such a signal aims to learn celf-attention-coniitte 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, base CoBERL on Gated Transformer-Xi XL (GT:XL; Parisotto et al. 2020) and Long-Shor Term Memories (LSTMs; Hochreiter & Schmidhuber. 1997). GTxXL is an adaptation of transformer architecture specific for RI. domains. We combine GTXXL with LSTMs using a gate trained via an RI. loss. This allows the agent to learn to exploit the representations
Mend London, Equal Contribution, Corresponding author: abaní no@deepmind.

---

Published as a conference paper at ICLR 2022
LETTO
GATE
IMASSOORBEE
TRAB-FOOBUEE INGeSroeeee
...
LOSS
ENCODER
CROSS
Figure 1: COBERL. A) General Architecture. We use a residual network to encode observations into embeddings Y.. We feed Y, through causally masked GTrXL transformrr, which computes the predicted masked inputs X, and passes those together with Y, learnt gate. The output of the gate is passed through single LSTM layer produce compute contrastive loss using predicted masked inputs X, and Y; as targets. For this, we do not use the causal mask of the Transfomer. For details about the contrastive loss, please see Section 2.1 C) Computation of q and d (from Equation I) with the round brackets denoting the computation of the similarity between the entries D) Regularization terms from Eq. 3 which explicitly enforce self attention consistency.
values that we use computing the loss. B) Contrastive loss. We also
offered by the transformer only when an environment requires it, and avoid the extra complexity when not needed.
We extensively our proposed agent across 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 RL using contrastive objectives without the need of hand-engineered augmentations.
An improved architecture that, using gate, allows us to flexibly integrate knowledge from both the transformer and the LSTM.
Improved data efficiency (and in some cases also overall performance) varied set of environments. Also, we show that individually both our contrastive loss and the architecture improvements play role in improving performance.
2
METHOD
To tackle the problem of data efficiency in deep reinforcement learning, we propose two modifications the status quo. First, we introduce novel representation learning objective aimed 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 L.STMs and transformers.
REPRESENTATION LEARNING IN REINFORCEMENT LEARNING
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,
2

---

Published as a conference paper at ICLR 2022
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 supe-sueee 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 world model and uses this collect imagined rollouts, which then act as extra data to train the RL al 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 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 right-to-left, as is common with RNN- 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 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 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.] (02T). 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
explain the details of the auxiliary loss. In 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 T represent indices in the sequence that have been randomly masked and let t E TT For the i-th training sequence in the batch, for each index & E TT, let 2 be the output of the GTrXL and yt the corresponding input to the GTrXL from the encoder (see Fig 18).
Let ( )) be the inner product defined on the space of critic embeddings. i.e. o0x, v) = g(x)² π()).
where g is a critic function. The critic separates the embeddings used for the contrastive proxy
task and the downstream RI. task. Details of the critic function are in App. [ 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 side benefit, the critic can be used to reduce the dimensionality of the dot-product. To learn the embedding of at mask locations f, we use of as positive example and the sets number of sequences in the minibatch. We model dh as
{ 100,000 and {x 10 as the negative examples with B the
exp(ø(e), v() E (xp()(ai, vtt) + (xp()(xt,=7))
with of denoting v ( ( oz : o is computed analogously (see Fig 1 IC). In order to
enforce self- attentitn consistency in the learned representations, we explicitly regularize the similarity between the pairs of transformer embeddings and inputs through Kullback-Leibe 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:
3

---

Published as a conference paper at ICLR 2022
exp(o(a), vi)) B
of explo(xi, x;))
P/
Le 2o exp(o(x) v())
Lr exp(o(x),,
B
with pl and s/ shorthand for pl (yt and s (1 (atl) respectively; Py (w) and
Syill )o defined analogously (see Fig ID). All together, the final objective takes the form:
((X,Y)= = Σ TE
(log d log (.y) +a
Σ 1ET
K KL sg(s;)) K
(3) with sg(-) indicating a ttop-gradiens As in our RL objective, we use the full batch of sequences that sampled from the replay buffer to optimize this contrastive objective. Finally, optimize weighted sum of the RL objective and L(X, (see App. C r the details on the weighting).
KL sg($;)) + K sg( (s)))
2.2 ARCHITECTURE OF COBERL.
While transformers have proven very effective at connecting long-range data dependencies natural language processing ( (Vawwani et al. 2017 Brown et al. 2020 Devlin et al., 2019) and computer vision (Carion et al., 2020 Dosovitskiy
(021), ) in the RL setting they are difficult to train and are prone to overlitting (Parisotto et al. (020). In contrast, LSTMs have long been demonstrated to be useful in RI. 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 [TA. Finally we also have skip connection from the transformer input to the LSTM output.
More concretely, let Y, be the output of the encoder network time t, then the additional module can be defined by the following equations (see Fig [A Gate(), below, has the same form as other gates internal to GTrXL):
X, GTrXL(YY): Z; Gate X,);
Output, concatenate(LSTTM())) Y))
(4)
The architecture of COBERL based on the idea that LSTMs and Transformer can complementary. In particular, the LSTM is known for having short contextual memory. However, by putting 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 that adding recurrent networks in architectures has the ability to closely emulate the behavior of non-recurrent but deeper models, but does so with far fewer parameters (Schwarzschild Having fewer parameters is an aspect of particular interest in RL, where gradients are noisy and hence training larger model more complicated than supervised learning (McCandlish et al., 2018). Both hypotheses have been empirically confirmed by our ablations in section 4.2
loss in performance. This hypothesis come from studies showing
(021).
The learnt gate (eq. and Appendix K). was done 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 tune the weights the gate to learn from them. We hypothesize that this would also help in terms of data efficiency 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 the output of the LSTM taken from Kapturowski et al. (018))
Since the architecture is agnostic to the choice of RL regime evaluate it in both on-policy and off off-policy settings. For on-policy, we use V-MPO (Song et al. 2219)), and for off-policy we use R2D2 (Kapturowski et al. 2018).
4

---

Published as a conference paper at ICLR 2022
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 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" portion of the unrolled network with replayed sequences during training.
CoBERL architecture
V-MPO Agent Given V-MPO's strong performance on DMI abb30,, in particular in conjunction with the GTrXI, architecture (Parisotto et al., 2020) which a key component of CoBERI., 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 target distribution for policy updates, subject to a sample-based based KL constraint, and gradients are calculated to partially move the parameters towards the target, again subject KL constraint. Unlike MPO, V-MPO uses a learned saate-value function (s) instead of a state-action value function.
3 RELATED WORK
Transformers in RL. The transformer architecture (Vaswani et al. 2017) has recently emerged as one of the best performing approaches in language modelling (Dai al. 2019 Brown et al. 2020) and question answering (Dehghani et al. 20188] Yang et al. 2019). More recently it has also been successfully applied to computer vision (Dosovitskiy et al. (021). 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 the high variance of the gradients in RL relative to that of (un)supervised learning problems. In this work, we build upon GTrXI. and demonstrate that, perhaps for RL: attention is not all you need, and by combining GTyXI. 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 simple gating is most effective in our experiments. Also, unlike GTrXL, we show that using unsupervised representation learning loss that enforces self-attention consistency is an effective way to enhance data efficiency 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 al. 2020 Caron al.) 2020 Mitrovic et al. (021). These methods have also been leveraged in an RL setting with the hope of improving performance. Apart from MRA (Fortunato al., 2019) mentioned above, one of the early examples of this is CURL (Srinivas et al. (020aa) which combines Q-Learning with separate encoder used for representation learning with the InfoNCE loss from CPC (Oord 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 | (021) and learning abstract representations of state-action pairs (Liu et al. 202II. The closest work to our is M-CURL (Zhu et al. 2020). Like our work, it combines mask prediction, transformers, and contrastive learning, few key differences. First, unlike M-CURL who use 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. in App. Kto produce the input for the Q-network, while M-CURI 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 al. 2020a) and compares encodings from the transformer with the separate momentum encoder.
but there
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
5

---

Published as a conference paper at ICLR 2022
experiments, we demonstrate performance Atari57 (Bellemare al. 2013), the DeepMind Control Suite Tassa et al., 2018), and the DMLab-30 (Beattie et al. 2016). Recently, Dreamer V2 (Hatner et al. 2020b) has emerged as 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 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
the distributed setup used see App. A. 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-parametter performance (see App.[C.5|rr for the detailed procedure). We use a ResNet as the encoder for COBERL. We use Peng's s (()) as our loss (Peng & Williams 1994). To ensure that this is comparable R2D2, we also run an R2D2 baseline with this loss. For thorough description of the architectures for all environments see App.
of all the baselines are tuned individually to maximise
For DMLab-30 30 use V-MPO (Song al. 2019) to directly compare COBERL with (Parisotto cet al. 2020) and also demonstrate how CoBERL may be applied to both on and off-policy learning. The experiments were run using Podracer setup (Hessel et al. (021). details of which may be found in 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 al. 2020). Compared to (Parisotto et al. 2020) we have two differences. Firstly, all the networks run without pixel control loss (Jaderberg al. 2016) as not 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.
4.1 TESTING COBERL 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 al. 2013). DeepMind Lab (Beattic et al. 20166) as well as the DeepMind Control Suite (Tassa et al. 2018). 8) To help with comparisons, 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. 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. 1.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.
has
> h.
Mean
Median
25thh Pet
5th Pct
Nrca Undertrc MUC
COBERL
49
1368.37%) + 34.10%
296.22%
13.65% 150.56%
R2D2-GTTXI.
48
1088.90% + 53.45%
273.83%
141.66%
5.94%
R2D2
47
1017.18% + 33.85%
269.26%
118.96%
4.18%
Rainbow
43
874.0%
231.0%
101.7%
4.9%
Dreamer V2*
37
631.1%
162.0%
76.6%
2.5%
Table 1: The human normalized scores on Atari-57. > h indicates the number of tasks for which performance above average human achieved. indicates that was run on 55 games with sticky actions; Pet refers to percentile.
gTTXL
Figure 2: Area under the curve (AUC) for Atari on all 57 levels (higher is better).
2015 Hessel et al. 2018 Machado et al. 2018 Hafner al., (020b),, 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 84x84 gray-scaled frames that are used as input to our agent. We do not use frame stacking.
Atari As commonly done in literature (Mnih et al.
Table [ 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-GTiXL shows that the addition of GTxXL is not sufficient to obtain the improvement in performance that CoBERL exhibits (in 4.2 we will demonstrate that both the contrastive loss and L.STM contribute this improvement). R2D2 R2X2-GT-XL also exhibits slightly

---

Published as a conference paper at ICLR 2022
better median than COBERL, showing that R2D2-GT:XL is indeed powerful variant Atari. Additionally, we observe that the difference in performance in CoBERL is higher when examining the lower percentiles. This suggests that CoBERL causes improvement in data efficiency. To confirm this tendency we calculated the area under the curve (AUC) for the learning curves presented in Appendix and we perform a I-fest analysis it (please see App details on how AUC was calculated). Figure 2 present this analysis, and it confirms a sigmificant difference with respect to the model used, t(1 14)=3.433, p=.008, with COBERL(M-7380.23 SD=942.299 being better than GTrXL (M=5192.76, S01776,,DD94444 SD=942.01), thus showing that COBERL is more data efficient.
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 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 App. C.4 We do not use pre- processcsssng on the environment frames. Finally, COBERL is trained only from pixels without state information.
Suiie
COBERL
ROD--GTINL
R3D2
D4FG-Pruel
acrohor swingup
355.25 + 332
265.00 + 113.27
327.16 16 5,35
81.7 = 4,4
swim
597.66
60.09
82,30
345.63
.44
upright
95260
5.16
21.70
936.09
11.58
408.7
19.6
land swingup
835.06 + 9.38
775.72
50.06
831,86
= 61.54
690.9
+ 41.9
swimmer swimmerts
419.26
48.37
200.55
+ 58.37
329.61
26.77
194.7 - 15.9
finger spun
965.96
1.69
983.74
10.25
980.55
965.7 0.6
rescher easy
984.00
2.76
+: 2.24
967.4 4.1
Anaa Undd
ehectah zun
$23.36 +: 41.38
105.23 +: 144.82
365.45 =: 30,40
523.8 6.8
walker walk
781.
26.16
50 70.78
687.18 = 18.15
968.3 1.8
ball cup
979.10
980.5
cattpole swingup
812.01 7.61
83601
4.37
816.23
$62.0
compole swingup sparse
714.80 + 16.18
743.71
+ 9.27
762.57 2: 6.71
480.0 - 56.6
Table 2: Results tasks in the DeepMind Control Suite. CoBERL, R2D2-GTrxL, R2D2, and D4PG-Pixels are trained on 100M frames. the appendix (Table 17). we show three other approaches as reference and not as directly comparable baseline.)
Coberl
R2D2-gTiXL
Figure 3: Area under the curve (AUC) for DmControl (higher is better).
We include six tasks popular in current literature: ball_in
catch, cartpole swingup,
cheetah run, finger spin, reacher easy, and wa 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, fish swim, fish upright, pendulum swingup, and
cartpole swingup_sparee 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 (Haamoja et al. 2018) pixels demonstrated in (Srinivas al., 2020a), and D4PG-Pixels (Tassa 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 much larger number of available environment steps per gradient update. We run for 100M environment steps as with D4PG-Pixel and we compute performance for our approaches by taking the evaluation performance of the final 10% of steps. Across the majority of tasks, COBERI. outperforms D4PG-Pixels. The increase in perfor- mance is especially apparent for the more difficult tasks. For most of the casier tasks, the performance difference between the COBERL, R2D2-GTIXL, and R2D2 is negligible. For ball 1elliin_on oup catch, cartpole swingup, finger spin and reacher easy, even the original R2D2 agent performs on par with the D4PG-Pixel baseline. On more difficult tasks such fish swim, and swimmer swimmers r 6, there a very large, appreciable difference between COBERL, R2D2, and R2D2-GTYXL. The combination of the LSTM and transformer specifically makes a large differ- ence 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. we report learning curves for each game). Figure shows that also this domain, we see a significant AUC improvement for COBERL when compared to R2D2-gTrXL, 1(24)=1.609, p=.03.
DMLab-30. To test CoBERL in a challenging 3 dimensional environment we run it in DmLab- 30 (Beattie al. 2016). The agent was trained at the same time all the 30 tasks, following
7

---

Published as a conference paper at ICLR 2022
Anew Undeer the Curve (AUC)
1500
1:3
1000
vin
GTXL
Cobori
Figure 4: DMLab-30 experiments, a) Human normalised returns in DMLab 30 across all the 30 levels (higher is better). b) Area under the curve for all the 30 DMLab levels. Results are over 5 seeds and the final 5% of training (lower) c) Area under the curve (AUC) for DML ab- across all the 30 levels. (higher is better)
the setup of GTrXI (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.477) + 4.21%, G1xXL=101554% + 0.50%, 1(60))4.37, 1=1.14...00.1 1e-5, Figure A). In Figure 4 we analysed data efficiency by computing the AUC and associated average statistics. In DMLab the difference between models is more significant than Atari, 1(60)=6.09,, p=9.39e-09, with COBERL (M=1193.28, SD=383.22 having better average AUC than GTrXL (M=764.09, SD=224.51), hence showing that our methods scales well more complex domains. (see Appendix for learning curves)
4.2 ABLATIONS
In Sec. 2, explained contributions that 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 the results. Ablations are Atari games chosen match the ones in the original DQN publication (Mnih et al. 2013), and on all the 30 DMLab games.
4.2. IMPACT OF AUXILIARY LOSSES
In Table 3 we show that our contrastive loss contributes significant gain in performance, both in Atari and D0, 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 we consider the median score on the reduced ablation set of Atari games. However the case of the DmLab-30, where we consider larger set of levels (7 vs. 30), there is a clear benefit of the auxiliary loss.
Moreover, Table 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 close possible to (Mitrovic the column CoBERL w/o masking shows that removing the masking technique derived from the language literature makes the results substantially worse
(ct al. (021). As seen on Table 4
COBERL
COBERL w/o aux loss
GTrXL baseline*
Mean
115.47% 1 4.21%
105.29%
1 2.02%
101.54% 1 0.50%
DMLab-30
Median
110.86%
104.84%
101.35%
Mean
726.34%
44.06%
490.61% 22.48%
605.62%
47.46%
Atari
Median
270.03%
357. 757.96%
394.53%
Table 3: Impact of contrastive loss. Human normalized scores on Atari-57 ablation tasks and DMLab-30 tasks. for DMLab the baseline is GTxXL trained with VMPO, for Atari the baseline is GTxXL trained with R2D2.
CoBERL
CoBERL CURL
with CoBERL with Sim- CLR
CoBERL w/o mask-
ing 331.05% 59.05%
Mean
726.34% 270.03%
44.06%
608.58 53% 62.93%
347.83 447.83% 46.74% 265.00%
Atari
Median
272.34%
233.63%
Table 4: Comparison with alternative auxiliary losses.
8

---

Published as a conference paper at ICLR 2022
4.2.2 IMPACT OF ARCHITECTURAL CHANGES
shows the effects of removing the LSTM from CoBERL (column "wo 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 [ 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.
CoBERL
w/o LSTM
w/o Gate
Mcan
115.47% + 4.21%
96.29%
+ 3.39
86.42% + 7.25%
DMLab-30
Median
110.86%
99.83%
95.21%
Mean
726.3 34% + 44.06
504.42%
I 48.51%
528 528.61% 73.21%
Atari
Median
270.03%
249.62%
283. 16%
Table 5: Impact of Architectural changes. Human normalized scores Atari ablation tasks and DMLab-30
4.2.3 IMPACT OF NUMBER OF PARAMETERS
Table 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 transforner 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 COBERL. This goes in the direction of hypothesis that having the LSTM after the Transformer could allow the former to exploit the extra context provided by the latter. For DMLab-30, is worth noting that CoBERL has 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 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 reduce computation complexity (given fewer parameters).
COBERL
GTxXL
COBERL with R2D2 LSTM
LSTM before
Mean
115.47% + 4.21%
101.5%% + 0.30%
101 101.99% + 180%%
N/A
DMLab-30
Median
110.86%
103.82%
100.0%
N/A
Num. Params.
17M
66M
47M
N/A
Mean
726.34% 44.06
605.62%=444 47.46%
604.36% 78.60%
362.38%
20.28%
Atari
Median
270.03%
394.53% 53%
296.72%
202.71%
Num. Params.
46M
42M
46M
18M
Table 6: Effeec of number of parameters. Human normalized scores Atari-57 ablation tasks and DMLab-30 tasks. *for DMLab-30 the baseline is GTrXL trained with VMPO with a memory size 512, for Atari the baseline is GTrXL trained with R2D2 with memory size of 64.
5 CONCLUSIONS
We proposed novel RL agent, Contrastive BERT for RL (COBERL), which introduces new contrastive representation learning loss that enables the agent to efficiently learn consistent represen- tations. This, paired with an improved architecture, resulted in better data efficiency on 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 g gTxXL. 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

---

Published as a conference paper at ICLR 2022
OPTIONAL 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
We also report the pseudo-code for the algorithm and the auxiliary loss in Appendix H
10

---

Published as a conference paper at ICLR 2022
REFERENCES
Abbas Abdolmalcki, Jost Tobias Springenberg, Jonas Degrave, Steven Bohez, Yuval Tassa, Dan Belov, Nicolas Heess, and Martin Riedmiller. Relative entropy regularized policy iteration. arXiv preprint /rXiv:/8/2.022256 2018.
Rishabh Agarwal, Marlos C. Machado, S. Castro, and Marc G. Bellemare. Contrastive behavioral similarity embeddings for generalization reinforcement learning. ArXiv, abs/2101.05265, 2021.
Charles Beattie, Joel Z. Leibo, Denis Teplyashin, Tom Ward, Mareus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Victor Valdés, Amir Sadik, et al. Deepmind lab. arXiv preprint arXiv:1612.03801. 2016.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. NeurIPS, 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In ECCV, 2020.
M. Caron, I. Misra, J. Mairal, Priya Goyal, P. Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. ArXiv, abs/2006. 00882. 2020.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, pp. 1597-1607, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoe Le, and Ruslan Salakhutdi- Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv:1901.02860, 2019.
Mostafa Dehghani. Stephan Gouws, Oriol Vinyals. Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. arXiv preprint arXiv.1807.038/1 2018
Jacob Devlin, Ming-Weieeeeee Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre- Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume 1 (Long and Short Papers), pp. 4171-4186, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2021.
Meire Fortunato, Melissa Tan, Ryan Faulkner, Steven Hansen, Adrià Puigdomènech Badia, Gavin Buttimore, Charlie Deck, Joel Z Leibo, and Charles Blundell. Generalization of reinforcement learners with working and episodic memory. NeurIPS, 2019.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pp. 297-304, 2010.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.10122. 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In ICMI, 2018.
11

---

Published as a conference paper at ICLR 2022
Raia Hadsell, Sumit Chopra, and Yann LeCun. Dimensionality reduction by learning invariant mapping. In 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR 066,, volume 2. pp. 1735-1742. IEEE, 2006.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In ICLR, 2020a.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020b.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence. volume 32, 2018
Matteo Hessel, Manuel Kroiss, Aidan Clark, lurii Kemaev, John Quan, Thomas Keck, Fabio Viola, and Hado van Hasselt. Podracer architectures for scalable reinforcement learning, 2021.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1780, 1997.
Max Jaderberg, Volodymyr Mnih, Wojeiech Marian Czarnecki, Tom Schaul, Joel Z Leibo, David Silver, and Koray Kavukcuoglu. Reinforcement learning with unsupervised auxiliary tasks. arXiv preprint ar (0..61...5999 .0539, 2016.
Lukasz Kaiser, Mohammad Babaeizadeh, Piotr Milos, Blazej Osinski, Roy H Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, al. Model-based reinforcement learning for atari. arXiv preprini D00i:v:/900000000 2019.
Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations, 2018.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
Brenden M Lake, Tomer D Ullman, Joshua Tenenbaum, and Samuel Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
Guoqing Liu, Chuheng Zhang, Li Zhao, Tao Qin, Jinhua Zhu, Li Jian, Nenghai Yu, and Tie- Yan Liu. Return-based contrastive representation learning for reinforcement learning. In ICLR 2021, January 2021.
Marlos C Machado, Mare G Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 61 61:5255,,, 2018
Bogdan Mazoure, R' eim Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjelm. Deep reinforeement and infomax learning. ArXiv, abs/2006.07217, 2020.
Sam McCandlish, Jared Kaplan, Dario Amodei, and OpenAl Dota Team. An empirical model of large-batch training. arXiv preprint arXiv: arXiv:18/2.06162,, 2018
Jovana Mitrovic, Brian McWilliams, Jacob Walker, Lars Buesing, and Charles Blundell. Repre- sentation learning via invariant causal mechanisms. In International conference on learning representations, 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Bellemare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 518(7540::59-5533 5333, 2015.
12

---

Published as a conference paper at ICLR 2022
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1rXiv:180.033748, 2018.
Emilio Parisotto, Francis Song. Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International Conference Machine Learning, pp. 7487-7498. PMLR, 2020.
Jing Peng and Ronald J Williams. Incremental multi-stey q-leariin.. In Machine Learning Proceed- ings 1994, pp. 222.22..22 Elsevier, 1994.
Jürgen Schmidhuber. Making the world differentiable: On using self-superviser fully recurrent eu al networks for dynamic reinforcement learning and planning in non non-stationary environm nts. 1990.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839)6604 2020.
Max Schwarzer, Ankesh Anand, R. Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bach- man. Data-efficient reinforcement learning with momentum predictive representations. ArXiv, abs/2007.05929 2020.
Avi Schwarzschild, Arjun Gupta, Amin Ghiasi, Micah Goldblum, and Tom Goldstein. The uncanny similarity of recurrence and depth. arXiv preprint arXiv:2102.11011, 2021.
H Francis Song, Abbas Abdolmalcki, Jost Tobias Springenberg. Aidan Clark, Hubert Soyer, Jack W Rae, Seb Noury, Arun Ahuja, Siqi Liu, Dhruva Tirumala, et al. V-mpo: On-policy maximum posteriori policy optimization for discrete and continuous control. arXiv preprint Niti:1999..12388.. 2019.
A. Srinivas, M. Laskin, and P. Abbeel. Curl: Contrastive unsupervised representations for reinforce- ment learning. In ICML, 2020a.
Aravind Srinivas, Michael Laskin, and Pieter Abbeel. Curl: Contrastive unsupervised representations for reinforcement learning. arXiv preprint arXiv:2004.044 041366 2020b.
Richard S Sutton. Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. In Machine learning proceedings 1990, pp. 216-224. Elsevier, 1990.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXivv 2018
Hado van Hasselt, Arthur Guez, Matteo Hessel, Volodymyr Mnih, and David Silver. Learning values across many orders of magnitude. arXiv preprint arXiv:1602.0777: 2016.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. NeurlPS, 2017.
Pauli Virtanen, Ralf Gommers, Travis F. Oliphant, Matt Haberland, Tyler Reddy, David Cournapeau, Evgeni Burovski, Pearu Peterson, Warren Weckesser, Jonathan Bright, Stéfan J. van der Walt, Matthew Brett, Joshua Wilson, K. Jarrod Millman, Nikolay Mayorov, Andrew R. J. Nelson, Eric Jones, Robert Kern, Eric Larson, C J Carey, ÎIhan Polat, Yu Feng, Eric W. Moore, Jake VanderPlas, Denis Laxalde, Josef Perkiold, Robert Cimrman, Ian Henriksen, E. A. Quintero, Charles R. Harris, Anne M. Archibald, Antônio H. Ribeiro, Fabian Pedregosa, Paul van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, -722, 2020. doi: 1031044//4 592-019-0666-2
Elena Voita, Rico Sennrich, and Ivan Titov. The bottom-up evolution of representations in the transformer: A study with machine translation and language modeling objectives. arXiv preprint arXiv:1909.011380. 2019.
13

---

Published as a conference paper at ICLR 2022
Zhilin Yang, Zihang Dai, Yiming Yang. Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V Le. Xlnet: Generalized autoregressive pretraining for language understanding. arXiv preprint arXi::1906.08237, 2019.
Jinhua Zhu, Yingce Xia, Lijun Wu, Jiajun Deng, W. Zhou, Tao Qin, and H. Li. Masked contrastive representation learning for reinforcement learning. ArXiv, abs/2010,07470, 2020.
14

---

Published as a conference paper at ICLR 2022
A SETUP DETAILS
A.1 R2D2 DISTRIBUTED SYSTEM SETUP
Following R2D2, the distributed system consists of several parts: actors, a replay buffer, learner, and an evaluator. Additionally, we introduce 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 central batch inference process:
Fy the observation time I.
Tt-1 the reward at the previous time, initialized with r-1 O.
d the action at the previous time, is initialized to O.
h recurrent state at the previous time, is initialized with h = O.
They block until they receive Q(2,, a; 0). The I-th actor picks a using an e--greedy policy. As R2D2, the value of computed following:
where 0.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 sufficient amount of actors have sent inputs, forming a batch. We use batch size of 64 our experiments. After a batch formed, the neural network of the agent is run compute ( a, 0) for the whole batch, and these values sent to their corresponding actors.
Replay buffer: stores fixed-length sequences of transitions
(wa) +--11 along with A transition is of the form w,
their priorities pr, where T. is the trace length (rs-1, O0-10 1,- Rs, as hg, Tsy I1+1). Concretely, this consists of the following elements:
r reward at the previous time.
a=-1: 1 action done by the agent at the previous time.
hg- 1 recurrent state (in our case hidden state of the LSTM) the previous time.
Zgi observation provided by the environment at the current time.
a,: action done by the agent at the current time.
hg: recurrent state (in our case hidden state of the LSTM) at the current time.
73 reward at the current time.
T++1: observation provided by the environment at the next time.
The sequences never cross episode boundaries and they are stored into the buffer in overlapping fashion, by 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 7/ 0.9.
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 not sent to the replay buffer. The evaluator acts 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 In this paper 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).
Evaluator:
episodes the results of those 5 episodes are average and reported.
15

---

Published as a conference paper at ICLR 2022
The learner contains two identical networks called the online and target networks with different weights 9 and respectively (Mnih et al. 2015). The target network's weights &' are updated to 8 every 400 optimization steps. is updated by executing the following sequence of instructions:
Learner:
First, the learner samples batch of size (batch size) of fixed-length sequences of transitions from the replay buffer, with each transition being of length L: T, (w()
Then, forward pass is done
the online network and the target with inputs (xt, 1. a. " 7 in order to obtain the state-action values
{(Q(x), 0)0), (1) ()))
With ((Q(2), a: 0), (2 8)), the Q(A) 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 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, used Podracer setup similar to (Hessel al. 2021) for fast usage of experience from actors by learners.
TPU learning and acting: As in the Scbulba setup of (Hessel et 202 acting and learning
network computations were co-located on a set of TPU chips, split into ratio of cores used for learning for every core used for inference. This ratio then scales with the total number of chips used.
Environment execution: Due 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 the agent states on the host CPUs.
A.3
COMPUTATION USED
R2D2 We train the agent with single TPU v2-based learner, performing approximately network updates second (each update on mini-batch of 64 sequences of length 80 for Atari and 120 for Control). We use 512 actors, using actors per CPU core, with each one performing environment steps per second on Atari. Finally for the batch inference process TPU v2, which allows all actors to achieve the speed we have described. particular, we used 8 TPU cores for learning and 2 for inference.
64
train the agent with hosts each with 8 TPU v2 cores. Each of the cores per host 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 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.
V-MPO
A.4 COMPLEXITY ANALYSIS
As stated, the agent consists of layers of convolutions, transformer layers, and linear layers. Therefore the complexity is mar (O(n² dd) O(k n--nn-d- n-d)), where is the kernel size in the case of convolutions, mar(O(n²
n is the size of trajectories, and is the size of hidden layers.
16

---

Published as a conference paper at ICLR 2022
B ARCHITECTURE DESCRIPTION
B.1 ENCODER
As shown Fig. observations O; are encoded using an encoder. In this work, the encoder we have used is ResNet-47 encoder. Those 47 layers are divided in 4 groups which have the following characteristics:
An initial 2 convolution with filter size 3x3 (1 layers).
Number of residual bottleneck blocks (in order): (2, 4, 6, 2). Each block has convolutional
layers with ReLU activations, with filter sizes 1x1, 3x3, and 1x1 respectively
layers).
Number of channels for the last convolution in each block: (64, 128, 256, 512).
Number of channels for non-last convolutions in each block: (16, 32, 64, 128).
Group norm is applied after each group, with group size of 8.
After this observation encoding step, final 2-1ayer MLP with ReLU activations of sizes (512, 448) applied. The previous reward and one-hot encoded action are concatenated and projected with linear layer into 64-dimensional vector. This 64-dimensional vector is concatenated with the 448-dimensional encoded input to have final 312-dimensional output.
B.2 TRANSFORMER
As described in Section 2 the output the encoder is fed to a Gated Transformer XL. For Atari and Control, the transformer has the following characteristics:
Number of layers:
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 CoBERL 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 Memory size of 512.
B.3 LSTM AND VALUE HEAD
For both R2D2 and V-MPO the outputs of the transformer and encoder passed through GRU transform to obtain 512-dimensionall vector. After that, an LSTM with 512 hidden units is applied. The the value function is estimated differently depending the RI. algorithm used.
17

---

Published as a conference paper at ICLR 2022
Following the LSTM, Linear layer of size 512 is used, followed by ReLU activation. Finally, to compute the Q values from that 512 vector a dueling head is used, as in Kapturrowkki et al.) (2018), dueling head is is used, which requires linear projection to the number of actions of the task, and another projection to a unidimensional vector.
R2D2
V-MPO Following the LSTM, 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 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.4
CRITIC FUNCTION
For DmLab-30 (V-MPO), we used a layer MLP with hidden sizes 512 and 128. For Atari and Control Suite (R2D2) used single layer of size 512.
18

---

Published as a conference paper at ICLR 2022
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 separate evaluator process that shares weights with our learner in order to measure the performance of our agents. We report scores at the end 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 120) as the environment does not require a a squashing function Atari and the Control Suite
long context inform decisions, and ii) to transform our Q values (as done in (Kapturowski et al. 2018) since reward structures vary highly in magnitude between tasks.
C.I ATARI AND DMLAB PRE-PROCESSING
We the commonly used input pre-processing Atari and DMLab frames, shown Tab. One difference with the original work of Mnih et al. (2015), is that we do not use frame stacking, as we rely our memory systems be able integrate information from the past, done in Kap- turowski et al. (018). ALE is publicly available at wttpps:// / //1i thub. com/ /mgbe//bee./ 1 lemare/ Arcade-Learning Envi ronment
Hyperparameter
Value
Max episode length
30 man
Num. action repeats
Num. stacked frames
Zero discount on life loss
false
Random noops range
30
Sticky actions
false
Frames max pooled
3 and 4
Grayscaled/RGB
Grayscaled
Action set
Full
Table 7: Atari pre-processing hyperparameters.
C.2 CONTROL SUITE PRE-PROCESSING
As mentioned in
we use no pre-processing on the frames received from the control environment.
C.3 DMLAB PRE-PROCESSING
Hyperparameter
Value
Num. action repeats
4
Num. stacked frames
Grayscaled/RGB
RGB
Image width
96
Image height
72
Action set
as in Parisotto al. (020)
Table DmLab pre-processing hyperparameters.
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" possible actions. For the ca cartpol tasks, we take diagonal approach, utilizing each unit vector in the action space and
19

---

Published as a conference paper at ICLR 2022
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.
Task
Action space size
Total amount of actions
Acrobot
2
Cartpole
5
Cup
2
4
Cheetah
64
Finger
4
Fish
5
32
Pendulum
2
Reacher
Swimmer
5
32
Walker
64
Table 9: Control discretization action spaces.
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 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 fixed throughout all the experiments.
Table 10 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 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 algorithm were chose in accordance to Kapturowski al. (2018) and the related sweep. And for the architecture typer-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 I1] reports the chosen hyperparameters that we found to optimize the performance of GTrXL on the TO Atari games. We then moved to CoBERL. CoBERL, 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 fixed all the previous hyperparameters optimized for GTrXL (see Tab. [T] 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.
(both for the architecture and
over which we did the sweep for the
We ended-up up picking 1, although the difference between 0.1 and 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 13 reports both the fixed ones and the ones with the ranges over which. The fixed parameearseeters were taken directly from Parisotto et al. (2020) and we sweep over "Epsilon Alpha", "Target Update Period" and 'Memory size' to make sure maximised performance of this baseline. Table [4] 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. [55 the reduced memory size, in accordance with our hypothesis that the LSTM on top of Transformer would help reducing the size the memory especially in last.
the extra parameters used for CoBERL and
20

---

Published as a conference paper at ICLR 2022
Hyperparameter
Value
Optimizer
Adam
Learning rate
{0.0001, 0.0003)
Q's A
(0.8, 0.9)
Adam epsilon
10
Adam betal
0.9
Adam beta2
0.999
Adam clip norm
40
Q-value transform (non-DMLai) DMLab)
h(x) sign(a V [2||
+
1
1)
+ c
Q-value transform (DMLab)
h((z)
Discount factor
0.997
Batch size
32
Trace length (Atari)
80
Trace length (non-Atari)
120
Replay period (Atari)
40
Replay period (non-Atari)
60
Replay capacity
80000 sequences
Replay priority exponent
0.9
Importance sampling exponent
0.6
Minimum sequences to start replay
5000
Target Q-network update period
400
Evaluation
0.01
Target e
0.01
Number of layers
[6, 8, 12)
Memory size
(64, 128)
Number of heads
{4, 8}
Attention size
644, 128}
Table 10: GTxXL Hyperpraamtter used in all the R2D2 experiments with range
of sweep.
Hyperparameter Optimizer
Value
Adam
Learning rate
0.0003
Q
0.8
Adam epsilon
10
Adam betal
0.9
Adam beta2
0.999
Adam clip norm
40
Q-value transform (non-DMLab)
h(x)= sign( (z)
/ [2]]
+
1)
+
ex
Q-valuc transform (DMLab)
h(z))
Discount factor
0.997
Batch size
32
Trace length (Atari)
80
Trace length (non-Atari)
120
Replay period (Atari)
40
Replay period (non-Atari)
60
Replay capacity
80000 sequences
Replay priority exponent
0.9
Importance sampling exponent
0.6
Minimum sequences to start replay
5000
Target --ntwwork update period
400
Evaluation
0.01
Target e
0.01
Number of layers
8
Memory size
64
Number of heads
8
Attention size
64
Table 11: GTXXL Hyperparameters choosen for all the R2D2 experiments.
21

---

Published as a conference paper at ICLR 2022
Hyperparameter
Value
Contrastive loss weight
1.0
Contrastive loss mask rate
0.13
Table 12: Extra hyperparameters for CoBERL for the R2D2 experiments
22

---

Published as a conference paper at ICLR 2022
Hyperparameter
Value
Batch Size
120
Unroll Length
95
Discount
0.99
Target Update Period
{10, 20, 50)
Action Repeat
4
Initial 7
1.0
Initial a
5.0
0.1
Ca
{0.001, 0.002)
Popart Step Size
0.001
Memory size
[256, 512]
Table 13: GTXXL Hyperparameters used in all the VMPO experiments with range of sweep.
Hyperparameter
Value
Batch Size
120
Unroll Length
95
Discount
0.99
Target Update Period
50
Action Repeat
4
Initial 17
1.0
Initial
5.0
0.1
Cor
0.002
Popart Step Size
0.001
Memory size
512
Table 14: GTrXL Hyperparameters choosen for all the VMPO experi- ments.
Hyperparameter
Value
Contrastive loss weight
1.0
Contrastive loss mask rate
0.15
Memory size
256
Table 15: Extra hyperparameters for CoBERL for the VMPO experiments
23

---

Published as a conference paper at ICLR 2022
D ADDITIONAL ABLATIONS
Table T6] shows the results of several gating mechanisms that we have investigated. As we can observe the GRU gate is clear improvement especially DMLab, only being harmful in median on the reduced ablation set of Atari games.
COBERL
'Sum' gate
'Concat' gate
w/o Gate
Mean
115.47% + 4.21%
108. 92% + 4.56%
105.93% + 4.82%
86. 36.42% +: 7.25%
DMLab
Median
110.86%
108. 12%
107,82%
96.21%
Mean
726.34% 44.06%
548, 188.666%-5 = 16.16%
653.20% 59.1%%
591.33%.111..557
Atari
Median
270.03%
23%
325.96%
320.09%
Table 16: Gate ablations. Human normalized scores Atari-57 ablation tasks and DMLab tasks. For the mean we include standard error over seeds.
DM Suaie acroloo swingup
COBERL
R2D2-GTFXL
K2D2
DrPC-PFuels
CURL
Dreamer
Fixel SAC
355.25 + 3
265.00 + 113.27
327.16 5.35
81.7 4.4
rish swim
60.09
265.07
upright
5.16
21.70
penduluam swingup
835.06 1: 9.38
775.72
+ 50.06
831 86
61.54
680.9
41.9
swimmer swimmzr6
419.26
48.37
206-95 + 58.37
329.61
26.77
194.7
15.9
spin
1.69
166
ensy
984,80 2.76
chectah run
523.36 41.38
105.23 + 144.82
365.45 = 50,40
523.8
6.8
518
570
253
197
walker walk
781 1: 26.266111 26.16
584 50
70.28
687
18.15
968.3 1.8
907
43
897
49
42 4:
cap
1.63
compole swingup
762
compole swingup sparse
714.80 80 16.18
743.71
9:27
762.57 6.71
462.0
56.6
Table 17: Results on tasks in the DeepMind Control Suite. CoBERL, R2D2-GTrXL, R2D2, and D4PG-Pixels are trained on 100M frames, while CURL, Dreamer, and Pixel SAC trained 500k frames. We show these three other approaches as reference and not a directly comparable baseline.
24

---

Published as a conference paper at ICLR 2022
E
GAME SCORES
(ablation gaaees)
COBERL
K2D2-gTXXL
R2D2
COBERL-auxilary
beam rider
29601.17 I1973.77
65709.85 = 29911.8
3 11340.33
enduro
2323.66 31.42
2247.76 139.16
2338.98 = 13.23
2250.01 69.20
breakout
409.87 11.33
379.23 32.37
363.43 98.32
420.72 9.86
pong
21.00 1: 0.00
20.95 + 0.09
21.00 = 0.00
21.00 = 0.00
qbert
34000.79 6074.45
22444, 4 4270.57
25375 7451.33
47741.30 + 5288.90
seaquest
164587.95 + 122633.80
310745.77 67781.42
83490.38
121434.46
174867.68 123876.36 30
space invaders
37497.11 + 9973.19
19113.28 8512.98
4704.35 # 2763.78
8783.28 6430.09
Atari (ablation games)
COBERI
COBERL with LSTM before
COBERL, w/ LSTM
beam rider
29601.17 + 11973.77
35188.16 23357.60
62234.03 = 12764.95
enduro
2323.66 + 31.42
2321.02 + 43.05
2286.74 112.37
breakout
409.87 = 11.33
418.13 6.77
411.68 13.62
pong
21.00 + 0.00
21.00 0.00
21.00 0.00
ghert
34000.79 6074.45
30055.59 9892.69
31198.54 5557.65
seaquest
164587.95
122633.80
172136.83 119874.67
19216144 94730.21
space invaders
37497.11
9973.19
23746.32
18961.05
10169.22 12746.13
games)
CoBERL
Skip
gate
Concal
beam rider
29601.17 + 11973.77
52498. 32 + 2934.21
72882.42 + 21239.63
51371.38 12560.94
enduro
2323.66 = 31.42
2168.37 + 362,99
2247 23 = 82.22
2288. 13 + 58 59
409.87
403.36
357.52
21.00 + 0.00
20.88 + 0.18
21.00 +: 0.00
21.00 0.00
qhert
34000.79 6074.45
27660.71 71 + 6830.06
33807.00 4294.99 93
43487_35
14518.46
seaquest
164587.95 + 122633.80 80
118960.81 112462.18
294976.55 41827.55
399817.75
36729.78
space invaders
37497.11 1: 9973.19
22347.29 15281.67
10387.59 1: 2858.63
20933.98 +: 13072.95
25

---

Published as a conference paper at ICLR 2022
Atari (ablation games)
CoBERL with CURL
CoBERL with SimCLR
beam rider
25605.45 + 12772.65
36364.92 27431.91
enduro
2280.38 + 152.75
2343.53 + 5.10
breakout
337.39 + 53.97
397.40 + 40.07
pong
19.99 - 2.02
20.94 + 0.11
qbert
15935.73 + 580.68
24491.28 + 14491.24
seaquest
163715.19 + 80515.52
273070.09 + 124529.37
space invaders
14038.27 + 16158.41
22120.53 + 17354.48
26

---

Published as a conference paper at ICLR 2022
Atari-57
COBERL
R2D2-gTrXL
R2D2
alien
10728.64 + 6771.77
9593.30 :: 32266.20
8960.97 4233.52
amidar
3089.76 + 1028.80
2888.93 1324.27
1980.60 + 181.03
assault
10055.81 + 7856.45
7569.86 1 4705588
18601.87 + 6948.52
asterix
732004.08 3422955.21
535334.42 = =400048.35
777461.28 251661.19
asteroids
138026.82 63177.27
108693.26 + 22302.14
56238.62 = 44595.74
atlantis
1062735.42 47026.77
1005098.28 + 36831.12
994240.79 42158.91
bank heist
1126.69 + 620.85
1131.62 + 272.21
1110.12 442.89
battle zone
92092.43 #39167.40
77051.40 + 45255.27
94903.82 = 20522.95
beam rider
29601.17 + 11973.77
65709.85 29911.83
33938.58 = 11340.33
berzerk
1288.99 + 615.45
594.60 + 117.78
1314.18 1 444.67
bowling
160.94 + 53.41
49.23 47.13
86.34 = 27.65
boxing
51.78
100.00 + 0.00
99.28 0.79
breakout
409.87 + 11.33
379.23 32.37
363.43 + 98.32
centipede
62816.39 + 13516.24
98575.88 + 43443.70
75733.83 + 199166.9
chopper command
614750.54 283486.50
53852.34 13906.17
26147.63 = 14180.38
crazy climber
1 49 50591.49
117777.36 + 18382.23
134610.46 35360.45
defender
3 90741.03
304551.03 75288 86
383931.18 151219.51
demon attack
+ 12721.85
115010.85 + 26125.20
96371.00 = 51049.84
double dunk
24.00 = 0.00
21.25 2.76
23.77 = 0.33
enduro
2323.66 + 31,42
2247.76 139.16
2338.98 + 13.23
fishing derby
45.74 + 9.05
39.83 4.85
47.47 + 8.12
freeway
34.00 + 0.00
34.00 0.00
33.67 + 0.42
frostbite
7877.53 + 1919.45
6758.04 + 2227.07
5850.66 = 1443.76
gopher
98169.07 + 12126.03
84697.17 26263393
106325.84 17858.36
gravitar
6219.65 + 484.85
5181.25 + 1615.50
5186.02 + 808.22
hero
20487.23 288.81
16544.03 31066.3
16764.65 = 2575.03
ice hockey
24.64 19.00
9.94 = 12.69
21.38 = 15.82
jamesbond
5589.27 + 4360.99
7149.04 3244..00
5872.92 + 3432.57
kangaroo
11703.01 + 2892.01
11760.61 3657.70
13223.97 = 2123.85
krull kung fu master
27738.24 = 19282,30
42886.34 35329.557
19821.18 = 17250.91
128429.78 + 12649.00
64979.95 25384.52
99741.61 + 14185.70
montezuma revenge
+ + 998.87
0.00 = 0.00
240.00 + 195.96
ms paeman name this game
10567.67 + 3692.12
9905.93 884.03
9576.99 = 2286.69
26311.25 + 6607.53
26669.02 + 3614.10
23889. 30 2608.30
phoenix
424708.26 + 173336.21
283179.23 94106.97
138558.03 = 154881.71
pitfall
0.00 + 0.00
0.00 + 0.00
0.00 = 0.00
pong
21.00 = 0.00
20.95 =
21.00 0.00
private eye
12034.62 + 5963.85
12690.27 + 4372.38
3065.63 + 5867.30
qbert
34000.79 + 6074.45
22444.14 4270.57
25375.24 + 7451.33
riverraid
22623.35 3512.77
23692.95 + 2100.11
28147.78 678.51
road runner
212345.12 192339.01
542775.42 58830.61
0.00 = 0.00
robotank
80.69 + 7.32
50.75 + 31.02
63.82 - 20.55
seaquest
164587.95 = 122633.80
310745.77 67781.42
83490.38 = 121434.46
skiing
-29958.20 20:20:::4 4.99
-21898.19 9898.02
-28243.31 + 3453.55
solaris
6037.15 + 3103.57
3675.08 + 5575.93
2512.84 + 1288.33
space invaders
37497.11 + 9973.19
19113.28 8512.98
4704.35 2763.78
star gunner
103870.95 + 23124.17
78559.08 + 421100881
88169.92 = 13783.74
surround
9.37 + 0.45
9.40 + 0.51
8.60 - 1.00
tennis
14.08 + 11.58
9.20 = 11.29
-0.04 0.08
time pilot
36085.82 + 6934.00
14250.32 1651.02
21051.93 4156.28
tutankham
33.93 + 28.18
21.66 + 13.11
79,51 = 40.10
up down
407873.64 + 96615.12
237945.96 + 101020.61
438679.20 + 38137.36
venture video pinball
1754.74 + 229.65
1632.54 +: 247.91
1624.82 + 100.65
248629.59 = 231446.28
48305.17 = 444432.44
339391.16 194812.76
wizard of wor yars revenge
19454.60 + 7061.60
14127.73 9153.21
18270.16 12116.04
385904.98 264624.16
217128.25 28248.52
248086.47
-162482.95
zaxxon
20411.27 + 12535.64
25063.44 + 10385.30
31679.84 10044..57
27

---

Published as a conference paper at ICLR 2022
DmLab Levels
coberl
gtrxl
rooms collect good objects test rooms exploit deferred effects test
0.14
9.69 + 0.23 0.23
56.57
54.87
rooms select nonmatching object
48.48:4.33
rooms watermaze
56.61
3.44
rooms keys doors puzzle
33.99 5.44
35.53 8.33
language select described object
611.02 1.02
605.59 14.29
language select located object
606.05 3.73
584.23 +
language execute random task
2 +:
+
language answer quantitative question
4.22
297.57
lasertag one opponent large
15.47
11.65
lasertag three opponents large
22.99 1009::3
28.46 #
lasertag one opponent small lasertag three opponents small
+.550
2
42.96 4.100
0.447
natlab fixed large map natlab varying map regrowth natlab varying map randomized
46.03 8.85
43.46 14.38
3.70
27.62
48.97
32.79 = 11.30
platforms hard platforms random
7.97
35.67 = 23.15
1.55
69.01 + 1.04
psychlab continuous recognition
58.91
54.96 + 3.87
psychlab arbitrary visuomotor mapping
51.56 0.44
59.96 + 0.33
psychlab sequential comparison
33.56 0.84
+ 3 0.64
psychlab visual search
77.96 1.33
7.882
explore object locations small
3..07:5.66
71.85 6.22
explore object locations large explore obstructed goals small
62.57 + 5.86
60.61 2.72
2.85 + 5.35
245.40 + 10.15
explore obstructed goals large explore goal locations small explore goal locations large explore object rewards few explore object rewards many
109.61 1.20
+
369.20 3.54
340.19
139.741
123.20 14.65
- 14.95
69.02 10.02#0.00 0.95
104.62 2.63
94.
28

---

Published as a conference paper at ICLR 2022
F LEARNING CURVES
F.1
ATARI LEARNING CURVES
125
-075
23
0.00
-25
-044
-3.8
7.5
-222
1531
123g22a
MX21
Human Human Normateed Score
D653
3.5
1.2
3eardfen
-23
000011
H130
1553
1023
3.5
Environment
Figure 5: Learning Curves for Atari. The x-axis represents number of enviromment steps millions. The y-axis represent the Human Normalised score. The error represents the 95% confidence interval. Red is COBERL, BLUE is GTxXL
29

---

Published as a conference paper at ICLR 2022
F.2
DMCONTROL LEARNING CURVES
359
275
3330
700 20022OOP
130
Eplocce
DOCEESS EVITERIS
1002
120
200
Environment steps
Figure 6: Learning Curves for DMControl. The x-axis represents number of environment steps millions. The y-axis represent the Episode return. The error represents the 95% confidence interval. Red is COBERI, BLUE is GTrXI.
30

---

Published as a conference paper at ICLR 2022
F.3 DMLAB LEARNING CURVES
0.4
6.3
6.2
6.2
8.
15
EJOCOBER TIMIS
KOROOIOSOO SHAEE TIME
0.5
0.75
575
025
335
0.2
2.15
335
335
1.2
4.39
352
335
Figure 7: Learning Curves for DMLab. The x-axis represents number of environment steps in billions. The axis represent Human Normalised score. The error represents the 95% confidence interval. Red COBERL, BLUE is GTrXL
31

---

Published as a conference paper at ICLR 2022
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://arcadeeaarigeemmmmmmmem
DeepMind Control Suite Tassa al. 2018 is released free, open-source software der the terms of Apache-2.0 License. The latest version of the source code is pub-
licly available at: hattps://gituu: com/ http://emmmoommmoooooooo//////// http://c.....o/ooooo///tt/eee/// cim
control/ su te REEAARRADEEE.. mc
DmLab Beattie al. 2016) is released as free, source software under the terms of Apache-2.0 License. The latest version of the source code is publicly available at: https://github. com/
H
PSEUDO-CODE
Algorithm Pseudo-code for CoBERL
training_iteratin iterations 0
initialise mnttialise_eeights(( www.wwewwtittwwuueeeeeee. Encoder TransformerStack. Gate, ValueNetwork) S max__rainingg... seiniinggtterationss do sampled_ batch t sample experience()
while beiining_terrttonn
encoded _aaddeesmeseeeee VisualEncodder(samplee batch)
encoded actions + mmeffor/cctionEnnoddrrssmpppppp batpeeedbaadeebhh encoded_action pambineed_ipuuss-e
combined inputs < concat(encoded, -ameddss,___,,,eeeee
transformer inputs,
teaaffommer_arrgette mask +- create create_masks maasss_taggessccomiine
output_tran transformerrrrr..... cotpressttotttttveeeeee +- TransfornerrSucktaaasffoomee titprfforrrrrrppppssestos inputs,, causal mask=Falsc)
contrastive Joss compute caattte_axx__leee wow._lossoouuppttoooo httprppptttttrsssooeeeeeoooooo contresttttottttteeeeeeeoee transformer caanformer_targetts. mask)
output transformer RL Transformerrscckktaaffomeer intputs,,p,ttpp causal mask=TTrue)
gated_output + Gate<(ombiied_iinpgl outpur_trannformme__RL]
lstm_output LSTM(gatee_ouutpu))
combined combined_inputs__o1 lacl__forrrruuueaee value_net < concat(Istm_ _ottput,, combined_innuus
value estimation +- ValueNeetworkcoombiedd inputs for for_value value_ _nteee rl loss + compute rl losoosvalle seto(value_sstmationn... estimation, extra _rggs)
total loss <- loss contrastive_loss
end while
32

---

Published as a conference paper at ICLR 2022
Pseudo-code for the auxiliary loss calculation
def compute_aux__es (inputl,, input2, mask_ext);
batch_size, seq_dim, Ceat_ding = input shape
inputl = reshape (inputl, [-1, feat_dim])
input2 = reshape (input2, [-1, feat_dim])
inputl Lz_e (inputl, axis=-1)
input2 = 12_normalize (input2, axis=-1)
# Compute labels index
labels_idx arange arange (input shaape [0])
12 labels_idx - labels_idx. astype
13 # Compute pseudo-labels for contrastive loss.
labels = one hot (labels_idx, input1. [0] * 2)
15 # Mask
the
image pair.
16 mask ene_hot (labels_idx, ayputi.shapere [0])
17 # Compute logits.
18 logits_ll = matmul (inputl, jnp.transpose (input
logits_22 matmul (input2, inp.transpose (input)))
logits_12 matmul (inputl, jnp.transpose (inputt)))))) (input2))
21
logits_21 matmul (input2, jmp.transpose (input1)
22
# Calculate invariance penalty.
inv_penalty = kl logiitth_ooiiss( # BB T
Logits 22, axis=-1)
25
inv Liv_penaltyylo += kl with Logth
26
stog_ggradieen((oggits___])))) logits_22, axis--1)
27
+- (
[cce_ggradient((lsits__]]]]]] logits_11, axis=-1)
inv_penalty kl_with_logitss
30
stop_ stcp_ggradeett/ (logits__]) legits_21 _2_, axis=-1)
31
inv inv_penalty = inv_penalty / 4. [B T]
32
logits_ logits_1l mask 1e9
34
22 logits_22 mask leg
35
36
1ogits_1211 concatenate concatenatee(looiitts ([logits_12, [ogits]l]], axis=-1) 1ogits_2122
concttenatee1looitaaaa
sigits_221, axis=-1)
39
loss 12 -sum (labels log softmax (oognas 1211), axis=-1)
40
loss 21 = -sum -sum (labels log soffmax(ooggtt _og_ls____2i_,, axis=-1)
41
loss reshape (loss_12 + loss_21, [batch_size, sen_din]) mask_ext inv_penalty reshape (inv_penalty, [batch_size, seg_din]) -xask_xxt
43
44
45
loss mean (loss inv_penalty)
return loss
I AREA UNDER THE CURVE
For all the levels we calculated the AUC by integrating composite Simpson's 's rule with a delta(x) of steps. We use the intergate package from scipy (Virtanen et al. 22020).
J LIMITATIONS AND FUTURE WORK
A limitation of our method is that it relies single time step information to compute its auxiliary objective. Such objective could naturally be adapted to operate on semporally-extent extended patches, and/or action-conditionel inputs. Also, as done in the original BERT (Devlin 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 directly the embeddings
33

---

Published as a conference paper at ICLR 2022
of the CLS token as sequence summary and hence provide more context to the policy estimation network. We regard those ideas as promising future research avenues.
K EXTRA INFORMATION ABOUT THE GATE EMPLOYED IN EQUATION
The gate we use in equation 4 is derived directly from the one used GTrXL PPrisotto et 2020 We report here the details for better clarity.
The gate is the defined the following way:
(V, = (1 - 00==000+=0)
(5)
where:
2 = o(W U2y - bg)
(6)
h tanh(W + (r(r
and:
r=o(W,r + U+y)
(8)
Wz, U2, Wg, Ug. W,andU,, are set of linear weights and by a bias.
In our case x is the output of the transformer network and y is the output of the encoder network.
34

---

Published as a conference paper at ICLR 2022
L
EXTRA ANALYSIS
coberi mas.inggggg...m 015
120
coberi masking 03
coberi masking 05
human nnrmalized
110
1.00
0.90
coberi mattri_masingg.om 015
masking, 00 model
coberi maskingg__.... 0_5
Figure 8: Human normalised
function of the masking pereentage the input sequence. coberl_masking relates to 15% masking of input sequence, coberl_man seterl_makking_Lne relates to 30% masking of the input sequence and @obber_maskkig_000 05 relates to 50% masking the input sequence.
We also attempted higher masking rate, but as shown figure iguur[[hhey they seem perform worse, probably because high level of masking are reducing much the number of frame present in the sequence, hence removing the information need to successfully perform the auxiliary task.
1200
model
1000
gtnx
600
400
cober!
gtni
model
Figure 9: AUC calculated at 100% human score.
Even when calulcate 100% human score, the AUC shows significant advantage of CoBERL over
(60)=266,) p=0.0011. CoBERL M=545.18, SD=282.58; gTrXL M=374.49, SD=282.58.
35


Reviews:
[{"review_id": "l90YdV6-mtA", "reviewer": "Reviewer_XmhC", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths\n\n* The problem of the growing size of transformers as we require larger time dependencies is of great concern, the trick of forwarding these representations to an LSTM layer is simple yet effective from the results provided\n* Avoiding the need of hand-engineered data augmentations is a big plus for representation learning subobjectives\n* The paper evaluates the proposed solution in multiple hard and well known benchmarks, authors also provide relevant ablation studies\n\n---\nWeaknesses \n\n* General comments:\n   * My biggest concerns are about clarity and presentation, the paper relies heavily on readers previous knowledge, which requires combined combined intuition from multiple disciplines. Additionally, several points of the architecture are never explained (neither given a intuition to the reader). Below I provide specific details\n   *  The number of random seeds in the empirical evaluation, only 3, is below current accepted standard for accepted RL works (which is a minimum of 5).\n---\nBelow I detail specific points that I believe need further clarification (in order of appearance):\n\n* I missed a brief background section explaining how representation learning is used in RL, also, this paper gets and improves ideas from ReLIC, we should be able to understand this work without requiring go to ReLIC's\n* Section 2.1, 3rd paragraph, you mention that 15% of the embedding is masked, it seems a very specific value, which ones did you explore? what happened with those different values? (some intuition)\n* Also there, and in section 2.1 in general, I missed to be given intuitions of why each component was designed in that specific way. E.g., starts directly explaining the masking, I would try to build into the reader the intuition of  why are we masking? before entering in the details. This is much better done in section 2.2 for instance.\n* What is exactly doing the \"Gate\" function? it is mentioned multiple times, and it is evaluated in the ablations, but I couldn't find where is it explained what is this function. It is mentioned that helps to skip the transformed output until the weights are warmed-up, how does it knows this?. I believe that the answer to this would also fit well in a background section.\n* Related work section, paragraph 2, says \"Second, CoBERL combines the transformer architecture with GRUs to...\" until this point GRUs are never mentioned, only LSTMs, so it is not evident if you use both indistinctively (if that is the case I would only say recurrent layer) or if not, where are the GRUs in the architecture and what is their role.\n* First paragraph page 6, second to last line, you say \"We use Peng's Q($\\lambda$) \" this is missing a proper citation.\n* Small typo, two lines below what I mentioned above, you are missing an space after V-MPO\n\n----------\n### Post Author's response and updated version\n\nI think that authors did a good work in the updated version, the work is much better self-contained now and easier to follow. Importantly, now all the experiments are the result of 5 independent runs instead of 3, which raises the reader's confidence in the empirical results.\n\nGiven the strengths I reported in my original review, I updated my scores.\n\n\n\n", "summary_of_the_paper": "The paper introduces a new transformer architecture, called CoBERL, targeting reinforcement learning. The key ingredients of the proposed architectures are the use of LSTMs to reduce the need of large transformers and a contrastive learning procedure that doesn't require of human data augmentation.", "summary_of_the_review": "The new architecture is interesting, proposes novel ideas that solve important issues of previous solutions and shows good empirical results. However, the paper has clarity issues and doesn't seem completely self-contained. Also the reduced number of random initializations goes below usual quality standards in published RL work. \n\nStill, I believe these points could be fixed through the rebuttal, so I am open to revising my score if the authors submit an updated version that tackles these points.\n\n----------\nAuthors correctly addressed my concerns in the updated version", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "XooojLHIBQK", "reviewer": "Reviewer_YoL5", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Normally when I hear \"LSTM\" I think \"hard to train, requires many samples\". When I hear \"Transformer\" I think the same thing. Yet in this work it is being claimed that a combination of these data-hungry methods results in higher data efficiency that with other methods. At which point I naturally get suspicious :) and want to investigate more.\n\nMajor concerns / criticisms / discussion\n----------------------------------------\n\n* The abstract promises that \"COBERL consistently improves data efficiency across the full Atari suite, a set of control tasks and a challenging 3D environment\", however most of the reported results are about performance improvement (after 200M steps (on Atari)). The results on improved data efficiency come in a for of AUC curves as compared to GTrXL. My understanding is that GTrXL was chosen as baseline because it demonstrated the highest human normalized score, but if the main question is about data efficiency and AUC, then a method that rises quickly (Rainbow?) and plateaus later would have a higher AUC, especially at an earlier cut-off points (for example 50M instead of 200M). What if we change the metric from \"AUC at 200M steps\" to \"AUC at the moment when an algorithms achieves 100% human performance\"? The motivation behind this line of thinking comes from domains where data samples are expensive (real robots, medical observations, etc) where what we care about is getting to good performance fast, and not so much about \"being more data efficient by the time we have collected 200M data samples\". Do we know that GTrXL is a good baseline for data efficiency even if we discard the arbitrarily chosen (by the community) mark of 200M samples?\n \n* The concept of \"self-attention consistency\" that seems to be a crucial idea driving this contribution seems underexplained to me: what do authors mean by self-attention being consistent, consistent across what, and why that consistency is a good thing?\n\n* A subjective note on readability: This paper was quite a hard read, perhaps because the explanation of the introduced concept and main ideas is mixed with implementational details and minor notes throughout the text. This makes it a bit hard to distill the core idea and important parts of the contribution from non-essential implementational details. \n\n\nQuestions\n---------\n\n* In the Rainbow paper (https://arxiv.org/pdf/1710.02298.pdf) the reported improvement over human is ~200%+ (Figure 1), in the current work it is reported as 874%, where does this difference come from?\n\nMinor remarks\n-------------\n\npage 3 \"augmentations make other methods less data efficient\" -- unclear what is meant by this statement, isn't the point of augmentations the opposite?\n\npage 7 par 1: typo \"Figure 2 presentt this analysis\"\n\nFigure 2 caption \"els.(higher is better)\" -- some formatting issue with spaces and the dot", "summary_of_the_paper": "The paper replaces the recurrent architecture of R2D2 with a Transformer-LSTM hybrid, and introduces a bi-directional contrastive loss that will be combined with the RL loss for overall optimization. The architecture is supplemented with a learnable gate between the Transformer and LSTM layers that should learn to regulate whether Transformer's output at time t should be fed into LSTM or it is not necessary, making the overall learning easier and leading to higher data efficiency and performance than the baseline (GTrXL).", "summary_of_the_review": "My evaluation largely depends on whether the chosen beaseline is consistent with the stated main contribution of this work: would AUC of other methods for RL would also be lower at earlier cut-off points. I see two potential possibilities:\n(a) The chosen baseline is actually good for data efficiency -> the main claim of the paper holds, all's good\n(b) The baseline was not good for data efficiency -> the paper should not claim data efficiency, but rather just go with performance increase (1425%, while baseline is 1202%).\nI am looking forward to the discussion to gain clarity on this point.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "N7QypRBUUXt", "reviewer": "Reviewer_gQY1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proposed approach seems like a novel combination of existing methods with the advantage of requiring less prior knowledge and higher sample efficiency.\nUsing time steps from other trajectories in the buffer seems well motivated and avoids the need for designing augmentations.\nExperiments are comprehensive and the ablation study is appreciated.\nA minor gripe is the use of Y for input and X for output of the transformer.", "summary_of_the_paper": "The paper introduces a contrastive learning objective to learn representations for an RL agent. The basic objective is taken from ReLIC, but instead of augmenting states (which requires prior knowledge of the problem domain) for each timestep the authors use timesteps of the other trajectories in a minibatch as negative examples.", "summary_of_the_review": "The proposed approach is a novel combination of existing methods that is well motivated. Experiments are comprehensive, including an ablation study.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduces a new transformer architecture for representation learning in RL. The key ingredients of the proposed architectures are a novel combination of existing methods: (1) the use of LSTMs to reduce the need for large transformers and (2) a contrastive learning procedure that doesn't require human data augmentation.  The resulting approach requires less prior knowledge and provides higher sample efficiency. The paper is convincing, with comprehensive experiments on multiple challenging and well-known benchmarks and an ablation study. The reviewers did expressed concerns that parts of the paper are a very difficult read and could use improvement, especially those relying on substantial external background. The intuition behind several components could be improved, and there are some clarity issues, as detailed in the individual reviews.
Decision: Accept (Spotlight)