An educational and commented implementation of MuZero, based on DeepMind's 2019 paper and original pseudocode, built with PyTorch and Ray.
MuZero builds on AlphaZero by learning the dynamics used for planning from experience. Its network represents observations as a latent state, predicts how actions change that state, and estimates rewards, values and policies. Monte Carlo tree search (MCTS) uses these predictions to choose actions.
For an introduction, watch the MuZero explanatory video or read How MuZero works and the project wiki.
The implementation supports fully connected and residual networks, parallel self-play, prioritized replay, reanalysis, and single-player or two-player turn-based games with discrete actions.
To adapt it to a new environment, add a game file containing its configuration and Game adapter. See CartPole for an example and Add an environment for the interface.
| CartPole | Lunar Lander | Breakout |
|---|---|---|
 |
 |
 |
| Learning curve · Checkpoint | Learning curve · Checkpoint | Learning curve · Checkpoint |
Use Python 3.12 or later. The included checkpoints were trained on Linux with Python 3.12.
git clone https://github.com/werner-duvaud/muzero-general.git
cd muzero-general
python3 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txtSee the official PyTorch installation guide for GPU setup and Gymnasium Box2D instructions for Lunar Lander build dependencies. MCTS tree visualization also needs the system Graphviz executable, in addition to its Python package.
Choose an environment:
python muzero.py cartpole
python muzero.py lunarlander
python muzero.py breakoutEach command starts a fresh model using that game's default config, runs the configured training budget, then evaluates the selected checkpoint. The defaults use CPU training for CartPole and Lunar Lander, and a GPU for Breakout when CUDA is available. Use --cpu to force CPU execution.
From Python:
from muzero import MuZero
agent = MuZero("cartpole")
agent.train()python muzero.py cartpole --evaluate
python muzero.py lunarlander --evaluate
python muzero.py breakout --evaluateThese commands load checkpoints/<env>/model.checkpoint and evaluate it on CPU, without training or rendering. The default is 100 episodes. Checkpoints contain their saved configuration; no JSON sidecar is needed.
Run TensorBoard in another terminal:
tensorboard --logdir ./resultsFor interactive rendering, human play, model diagnostics or hyperparameter search, open the menu:
python muzero.pyRun python muzero.py --help for all command-line options.
Each game file contains a complete, commented MuZeroConfig. All games share the same parameter names; values are tuned per environment. Edit the Python config to change the default recipe. config.py provides shared validation and temperature-schedule helpers.
You can also override fields for one run:
python muzero.py cartpole --config '{"num_workers": 4}' --output results/cartpole/example--config accepts a JSON object or a JSON file. --seed changes the training seed; evaluation seeds are configured separately.
By default, runs are saved under results/<env>/<timestamp>/:
| File | Contents |
|---|---|
model.checkpoint |
Latest saved training state, including weights and optimizer state |
model-best.checkpoint |
Best validation checkpoint, when validation is enabled |
replay_buffer.pkl |
Collected games for resuming training |
config.json, run.json |
Effective configuration and runtime/source metadata |
validation.json, evaluation.json |
Validation history and final evaluation results |
TensorBoard event files are saved in the same directory. With checkpoint_selection="best", the final evaluation uses the best validation checkpoint, which may precede the last gradient update.
Replay snapshots can be many gigabytes for image-based environments.
To evaluate your own run or resume interrupted training, replace RUN with its directory name:
python muzero.py cartpole --evaluate --checkpoint results/cartpole/RUN/model-best.checkpoint
python muzero.py cartpole --checkpoint results/cartpole/RUN/model.checkpoint --replay-buffer results/cartpole/RUN/replay_buffer.pklResume with both the checkpoint and replay buffer to retain collected experience. To train beyond the saved budget, increase training_steps through --config.
Additional adapters are available for:
- MiniGrid (
gridworld) and a simple grid world (simple_grid). - Tic-tac-toe, Connect4 and Gomoku.
- Twenty-One (
twentyone), a simplified card game. - OpenSpiel (
spiel), configured for tic-tac-toe by default. The adapter supports deterministic games with one player or two alternating players and 1D–3D observation tensors. Chance-based games such as backgammon and simultaneous-action games are unsupported. - An Atari template (
atari), configured for Breakout by default. Usebreakoutfor the tuned recipe and packaged checkpoint shown above.
Only CartPole, Lunar Lander and Breakout currently include packaged checkpoints.
- Copy games/cartpole.py to
games/my_game.py, keeping the common config fields. - Set the observation shape, action/player lists, network and training parameters. Observations must have shape
(channels, height, width); reshape a vector of lengthnto(1, 1, n). - Implement
Gameusing AbstractGame: initialization,reset(),step(action),legal_actions()andrender().step()returns(observation, reward, done). Implementto_play()for multiplayer games andclose()to release resources. - Run
python muzero.py my_game.
Actions and players use consecutive IDs starting at zero. Wrap Gymnasium's terminated and truncated flags into done. Trajectory indexing and training-target alignment are documented in game_history.py.
| File | Role |
|---|---|
| muzero.py | CLI, worker orchestration, training and checkpoint loading |
| self_play.py | MCTS, action selection and self-play |
| models.py | Fully connected and residual networks |
| trainer.py | Losses, recurrent unrolling and optimizer updates |
| replay_buffer.py | Replay sampling, eviction and reanalysis |
| game_history.py | Trajectories, observation stacks and training targets |
| evaluation.py | Validation, checkpoint selection and evaluation reports |
| random_bootstrap.py | Random-episode collection and replay initialization |
| diagnose_model.py | Compare learned dynamics with real environment transitions |
Format code with Ruff: ruff format --target-version py312 .
- MuZero paper and original pseudocode.
- Project wiki and explanatory video.
- Network diagram.
- EfficientZero and Sampled MuZero.
Werner Duvaud, Aurèle Hainaut, Paul Lenoir, and contributors. Released under the MIT license.
@misc{muzero-general,
author = {Werner Duvaud, Aurèle Hainaut},
title = {MuZero General: Open Reimplementation of MuZero},
year = {2019},
publisher = {GitHub},
url = {https://github.com/werner-duvaud/muzero-general},
}
