MuZero

The representation network encodes the raw observation (e.g., a game board image) into a hidden state s_t. This initial state serves as the starting point for all internal planning. Its key function is to compress the high-dimensional observation into a latent space that is both informative for prediction and stable for the dynamics model to process over many simulated steps. It is trained jointly with the other networks via backpropagation through the unrolled planning path.

The dynamics network is the learned internal model. It takes a hidden state s_{k-1} and a hypothetical action a_k and predicts the next hidden state s_k and an immediate reward r_k. Crucially, it operates entirely in the latent space, not in the pixel space. This allows MuZero to plan by repeatedly applying this network to simulate potential future trajectories, evaluating sequences of actions without interacting with the real environment. Its predictions are trained to match the targets generated by the representation and prediction networks on real data.

The prediction network takes a hidden state s_k and outputs two critical policy components:

• Policy (p_k): A probability distribution over possible actions from that state.
• Value (v_k): The expected future return (discounted sum of rewards) from that state. This network provides the search guidance for Monte Carlo Tree Search (MCTS) during planning by estimating which actions are promising and what the long-term value of a state is.

MuZero uses an MCTS variant as its planning algorithm. For each real environment step, it runs multiple simulations using its internal model (dynamics network) and predictions. The search builds a tree of possible futures by:

• Selecting actions using the Upper Confidence Bound for Trees (UCT) formula, balancing exploration and exploitation based on the prediction network's outputs.
• Expanding nodes using the dynamics and prediction networks.
• Backpropagating the estimated value up the tree. The final action is selected proportional to the visit counts at the root, resulting in a refined, search-based policy π_t that is stronger than the raw network policy p_t.

A FIFO replay buffer stores sequences of real experience: observations, MCTS search policies (π_t), actions taken, and observed rewards. Training is performed by uniformly sampling trajectories from this buffer. This is essential for:

• Breaking temporal correlations between consecutive states.
• Reusing data for greater sample efficiency.
• Stabilizing training by learning from a more stationary distribution of past experiences. The buffer is a core component that enables the off-policy training of all three networks.

All three networks are trained jointly to minimize a single, combined loss function on sampled trajectories from the replay buffer. For each step k in an unrolled sequence, the loss has three components:

• Policy Loss: Cross-entropy between the MCTS search policy π_{t+k} and the prediction network's policy p_k.
• Value Loss: Mean squared error between the observed return z_{t+k} (e.g., from n-step bootstrapping) and the predicted value v_k.
• Reward Loss: Mean squared error between the observed reward u_{t+k} and the reward predicted by the dynamics network r_k. The gradients flow back through the dynamics and representation networks, teaching the entire system a consistent latent model.

MuZero's architecture integrates several key mechanisms:

• Representation Network: Encodes the observation (e.g., game board) into a hidden state.
• Dynamics Network: Predicts the next hidden state and immediate reward given the current state and a planned action.
• Prediction Network: From a hidden state, outputs the policy (action probabilities) and value (expected return).
• Monte Carlo Tree Search (MCTS): Uses the learned model to simulate thousands of possible futures, selecting actions by balancing exploration and exploitation.

MuZero's experience replay buffer is critical for sample-efficient learning. It stores sequences of real interactions (trajectories) containing observations, actions, and rewards. During training, it samples these trajectories to update all networks. This breaks temporal correlations, reuses data, and is a form of off-policy learning, allowing the agent to learn from both recent and past experiences.

MuZero differs from earlier model-based RL like AlphaZero and Dreamer:

• vs. AlphaZero: AlphaZero requires a perfect, known simulator (game rules) for its MCTS. MuZero learns its own simulator from scratch, making it applicable to unknown environments.
• vs. Dreamer: Dreamer also learns a world model but uses it for latent imagination to train an actor-critic. MuZero uses its model exclusively for planning (MCTS) and learns value/policy functions directly from real and planned trajectories.

A paradigm where an agent learns an internal model of the environment's dynamics (transition function) and reward function. This model is used for planning by simulating future trajectories to select optimal actions, rather than relying solely on direct experience. MuZero is a model-based algorithm, but it learns a model that is implicit and latent, optimized solely for accurate value and policy predictions, not for reconstructing the true environment state.

A neural network that predicts future states and rewards in a compressed latent space. It serves as the agent's internal simulation of the environment. In MuZero, the learned dynamics function acts as a world model, predicting the next latent state and immediate reward given a current latent state and action. This allows the agent to imagine or dream rollouts for planning without interacting with the real environment.

A heuristic search algorithm used for decision-making in large spaces. MuZero uses MCTS as its planning algorithm, guided by its learned model. The process involves:

• Selection: Traversing the tree using a policy (e.g., PUCT) that balances exploration and exploitation.
• Expansion: Adding new nodes for unexplored actions.
• Simulation (Rollout): Using the model to predict outcomes (in MuZero, this is a truncated rollout using the dynamics and prediction networks).
• Backpropagation: Updating node statistics (visit counts, value estimates) back up the tree. The action with the highest visit count from the root is then executed.

A compressed, abstract encoding of the environment state learned by the representation function. MuZero does not predict the pixel-level observation. Instead, it learns a latent representation that contains only the information necessary for accurate planning and value prediction. This representation is the input to the dynamics and prediction networks, enabling efficient planning in a lower-dimensional, semantically meaningful space.

Temporal-Difference learning methods that unify Monte Carlo (full returns) and one-step TD updates. MuZero uses N-step bootstrapped targets for training. The value target for a state is computed by unrolling the model K steps into the future, summing the predicted rewards, and adding the predicted value at the final step. This multi-step return reduces bias compared to one-step TD and variance compared to full Monte Carlo returns, leading to more stable and efficient learning.

A key distinction in how agents use experience:

• On-Policy: Learns from experience generated by the current policy being improved (e.g., A3C).
• Off-Policy: Can learn from experience generated by a different behavior policy, often stored in a replay buffer (e.g., DQN). MuZero is off-policy. It uses a replay buffer to store trajectories generated by past versions of its policy (guided by MCTS). This allows for sample-efficient learning by reusing old data and breaking temporal correlations.