Introduction to Modern Reinforcement Learning

AlphaZero plans by using the real rules of a game. It can ask, "If I make this move, what board position comes next?" because the simulator is given. MuZero removes that assumption. It learns the internal model it uses for planning, then combines that learned model with search.

The surprising part is what MuZero does not learn. It does not need to reconstruct the full observation, draw every future board, or explain the environment in human terms. It learns predictions that are useful for action selection: rewards, values, and policy priors in a latent state space. This is a practical lesson for world models generally. A model can be valuable because it supports decisions, even if it is not a faithful simulator of every detail.

MuZero is usually described with three learned functions:

$s_0=h_\theta(o_0),$

$s_{k+1}, r_k=g_\theta(s_k,a_k),$

$p_k, v_k=f_\theta(s_k)$

The representation function $h_\theta$ encodes the current observation into an internal state $s_0$ . The dynamics function $g_\theta$ takes a latent state and action, then predicts the next latent state and reward. The prediction function $f_\theta$ outputs a policy prior $p_k$ and value estimate $v_k$ for search. Monte Carlo tree search uses these predictions to allocate visits, compare branches, and produce an improved action target.

Equation walkthrough

MuZero learns three decision-useful functions

Click each function to see how MuZero plans without knowing the environment rules.

Representation

observation to hidden state

The representation function h turns the real observation into the first hidden state for search. This is the bridge from the environment into MuZero's internal planning space.

Subcomponents

o_0

Current observation

The real input from the environment, such as a board position, Atari frame stack, or sensor reading.

h

Representation function

A neural network that compresses the observation into a latent state useful for planning.

s_0

Initial hidden state

The starting latent state for MCTS. It does not need to reconstruct the full observation; it needs to support good decisions.

MCTS operates in latent states: use g to imagine, use f to evaluate, and train h, g, and f end to end.

At acting time, MuZero performs MCTS at each timestep. An action is sampled from the search policy, the environment returns a new observation and reward, and the episode trajectory is stored in replay. At training time, a trajectory is sampled from replay, the representation function encodes past observations into $s_0$ , and the dynamics function is unrolled for $K$ steps using the real actions from that trajectory.

The model is trained end to end through time to predict policy targets from search, observed rewards, and value targets such as final rewards or n-step returns. Search improves the policy target; learning improves the model that future search will use.

MuZero process

Planning, acting, and training

Choose a stage to see how MuZero uses its learned model, real environment feedback, and MCTS targets.

Tree search over imagined futures

MuZero plans in latent space. It starts from observations, unrolls candidate actions with learned dynamics, then evaluates each imagined state with policy and value predictions.

Step 1Represent history

$s_0=h(o_1,\ldots,o_t)$

Past observations are encoded into the initial hidden state for search.

Step 2Unroll an action

$s_k,r_k=g(s_{k-1},a_k)$

Given a previous hidden state and candidate action, the dynamics function predicts a new hidden state and immediate reward.

Step 3Evaluate the node

$p_k,v_k=f(s_k)$

The prediction function computes a policy prior and value estimate from the hidden state.

Takeaway

MCTS combines imagined rewards, values, and policy priors to compare action branches before any real action is taken.

MuZero learns only what matters for decision-making: latent dynamics, expected rewards, policy improvements, and long-horizon values.

MuZero planning, acting, and training process — a, MuZero plans by combining representation, dynamics, and prediction functions in latent space. b, At acting time, MCTS produces a search policy, the agent samples an action, and the environment returns the next observation and reward. c, During training, replayed trajectories are unrolled through the model so value, reward, and policy predictions can be improved end to end. [Source]

ℹ

MuZero vs Dreamer

Both use learned latent dynamics. Dreamer mainly trains actor and critic through imagined rollouts. MuZero mainly uses the learned dynamics inside search, producing stronger planning targets from each real state. The shared idea is latent prediction; the difference is how the imagined future is used.