Quantile Regression DQN (QR-DQN)

QR-DQN learns a return distribution by fixing evenly spaced quantile levels and learning where each one sits on the return axis, so there is no fixed grid to clip tails and no projection step to smear them.

QR-DQN fixes quantile fractions at one sixth, one half and five sixths and learns return locations, turning reward 1.0 into targets 2.8, 4.6 and 6.4

Why Does This Exist?

Categorical DQN showed that learning the full return distribution beats learning one average. But it nails the distribution to a fixed grid: 51 buckets between a chosen minimum and maximum. Returns outside the range get clipped, and every Bellman update lands between buckets, so probability mass gets smeared across neighbors by a projection heuristic. The grid is a cage, and the projection is noise.

QR-DQN (Dabney and coauthors, 2018) transposes the design. Fix the probabilities at evenly spaced quantile fractions $\hat{\tau}_i = (i - 0.5) / N$ , and learn where each quantile sits on the return axis: $\theta_1(s, a)$ through $\theta_N(s, a)$ . No bounds to choose, nothing to clip, no projection. The expected value is just the mean of the locations, $Q(s, a) = \frac{1}{N} \sum_i \theta_i(s, a)$ , and the Atari version used 200 quantiles with Huber threshold $\kappa = 1$ .

This page covers the transpose, the loss, and one fully traced 3-quantile update. For the fixed-grid baseline it replaces, see categorical DQN. For the mean-only starting point, see deep Q-networks.

Think of It Like This

Cutting a loaf into equal-weight slices

You must describe an oddly shaped loaf, dense in the middle and airy at the crust.

The old way bolts 51 fixed baskets to the counter and crumbles each slice across the two nearest baskets, throwing away any overhang past the last basket. Slice positions are fixed, weights are learned, and the edges get mangled.

The new way cuts the loaf into slices of exactly equal weight and measures where the cuts land. Dense regions get tightly packed cuts, airy regions get spread-out cuts, the knife roams wherever it must, and nothing is crumbled or thrown away. Fixed weights, learned positions.

Where it stops: bread is positive and finite, while returns can be negative and unbounded. That is precisely why the knife needs the whole real line, and why the fixed-basket version clips exactly where the danger lives.

How It Actually Works

Write $Z(s, a)$ for the random return and $F_Z^{-1}(\tau)$ for its quantile function: the return value below which fraction $\tau$ of outcomes fall. QR-DQN learns $N$ locations $\theta_i(s, a)$ sitting at the fixed midpoints $\hat{\tau}_i = (i - 0.5)/N$ . In statistics, the $\tau$ -th quantile is whatever minimizes the pinball loss $\rho_\tau(u) = u(\tau - \mathbb{I}(u < 0))$ , where $u$ is target minus prediction and $\mathbb{I}$ is one when $u$ is negative. Underestimate and you pay $\tau$ per unit; overestimate and you pay $1 - \tau$ .

The pinball loss has a kink at zero that rattles neural gradients, so the absolute error is swapped for the Huber loss with threshold $\kappa$ : quadratic within $\kappa$ , linear beyond it. The quantile Huber loss is therefore:

\rho_\tau^\kappa(u) = \lvert \tau - \mathbb{I}(u < 0) \rvert \, \mathcal{L}_\kappa(u)

Each transition builds Bellman targets $\mathcal{T}\theta_j = R + \gamma \, \theta_j(S', a^*)$ from the greedy next action, then scores every prediction against every target, all $N \times N$ pairs:

\mathcal{L}(\theta) = \frac{1}{N^2} \sum_{i=1}^{N} \sum_{j=1}^{N} \rho_{\hat{\tau}_i}^\kappa(\mathcal{T}\theta_j - \theta_i(S, A))

Because quantile regression minimizes the 1-Wasserstein distance directly, the contraction guarantees of dynamic programming survive. The categorical projection has no such promise.

Worked example: a 3 by 3 loss grid by hand

Take $N = 3$ with $\kappa = 1.0$ . The fixed fractions are $\hat{\tau} = [1/6, 1/2, 5/6]$ . Current predictions are $\theta = [1.0, 3.0, 5.0]$ , next-state quantiles are $[2.0, 4.0, 6.0]$ , reward $R = 1.0$ , discount $\gamma = 0.9$ . Targets:

\mathcal{T}\theta = 1.0 + 0.9 \times [2.0, 4.0, 6.0] = [2.8, 4.6, 6.4]

Now three representative pairs. First, prediction $1.0$ against target $2.8$ at $\tau = 1/6$ : error $u = +1.8$ exceeds $\kappa$ , so Huber gives $1.8 - 0.5 = 1.3000$ , weighted by $1/6$ for a loss of $0.2167$ . Second, prediction $3.0$ against target $2.8$ at $\tau = 1/2$ : error $u = -0.2$ sits inside $\kappa$ , so Huber gives $0.5 \times 0.04 = 0.0200$ , weighted by $1/2$ for $0.0100$ . Third, prediction $5.0$ against target $2.8$ at $\tau = 5/6$ : error $u = -2.2$ exceeds $\kappa$ , so Huber gives $2.2 - 0.5 = 1.7000$ , weighted by $1/6$ for $0.2833$ .

The full pairwise grid is:

\begin{bmatrix} 0.2167 & 0.5167 & 0.8167 \\ 0.0100 & 0.5500 & 1.4500 \\ 0.2833 & 0.0133 & 0.7500 \end{bmatrix}

Rows are predictions $[1.0, 3.0, 5.0]$ , columns are targets $[2.8, 4.6, 6.4]$ . The nine entries sum to $4.6067$ , so the mean loss is $0.5119$ . The diagram above plots these same fractions, targets, and the $0.5119$ result.

Reading the grid

Each row answers one question: how should this quantile location move given everything the target distribution says? The diagonal stays small because each prediction sits near its own target, while the largest entry, row 2 column 3 at $1.4500$ , is the median prediction facing the top target. For risk-sensitive control that reads the lower quantiles, see risk-sensitive RL.

Code

This snippet rebuilds the grid with plain lists, no installs needed:

1def huber(u: float, kappa: float = 1.0) -> float:2    au = abs(u)3    return 0.5 * u * u if au <= kappa else kappa * (au - 0.5 * kappa)4
5
6def qr_loss(pred: list[float], tgt: list[float], taus: list[float], kappa: float = 1.0) -> tuple[list[list[float]], float]:7    grid = []8    for i, p in enumerate(pred):9        row = []10        for t in tgt:11            u = t - p12            w = abs(taus[i] - (1.0 if u < 0 else 0.0))13            row.append(w * huber(u, kappa))14        grid.append(row)15    mean_loss = sum(sum(row) for row in grid) / (len(pred) * len(tgt))16    return grid, mean_loss17
18
19taus = [(i - 0.5) / 3 for i in (1, 2, 3)]20grid, mean_loss = qr_loss([1.0, 3.0, 5.0], [2.8, 4.6, 6.4], taus)21
22print([round(v, 4) for v in grid[0]])  # -> [0.2167, 0.5167, 0.8167]23print([round(v, 4) for v in grid[1]])  # -> [0.01, 0.55, 1.45]24print([round(v, 4) for v in grid[2]])  # -> [0.2833, 0.0133, 0.75]25print(round(mean_loss, 4))             # -> 0.5119

The prints match the hand-built grid entry for entry, including the $0.5119$ mean.

Watch Out For

Trusting unsorted quantiles for risk decisions

Nothing in the network forces $\theta_1 \leq \theta_2 \leq \dots \leq \theta_N$ . Early in training, or on strange states, the heads cross and the lower quantile reads above the median. If your policy avoids actions by their worst decile, crossed quantiles corrupt the decision silently. Always sort a copy before computing risk metrics: sorted_theta = sorted(theta). The loss tolerates crossing during training; your deployment must not.

Setting the Huber threshold by gut feel

Push $\kappa$ toward zero and the loss reverts to the kinked pinball form, so gradients jitter around every target. Push it far too high and the loss turns into symmetric squared error, erasing the asymmetry that pins each head to its own quantile. The symptom in both cases is quantiles that bunch together instead of spreading across the distribution. Keep $\kappa = 1.0$ unless your reward scale demands otherwise, and scale it with your rewards if you change that scale.

The Quick Version

QR-DQN fixes quantile fractions at $(i - 0.5)/N$ and learns return locations, transposing the fixed C51 grid.
No return bounds to tune, no clipped tails, no projection smear between buckets.
The quantile Huber loss pairs pinball asymmetry with Huber smoothness for stable neural gradients.
One traced update gives targets $[2.8, 4.6, 6.4]$ and mean pairwise loss $0.5119$ .
Minimizing the 1-Wasserstein distance keeps the Bellman contraction guarantees that projection breaks.