Improving Sample Efficiency of Value Based Models Using Attention and Vision Transformers

TL;DR

This paper presents a novel deep reinforcement learning architecture that leverages attention mechanisms and vision transformers to enhance sample efficiency and performance in Atari games.

Contribution

It introduces a visually attentive model using transformers on feature maps, combining advances from NLP and computer vision for improved RL sample efficiency.

Findings

Improves sample complexity in Atari environments

Achieves better performance in some games

Demonstrates the effectiveness of attention-based architectures in RL

Abstract

Much of recent Deep Reinforcement Learning success is owed to the neural architecture's potential to learn and use effective internal representations of the world. While many current algorithms access a simulator to train with a large amount of data, in realistic settings, including while playing games that may be played against people, collecting experience can be quite costly. In this paper, we introduce a deep reinforcement learning architecture whose purpose is to increase sample efficiency without sacrificing performance. We design this architecture by incorporating advances achieved in recent years in the field of Natural Language Processing and Computer Vision. Specifically, we propose a visually attentive model that uses transformers to learn a self-attention mechanism on the feature maps of the state representation, while simultaneously optimizing return. We demonstrate empirically that this architecture improves sample complexity for several Atari environments, while also achieving better performance in some of the games.