Classical TD3 trains a critic that predicts a single number: the expected return Q(s, a). Distributional QR-TD3 replaces that scalar with a whole distribution of the return, approximated by a fixed set of quantiles — and then goes one step further, estimating that distribution at several bootstrap horizons at once.
The result is a critic organised around two axes that never interact. They are orthogonal — they simply multiply.
Two orthogonal axes
A critic in QR-TD3 is parameterised by two independent settings:
- num_quantiles — how finely the return distribution is described inside a single estimator. With num_quantiles = 32, the critic emits 32 quantiles instead of one mean.
- nstep_horizons — how many different n-step bootstrap targets the same optimal Q* is estimated against. With nstep_horizons = [4, 8, 12, 16] there are 4 horizon heads.
Quantile regression: distribution, not mean
Quantile regression (QR) makes each estimator distributional. Instead of predicting E[return], the head predicts the value of the return at 32 evenly spaced probability levels — the τ-fractions.
The quantile positions are fixed and shared:
τ_i = (i + 0.5) / 32, i = 0 … 31
n-step horizons: the same Q*, different bias / variance
Each of the 4 horizon heads estimates the same optimal Q* — but bootstraps over a different number of steps before falling back on the critic's own prediction.
Short and long horizons trade bias against variance:
- A 4-step head reacts quickly and has low variance, but inherits more bootstrap bias from the critic.
- A 16-step head is closer to the true Monte-Carlo return — less biased — but noisier.
- Aggregating across all four gives the actor a bias / variance spectrum of the same target, which stabilises the policy gradient.
The critic output tensor
Because the axes multiply, one critic per symbol now produces a 4-dimensional tensor of shape (B, N, H, nq) = (B, N, 4, 32): for every item in the batch and every symbol, 4 return distributions (one per horizon), each described by 32 quantiles.
That is 4 × 32 = 128 values per symbol per critic, where a plain distributional critic stored 32 and a scalar critic stored 1.
sa → shared trunk → [ head n=4 | n=8 | n=12 | n=16 ] → (B, N, 4, 32)
How it looks in code
Implementation is a shared trunk plus one lightweight head per horizon. The trunk (input projection + residual blocks) does the heavy lifting; each head is a single linear layer mapping the hidden state to num_quantiles outputs.
head_norms = [LayerNorm(hidden) × 4] heads = [Linear(hidden, 32) × 4]
Where the axes meet — the loss
Training is the only place the two axes come together, and even there they stay separated. The quantile Huber loss is computed head-by-head; the quantile axis is identical for every head.
- pred[:, h] is (B, N, 32) — the prediction of head h.
- backup[:, h] is the n_h-step target for that head's horizon.
- The τ-grid (i + 0.5) / 32 is the same for all 4 heads.
- twin-min, the q_crossing penalty and CVaR all operate along the 32-quantile axis — separately within each horizon.
for h in range(H):
loss_h = quantile_huber(pred[:, h], backup[:, h])The equal-nq invariant
Every head must carry the same num_quantiles — this is enforced by an assert. The reason is the per-quantile twin-min: TD3 takes the minimum of two critics to fight overestimation, and that minimum is taken quantile-by-quantile. If two heads disagreed on the number of quantiles, the quantile axes would not line up.
Summary
nstep_horizons sets how many heads there are and which n-step target each one learns — the count of distinct estimates of the same Q*. num_quantiles sets how finely each of those heads describes its return distribution.
One axis is "how many different views of the target", the other is "how detailed each view is". They multiply into a (B, N, 4, 32) critic — and never interfere.