PaperlayerView exact revision on arXiv

Source-linked AI summary

Mastering Atari Games with Limited Data

Weirui Ye, Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, Yang Gao

arXiv:2111.00210v2cs.LGcs.AIcs.CVcs.RO

TL;DR

Reinforcement learning still faces a sample-efficiency challenge, especially when image-based methods must reach human-level Atari performance with limited data. EfficientZero builds a model-based visual RL agent on MuZero using self-supervised consistency, end-to-end value-prefix learning, and off-policy value-target correction. It achieves super-human Atari 100k performance with two hours of gameplay and state-of-the-art performance on some DMControl tasks.

  • Problem

    Image-based reinforcement learning had not consistently reached human-level Atari performance under limited data, despite progress in sample-efficient methods.

  • Method

    EfficientZero is a MuZero-based model-based visual RL algorithm using a self-supervised environment model, end-to-end value-prefix learning, and model-corrected off-policy value targets.

  • Results

    194.3% mean and 109.0% median human normalized performance are achieved on Atari 100k with only two hours of real-time gameplay, alongside state-of-the-art performance on some DMControl tasks.

  • Takeaways & Limitations

    EfficientZero demonstrates that model-based visual reinforcement learning can achieve super-human Atari performance with very limited gameplay data.

  • Takeaways & Limitations

    MuZero’s environment model can suffer from insufficient supervision because reward signals may be sparse and value targets are noisy, causing prediction errors to accumulate during deep MCTS expansions.

Abstract

from arXiv · show

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.

1 Introduction

EfficientZero addresses the challenge of achieving strong image-based reinforcement learning performance with limited environmental data. Built on MuZero, it combines three changes and reaches super-human Atari 100k performance with only two hours of gameplay.

  • Motivation: Sample complexity remains a major challenge because successful reinforcement learning methods can require millions or billions of environmental interactions.AlphaZero, for example, needs 21 million games during training.
  • Motivation: High performance with low sample complexity is important for applying reinforcement learning to real-world problems such as robotics, healthcare, and recommendation systems.
  • Motivation: Model-based methods can improve sample efficiency by training policies with both real-environment data and model-generated imagined data, but many successes remain state-based.
  • Approach: EfficientZero modifies MuZero with a self-supervised temporally consistent environment model, end-to-end value-prefix learning, and model-based correction of off-policy value targets.
  • Results: 194.3% mean and 109.0% median human normalized performance are achieved on Atari 100k with only two hours of real-time gameplay.

2 Related Work

Prior work improved sample-efficient reinforcement learning, but limited-data image-input settings still lacked consistent human-level performance. EfficientZero targets this gap with model-based learning while addressing multi-step value estimation and off-policy concerns.

  • Sample-Efficient Reinforcement Learning: Sample-efficient reinforcement learning has advanced, including image-based methods for simulated robotic control and Atari performance with as little as 400k frames.
  • Sample-Efficient Reinforcement Learning: Despite recent progress, limited-data reinforcement learning algorithms remained well behind human performance.
  • Sample-Efficient Reinforcement Learning: Traditional model-based reinforcement learning is considered more sample efficient, while model-free methods dominated performance in image-input settings.
  • MCTS-Based Reinforcement Learning: MCTS-based reinforcement learning combines model-based and model-free learning and has been extended from board games to learned world models, continuous actions, and offline data.
  • Multi-Step Value Estimation: Multi-step value targets can accelerate convergence but introduce off-policy issues because replay-buffer rewards may come from a different policy.
  • Multi-Step Value Estimation: EfficientZero uses model-based value estimates to bootstrap from real-world rewards, balancing off-policy correction against the risk of model exploitation.

3 Background

MuZero uses learned latent representations, environment dynamics, policy priors, and value estimates within MCTS to improve action selection. Training collects MCTS-guided experience and optimizes predictions from replayed data.

  • MuZero: MuZero learns an environment model with reward and dynamics functions, alongside policy and value functions used by MCTS.
  • MuZero: MCTS expands a search tree using the learned model, policy prior, and value function, then returns a root action-visit distribution as an improved policy.
  • MuZero: The functions operate on hidden abstract states extracted from observations by a representation function for computational efficiency and easier environment modeling.
  • Training: Training stores MCTS-generated rollout data in a replay buffer and minimizes a loss over sampled data.
  • Training: The training targets include environment rewards, predicted rewards, MCTS visit distributions, predicted policies, and bootstrapped value targets.
  • Monte-Carlo Tree Search: MCTS balances exploration and exploitation by combining Q-values, policy priors, and visit counts, with training-time noise encouraging exploration.

4 EfficientZero

EfficientZero addresses three obstacles that limit MuZero-style model-based RL with scarce data: insufficient model supervision, compounding reward-prediction errors, and off-policy multi-step targets. It introduces self-supervised consistency learning, end-to-end value-prefix prediction, and model-based off-policy correction.

  • 4 EfficientZero: EfficientZero targets limited-data weaknesses in MuZero: inadequate environment-model supervision, aleatoric reward uncertainty, and off-policy multi-step values.These issues respectively weaken model learning, accumulate errors during deep MCTS expansion, and hinder convergence.
  • 4 EfficientZero: The three modifications are proposed as critical changes intended to improve performance when samples are limited.EfficientZero is presented as a model-based algorithm built on MCTS that achieves super-human performance on Atari 100k.
  • 4.1 Self-Supervised Consistency Loss: Self-supervised consistency learning trains predicted hidden states to match representations of subsequent observations, providing richer signals than scalar rewards and values.The method shapes hidden-state and predicted-state consistency through the dynamics function and extends representation learning to adjacent observations and transitions.
  • 4.2 End-To-End Prediction of the Value Prefix: Longer model rollouts compound prediction errors, making exact future rewards difficult to predict from aliased states and degrading MCTS expansion and action search.The paper motivates longer-horizon prediction because broad outcomes can be easier to forecast than the exact timestep of a reward.
  • 4.2 End-To-End Prediction of the Value Prefix: Value-prefix prediction estimates discounted rewards from unrolled states, allowing end-to-end and per-step supervision through an LSTM.The value prefix is used as a prefix in later Q-value computation and can be supervised whenever a new state arrives.
  • 4.3 Model-Based Off-Policy Correction: Model-based off-policy correction uses shorter reward horizons for older trajectories, then redoes MCTS with the current policy at the resulting state.This imagined correction reduces policy divergence while addressing bias introduced by shortening the horizon.

5 Experiments

Experiments evaluate EfficientZero on Atari 100k and DMControl 100k, then analyze how its components affect performance, latent-state consistency, value prediction, and off-policy correction.

  • Atari 100k: 100k environment steps correspond to 400k frames and roughly two hours of gameplay, whereas DQN uses 200 million frames or about 925 hours.The benchmark contains 26 Atari games and reports raw scores plus mean and median human-normalized scores.
  • Atari 100k: EfficientZero achieves a 1.904 mean and 1.160 median human-normalized score on Atari 100k, outperforming DQN's 0.959 median with 500 times less data.It outperforms human performance in 14 of 26 games and is reported as 170% and 180% better than SPR in mean and median score, respectively.
  • DMControl 100k: EfficientZero achieves state-of-the-art performance on selected DMControl 100k tasks, with results comparable to state SAC and lower variance than CURL.MuZero does not work well on these simulated tasks, while state SAC uses ground-truth states.
  • Component ablations: Removing self-supervised consistency, the end-to-end value prefix, or model-based off-policy correction reduces performance, with consistency producing the largest drop.The ablation evaluates each removed component across 26 Atari games.
  • Temporal consistency: Consistency enables predicted next latent states to reconstruct observations, while states from the version without consistency cannot be reconstructed reliably.The analysis attributes this to distributional shifts between representation states and dynamics-function states that consistency reduces.
  • Value prefix: The value prefix has lower validation error after five-step unrolling despite direct reward prediction having lower training loss, indicating less overfitting.The authors connect this behavior to reduced state aliasing and better generalization.
  • Off-policy correction: Value-target error decreases for fresher trajectories, while off-policy correction yields smaller errors across trajectories of different ages.The result indicates that stale replay data creates a substantial off-policy issue under limited data.

6 Discussion

EfficientZero is presented as a sample-efficient model-based method that reaches super-human Atari performance with two hours of gameplay and state-of-the-art performance on some DMControl tasks. The authors identify future work in continuous-action design, MCTS acceleration, and lifelong learning.

  • Discussion: EfficientZero achieves super-human Atari performance with two hours of gameplay and state-of-the-art performance on some DMControl tasks.The paper frames this result as one step toward running reinforcement learning with complex sensory inputs in the physical world.
  • Discussion: Future work targets better continuous-action design, faster MCTS, and combining the framework with lifelong learning.

A.1 Models and Hyper-parameters

EfficientZero uses a compact MuZero-style network pipeline with representation, dynamics, and prediction components, tailored for limited-data visual RL. Training uses self-supervised consistency, recurrent prediction, stacked image histories, and 120k optimization steps with data collected during the first 100k.

  • Architecture: The model has representation, dynamics, and prediction components, following a MuZero-style pipeline.The dynamics network reduces MuZero’s residual blocks from 16 to 1 and adds a residual link for historical hidden states.
  • Architecture: The dynamics network concatenates states and actions into 65 planes, then applies downsampling, a residual link, and one 64-plane residual block.The input action-state tensor is followed by a stride-2 convolution and batch normalization.
  • Prediction: Prediction networks use two-layer MLPs with batch normalization to predict reward, value, or policy.The reward network predicts a value prefix using an LSTM with hidden size 512; output dimensions differ between reward and policy prediction.
  • Prediction: The LSTM is trained for 5 unrolled steps and reset after ζ = 5 recurrent-inference steps.MCTS can use deeper dynamics rollouts than the training unroll length.
  • Inputs: Inputs stack 4 historical frames with a 4-frame skip, producing a 96 × 96 × 12 tensor covering 16 frames.The method uses batch normalization, reward clipping, and no additional state normalization.
  • Training schedule: EfficientZero trains for 120k steps while collecting data only during the first 100k steps, then lowers the learning rate from 0.2 to 0.02.The later trajectories can therefore be fully used during training.

A.2 More Ablations

Ablations show that EfficientZero’s full combination performs best, with self-supervised consistency especially important across environments. Within its components, value prefix helps sparse-reward settings, while dynamic horizon is more important than MCTS root value for off-policy correction.

  • Component ablations: The full EfficientZero version outperforms variants missing self-supervised consistency, value prefix, or off-policy correction.The importance of individual components varies across environments.
  • Component ablations: Removing self-supervised consistency produces much poorer results across all evaluated environments.This component is more consequential than the tested data augmentations.
  • Component ablations: Value prefix matters in sparse-reward environments such as Pong, whereas state aliasing has less effect on reward signals in dense-reward environments such as Qbert.The passage attributes the difference to reward availability in the two settings.
  • Data augmentation: Removing data augmentation has similar performance to the full consistency component, indicating that the self-supervised loss provides most of its improvement.The augmentation study uses random 0–4 pixel shifts and intensity changes.
  • Off-policy correction: Removing dynamic horizon performs worse than removing MCTS root value, making dynamic horizon the more important off-policy-correction technique.This comparison is based on several Atari games trained for 100k steps.

A.3 MCTS Details

EfficientZero uses MuZero-style MCTS with 50 simulations, UCT-based selection, expansion with predicted reward and policy, and backup through normalized Q-values. It additionally uses mean Q-values for unvisited nodes and root-level exploration noise during training-related searches.

  • Search procedure: Policy search repeats MuZero’s three-stage MCTS procedure for Nsim = 50 simulations.The stages are selection, expansion, and backup.
  • Selection: UCT selection balances Q-values, visit counts, and policy priors while traversing the tree from the root state.The usual constants are c1 = 1.25 and c2 = 19652.
  • Selection: A mean Q-value mechanism estimates unvisited nodes instead of assigning them the default value 0.The experiments report better exploration from the mean Q-value mechanism.
  • Expansion: Expansion assigns predicted reward and policy as priors, adding Dirichlet noise at the root during self-play and reanalysis.The noise is omitted for non-root nodes and evaluation, where ρ is set to 0.
  • Backup: Backup normalizes Q-values using tree-wide minimum and maximum values with ε = 0.01 to smooth the range.The normalization addresses overconfidence caused by small value differences under limited data.
  • Outputs: After search, the root’s average value supports off-policy correction and its visit-count distribution becomes the target policy distribution.The MCTS output temperature is reduced to 0.5 and then 0.25 during training.

A.4 Training Details

Training combines replay, reanalysis, dynamic-horizon off-policy correction, and parallel self-play and learner workers. Targets become fresher through reanalysis, while the dynamic horizon shortens unrolling for stale trajectories.

  • Pipeline: The implementation uses a parallel architecture with double buffering in PyTorch and Ray.The design shares replay, context, and batch queues across worker types.
  • Pipeline: Self-play actors collect trajectories into replay while CPU rollout workers prepare sampled transition contexts and GPU workers reanalyze batches.The learner, data workers, CPU workers, and GPU workers run in parallel through shared queues.
  • Replay: Prioritized replay uses value-training error, α = 0.6, and β annealed from 0.4 to 1.0, while collecting intermediate sequences of 400 moves.The priority mechanism improves performance only slightly with limited data.
  • Reanalysis: Reanalysis revisits past trajectories with the latest target model and MCTS to obtain fresher value and policy targets.Policy is reanalyzed for 99% of the data and value for 100%.
  • Off-policy correction: The dynamic horizon uses trajectory freshness to choose the unroll length, with k = 5, total collection duration 100k steps, and τ = 0.3.Stale trajectories are unrolled less to reduce gaps between current predictions and old rollouts.
  • Off-policy correction: The target value replaces the predicted v_t+k with the averaged MCTS root value to alleviate off-policy bias.This correction is paired with the dynamic horizon.

A.5 Evaluation

EfficientZero is evaluated on Atari 100k across 26 games using training curves, multiple evaluation seeds, and aggregate metrics with confidence intervals. It significantly outperforms other methods across the reported aggregate metrics.

  • Evaluation setup: Evaluation covers 26 Atari 100k games, with training curves tracking average rewards across 32 evaluation seeds for three runs.The curves span 120,000 total training steps.
  • Aggregate metrics: Aggregate evaluation uses mean, median, IQM, and optimality gap, with 95% confidence intervals estimated by stratified percentile bootstrap.Higher mean, median, and IQM scores and lower optimality gaps indicate better performance.
  • Evaluation setup: Scores are reported for three random seeds per game, with each run evaluated using 32 different seeds and summary means and standard deviations.Table 10 reports the per-game scores and aggregates across runs.
  • Aggregate metrics: EfficientZero significantly outperforms the other methods on all four reported aggregate metrics for Atari 100k.The comparison includes results from Agarwal et al., with most methods based on 10 runs per game, SimPLe on five, and EfficientZero on three.
Loading 2111.00210v2…