Beyond Expected Value: Distributions, Models, and Search

The return is a random variable

A classical value function reports a single number: the expected return from a state. But the return is genuinely random — the same state can lead to a jackpot or a disaster depending on stochastic dynamics and policy choices. Distributional RL learns the entire distribution of returns Z(s,a) instead of just its mean Q(s,a) = E[Z], and replaces the ordinary Bellman equation with a distributional Bellman equation that updates the whole distribution under the random reward and discounted next-state distribution.

The practical question is how to represent a distribution. Two answers dominate: a categorical one fixes a grid of return values and learns probabilities on them (the C51 algorithm, trained by projecting the Bellman target back onto the grid and minimizing cross-entropy); a quantile one fixes the probabilities and learns the return values at each quantile (trained with the quantile-regression loss). The quantile view is usually preferred because it sidesteps the awkward projection step.

Z(s,a) \stackrel{\mathcal{D}}{=} R(s,a) + \gamma\, Z\bigl(S', A'\bigr)

The distributional Bellman equation: the return Z is a random variable, equal in distribution to the reward plus the discounted next-state return.

Why modeling the distribution helps

If you only ever act on the mean, why carry the whole distribution? Empirically, distributional agents simply learn better representations: matching a full distribution is a richer, more demanding prediction target, so the network is forced to encode more about the dynamics. Beyond that, the distribution unlocks risk-sensitive control (optimize a conditional value-at-risk instead of the mean) and better exploration (the spread of Z is a signal of uncertainty about a state-action).

Learning a model and planning in imagination

Model-based RL learns the environment's dynamics — predict the next state and reward — and then exploits that model to generate extra experience without touching the real environment. The promise is sample efficiency: one expensive real step can seed many cheap imagined ones. The peril is model error: a learned model is wrong, and a long rollout compounds those errors until the agent is optimizing against a fantasy.

\hat{s}_{t+1},\, \hat{r}_t \sim p_\theta\!\left(s_{t+1}, r_t \mid s_t, a_t\right)

A learned dynamics model predicts the next state and reward, letting the agent generate imagined experience without touching the real environment.

Model-based policy optimization (MBPO) is the clean resolution. Instead of trusting long imagined trajectories, it generates short, branched rollouts: start from real states drawn from the replay buffer and roll the model forward only a handful of steps, then train an off-policy learner (like SAC) on the mix of real and imagined data. The supporting theory bounds the policy's true value gap in terms of the model error and the rollout length — making explicit why keeping rollouts short keeps the bias controlled.

Fit an ensemble of one-step dynamics models to the real-experience buffer (the ensemble flags where the model is unsure).
Sample real start states; roll the model forward k short steps to synthesize transitions.
Train the policy and critic on real + imagined data with an off-policy algorithm.
Keep k small early and grow it only as the model proves reliable.

Search as planning — MCTS

When you have a model — or the true rules of a game — you can plan at decision time by looking ahead. Monte Carlo tree search (MCTS) builds a search tree incrementally through four repeated phases: select a promising path using an upper-confidence rule that balances a node's estimated value against how rarely it has been visited; expand a new leaf; evaluate it (by a rollout or a learned value network); and back up the result to update value estimates along the path. Over many iterations the tree concentrates effort on the most promising lines.

# UCB-for-trees (UCT) selection score at a node
score(a) = Q(s,a) + c * sqrt( ln N(s) / N(s,a) )
#          exploit value      explore rarely-tried actions

The selection rule that makes MCTS focus without going blind.

This is the engine behind AlphaGo: a learned policy network proposes moves, a value network scores positions, MCTS stitches them into deep lookahead, and self-play generates the training data. Its successor MuZero went further — it learns its own latent dynamics model and runs MCTS entirely inside that learned model, dissolving the line between the model-based and search families in this guide.