Abstract

Many reinforcement learning (RL) agents require a large amount of experienceto solve tasks. We propose Contrastive BERT for RL (CoBERL), an agent thatcombines a new contrastive loss and a hybrid LSTM-transformer architecture totackle the challenge of improving data efficiency. CoBERL enables efficient,robust learning from pixels across a wide range of domains. We usebidirectional masked prediction in combination with a generalization of recentcontrastive methods to learn better representations for transformers in RL,without the need of hand engineered data augmentations. We find that CoBERLconsistently improves performance across the full Atari suite, a set of controltasks and a challenging 3D environment.

 

Quick Read (beta)