Distributional RL and Rainbow

This builds on Sharpening DQN: Improvements and the Distributional Idea, which recapped the five improvements that keep the scalar $Q$-value and then developed the sixth — distributional RL — through the distributional Bellman equation and the C51 algorithm. Recall the shift: the return $Z(s,a)$ is a random variable, $Q(s,a)=\mathbb{E}[Z(s,a)]$ is only its mean, and C51 learns the whole distribution by fixing a grid of return atoms $z_1<\cdots <z_N$ and learning a probability for each, snapping every Bellman-updated distribution back onto the grid with a projection $\Phi$.

This lesson picks up the complement to C51, folds all six improvements into one agent — Rainbow — and follows the distributional line to its modern endpoint.

Quantile methods: QR-DQN

C51 fixes a set of return values and learns how much probability sits on each. The equally natural dual is to fix the probabilities and learn where the return values sit. That swap is QR-DQN, and it removes the projection step.

QR-DQN: learn the quantiles instead

C51 fixes the return values and learns their probabilities. QR-DQN (Dabney et al., 2018) does the reverse: fix the probabilities and learn the return values.¹ It represents $Z(s,a)$ by $N$ quantiles — the network outputs $N$ values ${\theta }_1(s,a),\dots ,{\theta }_N(s,a)$, one for each of the fixed cumulative probabilities ${\tau }_i=\frac{2i-1}{2N}$ (the midpoints of $N$ equal-probability bins). The distribution is a mixture of point masses at the learned locations,

Swapping which axis is fixed removes the projection entirely: because the support is learned, the Bellman-updated atoms need not be snapped back onto a grid. In its place QR-DQN minimizes the quantile (Huber) regression loss, which trains each ${\theta }_i$ toward the ${\tau }_i$-quantile of the target return distribution. The trade is clean: C51 carries a fixed support and a projection but a simple cross-entropy; QR-DQN drops the support constraint and the projection but pays with an asymmetric quantile loss. Both learn a distribution; they differ only in which coordinate of the histogram they let the network move.

The asymmetry is what makes a regression target a quantile estimate. Ordinary squared error, minimized, returns the mean; the pinball loss at level $\tau$ weights over- and under-shoots differently,

and its minimizer is the $\tau$-quantile of the target. Here $u=\hat{T}{\theta }_j-{\theta }_i$ is the target minus the estimate, so $u>0$ means the estimate sat below the target (an under-estimate). If $u>0$ the loss is $\tau u$; if $u<0$ it is $(\tau -1)u=(1-\tau )|u|$. For $\tau =0.9$ an under-estimate ($u>0$) is penalized $9\times$ as hard as an over-estimate, so the estimate is dragged upward until only $10\%$ of the target mass remains above it — the definition of the $0.9$-quantile. QR-DQN replaces the kink at $u=0$ with a Huber softening (the quantile Huber loss) for a smoother gradient near zero. For example, to estimate the median ($\tau =0.5$) of target samples $\{1,2,6\}$, the pinball loss weights overshoot and undershoot equally, so it is minimized at the middle value $2$ — the sample median — whereas squared error would return the mean $3$, pulled up by the outlier. Fixing $\tau =0.5$ recovers the median; sweeping ${\tau }_i=(2i-1)/2N$ across the $N$ heads recovers the whole distribution.

Why the distribution helps even when you act on the mean

C51 and QR-DQN still select actions by $\arg {\max }_{a}Q(s,a)=\arg {\max }_a\mathbb{E}[Z(s,a)]$ — the extra information is thrown away at decision time. So why does modeling the full distribution improve the resulting policy? Three reasons the papers give.² First, the distributional target is a richer, denser learning signal: predicting fifty-one probabilities (or fifty-one quantiles) per action shapes the representation with far more constraints than a single scalar, and a better representation yields a better-estimated mean. Second, it interacts well with function approximation: a categorical target reframes the regression as a classification over atoms, whose cross-entropy loss is better-behaved for a neural network than chasing a moving scalar, and it keeps predictions bounded to $[V_{\min },V_{\max }]$. Third, learning the distribution appears to reduce the chattering and instability that plague approximate value iteration, giving more stable updates. The mean you act on is simply estimated better when it is estimated as one summary of a distribution the network was forced to get right.

Rainbow: all six at once

Each improvement was published on its own, measured against the DQN baseline. The natural question — do they add up, or do they overlap and cancel? — is what Rainbow (Hessel et al., 2018) answered by integrating all six into one agent and measuring the whole.³ The integration is mostly a matter of choosing compatible pieces:

• The value head is C51's distributional output, and the network uses the dueling decomposition — but applied to the atom logits: a shared state-value stream and a per-action advantage stream, combined before the softmax.
• The target is the multi-step ($n=3$) distributional target, run through Double-DQN action selection: the online network picks $a^{\ast }$, the target network supplies its distribution.
• Transitions are drawn by prioritized replay, with the priority set to the distributional loss (the KL / cross-entropy) rather than the absolute TD error, since there is no longer a scalar TD error to take the absolute value of.
• Exploration is NoisyNets throughout, so the $\epsilon$ schedule is dropped entirely.

Nothing here is a new algorithm; it is the same Q-learning skeleton — act, store, sample, form a bootstrapped target, regress toward it, periodically sync the target network — with a distributional head, a multi-step Double-DQN target, prioritized sampling, and learned exploration noise all slotted in at once. On the 57-game Atari benchmark Rainbow set a new state of the art in both final score and sample efficiency: it reached the previous best DQN's performance in a fraction of the frames.

The ablation

The paper's most cited figure is the ablation study: retrain Rainbow six times, each run with one component removed, and compare against the full agent. A component matters in proportion to how much removing it hurts. Some of the results were unexpected.³

Prioritized replay and multi-step returns were the two most important components: removing either produced the largest drop, and multi-step in particular sped up early learning across almost every game. The distributional head was next, its contribution growing over training and dominating in the later stages — evidence that the richer target matters most late in training. Dueling and NoisyNets helped moderately, each carrying a subset of games. Double DQN, on its own, barely mattered inside Rainbow — the distributional target already bounds the values and curbs the overestimation that Double DQN was invented to fix, so its separate correction becomes largely redundant. That last finding is the clearest illustration of why an integrated study was needed: a component's value is not a fixed property but depends on what else is in the agent.

Beyond Rainbow: implicit quantiles, full parametrization, and Agent57

Rainbow was not the end of the distributional line. Work after 2018 made the return distribution more expressive while keeping the same Q-learning skeleton, then folded the result into agents that finally cleared the whole Atari suite.

IQN (Dabney et al., 2018) turns QR-DQN's fixed quantile levels into a continuous sampling of them.⁴ QR-DQN learns $N$ heads at the fixed levels ${\tau }_i=(2i-1)/2N$. IQN instead learns a single network that, given a quantile level $\tau \sim U[0,1]$ sampled fresh, outputs the return at that level — an implicit quantile function $Z_{\tau }(s,a)$. Because $\tau$ is an input, the network represents the entire inverse CDF rather than a fixed histogram, and the action value is a Monte Carlo average $Q(s,a)\approx \frac{1}{K}{\sum }_k{Z}_{{\tau }_k}(s,a)$ over sampled levels. Two payoffs follow. First, expressiveness is no longer capped by a head count: more samples buy a finer distribution at test time without retraining. Second, the sampling of $\tau$ can be reshaped to encode risk sensitivity — the reason the distributional view matters even though the agents still act on the mean. Sample $\tau$ from the low end and the induced policy is risk-averse (it optimizes a conservative quantile, a CVaR-like objective); sample from the high end and it is risk-seeking. On the 57-game Atari benchmark IQN surpassed Rainbow's distributional component and closed much of the gap to a full Rainbow with far less machinery.

FQF (Yang et al., 2019) takes the last step: it also learns where to place the quantile levels.⁵ IQN samples $\tau$ uniformly; FQF adds a small fraction proposal network that outputs an adjusted set of $\tau$ values, trained to minimize the $1$-Wasserstein distance between the quantile approximation and the true distribution — so the atoms cluster where the distribution has structure and thin out where it is flat. It is the fully-parametrized quantile function: both the levels and the values are learned, generalizing C51 (fixed values, learned probabilities) and QR-DQN (fixed probabilities, learned values) at once.

The endpoint of the value-based line is Agent57 (Badia et al., 2020), the first agent to exceed the human baseline on all 57 Atari games — including the hard-exploration holdouts Montezuma's Revenge and Pitfall that had resisted every DQN-style method.⁶ It layers three things onto the distributional Q-learning core: a family of policies indexed by a discount and an exploration weight, so a single network holds both near-sighted exploiting policies and far-sighted exploring ones; a separate intrinsic-reward stream (the never-give-up novelty bonus) to drive the hard-exploration games; and a bandit that picks, per episode, which policy in the family to follow. The distributional head is still there — the ancestry runs straight from C51 through Rainbow and IQN — but the win on the last few games came from bolting on the directed exploration this course treats in the exploration lesson. Value-based deep RL, begun by DQN and consolidated by Rainbow, was completed when the distributional target was combined with a serious exploration mechanism.

Where this leaves value-based deep RL

Every improvement from DQN to Rainbow kept the Q-learning frame — act, store to a replay buffer, sample a mini-batch, build a bootstrapped target with a frozen target network, regress toward it — and changed exactly one slot: the target (Double DQN, multi-step), the network (dueling), the sampling (prioritized replay), the exploration (NoisyNets), or the thing being predicted (distributional). Because they touched different slots they composed, and Rainbow's ablation confirmed that the sum genuinely exceeds any part.

The deepest of the six, distributional RL, is the one that reframed the objective: learn the return distribution $Z(s,a)$ and recover $Q=\mathbb{E}[Z]$ as a by-product. That perspective — a value is a distribution, not a number — carries beyond discrete-action Atari. It reappears in the distributional critics of continuous-control agents and in risk-sensitive control, where an agent that knows the spread of returns, not just their mean, can prefer the safe action over the equal-mean gamble. Rainbow is the endpoint of the value-based line begun by DQN; the policy-based and model-based lines take up from actor-critic and PPO and model-based RL.