Auditing the Risk Claims of Distributional Reinforcement Learning

TL;DR: Distributional RL agents (QR-DQN, C51, IQN) learn full return distributions, and people increasingly read those distributions as meaningful — visualizing them for interpretability, thresholding them for risk-sensitive control via CVaR, monitoring them as safety signals. This paper builds a statistically careful audit that checks whether the risk trade-offs a trained agent claims actually exist in the environment, using snapshot-restart Monte Carlo as ground truth. The answer is brutal: 40–95% of the strongest claimed trade-offs are provably false, essentially none is confirmable, and where the agent places its most confident risk claims is statistically indistinguishable from random. The distributional shape is a training artifact — seed-idiosyncratic, formed early, uncorrelated with performance — and it survives at full Atari scale. The positive controls (environments with a known risk answer key, where the same audit confirms 96–100% of real claims) are what make this hard to dismiss.

The assumption being tested

Distributional RL replaces Q(s,a) = \mathbb{E}[Z(s,a)] — the expected discounted return — with the full distribution of Z(s,a). C51 learns probability mass on a fixed grid of return values; QR-DQN learns the locations of quantile atoms; IQN samples quantile levels implicitly. The original motivation was performance, and it delivered: distributional heads drove state-of-the-art Atari results.

But somewhere along the way, the distribution stopped being an implementation detail. If the head says action A yields a tight, safe return distribution while action B yields a wide gamble of equal mean, that reads like actionable risk information. Papers use it exactly that way: pick actions by CVaR instead of the mean, flag risky states in autonomous driving, plot the distributions to explain what the agent "believes."

Theory has quietly warned this might be unfounded. Quantile TD learning has non-unique fixed points that need not equal the true return distribution (Rowland et al. 2023), and the performance benefits of distributional methods show up even when only the mean is used — suggesting the distributions can help without being correct. Clements et al. (2019) argued architecturally that quantile spread conflates aleatoric and epistemic uncertainty. What nobody had done is measure it: take a trained agent, take the specific states where its distribution makes a decision-relevant claim, and check that claim against ground truth. That's the gap this paper fills, and the measurement comes back close to worst-case.

The aha: one number that isolates the claims users act on

Comparing full distributions against ground truth everywhere would drown you in errors nobody cares about. The clever move is a screening metric that isolates precisely the claims a risk-sensitive user would act on.

Take the top two actions by mean at a state s. Define the excess Wasserstein gap:

\Delta_{\mathrm{x}}(s) = W_1\big(Z(s,a_1), Z(s,a_2)\big) - \big|\mu_1(s) - \mu_2(s)\big| \geq 0,

where W_1 is the 1-Wasserstein distance between the two return distributions (for quantile representations, just the mean absolute difference between corresponding atoms) and \mu_j are the means. The paper's Proposition 1 gives this quantity a sharp meaning: since W_1(F,G) = \int |F - G|\,dx and the mean gap is |\int (F-G)\,dx|, we have \Delta_{\mathrm{x}} = 0 if and only if F - G never changes sign — i.e., one action first-order stochastically dominates the other. And when \Delta_{\mathrm{x}} > 0, it equals exactly twice the mass of the dominance violation.

The audit metric in one picture. When one action's return CDF sits entirely to one side of the other's (left), every monotone risk functional ranks the actions the same way — there is no risk decision to make. When the CDFs cross (right), the agent is claiming a genuine risk trade-off: one action is safer in the tail, the other better on average. Δx measures exactly the crossing mass, so screening on Δx > 0 finds precisely the states where a risk-sensitive user would deviate from a mean-greedy policy.

This is the audit's conceptual core. \Delta_{\mathrm{x}}(s) > 0 is exactly the agent's claim that the two actions trade off risk — that a CVaR policy and a mean-greedy policy could disagree at s. Errors that both actions' distributions share cancel out; only relative risk claims, the kind people act on, get tested. A distribution can be arbitrarily wrong in ways nobody consumes; this metric doesn't care. It cares only about the states a practitioner consulting the head would flag first.

Getting ground truth without fooling yourself

The audit's mechanics: roll out the trained agent, collect 2,000 states, rank them by \Delta_{\mathrm{x}}, and sample 49 per run stratified across the ranking (7 from the top 0.5%, 7 from 0.5–2%, etc.). For each audited state, restore the environment from a stored snapshot, force each of the top-2 actions, then run 2,000 independent rollouts following the greedy policy. Environment stochasticity (MinAtar's sticky actions, random spawns) is the only noise source — matching what Z^\pi is supposed to capture. Compute the empirical \hat{\Delta}_{\mathrm{x}}^{\mathrm{MC}} from those rollouts and compare.

The refreshingly honest part is the section on two pitfalls that produced convincing but wrong audits of the authors' own:

P1 — RNG cloning. deepcopy of an environment clones its internal random generator, so every "random" restart replays the identical future. The first version of this audit concluded all risk was artifact for this silent, directionally plausible reason. Anyone doing restart-based evaluation should check for this.

P2 — the empirical Wasserstein noise floor. For two n-samples from the same distribution with spread \sigma, the empirical \Delta_{\mathrm{x}} concentrates around c\,\sigma n^{-1/2} > 0, not zero — a classical result (Fournier & Guillin 2015). At the pilot n=200, the floor's 95th percentile was ≈0.22, dwarfing a naive 0.05 threshold; a second wrong audit "confirmed" two-thirds of claims that were pure noise. The fix is per-state permutation nulls (split one action's rollouts in half; any excess between the halves is estimation noise by construction), bootstrap refutation tests, and Benjamini–Hochberg FDR control at 0.10. A claim is confirmed only if the true excess exceeds the permutation null, and refuted only if the learned claim exceeds the 97.5th bootstrap percentile of the MC excess — a conservative bar, since the estimator's finite-sample bias inflates that percentile in the agent's favor.

The findings

Across QR-DQN on three MinAtar games and five seeds (735 audited states):

Three facts stack up:

The strongest claims are the falsest. 66–84% of top-2% claims are refuted; the rate rises to ≈94% in the top half-percent for Breakout and Seaquest. Among low-alias states (where the observation approximately pins down the true state, ruling out partial observability as an excuse), refutation is equal or higher — 80%, 100%, 91%.

Placement is truth-blind. The declining refutation curve down the ranking could be a test-power artifact (small claims can't be refuted). So the paper runs a truth-blind null: permute each run's learned claims across its audited states, preserving magnitudes, randomizing placement relative to ground truth. The null head gets 81% top-strata refutation; the actual heads get 79% (p = 0.22). The agents place their most confident risk claims no better than chance. Pooled learned-vs-true correlation: 0.02, −0.01, −0.04.

Nothing confirms. With power to detect true trade-offs above ≈0.07, pooled FDR-corrected confirmations number zero across all 735 QR-DQN states, all 441 C51 states, all IQN states. MinAtar's aleatoric noise simply doesn't produce action-conditional dominance violations of the size the heads routinely report. The agents aren't exaggerating real trade-offs; they're inventing them.

The most vivid specimens are doomed Breakout states: the ball is irrecoverably lost, every one of 2,000 futures returns the identical value (the true distribution is a point mass), yet the head reports wide, crossing distributions with excess up to 0.44. A fifth of top-stratum Breakout states are this type.

Why the positive controls matter

A zero-confirmation result should make you suspect the audit, not the agents. The paper's answer is riskygrid: a ten-step corridor where each column offers "safe" (deterministic 0.5) or "gamble" (0.5 ± g, fair coin), with per-column gamble scales rendered into the observation. Means are identical, so every g > 0 column is a pure risk trade-off of known magnitude, and zero aliasing by construction. Same architecture, same training recipe, same audit.

On the control, everything inverts: 143/147 = 97% confirmations, zero false refutations in the top strata, learned–true correlation 0.92, and the head's CVaR choice picks the truly safer action 99–100% of the time. The obvious counter — "MinAtar's real trade-offs are just faint" — is preempted by a faint variant with gambles scaled down to g \in \{0.05, \dots, 0.4\}, straddling the audit's own detection floor: still 96–100% confirmed at every g > 0, with g = 0 states at the FDR budget. The same architecture and audit that report near-perfect risk here report zero on MinAtar. The reading measures the agents.

Does it matter for decisions?

Yes, and in the worst possible way: inconsistently. At the heads' own most-flagged states, comparing the head's CVaR ranking against a risk-blind mean-greedy baseline:

In Breakout the head is genuinely informative. In Seaquest it is anti-predictive — significantly worse than a coin flip at exactly the states it flags as most risk-critical. In Asterix it's chance. And a probe over 13 head-shape features finds no signal that tells you which regime you're in. Advice that's sometimes right and sometimes inverted, with no tell, is arguably worse than uniformly useless advice.

Structural, and not repairable by the obvious fixes

The artifact isn't a symptom of weak agents. It's fully formed by 500k steps and unchanged at 3M while score doubles; across 33 runs, score and top-strata refutation are uncorrelated (Spearman −0.07); sibling seeds reproduce a median 3.8% of each other's claim magnitudes — each seed invents its own spurious risk map. It replicates under C51's categorical projection and IQN's implicit quantiles, which implicates the shared ingredient: bootstrapped distributional targets from a changing policy under function approximation, not any particular representation. And a pretrained 10M-frame SB3-zoo QR-DQN (Breakout score ≈400) has all 14 of its top-strata claims refuted — mean claimed excess 0.184 against ground-truth 0.000.

The repairs section is where the paper earns its strongest conclusion. Training a CVaR-greedy agent doesn't help (88–98% top-strata refutation). Five-member ensembling attenuates but doesn't calibrate (recovers 8 Breakout confirmations; Seaquest and Asterix stay at zero). Most tellingly, the best possible monotone recalibration — an isotonic map from learned to true excess, fit on held-out data — rescues the positive-control head completely but passes MinAtar only by collapsing all claims to 12–22% of their magnitude and asserting nothing. The distinction matters: a miscalibrated head has signal at the wrong scale and is fixable post hoc; an uninformative head has no signal for any correction to work with. A distillation control confirms the architecture can fit true distributions when directly supervised — the information is missing from what TD learning produces, not from what the network can express.

What to make of it, and what to be skeptical about

If this holds, the practical implication is blunt: face-value reads of distributional heads are least reliable exactly where they'd be consulted — the top of the head's own risk ranking is a winner's-curse selection over noise. Interpretability work that visualizes return distributions needs ground-truth validation, just as saliency maps did after Atrey et al. (2020). Risk-sensitive policies built on raw quantile heads (Keramati et al. 2020; Lim & Malik 2022) are acting on spurious risk at critical states. And the growing family of calibration wrappers now has a concrete standard: an audit-passing head confirms its top claims. None tested here does.

Caveats worth holding onto. The environments are mostly MinAtar — small, and its aleatoric structure (sticky actions, spawns) may genuinely produce few large action-conditional trade-offs, so "the head invents risk" and "there was little risk to learn" are entangled; the positive controls address the audit's power but not whether richer environments would give the heads more to latch onto. The full-ALE result has an acknowledged asterisk: the zoo agents were trained without sticky actions, so their claimed spread couldn't reflect environment stochasticity by construction — that audit quantifies the artifact's size, not an independent replication; an ALE agent trained under sticky actions is the missing test at scale. The audit requires snapshot-restart access, which limits it to simulators. And newer calibrated variants (evidential, conformal) weren't tested — though the "uninformative, not miscalibrated" result suggests post-hoc wrappers around the same head face a hard ceiling.

Where to spend your reading time: the Pitfalls section (P1/P2) is short and valuable to anyone doing restart-based evaluation of anything, not just distributional RL. Then the positive-control subsection ("The Audit Passes Its Positive Control") — the selection note for this paper was right that this is the load-bearing evidence — and the "Miscalibrated or uninformative?" analysis, which is the sharpest single result. Figure 6's doomed-state specimens are the picture worth a thousand refutation percentages: a point-mass ground truth against a confidently wide, crossing pair of learned distributions.