Source-linked AI summary
Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, Timothy Lillicrap, David Silver
TL;DR
Planning methods typically require known environment dynamics, while model-free methods struggle with domains requiring precise lookahead. MuZero combines tree search with a learned model that predicts planning-relevant rewards, policies, and values. It reaches state-of-the-art Atari performance and matches superhuman AlphaZero performance in Go, chess, and shogi without game rules.
Problem
Planning algorithms require known dynamics, while prior model-based and model-free RL methods each remain limited in domains where the other approach excels.
Method
MuZero combines AlphaZero-style tree search with a recurrent learned model that predicts reward, policy, and value from hidden states and hypothetical actions.
Results
MuZero achieves state-of-the-art performance across 57 Atari games and matches superhuman performance in Go, chess, and shogi without game rules.
Takeaways & Limitations
A learned model can support powerful planning across visually complex Atari games and precision board games without supplied environment dynamics.
Takeaways & Limitations
Imperfect-information games such as Poker are not directly addressed by MuZero.
Abstract
from arXiv · showhide
Constructing agents with planning capabilities has long been one of the main challenges in the pursuit of artificial intelligence. Tree-based planning methods have enjoyed huge success in challenging domains, such as chess and Go, where a perfect simulator is available. However, in real-world problems the dynamics governing the environment are often complex and unknown. In this work we present the MuZero algorithm which, by combining a tree-based search with a learned model, achieves superhuman performance in a range of challenging and visually complex domains, without any knowledge of their underlying dynamics. MuZero learns a model that, when applied iteratively, predicts the quantities most directly relevant to planning: the reward, the action-selection policy, and the value function. When evaluated on 57 different Atari games - the canonical video game environment for testing AI techniques, in which model-based planning approaches have historically struggled - our new algorithm achieved a new state of the art. When evaluated on Go, chess and shogi, without any knowledge of the game rules, MuZero matched the superhuman performance of the AlphaZero algorithm that was supplied with the game rules.
1 Introduction
MuZero addresses the tension between planning methods that require known dynamics and model-free methods that struggle with precise lookahead. It combines tree search with a learned model focused on planning-relevant predictions.
- Motivation: Planning algorithms succeed in games and real-world applications but require known environment dynamics, limiting direct use in robotics and other complex domains.Such knowledge may be game rules or an accurate simulator.
- Motivation: Prior model-based RL methods remain below state of the art in visually rich Atari games, while model-free RL struggles in chess and Go.The methods differ in whether they reconstruct observations or directly estimate policy and value.
- Contribution: MuZero combines AlphaZero-style search with a learned model to achieve state-of-the-art Atari performance and maintain superhuman chess, shogi, and Go performance.The approach also extends AlphaZero to single-agent domains and intermediate rewards.
- Contribution: MuZero iteratively transforms observations into hidden states and predicts policy, value, and immediate reward for hypothetical actions.The model is trained end-to-end to estimate these quantities for planning.
2 Prior Work
Prior model-based RL often models environmental states or observations, whereas value-equivalent approaches focus on planning-relevant outcomes. MuZero builds on this trajectory with a hidden model that supports policy, value, and reward prediction.
- Model-based RL: Model-based RL learns an environment model before planning, traditionally representing state transitions and expected rewards conditioned on actions.This model is commonly formalized as a Markov decision process.
- Observation models: Pixel-level observation modeling is computationally difficult for large-scale planning, while latent models still devote substantial capacity to predicting observations or latent states.These approaches target reconstruction or prediction rather than only planning-relevant quantities.
- Value-equivalent models: Value-equivalent models construct abstract dynamics whose cumulative rewards match those of trajectories in the real environment.Their transition model need not reproduce real environment states.
- MuZero’s relation to prior work: MuZero’s model uses representation, dynamics, and prediction components to process observations, actions, rewards, policies, and values.Its recurrent hidden-state structure supports iterative planning without requiring a pixel-level reconstruction.
- Value-based predecessors: TreeQN, value iteration networks, and value prediction networks learn abstract or local models for value optimization, with value prediction networks the closest precursor to MuZero.These methods ground learned models in planning or real actions to approximate value functions.
3 MuZero Algorithm
MuZero predicts planning-relevant quantities by recurrently applying a learned model to past observations and hypothetical actions. Its jointly trained components support MCTS while avoiding semantic dependence on true environment states.
- Model predictions: For each future step k = 1...K, MuZero predicts policy, value, and reward from past observations and future actions.The model is conditioned on observations o1, ..., ot and actions at+1, ..., at+k.
- Model structure: The dynamics function recurrently maps a hidden state and hypothetical action to an internal state and immediate reward.This creates the model’s internal transition process.
- Model structure: Unlike traditional environment models, MuZero’s hidden state has no environment-state semantics and exists solely to predict policies, values, and rewards.The representation, dynamics, and prediction functions jointly form the internal model.
- Planning: MuZero searches hypothetical action trajectories with an MCTS algorithm similar to AlphaZero’s, generalized to single-agent domains and intermediate rewards.The search operates over internal rewards and states induced by the learned dynamics function.
- Training: The model is trained end-to-end to match predicted policies, values, and rewards with search-generated or observed targets at corresponding future steps.Policy targets come from MCTS, value targets come from game or MDP outcomes, and reward targets are observed rewards.
- Training: The value target may use discounted n-step bootstrapping, allowing long episodes and intermediate rewards beyond board-game terminal outcomes.Board-game outcomes are represented as −1, 0, or +1 rewards.
- Training: The overall loss combines reward, value, and policy losses with L2 regularization.The component losses are denoted lr, lv, and lp.
4 Results
MuZero was evaluated across board games and visually complex Atari domains, including scalability, model-free comparisons, and sample-efficient training. It matched or exceeded strong baselines while using learned models for planning.
- Evaluation scope: MuZero was evaluated on Go, chess, shogi, and all 57 Atari games.Board games benchmark challenging planning, while Atari benchmarks visually complex reinforcement-learning domains.
- Board-game performance: MuZero slightly exceeded AlphaZero in Go despite using less computation per search-tree node.MuZero used 16 residual blocks per evaluation compared with AlphaZero’s 20, suggesting possible computation caching in the search tree.
- Atari performance: MuZero achieved a new state of the art in Atari mean and median normalized scores across 57 games.It outperformed R2D2 in 42 of 57 games and surpassed SimPLe in all games.
- Sample efficiency: 731% median normalized score was achieved by MuZero Reanalyze using 200 million frames per Atari game.This exceeded reported scores of 192%, 231%, and 431% for IMPALA, Rainbow, and LASER, respectively.
- Planning scalability: MuZero matched a perfect model in Go while searching up to 10 seconds, despite training with approximately 0.1-second searches.The learned model therefore remained effective when search time greatly exceeded the training setting.
- Planning scalability: Atari performance plateaued around 100 simulations, while a single simulation still produced strong performance.The weaker scaling than Go was attributed to greater model inaccuracy in Atari; the raw policy appeared to internalize search benefits.
- Learning comparison: On Ms. Pacman, MuZero learned faster and reached a higher final score than a comparable Q-learning system.The Q-learning implementation matched R2D2 but learned significantly slower and converged to a much lower final score.
- Training scalability: MuZero learned an effective Ms. Pacman policy with only 6 simulations per move during training.This was fewer simulations than the number of available actions.
5 Conclusions
MuZero combines high-performance planning with model-free reinforcement-learning strengths through a learned model that requires no environment dynamics. It matches superhuman board-game planning performance and outperforms state-of-the-art model-free methods in Atari.
- Conclusion: MuZero combines tree-based planning with learned models without requiring game rules or environment dynamics.The paper presents this combination as addressing both logically complex board games and visually complex Atari games.
- Conclusion: MuZero matched superhuman planning performance in chess and Go while outperforming state-of-the-art model-free reinforcement-learning algorithms in Atari.The conclusion frames these as the favored domains of high-performance planning and model-free reinforcement learning, respectively.
- Implications: The method may support applying learning and planning to real-world domains without perfect simulators.The paper identifies robotics, industrial control, and intelligent assistants as examples of such domains.
Supplementary Materials
The supplementary materials provide algorithm pseudocode, machine-readable data for the main and supplementary figures and tables, and related background on MuZero’s generalization beyond AlphaZero.
- Supplementary materials: The supplementary materials include a pseudocode description of the MuZero algorithm.
- Supplementary materials: Machine-readable JSON data are provided for Figures 2, 3, S2, S3, and S4 and Tables 1, S1, and S2.
- Supplementary materials: Supplementary materials are accessible through the ancillary files of the arXiv submission.
- Method context: MuZero is designed for a more general setting than AlphaGo Zero and AlphaZero.The surrounding comparison describes how MuZero replaces game-rule components with a learned implicit model and supports broader reinforcement-learning settings.
- Method context: MuZero replaces AlphaZero’s simulator-based state transitions, legal-action masking, and terminal-node handling with a learned implicit model.The learned model represents tree nodes with hidden states and predicts the quantities used for search.
- Method context: MuZero extends beyond two-player games with terminal ±1 rewards to single-agent domains and discounted intermediate rewards of arbitrary magnitude.
Appendix B Search
MuZero uses Monte-Carlo tree search with upper confidence bounds over learned internal states, adapting search to intermediate rewards, discounting, and unbounded values.
- MuZero’s search uses Monte-Carlo tree search with upper confidence bounds, converging asymptotically to optimal policies in single-agent domains and minimax values in zero-sum games.
- Each search-tree edge stores visit count, mean value, policy, reward, and state-transition statistics.
- Selection: Selection repeatedly chooses actions by maximizing an upper confidence bound over stored statistics until reaching a leaf node.
- Selection: Before the final simulation step, next states and rewards are retrieved from the stored transition and reward tables.
- Expansion: At expansion, the dynamics function computes the next hidden state and reward, while the prediction function supplies policy and value estimates for initializing a new node.
- Backup: Backup supports intermediate rewards, discounts γ different from 1, and unbounded value estimates through bootstrapped discounted-return estimates.
- MuZero normalizes values during search to avoid game-specific prior knowledge when values are unbounded.
- Board games use 800 simulations per search, whereas Atari uses 50 because of its smaller branching factor and simpler policies.
Appendix D Data Generation
Data generation uses MCTS self-play with domain-specific simulation budgets, replay organization, exploration schedules, and action-history representations tailored to board games and Atari.
- Training data comes from games played by the latest network checkpoint with MCTS, using 800 simulations per board-game move and 50 per Atari move.
- Board-game data is submitted after games finish, while Atari sequences are submitted every 200 moves because Atari episodes can reach 108,000 frames.
- Atari samples actions from visit-count distributions throughout games, unlike board-game exploration that follows AlphaZero’s scheme.
- The Atari temperature decays from 1 to 0.5 to 0.25 across successive training intervals, making action selection greedier over training.
- Board-game inputs encode recent board histories, using the last 8 states in Go and shogi and the last 100 in chess to help predict draws.
- Atari inputs contain the last 32 RGB frames at 96x96 resolution and the 32 actions that produced them.
- The dynamics function receives the hidden state plus an action representation, with action planes matching the hidden-state or board resolution.
Appendix F Network Architecture
The network uses AlphaZero-like prediction and residual architectures, with domain-specific downsampling and action encoding; Atari value and reward targets use transformed categorical supports.
- The prediction function uses AlphaZero’s architecture: convolutional layers reduce planes while preserving resolution, followed by a fully connected output layer.
- Target representation: Atari value and reward targets are scaled with an invertible transform and represented categorically over 601 supports from −300 to 300.
- Representation and dynamics functions use 16 residual blocks with 256 hidden planes and 3x3 kernels, while preserving AlphaZero-like design choices.
- Prediction architecture: The network uses three residual blocks with 256 planes followed by average pooling to 6x6 in the prediction pathway.
- Dynamics architecture: The dynamics function encodes each action as an image and stacks it with the previous hidden state along the plane dimension.
Appendix G Training
Training unrolls MuZero from replayed trajectory states, jointly fitting search-derived policy and value targets plus rewards, with prioritized Atari replay and gradient scaling.
- Training unrolls the network for K hypothetical steps from states sampled from replay-buffer trajectories.
- Atari sequences use prioritized replay based on absolute search-value versus observed-return error, while board-game states are sampled uniformly.
- Each unrolled step predicts policy, value, and reward targets, with reward loss omitted for board games without intermediate rewards.
- Losses are jointly trained across representation, dynamics, and prediction functions using targets aligned to corresponding future real-time observations.
- Each prediction head’s loss is scaled by 1/K so total gradient magnitude remains similar across different unroll lengths.
- The dynamics-function gradient is additionally scaled by 1/2 to keep its total applied gradient constant.
- Experiments always use K = 5 unrolled steps.
- Training uses 16 TPUs and 1000 self-play TPUs per board game, versus 8 training and 32 self-play TPUs per Atari game.
Appendix H Reanalyze
MuZero Reanalyze improves sample efficiency by revisiting past trajectories with the latest model and search, while stabilizing value targets and increasing sample reuse.
- Reanalysis: MuZero Reanalyze re-executes search on past time-steps using the latest model parameters to produce potentially better policy targets.The fresh policy supplies the policy target for 80% of training updates.
- Value targets: A target network based on recent parameters provides a fresher, stable n-step bootstrapped target for the value function.The target is formed from future rewards and the target network’s value estimate.
- Training adjustments: Sample reuse was increased by drawing 2.0 samples per state instead of 0.1, while the value-target weight was reduced to 0.25 and n-step returns to n = 5.Policy and reward targets retained weights of 1.0; the earlier n-step setting was n = 10.
Appendix I Evaluation
The appendix describes evaluation protocols for MuZero across board games and Atari, including Elo-based board-game comparisons, repeated Atari training runs, and supplementary algorithm and learning-curve analyses.
- Board-game evaluation: Board-game strength was measured with Elo ratings estimated by Bayesian logistic regression from tournament outcomes.The win probability used p(a defeats b) = (1+10(celo(e(b)−e(a))))−1 with celo = 1/400.
- Board-game evaluation: MuZero’s board-game Elo was computed from tournaments using 800 simulations per move against Stockfish, Elmo, or AlphaZero baselines.Baseline players used equivalent 100ms-per-move search time, with ratings anchored to publicly available values.
- Supplementary analyses: The supplementary evaluations report Atari training repeatability and learning curves using total reward against millions of training steps, with medians, individual runs, means, and variability bands.Figure S1 summarizes five games across 10 runs; Figure S5 reports individual-game means across 1000 evaluation games.
- Atari evaluation: Atari evaluation used mean reward over 1000 episodes per game, with 30-minute or 108,000-frame episode limits and 50 simulations per move unless otherwise indicated.Two evaluation strategies mitigated deterministic-simulator effects: 30 random no-op starts and human starts.
- Supplementary analyses: Supplementary materials also summarize MuZero’s planning and learning equations and analyze evaluation depth and policy-improvement ablations.Figure S3 distinguishes learned-model evaluation depth from policy-improvement analyses.