A Distributional Perspective on Reinforcement Learning (C51)

Prerequisites: Deep Q-Networks, DQN Improvements

The Problem: Averages Hide Information

Every value-based method you’ve seen so far—from tabular Q-learning to DQN—learns a single number per state-action pair: the expected return $Q(s, a) = \mathbb{E}[G_t]$ . But the return is a random variable. Averaging it throws information away.

Consider two actions with identical expected value:

Action A: Always returns exactly 10
Action B: Returns 0 or 20 with equal probability

Both have $Q(s, a) = 10$ , yet they describe very different situations. A network forced to predict only the mean must reconcile the conflicting outcomes of Action B into one scalar—and in deep RL, where the same network weights serve many states, that averaging can blur the learning signal itself.

Bellemare, Dabney, and Munos asked: what happens if we model the whole return distribution $Z(s, a)$ , with $Q(s, a) = \mathbb{E}[Z(s, a)]$ ? The answer became C51, and later a core component of Rainbow.

The Distributional Bellman Equation

The classic Bellman equation says the value of a state-action pair equals the immediate reward plus the discounted value of what comes next. The distributional version makes the same statement about distributions: the distribution of your return is the distribution of your immediate reward, shifted and scaled by the discounted distribution of future returns.

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

where $\stackrel{D}{=}$ means “equal in distribution.” Randomness now enters from three sources: the reward, the transition to $S'$ , and the future return $Z(S', A')$ itself.

∑Mathematical Details

Define the distributional Bellman operator $\mathcal{T}^\pi$ acting on return distributions. The paper’s central theoretical result concerns policy evaluation:

Two caveats the paper is careful about, and that are easy to miss:

Control is not covered. For the Bellman optimality operator, the distributional operator is not a contraction in any metric the paper considers; only the means are guaranteed to behave. C51’s strong control results are empirical, not a corollary of the theorem.
The theory and the algorithm use different metrics. The contraction holds in Wasserstein distance, but Wasserstein cannot be minimized directly with unbiased stochastic gradients from samples. C51 instead minimizes a KL divergence (cross-entropy) after projecting onto fixed atoms—a gap between theory and practice that later work (quantile regression methods such as QR-DQN) addressed.

The C51 Algorithm

To make distributions learnable by a network, C51 represents $Z(s, a)$ as a categorical distribution over 51 fixed atoms—evenly spaced return values between $V_{\min}$ and $V_{\max}$ . The network’s job for each action is to output a probability for each atom: “how likely is the return to be near this value?”

Training then looks a lot like classification:

Compute the target distribution $r + \gamma Z(s', a^*)$ using the target network
The shifted atoms $r + \gamma z_j$ generally fall between the fixed atom locations, so project each one onto its two nearest atoms, splitting probability mass proportionally
Minimize the cross-entropy between the projected target and the predicted distribution

Action selection stays simple: compute $Q(s, a) = \sum_i z_i \, p_i(s, a)$ and act greedily on those means (plus whatever exploration you use).

∑Mathematical Details

Parameterization. Fix $N = 51$ atoms:

$z_i = V_{\min} + i \cdot \frac{V_{\max} - V_{\min}}{N-1}, \quad i = 0, ..., N-1$

with $V_{\min} = -10$ , $V_{\max} = 10$ for Atari. Note this is a fixed design choice, not the true return range: with clipped per-step rewards and $\gamma = 0.99$ , discounted returns can reach roughly $\pm 100$ . Return mass that falls outside the support is projected onto the boundary atoms — the paper found this truncated support works well in practice. The network outputs logits $\ell(s, a) \in \mathbb{R}^{N}$ per action, converted to probabilities by a softmax:

$p_i(s, a) = \frac{e^{\ell_i(s, a)}}{\sum_j e^{\ell_j(s, a)}}$

Projection. For each target atom $\hat{z}_j = r + \gamma z_j$ :

Clip $\hat{z}_j$ to $[V_{\min}, V_{\max}]$
Find its position $b_j = (\hat{z}_j - V_{\min}) / \Delta z$ between the neighboring atoms $\lfloor b_j \rfloor$ and $\lceil b_j \rceil$
Split the probability $p_j(s', a^*)$ between those two atoms in proportion to proximity

Loss. With projected target probabilities $m_i$ , minimize the cross-entropy:

$L = -\sum_i m_i \log p_i(s, a)$

This is exactly the KL divergence between target and prediction (up to the target’s entropy), which is why C51’s loss behaves like a well-conditioned classification loss rather than a regression loss.

</>Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistributionalDQN(nn.Module):
    """
    C51 Distributional DQN.

    Outputs a probability distribution over returns for each action.
    """

    def __init__(self, state_dim: int, n_actions: int,
                 n_atoms: int = 51, v_min: float = -10.0, v_max: float = 10.0):
        super().__init__()

        self.n_actions = n_actions
        self.n_atoms = n_atoms
        self.v_min = v_min
        self.v_max = v_max

        # Support for the distribution
        self.register_buffer(
            'support',
            torch.linspace(v_min, v_max, n_atoms)
        )
        self.delta_z = (v_max - v_min) / (n_atoms - 1)

        # Network outputs n_actions * n_atoms values
        self.network = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, 128),
            nn.ReLU(),
            nn.Linear(128, n_actions * n_atoms)
        )

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """
        Returns log-probabilities over atoms for each action.

        Shape: [batch_size, n_actions, n_atoms]
        """
        batch_size = state.size(0)

        # Get raw outputs
        logits = self.network(state)

        # Reshape: [batch, n_actions * n_atoms] -> [batch, n_actions, n_atoms]
        logits = logits.view(batch_size, self.n_actions, self.n_atoms)

        # Apply log-softmax over atoms dimension
        log_probs = F.log_softmax(logits, dim=2)

        return log_probs

    def get_q_values(self, state: torch.Tensor) -> torch.Tensor:
        """Compute Q-values as expected value under distribution."""

        log_probs = self.forward(state)
        probs = log_probs.exp()  # [batch, n_actions, n_atoms]

        # Q(s,a) = sum_i p_i * z_i
        q_values = (probs * self.support).sum(dim=2)  # [batch, n_actions]

        return q_values


def project_distribution(next_dist: torch.Tensor,
                        rewards: torch.Tensor,
                        dones: torch.Tensor,
                        gamma: float,
                        support: torch.Tensor,
                        v_min: float,
                        v_max: float,
                        n_atoms: int) -> torch.Tensor:
    """
    Project the target distribution onto the fixed support.

    This handles the Bellman update: Z' = r + gamma * Z

    Args:
        next_dist: Probabilities over atoms for next state [batch, n_atoms]
        rewards: Rewards [batch]
        dones: Done flags [batch]
        gamma: Discount factor
        support: The fixed atoms [n_atoms]
        v_min, v_max: Support bounds
        n_atoms: Number of atoms

    Returns:
        Projected distribution [batch, n_atoms]
    """
    batch_size = rewards.size(0)
    delta_z = (v_max - v_min) / (n_atoms - 1)

    # Compute target support: T_z = r + gamma * z (clipped)
    # Shape: [batch, n_atoms]
    target_support = rewards.unsqueeze(1) + gamma * (1 - dones.unsqueeze(1)) * support.unsqueeze(0)
    target_support = target_support.clamp(v_min, v_max)

    # Compute the projection
    # b = (T_z - v_min) / delta_z gives the float index
    b = (target_support - v_min) / delta_z

    # Lower and upper atom indices
    lower = b.floor().long()
    upper = b.ceil().long()

    # Handle edge case where b is exactly an integer
    lower = lower.clamp(0, n_atoms - 1)
    upper = upper.clamp(0, n_atoms - 1)

    # Distribute probability proportionally
    # m_l = p * (u - b), m_u = p * (b - l)
    m = torch.zeros(batch_size, n_atoms, device=rewards.device)

    # Upper and lower proportions
    upper_prop = (b - lower.float())
    lower_prop = 1 - upper_prop

    # Add probability mass to lower atoms
    m.scatter_add_(1, lower, next_dist * lower_prop)
    # Add probability mass to upper atoms
    m.scatter_add_(1, upper, next_dist * upper_prop)

    return m


def compute_distributional_loss(online_net: DistributionalDQN,
                                target_net: DistributionalDQN,
                                batch: dict,
                                gamma: float) -> torch.Tensor:
    """
    Compute the distributional RL loss (cross-entropy).
    """
    states = batch['states']
    actions = batch['actions']
    rewards = batch['rewards']
    next_states = batch['next_states']
    dones = batch['dones']

    # Get current distribution (log probabilities)
    log_probs = online_net(states)  # [batch, n_actions, n_atoms]

    # Select distribution for taken action
    actions_expanded = actions.unsqueeze(1).unsqueeze(2).expand(-1, -1, online_net.n_atoms)
    current_log_probs = log_probs.gather(1, actions_expanded).squeeze(1)  # [batch, n_atoms]

    with torch.no_grad():
        # Select best next action by expected value
        next_q_values = target_net.get_q_values(next_states)
        best_actions = next_q_values.argmax(dim=1)

        # Get the target distribution for that action
        next_log_probs = target_net(next_states)
        best_actions_expanded = best_actions.unsqueeze(1).unsqueeze(2).expand(-1, -1, target_net.n_atoms)
        next_dist = next_log_probs.gather(1, best_actions_expanded).squeeze(1).exp()  # [batch, n_atoms]

        # Project target distribution
        target_dist = project_distribution(
            next_dist, rewards, dones, gamma,
            online_net.support, online_net.v_min, online_net.v_max, online_net.n_atoms
        )

    # Cross-entropy loss: -sum(target * log(current))
    loss = -(target_dist * current_log_probs).sum(dim=1).mean()

    return loss

Practical notes:

v_min and v_max are a design choice, and returns outside them get projected onto the boundary atoms. A little boundary pile-up is tolerable (C51’s own Atari setting of ±10 does not bracket the achievable ±100 discounted-return range), but if most of the mass saturates at a boundary the distribution stops being informative — in your own environment, start from an estimate of the typical discounted return range and widen if you see saturation.
The original C51 selects the greedy next action with the target network (as above). Rainbow combines C51 with Double DQN, selecting with the online network instead.
51 atoms is an empirical sweet spot from the paper’s atom-count sweep, not a magic number—hence the name.

Why It Works

The paper’s expected-value theory doesn’t fully explain C51’s gains—the means would converge under ordinary Q-learning too. The authors point to more practical mechanisms:

Richer learning signal. Each update matches a full distribution rather than a single scalar, giving the network many soft targets per transition.
Better-behaved gradients. Cross-entropy over a categorical output is a well-conditioned loss that deep learning tooling is highly optimized for, compared to regression on a moving scalar target.
Preserved multimodality. When outcomes genuinely diverge (a jump that either clears the gap or doesn’t), the model can represent both outcomes instead of splitting the difference, which reduces harmful averaging in the learned representation.

Which of these mechanisms matters most is still an active research question—the follow-up literature on distributional RL is partly an attempt to answer it.

Results

Limitations and Extensions

Fixed support is restrictive. You must choose $V_{\min}$ , $V_{\max}$ , and the atom count up front. QR-DQN (Dabney et al., 2018) flips the parameterization—fixed probabilities, learned atom locations—via quantile regression, and IQN (Implicit Quantile Networks) learns the full quantile function.
Theory-practice gap. The Wasserstein contraction result doesn’t cover the projected KL update; quantile-regression methods close part of this gap.
In Rainbow’s ablations, removing distributional learning barely affected early training but clearly hurt final performance—evidence that the distributional signal matters most once the easy gains are learned. See Rainbow for the full ablation picture.
The distributional view also opened the door to risk-sensitive RL: with the whole distribution in hand, you can optimize quantiles or conditional value-at-risk rather than the mean.

Key Takeaways

Model the distribution, act on the mean. C51 learns $Z(s, a)$ but still selects actions by expected value—the policy interface is unchanged from DQN.
The Bellman equation lifts to distributions, and for a fixed policy the distributional operator is a contraction in Wasserstein distance.
C51 is a classification-flavored algorithm: fixed atoms, softmax outputs, projected targets, cross-entropy loss.
The gains are empirical. The control setting has no contraction guarantee, yet C51 worked well enough to become a pillar of Rainbow.

Discussion Questions

Two actions have the same mean return but different variances. C51 can distinguish them—but its greedy policy still acts on the mean. When would you want the policy itself to be risk-sensitive, and how could you get that from $Z(s, a)$ ?
Why can’t the Wasserstein distance be minimized directly with sampled transitions and stochastic gradient descent? What does the KL projection sacrifice?
How would you set $V_{\min}$ and $V_{\max}$ for an environment without reward clipping? What failure mode appears if the range is too narrow? Too wide?
Rainbow’s ablation found distributional RL mattered mainly late in training. What might explain that timing?