World Models That Speak in Atari Frame Tokens

What it does

IRIS is a data-efficient reinforcement learning agent for Atari games. It first learns a world model from a small amount of real environment interaction, then trains its policy on millions of additional trajectories imagined by that model. The world model itself pairs a discrete autoencoder with an autoregressive Transformer: the autoencoder compresses visual frames into a compact token vocabulary, and the Transformer learns the grammar of how those tokens evolve.

The interesting bit

By casting world modeling as sequence prediction, IRIS turns Atari dynamics into a language problem—frames become words, rollouts become sentences. The repository includes scripts that let you watch the agent play, but more unusually, you can provide keyboard inputs to unroll trajectories inside the imagined world model itself, effectively stepping through the agent’s own hallucinated environment.

Key highlights

• Published at ICLR 2023 and ranked in the notable top 5%
• Discrete autoencoder builds a token vocabulary from Atari frames; the Transformer composes future sequences
• Pretrained checkpoints hosted on Hugging Face for direct evaluation
• Structured run folders automatically store checkpoints, reconstructions, and episode visualizations
• Includes a results notebook to reproduce the paper’s figures

Caveats

• Atari ROMs ship with the dependencies, and the README explicitly warns that you must acknowledge licensing compliance to use them
• Interactive world-model play flushes the Transformer’s memory every 20 frames to keep latency low, which truncates long-horizon imagination during human interaction

Verdict

Grab it if you are researching model-based RL or want to experiment with transformer-based world models in a concrete, reproducible setting. Look elsewhere if you need a general-purpose RL framework; the token vocabulary and architecture are tightly bound to Atari’s visual domain.