# CoBERL: Contrastive BERT for Reinforcement Learning

arXiv: https://arxiv.org/abs/2107.05431 · HTML: https://mlsanity.com/paper/2107.05431/

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 57 references that are in the map.
- First posted on arXiv 12 Jul 2021.

## Builds on (most cited first, 30 of 38)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Unsupervised Learning of Visual Features by Contrasting Cluster Assignments](https://mlsanity.com/paper/2006.09882/) (2020)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Rainbow: Combining Improvements in Deep Reinforcement Learning](https://mlsanity.com/paper/1710.02298/) (2017)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Reinforcement Learning with Unsupervised Auxiliary Tasks](https://mlsanity.com/paper/1611.05397/) (2016)
- [An Empirical Model of Large-Batch Training](https://mlsanity.com/paper/1812.06162/) (2018)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents](https://mlsanity.com/paper/1709.06009/) (2017)
- [Representation Learning via Invariant Causal Mechanisms](https://mlsanity.com/paper/2010.07922/) (2020)
- [The Bottom-up Evolution of Representations in the Transformer: A Study with Machine Translation and Language Modeling Objectives](https://mlsanity.com/paper/1909.01380/) (2019)
- [Contrastive Behavioral Similarity Embeddings for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2101.05265/) (2021)
- [Learning values across many orders of magnitude](https://mlsanity.com/paper/1602.07714/) (2016)
- [V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control](https://mlsanity.com/paper/1909.12238/) (2019)

## Built on by (most cited first, 21 of 21)

- [Retrieval-Augmented Reinforcement Learning](https://mlsanity.com/paper/2202.08417/) (2022)
- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)
- [SMART: Self-supervised Multi-task pretrAining with contRol Transformers](https://mlsanity.com/paper/2301.09816/) (2023)
- [On Transforming Reinforcement Learning by Transformer: The Development Trajectory](https://mlsanity.com/paper/2212.14164/) (2022)
- [Procedural Generalization by Planning with Self-Supervised World Models](https://mlsanity.com/paper/2111.01587/) (2021)
- [Abstraction for Deep Reinforcement Learning](https://mlsanity.com/paper/2202.05839/) (2022)
- [Light-weight probing of unsupervised representations for Reinforcement Learning](https://mlsanity.com/paper/2208.12345/) (2022)
- [Transformers in Reinforcement Learning: A Survey](https://mlsanity.com/paper/2307.05979/) (2023)
- [CORN: Contact-based Object Representation for Nonprehensile Manipulation of General Unseen Objects](https://mlsanity.com/paper/2403.10760/) (2024)
- [Transformer in Transformer as Backbone for Deep Reinforcement Learning](https://mlsanity.com/paper/2212.14538/) (2022)
- [PASTA: Pretrained Action-State Transformer Agents](https://mlsanity.com/paper/2307.10936/) (2023)
- [RePreM: Representation Pre-training with Masked Model for Reinforcement Learning](https://mlsanity.com/paper/2303.01668/) (2023)
- [Learning Goal-Conditioned Representations for Language Reward Models](https://mlsanity.com/paper/2407.13887/) (2024)
- [Contrastive Abstraction for Reinforcement Learning](https://mlsanity.com/paper/2410.00704/) (2024)
- [MarineFormer: A Spatio-Temporal Attention Model for USV Navigation in Dynamic Marine Environments](https://mlsanity.com/paper/2410.13973/) (2024)
- [Stackelberg Coupling of Online Representation Learning and Reinforcement Learning](https://mlsanity.com/paper/2508.07452/) (2025)
- [Using Forwards-Backwards Models to Approximate MDP Homomorphisms](https://mlsanity.com/paper/2209.06356/) (2022)
- [Differentially Private Optimization for Non-Decomposable Objective Functions](https://mlsanity.com/paper/2310.03104/) (2023)
- [Transformer-Enhanced Reinforcement Learning: Fundamentals and Applications in Communication Networks](https://mlsanity.com/paper/2606.05208/) (2026)
- [Statistical Convergence of Transformer Encoder-Accelerated Robust Reinforcement Learning](https://mlsanity.com/paper/2609.23775/) (2026)
- [Does Learning Recurrent Networks Matter in Partially Observable RL?](https://mlsanity.com/paper/2610.13197/) (2026)

## Abstract (arXiv)

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, robust learning from pixels across a wide range of domains. We use bidirectional masked prediction in combination with a generalization of recent contrastive methods to learn better representations for transformers in RL, without the need of hand engineered data augmentations. We find that CoBERL consistently improves performance across the full Atari suite, a set of control tasks and a challenging 3D environment.