Mastering Atari Games with Limited Data

TL;DR

EfficientZero targets the challenge of data-efficient reinforcement learning in image-based tasks by building on MuZero and adding three innovations: a self-supervised temporal consistency loss to better learn the environment model, an end-to-end value-prefix predictor to combat state aliasing, and a model-based off-policy correction to stabilize learning with limited data. The method achieves super-human performance on Atari 100k with only 2 hours of gameplay and reaches competitive results on DMControl 100k, outperforming several prior baselines and demonstrating strong practical potential for real-world applications. Through targeted ablations, the paper shows that the consistency signal is the most impactful component, with the value-prefix and off-policy correction providing additional gains in early learning and data-limited regimes. Together, these contributions advance sample-efficient, model-based RL for high-dimensional visual control and open avenues for lighter, more deployable planning-based agents.

Abstract

Reinforcement learning has achieved great success in many applications. However, sample efficiency remains a key challenge, with prominent methods requiring millions (or even billions) of environment steps to train. Recently, there has been significant progress in sample efficient image-based RL algorithms; however, consistent human-level performance on the Atari game benchmark remains an elusive goal. We propose a sample efficient model-based visual RL algorithm built on MuZero, which we name EfficientZero. Our method achieves 194.3% mean human performance and 109.0% median performance on the Atari 100k benchmark with only two hours of real-time game experience and outperforms the state SAC in some tasks on the DMControl 100k benchmark. This is the first time an algorithm achieves super-human performance on Atari games with such little data. EfficientZero's performance is also close to DQN's performance at 200 million frames while we consume 500 times less data. EfficientZero's low sample complexity and high performance can bring RL closer to real-world applicability. We implement our algorithm in an easy-to-understand manner and it is available at https://github.com/YeWR/EfficientZero. We hope it will accelerate the research of MCTS-based RL algorithms in the wider community.

Figures (9)

• Figure 1: Our proposed method EfficientZero is 176% and 163% better than the previous SoTA performance in mean and median human normalized score and is the first to outperform the average human performance on the Atari 100k benchmark. The high sample efficiency and performance of EfficientZero can bring RL closer to the real-world applications.
• Figure 2: The self-supervised consistency loss.
• Figure 3: A sample trajectory from the Atari Pong game. In this case, the right player didn't move and missed the ball.
• Figure 4: Evaluations of image reconstructions based on latent states extracted from the model with or without self-supervised consistency. The predicted next states with consistency can basically be reconstructed into observations while the ones without consistency cannot.
• Figure 5: Training and validation losses of direct reward prediction method and the value prefix method.