Learning to Play Atari in a World of Tokens

TL;DR

This paper introduces DART, a discrete representation-based transformer method for sample-efficient model-based reinforcement learning in Atari games, improving performance and surpassing previous methods and human scores in many games.

Contribution

DART is the first to utilize discrete abstract representations in transformer-based reinforcement learning, enhancing modeling of complex, discrete world properties and improving sample efficiency.

Findings

DART achieves a median human-normalized score of 0.790 on Atari 100k.

DART outperforms previous state-of-the-art methods without look-ahead search.

DART surpasses human performance in 9 out of 26 Atari games.

Abstract

Model-based reinforcement learning agents utilizing transformers have shown improved sample efficiency due to their ability to model extended context, resulting in more accurate world models. However, for complex reasoning and planning tasks, these methods primarily rely on continuous representations. This complicates modeling of discrete properties of the real world such as disjoint object classes between which interpolation is not plausible. In this work, we introduce discrete abstract representations for transformer-based learning (DART), a sample-efficient method utilizing discrete representations for modeling both the world and learning behavior. We incorporate a transformer-decoder for auto-regressive world modeling and a transformer-encoder for learning behavior by attending to task-relevant cues in the discrete representation of the world model. For handling partial observability, we aggregate information from past time steps as memory tokens. DART outperforms previous state-of-the-art methods that do not use look-ahead search on the Atari 100k sample efficiency benchmark with a median human-normalized score of 0.790 and beats humans in 9 out of 26 games. We release our code at https://pranaval.github.io/DART/.