Building my_zero: MuZero and Sampled MuZero from Scratch

Notes on building a from-scratch implementation of MuZero and Sampled MuZero for discrete and continuous control.

my_zero is my attempt to implement MuZero and Sampled MuZero. I built it as a learning and research project: the goal was to understand how the learned world model, tree search, self-play, replay, and recurrent optimization connect in a complete system rather than as isolated equations.

The repository currently contains end-to-end examples for both discrete and continuous control, along with a focused test suite. It is released under the MIT license.

Why build it from scratch?

A usable implementation must make several subsystems agree on data shapes, targets, conventions, and timing:

• The representation network maps an observation to a latent state.
• The dynamics network advances that state under a proposed action and predicts reward.
• The prediction network returns a policy and value for latent states.
• MCTS turns network predictions into improved action targets for self-play.
• Replay and recurrent unrolling convert collected trajectories into supervised training targets.

Implementing those pieces directly made the interface between planning and learning much more concrete. It also made it easier to examine practical choices such as value support, search exploration, action sampling, replay prioritization, and checkpointing.

What the repository implements

The codebase is organized around a configurable MuZeroNet with representation, dynamics, and prediction components. The network supports both scalar outputs and categorical-support value/reward targets, and it uses configurable MLP bodies with optional skip connections and latent normalization.

On top of the network, the repository includes:

• PUCT-based MuZero search for discrete action spaces, including root Dirichlet noise and legal-action masking.
• Sampled MuZero search for continuous actions, with policy-proposed and uniform candidate actions.
• Multi-step recurrent training of policy, value, and reward predictions.
• Standard and prioritized experience replay.
• Parallel self-play workers, checkpoints, JSONL metrics, TensorBoard logging, and environment callbacks.

The discrete and sampled search implementations deliberately live in separate modules. That keeps the core distinction visible: standard MuZero can expand every available action, while Sampled MuZero must choose a tractable candidate set before allocating search effort.

Two small but complete experiments

CartPole with MuZero

The CartPole example uses standard MuZero with a discrete two-action policy. It is intentionally small enough to make the training and search configuration easy to inspect: the example defines the latent-network dimensions, value support, MCTS budget, replay settings, and self-play worker count in one place.

In a representative run, the average episodic return reaches CartPole’s maximum return of 500. This is a smoke test for end-to-end behavior rather than an aggregate benchmark or a claim of state-of-the-art performance.

Figure 1. Representative CartPole training run with MuZero. Average reward reaches the environment maximum of 500.

Pendulum with Sampled MuZero

Pendulum provides the complementary continuous-control case. The policy predicts parameters for a bounded continuous action distribution, and the sampled search considers a limited set of candidate actions at each node rather than enumerating an impossible action space.

The example mixes policy-proposed actions with uniformly sampled candidates at the root, uses prioritized replay, and logs a representative return improvement from approximately -1,250 to approximately -200. The configuration is deliberately more expensive than CartPole and is intended as a substantial experiment rather than a quick smoke test.

Figure 2. Representative Pendulum training run with Sampled MuZero. Average reward improves from approximately -1,250 to approximately -200.

Training artifacts and reproducibility

Each example writes structured artifacts rather than only printing progress:

• logs/<environment>/ stores JSONL metrics and TensorBoard events.
• checkpoints/<environment>/ stores model and optimizer checkpoints.
• <environment>_train_log.json stores the returned training history.

This matters for reinforcement-learning experiments. Search-heavy training is sensitive to configuration changes, so explicit artifacts make it easier to compare runs, inspect failures, and resume work. The examples run on CPU by default, with optional Apple Metal support where available.

Try it

The repository requires Python 3.10 or newer, PyTorch, Gymnasium, NumPy, and TensorBoard:

git clone https://github.com/LorenzzoQM/my_zero.git
cd my_zero
python -m pip install -e ".[dev]"
python -m pytest -q
python examples/cart_pole_train.py

For the continuous-action example, run:

python examples/pendulum_train.py

The full source, examples, and test suite are available at github.com/LorenzzoQM/my_zero. For a quick summary of the theory, see Zero to Reality: AlphaGo, MuZero and the Road to Real-World AI.