Distributional RL and C51

Distributional RL learns the full distribution of returns, not just its mean, capturing risk and multi-modality.

Modeling the whole return

Classic value-based RL learns the expected return Q(s,a). Distributional RL instead learns the entire distribution Z(s,a) of the random return, of which Q is the mean. Two states can share the same expected value but differ sharply in risk; the distribution captures that, and learning it often stabilizes training.

The distributional Bellman equation

The Bellman operator generalizes to distributions: the return distribution at (s,a) equals the reward plus a discounted, shifted-and-scaled version of the next state's return distribution. Convergence is analyzed in a metric over distributions (the Wasserstein distance), which the exact operator contracts.

C51: categorical distributions

C51 (the categorical algorithm, so named for its 51 atoms) represents Z as a discrete distribution over a fixed grid of return values, learning a probability for each atom. Applying the Bellman update shifts the support by reward and discount, which generally lands between atoms, so the result is projected back onto the fixed grid before a cross-entropy loss is applied.

# C51 essentials
# atoms z_i evenly spaced in [V_min, V_max]
# network outputs p_i(s,a) = softmax logits over atoms
# Bellman: shift atoms by r + gamma*z_i, clip, project onto grid,
# then minimize cross-entropy( projected_target , p(s,a) )

Why it helps

Beyond enabling risk-sensitive policies, learning the distribution provides a richer, auxiliary-like training signal that empirically improves sample efficiency and final performance, which is why C51 became a component of Rainbow DQN. Its main limitation is the fixed support: values outside [V_min, V_max] cannot be represented, motivating the quantile-based successor QR-DQN which learns the support instead of fixing it.