Categorical DQN (C51)

C51 predicts a full histogram over future returns instead of one average Q-value, using 51 fixed return buckets. You keep the spread and the bad tail instead of averaging them away.

C51 shifts each return atom by reward plus discount, then splits its probability across the two nearest fixed buckets

Why Does This Exist?

Standard Q-learning compresses the future into one number per action:

Q(s, a) = E\left[ \sum_{t=0}^{\infty} \gamma^t R_{t+1} \;\middle|\; S_0 = s, A_0 = a \right]

One number hides the shape. An action that always pays +10 and a coin flip paying +100 or -80 can share the same mean, yet you would never treat them the same when hardware or money is on the line. Stochastic environments also branch into distinct modes, and averaging modes can invent a middle outcome that never happens.

Bellemare, Dabney, and Munos introduced C51 in 2017 as the first distributional agent: learn the whole return distribution $Z(s, a)$ instead of only its mean. The richer target steadies deep-network features because every bucket supplies gradient, not just one scalar error.

This page covers the fixed-grid mechanism: atoms, the Bellman shift, the projection back onto the grid, and the cross-entropy loss. For the mean-only baseline see deep Q-networks. For movable quantile locations instead of a fixed grid see quantile regression DQN. For C51 bundled with five other upgrades see Rainbow DQN.

You only need one prerequisite: a Q-value is the expected discounted return from a state and action, and the greedy action is the one with the highest Q.

Think of It Like This

A forecast that shows the storm, not just the average rain

You plan an outdoor wedding and ask two forecasters about rainfall.

The first says the expected rainfall is 13 mm. Useless. Is that a gentle all-day drizzle or blue skies plus a one-in-five flash flood? One number masks the danger.

The second hands you a histogram: 0 mm with 70 percent chance, 5 mm with 10 percent, 50 mm with 20 percent. Now you can compute the mean (10.5 mm) and the disaster chance (20 percent), and you rent the tent.

Map it point for point: days are episodes, rainfall amounts are returns, histogram bars are C51 atoms, and renting the tent is picking the action with the safer distribution.

Where it stops: forecasts do not change the sky. In C51 the distribution drives the action, and the Bellman update then shifts and re-projects that same distribution to train earlier states.

How It Actually Works

Let $Z(s, a)$ be the random return and $Q(s, a) = E[Z(s, a)]$ . The distributional Bellman equation says the return equals the reward plus the discounted next return, in distribution:

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

where $A^*$ is the greedy next action picked by mean value. C51 stores each $Z(s, a)$ as a categorical distribution over a fixed grid.

Fixed atoms and probabilities

Pick bounds $V_{\min}$ and $V_{\max}$ and $N = 51$ evenly spaced atoms:

z_i = V_{\min} + i \Delta z, \quad \Delta z = \frac{V_{\max} - V_{\min}}{N - 1}

The network outputs logits $\psi_i(s, a)$ per action, and a softmax turns them into probabilities $p_i(s, a)$ that sum to 1. The reported Q-value is the dot product:

Q(s, a) = \sum_{i=0}^{N-1} z_i p_i(s, a)

Action selection still uses the mean. The distribution is there to train better and to expose risk.

Shift, clamp, and project

Given a transition $(s, a, r, s')$ , each atom of the next-state distribution moves:

\hat{T}_j = \mathrm{clip}(r + \gamma z_j, V_{\min}, V_{\max})

The clipped position rarely lands on the grid, so express it in grid units:

b_j = \frac{\hat{T}_j - V_{\min}}{\Delta z}, \quad l = \lfloor b_j \rfloor, \quad u = \lceil b_j \rceil

Then split the probability $p_j(s', a^*)$ across the two neighbours in proportion to closeness:

m_l \mathrel{+}= p_j (u - b_j), \qquad m_u \mathrel{+}= p_j (b_j - l)

If $b_j$ is exactly an integer, all of $p_j$ lands on that atom. The vector $m$ is the projected Bellman target on the same grid.

Worked example with 5 atoms

Use a small grid so every step fits on the page. Bounds $V_{\min} = -2.0$ and $V_{\max} = 2.0$ with $N = 5$ give $\Delta z = 1.0$ and atoms $[-2.0, -1.0, 0.0, 1.0, 2.0]$ . Next-state probabilities at the greedy action are $[0.10, 0.20, 0.40, 0.20, 0.10]$ , whose mean is exactly $0.00$ . Reward $R = 0.5$ and discount $\gamma = 0.9$ .

Shift each atom: $-1.3, -0.4, 0.5, 1.4, 2.3$ . The last one exceeds $V_{\max}$ , so clamp it to $2.0$ . In grid units that is $b = [0.70, 1.60, 2.50, 3.40, 4.00]$ .

Split each mass:

Atom 0 ( $p = 0.10$ , $b = 0.70$ ): $m_0$ gains $0.10 \times 0.30 = 0.03$ , $m_1$ gains $0.10 \times 0.70 = 0.07$ .
Atom 1 ( $p = 0.20$ , $b = 1.60$ ): $m_1$ gains $0.20 \times 0.40 = 0.08$ , $m_2$ gains $0.20 \times 0.60 = 0.12$ .
Atom 2 ( $p = 0.40$ , $b = 2.50$ ): $m_2$ and $m_3$ each gain $0.40 \times 0.50 = 0.20$ .
Atom 3 ( $p = 0.20$ , $b = 3.40$ ): $m_3$ gains $0.20 \times 0.60 = 0.12$ , $m_4$ gains $0.20 \times 0.40 = 0.08$ .
Atom 4 ( $p = 0.10$ , clamped to index 4): $m_4$ gains $0.10$ .

Totals: $m = [0.03, 0.15, 0.32, 0.32, 0.18]$ , which sums to $1.00$ . Its mean is $(-2)(0.03) + (-1)(0.15) + (0)(0.32) + (1)(0.32) + (2)(0.18) = 0.47$ . The unclamped math says $0.5 + 0.9 \times 0.0 = 0.50$ , so the $0.03$ gap is the clamped tail. The diagram at the top draws these exact bars.

Why the loss is cross-entropy

Target $m$ and prediction $p(s, a)$ share the grid, so training minimizes their KL divergence, which is the cross-entropy:

\mathcal{L} = - \sum_{i=0}^{N-1} m_i \log p_i(s, a)

Use a separate target network for $m$ , as in DQN. The projection itself carries no gradient.

Code

This snippet runs the projection from the worked example with only the standard library. It prints the target masses, their sum, and their mean.

1import math2from typing import List3
4
5def project_c51(6    p_next: List[float],7    reward: float,8    gamma: float,9    v_min: float,10    v_max: float,11) -> List[float]:12    """Project a next-state distribution onto the fixed C51 support."""13    n = len(p_next)14    dz = (v_max - v_min) / (n - 1)15    atoms = [v_min + i * dz for i in range(n)]16    m = [0.0 for _ in range(n)]17    for j in range(n):18        shifted = reward + gamma * atoms[j]19        clipped = min(max(shifted, v_min), v_max)20        b = (clipped - v_min) / dz21        lo = math.floor(b)22        hi = math.ceil(b)23        if lo == hi:24            m[lo] += p_next[j]25        else:26            m[lo] += p_next[j] * (hi - b)27            m[hi] += p_next[j] * (b - lo)28    return m29
30
31p_next = [0.10, 0.20, 0.40, 0.20, 0.10]32atoms = [-2.0, -1.0, 0.0, 1.0, 2.0]33m = project_c51(p_next, reward=0.5, gamma=0.9, v_min=-2.0, v_max=2.0)34print([round(v, 2) for v in m])35# -> [0.03, 0.15, 0.32, 0.32, 0.18]36print(round(sum(m), 6))37# -> 1.038mean = sum(mi * zi for mi, zi in zip(m, atoms))39print(round(mean, 2))40# -> 0.47

Watch Out For

A support grid that cannot hold your returns

C51 clips everything outside $[V_{\min}, V_{\max}]$ into the edge atoms. Set the range too narrow and the whole right tail piles onto the last atom, so high-reward plans look identical and the policy undervalues them. Set it far too wide and neighbouring atoms sit dozens of points apart, so small but decisive value gaps vanish inside one bin.

Symptom: projected masses spike at an edge atom, or the agent cannot distinguish actions that differ by a few points. Fix: bracket returns from logged episodes before training, normalize rewards, and re-check the bracket when the environment changes. If you cannot bound returns at all, that is the signal to use quantile regression instead of a fixed grid.

Letting gradients flow into the projection target

The target $m$ must be treated as a constant built by the target network. If your code keeps the computation graph attached through the shift and split, gradients flow into the target and the loss chases a moving goalpost.

Symptom: loss falls but the policy oscillates or the distribution collapses toward uniform. Fix: build $m$ under no-gradient handling and confirm the optimizer only updates the online logits. The Bellman target is data, not a second prediction head.

The Quick Version

C51 stores each action value as a categorical distribution over 51 fixed return atoms instead of one expected number.
Q-values are the mean of that distribution, so action selection stays familiar while training sees the full shape.
The Bellman update shifts atoms by reward plus discount, clamps them into bounds, and splits each probability across its two nearest grid neighbours.
With atoms $[-2, -1, 0, 1, 2]$ , reward 0.5, and discount 0.9, the worked target is $[0.03, 0.15, 0.32, 0.32, 0.18]$ with mean 0.47.
Training minimizes cross-entropy between the projected target and the online distribution, with the target detached from gradients.