Sharpening DQN: Improvements and the Distributional Idea

Deep Q-networks made value-based deep RL work: approximate $q_{\ast }$ with a neural network, train it by bootstrapped semi-gradient regression, and stabilize the deadly triad with a replay buffer and a frozen target network. In the five years after the 2015 Nature paper, a stream of focused papers each fixed one weakness of that baseline without disturbing the frame. Six of them proved to matter enough to keep, and one — distributional reinforcement learning — turned out to be less a patch than a change of objective: instead of predicting the average return, predict the whole distribution of returns and act on its mean.

This lesson takes five of the six improvements — the ones that keep the scalar $Q$-value and adjust the target, the network, the data, or the exploration — and then the sixth and deepest, distributional RL, which changes what the value head predicts. We stop after the C51 algorithm and its projection; the quantile methods, the Rainbow agent that folds all six together, and the modern distributional line continue in a companion lesson.¹ The through-line is that every fix changes one slot of the same Q-learning loop, so they compose.

Six ways to sharpen DQN

Think of the DQN agent as a machine with a few interchangeable slots: which data it learns from, what its network looks like, what target it regresses toward, and how it explores. Each of the six improvements swaps out one slot and leaves the rest alone. Three of the six were introduced alongside the baseline in the DQN lesson and are recapped here only as a checklist; the other three are new. Each changes one slot of the agent, and — this is the point Rainbow rests on — they change different slots, so they compose.

Double DQN (van Hasselt et al., 2016) decouples action selection from action evaluation so that noise which inflates an action's estimate is not both chosen and trusted; the target becomes $r+\gamma Q(s^{\prime },\arg {\max }_{a^{\prime }}Q(s^{\prime },a^{\prime };\mathbf{w});{\mathbf{w}}^{-})$.²Dueling networks (Wang et al., 2016) restructure the head into a scalar $V(s;\beta )$ and a per-action advantage $A(s,a;\alpha )$ recombined with the mean advantage subtracted, so the state value is trained by every action's update.³Prioritized experience replay (Schaul et al., 2016) draws transition $i$ with probability $P(i)\propto p_i^{\alpha }$, $p_i=|{\delta }_i|+ϵ$, and corrects the induced bias with importance-sampling weights $w_i=(N\cdot P(i){)}^{-\beta }$.⁴ Those three were covered in full earlier. The remaining three deserve a paragraph each.

Multi-step returns

One-step Q-learning bootstraps immediately: the target uses one real reward $r$ and then trusts the network for everything after. Waiting for $n$ real rewards before bootstrapping trades a little variance for much less bias early in training, exactly as in $n$-step TD. The $n$-step target replaces the single reward with a truncated $n$-step return:

Reward information propagates $n$ times faster along a trajectory, so credit reaches early states in far fewer updates. Rainbow uses $n=3$. In principle a multi-step target off-policy needs importance-sampling corrections; in practice, with short $n$ and a replay buffer of recent-ish data, the uncorrected version works well and is what Rainbow ships.⁵

For example, suppose a transition earns rewards $r_t=0$, $r_{t+1}=0$, $r_{t+2}=1$ before the agent bootstraps, with $\gamma =0.99$, and the current network badly misestimates the state at horizon $h$ as $Q(s_{t+h},\cdot )=5$ when the true continuation value is $2$. The one-step target is $0+0.99\cdot 5=4.95$, dominated by the wrong bootstrap: almost the entire target is the network's own error. The three-step target is $0+0.99\cdot 0+0.99^2\cdot 1+0.99^3\cdot 5=0.980+4.851=5.831$, and the true three-step value is $0.980+0.99^3\cdot 2=0.980+1.940=2.920$. The absolute bootstrap error has shrunk from $|0.99\cdot (5-2)|=2.97$ in the one-step target to $|0.99^3(5-2)|=2.91$ — modest here, but the deeper point is that the two real rewards now anchor the target, so the network's mistake is diluted by real reward signal. The cost is variance: a three-step return sums three stochastic rewards, so its spread across episodes is larger. Short $n$ (Rainbow's $3$) sits near the sweet spot for Atari — enough real reward to cut early bias, not so many steps that variance swamps it.

NoisyNets for exploration

$\epsilon$-greedy explores by flipping a coin at every state and, with probability $\epsilon$, ignoring everything the network has learned. It is state-independent and clumsy: on games needing a long, specific sequence of actions to reach the first reward — Montezuma's Revenge is the standing example — uniform random jitter almost never stumbles onto that sequence. NoisyNets (Fortunato et al., 2018) replace the $\epsilon$ schedule with learnable noise injected into the network's linear layers.⁶ A noisy linear layer replaces the ordinary weight $\mathbf{w}$ and bias $\mathbf{b}$ with

where $\mathbf{\mu }$ and $\mathbf{\sigma }$ are learned parameters and $\mathbf{\epsilon }$ is sampled noise resampled each forward pass. The network learns how much noise to add, and where: in states where exploration still pays it can keep $\mathbf{\sigma }$ large, and in states it has mastered it can drive $\mathbf{\sigma }$ toward zero, annealing its own exploration per state rather than on a global clock. The extra parameters are trained by the ordinary gradient — no new loss term — and at evaluation the noise is switched off.

For a layer with $p$ inputs and $q$ outputs, independent noise would need $pq$ samples of ${\epsilon }_w$ per forward pass — expensive for the large linear layers of a value network. Fortunato et al. use factorized Gaussian noise to cut that to $p+q$: sample one vector ${\epsilon }^{\text{in}}\in {\mathbb{R}}^p$ and one ${\epsilon }^{\text{out}}\in {\mathbb{R}}^q$, pass each through $f(x)=\mathrm{sgn}(x)\sqrt{|x|}$, and set the weight noise to the outer product ${\epsilon }_{w,jk}=f({\epsilon }_j^{\text{out}})f({\epsilon }_k^{\text{in}})$ with bias noise ${\epsilon }_{b,j}=f({\epsilon }_j^{\text{out}})$. The rank-one structure trades a little noise diversity for a large saving in sampling cost, and it is what Rainbow ships. The initialization matters: $\mathbf{\sigma }$ starts near ${\sigma }_0/\sqrt{p}$ (with ${\sigma }_0=0.5$), so early training is noticeably stochastic and $\mathbf{\sigma }$ shrinks only where the return gradient indicates exploration no longer pays.

The distributional idea

The other five improvements adjust the target, the network, the data, or the exploration; distributional RL changes what the value head predicts. It is the central idea of this lesson, so we state the shift precisely before any algorithm.

The return from a state-action pair,

is a random variable. Reward is stochastic, transitions are stochastic, and the policy may be stochastic, so the same $(s,a)$ produces a spread of returns across episodes. The action-value function throws that spread away and keeps only its mean:

Distributional RL keeps the whole random variable $Z(s,a)$ and learns its distribution. A state where an action is a coin flip between a big win and a big loss and a state where it reliably yields the average of the two can share the same $Q$-value while having very different return distributions; the distributional agent distinguishes them.⁷

The distributional Bellman equation

The ordinary Bellman equation relates expected values: $Q(s,a)=\mathbb{E}[r+\gamma Q(s^{\prime },a^{\prime })]$. The distributional version relates the random variables themselves, before any expectation is taken:

where $\stackrel{D}{=}$ means equal in distribution and $s^{\prime },a^{\prime }$ are drawn from the transition and the policy. Read the right-hand side as an operation on distributions: sampling a next state-action pair $(s^{\prime },a^{\prime })$ selects one of the distributions $Z(s^{\prime },a^{\prime })$; multiplying by $\gamma$ scales it toward zero; adding the reward $r$ shifts it along the return axis; and mixing over the stochastic $(r,s^{\prime },a^{\prime })$ superposes the results into a new distribution. That composite map is the distributional Bellman operator ${\mathcal{T}}^{\pi }$, and its fixed point is $Z^{\pi }$.⁷ Taking expectations of both sides collapses it back to the familiar scalar Bellman equation, which is why any solution to the distributional problem also solves the ordinary one — distributional RL loses nothing and carries strictly more information.

C51: a categorical distribution on fixed atoms

To learn $Z$ with a network, you must choose how to represent a distribution. C51 (Bellemare, Dabney & Munos, 2017) takes the most direct route: fix a finite grid of $N$ possible return values — the atoms — and let the network output a probability for each.⁷ The atoms are spaced evenly on $[V_{\min },V_{\max }]$,

and for each state-action pair the network emits a softmax over the atoms, giving probabilities $p_i(s,a)$ that sum to one. The approximate return distribution is the categorical

with ${\delta }_{z_i}$ a point mass at atom $z_i$. The mean $Q$ used to act greedily is just the atom values weighted by their probabilities. The 51 is the atom count the paper found best; $[V_{\min },V_{\max }]=[-10,10]$ was standard for Atari with reward clipping.

The projection step

Fixing the atoms creates one difficulty, and resolving it is the whole technical content of C51. Apply the distributional Bellman operator to a categorical distribution and the atoms move: each atom $z_j$ of the next-state distribution maps to $r+\gamma z_j$, and in general that value is not one of the atoms$\{z_1,\dots ,z_N\}$. The scaled-and-shifted distribution lives on a displaced grid that no longer lines up with the fixed support the network can represent.

The fix is a projection $\Phi$ that redistributes each shifted atom's probability mass onto the two nearest fixed atoms, splitting it in inverse proportion to distance (and clipping to $[V_{\min },V_{\max }]$ at the ends). Formally, the mass that the Bellman update places at $\hat{T}{z}_j=r+\gamma z_j$ is spread onto atom $z_i$ by the amount

where $[\cdot {]}_a^b$ clips to $[a,b]$ and $a^{\ast }$ is the greedy next action. The inner bracket is the linear how close is the shifted atom to $z_i$ weight — one when they coincide, zero a full $\Delta z$ away — so each shifted atom's mass lands entirely within the grid. The result is a valid categorical distribution on the fixed atoms, and it becomes the regression target.

C51 then minimizes the cross-entropy between the network's predicted categorical $Z_{\theta }(s,a)$ and the projected target $\Phi \hat{\mathcal{T}}{Z}_{{\theta }^{-}}$ — a classification-style loss over atoms, not a squared TD error. Everything else is DQN: replay buffer, target network ${\theta }^{-}$, $\epsilon$-greedy (or NoisyNet) behavior, and greedy action selection by the mean $Q$. On the Atari suite C51 beat every prior single-improvement variant, and it did so while still acting on the mean — the richer target simply produces a better-shaped mean.⁷

A projection worked end to end

For example, take a small grid to keep the arithmetic legible: $N=5$ atoms on $[V_{\min },V_{\max }]=[0,4]$, so $\Delta z=(4-0)/(5-1)=1$ and the atoms are $z=(0,1,2,3,4)$. Suppose the greedy next-state distribution puts its mass on just two atoms, $p(s^{\prime },a^{\ast })=(0,0,0.6,0.4,0)$ — probability $0.6$ at $z_3=2$ and $0.4$ at $z_4=3$. Let the reward be $r=0.5$ and the discount $\gamma =0.9$.

Each atom is scaled and shifted by $\hat{T}{z}_j=r+\gamma z_j$, then clipped to $[0,4]$:

Neither $2.3$ nor $3.2$ is a grid atom, so each must be split. The fractional index of a shifted value is $b=(\hat{T}{z}_j-V_{\min })/\Delta z$; here $b_3=2.3$ and $b_4=3.2$. For the first, $l=\lfloor 2.3\rfloor =2$ and $u=\lceil 2.3\rceil =3$ (atoms $z_2=2$ and $z_3=3$ in index terms — indices $2$ and $3$ in the zero-based atom array). The mass $0.6$ splits by distance: $(u-b)=0.7$ to the lower neighbor and $(b-l)=0.3$ to the upper. So atom index $2$ (value $2$) receives $0.6\cdot 0.7=0.42$ and atom index $3$ (value $3$) receives $0.6\cdot 0.3=0.18$. For the second, $b_4=3.2$ gives $l=3$, $u=4$, splitting $0.4$ into $0.4\cdot 0.8=0.32$ at atom index $3$ and $0.4\cdot 0.2=0.08$ at atom index $4$. Collecting per atom:

The masses sum to $1.00$, as they must, and this vector $m$ — a valid categorical on the fixed grid — is the cross-entropy target for the current pair. Its mean is $0\cdot 0+1\cdot 0+2\cdot 0.42+3\cdot 0.50+4\cdot 0.08=0.84+1.50+0.32=2.66$, which equals $r+\gamma \mathbb{E}[Z(s^{\prime },a^{\ast })]=0.5+0.9\cdot (2\cdot 0.6+3\cdot 0.4)=0.5+0.9\cdot 2.4=2.66$ exactly. The projection moves mass around but preserves the mean — the scalar Bellman backup is a by-product of the distributional one.

Where this leaves us

Two threads are now in hand. Five improvements — Double DQN, multi-step returns, dueling networks, prioritized replay, and NoisyNets — each sharpen one slot of the Q-learning loop without touching the others, which is what lets them stack. And distributional RL reframes the objective: predict the whole return distribution $Z(s,a)$, act on its mean $Q=\mathbb{E}[Z]$, and get a richer, better-shaped estimate for free. C51 realizes that idea with a fixed grid of atoms and a projection that snaps each Bellman-updated distribution back onto the grid.

C51 fixes the return values and learns their probabilities. The natural complement — fix the probabilities and learn the values — gives QR-DQN, which drops the projection entirely. That method, the reasons the distribution helps even when you act on the mean, and the Rainbow agent that assembles all six improvements at once (with its component ablation) continue in Distributional RL and Rainbow.