Unsupervised State Representation Learning in Atari

TL;DR

This paper presents a novel unsupervised method for learning state representations in Atari games by maximizing mutual information, along with a new benchmark for evaluating how well these representations capture true game state variables.

Contribution

Introduces a mutual information-based approach for unsupervised state representation learning and a new Atari benchmark for evaluating the quality of learned representations.

Findings

Our method effectively captures ground truth state variables.

The proposed benchmark provides a new standard for evaluation.

Compared to other methods, our approach shows competitive performance.

Abstract

State representation learning, or the ability to capture latent generative factors of an environment, is crucial for building intelligent agents that can perform a wide variety of tasks. Learning such representations without supervision from rewards is a challenging open problem. We introduce a method that learns state representations by maximizing mutual information across spatially and temporally distinct features of a neural encoder of the observations. We also introduce a new benchmark based on Atari 2600 games where we evaluate representations based on how well they capture the ground truth state variables. We believe this new framework for evaluating representation learning models will be crucial for future representation learning research. Finally, we compare our technique with other state-of-the-art generative and contrastive representation learning methods. The code associated with this work is available at https://github.com/mila-iqia/atari-representation-learning