This repository provides code for unsupervised state representation learning in Atari games, targeting researchers and practitioners in reinforcement learning and representation learning. It enables the evaluation of learned representations by offering ground truth state variable annotations and a framework for probing these representations.

How It Works

The core innovation is the Atari Annotated RAM Interface (AtariARI), a Gym wrapper that exposes semantic state variable labels (e.g., player position, enemy locations) alongside standard Atari observations. This allows for direct evaluation of learned representations by training simple linear probes on these ground truth labels. The project also includes implementations of representation learning methods like Spatio-Temporal DeepInfoMax (ST-DIM).

Quick Start & Requirements

• Minimal Install: pip install 'gym[atari]' and pip install git+git://github.com/mila-iqia/atari-representation-learning.git
• Full Install: Requires PyTorch, scikit-learn, TensorFlow, and OpenAI baselines. Installation involves conda and pip commands as detailed in the README.
• Prerequisites: gym[atari], PyTorch, TensorFlow, scikit-learn.
• Usage: Examples for using the AtariARIWrapper and the ProbeTrainer are provided. Training ST-DIM can be initiated via python -m scripts.run_probe --method infonce-stdim --env-name {game}NoFrameskip-v4.

Highlighted Details

• AtariARI wrapper provides ground truth RAM annotations for state variables.
• ProbeTrainer facilitates evaluating learned representations with linear probes.
• Includes implementation of Spatio-Temporal DeepInfoMax (ST-DIM).
• Supports various Atari environments with detailed usage examples.

Maintenance & Community

The project originates from Mila and has contributions from authors affiliated with Mila and other institutions. Links to the paper, poster, and slides are provided. No specific community channels (Discord/Slack) or roadmap are mentioned.

Licensing & Compatibility

The repository does not explicitly state a license. However, it depends on libraries like gym, PyTorch, and TensorFlow, which have permissive licenses. Users should verify licensing for commercial or closed-source use.

Limitations & Caveats

RAM labels are intended for full-sized Atari observations (210x160); downsampling may affect probing accuracy. The project relies on specific versions of dependencies and may require careful environment setup.